Ceiba: An Efficient and Scalable DNN Scheduler for Spatial Accelerators

Spatial accelerators are domain-specific architectures to elevate performance and energy efficiency for deep neural networks (DNNs). They also bring a large number of schedule parameters to determine computation and data movement patterns of DNNs. Previous works formulate the schedule problem as design space exploration or integer linear programming. However, these advanced techniques face the challenge of efficiency or scalability. In this article, we propose Ceiba, which is a deep reinforcement learning-based DNN scheduler for spatial accelerators. Ceiba observes the running DNN computation as well as the spatial architecture to make schedule decisions. Then, Ceiba receives a reward to learn and produce the best-fit policy. To provide efficient and scalable scheduling, Ceiba constructs a DNN-architecture-specific action space. It is defined by upper and lower bounds to exclude invalid and sub-optimal schedule candidates. Extensive experiments demonstrate that Ceiba generally provides better performance for spatial accelerators under a fixed number of searching steps or a fixed amount of time. Specifically, Ceiba achieves an average 2.2× speedup for the Simba accelerator, compared with the state-of-the-art scheduler. When scaling the batch size and the hardware architecture up by 64×, the performance gains of Ceiba are 1.8× and 1.2× on average, respectively. Moreover, Ceiba exhibits better scalability for the Eyeriss accelerator.