TensorMap: A Deep RL-Based Tensor Mapping Framework for Spatial Accelerators
Fuyu Wang, Minghua Shen, Yutong Lu, Nong Xiao · IEEE Transactions on Computers · 2024
The mapping of tensor computation is a complex and important process for spatial accelerators. Today's mapping works depend on hand-tuned kernel libraries or search-based heuristics from human experts. The former is time-intensive while the latter easily leads to sub-optimal performance. In this paper, we propose TensorMap, a deep reinforcement learning (RL)-based mapping framework for tensor computations on spatial accelerators. We propose a sequential generation mode for mapping optimization and construct a coarse-grained action space to reduce the complexity of the mapping search space. An efficient policy network is devised to optimize mapping primitives in the RL-based search. We then propose a stop signal that is sampled fromBernoullidistribution to facilitate multi-level loop unrolling for spatial accelerators. Finally, a genetic algorithm is employed to further refine the optimized mappings. In the experiments, we demonstrate TensorMap's ability for different spatial accelerators with various tensor computations. On TPU, TensorMap provides 2.6$\times$, 2.7$\times$, and 2.4$\times$better energy-delay product (EDP) on average compared with FlexTensor, Ansor, and AMOS respectively. On Eyeriss, TensorMap provides 2.1$\times$, 1.8$\times$, and 1.7$\times$better EDP on average compared with FlexTensor, Ansor, and AMOS respectively.