Towards Efficient Scheduling of Transformer Neural Network Computation for Edge AI Deployment

transformer networks, edge AI, evolutionary algorithms

Transformer neural networks have gained popularity in recent years, demonstrating remarkable performance across many application domains. However, inference on resource-constrained embedded hardware remains challenging due to Transformers' substantial computational demands. We aim to address this problem by focusing on exploiting the inherent parallelism opportunities presented by the multi-head self attention operations of Transformers, to achieve a speedup in processing on embedded hardware. In this paper, we present an evolutionary-based scheduling approach for distribution and allocation of Transformer operations across systolic array-based hardware accelerators used for execution. Our methodology takes as input specifications of the Transformer workload and the target systolic array architecture and explores the large mapping space to identify an efficient plan of operation-to-array assignments. The plans are evaluated against a hardware-aware cost model, capturing the cost of computational cycles for a given operation and systolic array, with the objective to minimize the total sum across all operations. Through extensive experimental evaluations across diverse systolic array dimensions, we demonstrate that our evolutionary-based scheduler surpasses conventional heuristics and is able to find plans offering up to 33.8% average reduction in overall cycle count.