Fast and Fusiest: An Optimal Fusion-Aware Mapper for Accelerator Modeling and Evaluation

TL;DR

The paper tackles the problem of optimally mapping fused tensor-algebra workloads onto accelerators, where fusion reduces off-chip accesses but expands the search space exponentially with the number of Einsums. It introduces the Fast and Fusiest Mapper (FFM), which builds full fused mappings one Einsum at a time by pruning suboptimal partial mappings using compatibility and reservation criteria, guaranteeing optimality. Key innovations include a pmapping taxonomy (objective, compatibility, reservation), grouping for efficient skipping of incompatible joins, and reservation-based pruning that preserves optimality while scaling near-linearly with Einsums; empirically, FFM achieves over 1000× speedups over prior baselines and substantially improves Transformer-related mappings. The approach enables rapid, accurate accelerator modeling and evaluation, with practical impact for designing and evaluating fused-tensor accelerators in real workloads like GPT-3-scale transformers.

Abstract

The latency and energy of tensor algebra accelerators depend on how data movement and operations are scheduled (i.e., mapped) onto accelerators, so determining the potential of an accelerator architecture requires both a performance model and a mapper to search for the optimal mapping. A key optimization that the mapper must explore is fusion, meaning holding data on-chip between computation steps, which has been shown to reduce energy and latency by reducing DRAM accesses. However, prior mappers cannot find optimal mappings with fusion (i.e., fused mappings) in a feasible runtime because the number of fused mappings to search increases exponentially with the number of workload computation steps. In this paper, we introduce the Fast and Fusiest Mapper (FFM), the first mapper to quickly find optimal mappings in a comprehensive fused mapspace for tensor algebra workloads. FFM shrinks the search space by pruning subsets of mappings (i.e., partial mappings) that are shown to never be a part of optimal mappings, quickly eliminating all suboptimal mappings with those partial mappings as subsets. Then FFM joins partial mappings to construct optimal fused mappings. We evaluate FFM and show that, although the mapspace size grows exponentially with the number of computation steps, FFM's runtime scales approximately linearly. FFM is orders of magnitude faster ($>1000\times$) than prior state-of-the-art approaches at finding optimal mappings for Transformers.

Figures (11)

• Figure 1: An analogy of pmappings as puzzle pieces. (a) Tabs and sockets represent compatibility. Numbers represent a cost metric. (b) Pieces can be joined in several ways. (c) Pieces can be grouped by tabs/sockets and (d) pruned per-group.
• Figure 2: (a) Pseudocode of a mapping for two matrix-vector multiplications and (b) the equivalent LoopTree mapping.
• Figure 3: A two-step example of joining pmappings with puzzle analogy. Only storage nodes for shared tensors matter to joining compatibility, so other storage nodes are omitted for clarity. (a) Joining pmappings for Einsums $A$ and $B$. (b) Joining the pmapping for Einsums $A,B$ and the pmapping for Einsum $C$. The storage node of shared tensor (tensor $A$ between Einsums $A$ and $B$, and tensor $B$ between Einsums $B$ and $C$) and loops above the storage nodes must match, analogous to tabs/sockets in the puzzle.
• Figure 4: An example constrained optimization and (b) example puzzle pieces with two criteria (objective, reservation) to consider. Two groups of puzzle pieces, group G0 and group G1, are highlighted for discussion.
• Figure 5: (a) A mapping for two Einsums onto an accelerator with DRAM and GLB memory levels and (b) pmappings separated from the full mapping.