SAPar: A Surrogate-Assisted DNN Partitioner for Efficient Inferences on Edge TPU Pipelines

Pipelining deep neural networks (DNNs) across multiple Edge Tensor Processing Units (TPUs) can enhance on-device performance by increasing the capacity for DNN parameters caching and enabling pipeline parallelism. Effective deployment on pipelined Edge TPUs requires a partitioning tool to divide the DNN into segments, each assigned to a different Edge TPU in the pipeline. Achieving balanced workload distribution across these segments is crucial for optimal timing performance. However, workload balancing across Edge TPUs is challenging, as DNN execution time is influenced by proprietary hardware architecture and compiler internals, forming a black-box function inaccessible to partitioning tools. To address this challenge, this article introduces SAPar , a new surrogate-assisted DNN partitioner that integrates a neighborhood search engine with a surrogate-assisted evaluator for effective and efficient DNN partitioning. The neighborhood search engine systematically explores the decision space, guided by knowledge obtained from empirical insights and neighborhood evaluation feedback provided by the surrogate-assisted evaluator. The evaluator cooperatively applies an accurate yet time-consuming latency profiler and an efficient graph transformer-based surrogate model , achieving both precision and scalability. Experiments on real Edge TPU hardware demonstrate that SAPar achieves significantly better pipeline performance than Google’s current profiling-based partitioner with an 8.82× to 110× speedup in partitioning time. Moreover, SAPar reduces the bottleneck latency by 8.93% to 44.15% across five classic DNN models compared with a state-of-the-art reinforcement learning-based partitioner.