SparseMap: A Sparse Tensor Accelerator Framework Based on Evolution Strategy

Abstract

The growing demand for sparse tensor algebra (SpTA) in machine learning and big data has driven the development of various sparse tensor accelerators (SpTAs). However, most existing manually designed accelerators are limited to specific scenarios, and it is time-consuming and challenging to adjust a large number of design factors when scenarios change. Therefore, automating the design of SpTA accelerators is crucial. Nevertheless, previous works focus solely on either mapping (i.e., tiling communication and computation in space and time) or sparse strategy (i.e., bypassing zero elements for efficiency), leading to suboptimal designs due to the lack of consideration of both. A unified framework that jointly optimizes both is urgently needed. However, integrating mapping and sparse strategies leads to a combinatorial explosion in the design space(e.g., as large as $O(10^{41})$ for the workload $P_{32 \times 64} \times Q_{64 \times 48} = Z_{32 \times 48}$ ). This vast search space renders most conventional optimization methods [e.g., particle swarm optimization (PSO), reinforcement learning (RL), and Monte Carlo tree search (MCTS)] inefficient. To address this challenge, we propose an evolution strategy (ES)-based SpTA optimization framework, called SparseMap. SparseMap constructing a more essential-factor design space with the consideration of both mapping and sparse strategy. We introduce a series of enhancements to genetic encoding and evolutionary operators, enabling SparseMap to efficiently explore the vast and diverse design space. We quantitatively compare SparseMap with prior works and classical optimization methods, demonstrating that SparseMap consistently finds superior solutions.