A Fine-grained Workload Orchestration Framework for Processing-in-Memory Architecture

The Processing-in-Memory (PIM) architecture provides a promising option for DNN accelerators with high performance and energy efficiency. While computation and memory access have been the primary focus, its communication performance has not been fully optimized. This work proposes a fine-grained workload orchestration framework that enables computation-communication overlap to enhance the inference performance of DNN workloads on PIM. Firstly, we introduce a simulated annealing-based partitioning algorithm to convert the layer-level workload graph into a finer-grained operation-level graph. Second, we present an operation-level list scheduling algorithm to co-schedule computation and communication tasks. Finally, we develop an integer linear programming algorithm to map the scheduled tasks onto the PIM system. The experimental results reveal that we achieve a 34% and 17% improvement in inference performance compared to the baseline methods.