TileMap: Mapping Multi-Head Attention on Spatial Accelerators with Tile-based Analysis

Multi-head attention demonstrates enormous potential but incurs high costs (e.g., memory access) when deploying transformer models on spatial accelerators. Operator fusion is prevalent in conventional deep learning mappers to reduce off-chip memory access. However, designing operator-fusion mapping for multi-head attention is challenging. There exist strict data dependency between operators and hardware resource constraints of accelerators. In this paper, we propose TileMap, an operator-fusion mapping framework for multi-head attention on spatial accelerators. Central to this framework is tile-based analysis that can automatically satisfy both data dependency and resource constraints. Based on this analysis, we construct a mapping design space, which significantly prunes invalid candidates. We then propose an RL-based searching algorithm to explore the mapping space. The RL agent regards the mapping space as action space and further optimizes it to preserve all constraints. Experiments show TileMap achieves <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$1.8\times$</tex> to <tex xmlns:mml="http://www.w3.org/1998/Math/MathML" xmlns:xlink="http://www.w3.org/1999/xlink">$3.1\times$</tex> speedup on different spatial accelerators, relative to state-of-the-art mapping approaches.