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

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.