UniSparse: An Intermediate Language for General Sparse Format Customization

TL;DR

UniSparse is proposed, an intermediate language that provides a unified abstraction for representing and customizing sparse formats and develops a compiler leveraging the MLIR infrastructure, which supports adaptive customization of formats and automatic code generation of format conversion and compute operations for heterogeneous architectures.

Abstract

The ongoing trend of hardware specialization has led to a growing use of custom data formats when processing sparse workloads, which are typically memory-bound. These formats facilitate optimized software/hardware implementations by utilizing sparsity pattern- or target-aware data structures and layouts to enhance memory access latency and bandwidth utilization. However, existing sparse tensor programming models and compilers offer little or no support for productively customizing the sparse formats. Additionally, because these frameworks represent formats using a limited set of per-dimension attributes, they lack the flexibility to accommodate numerous new variations of custom sparse data structures and layouts. To overcome this deficiency, we propose UniSparse, an intermediate language that provides a unified abstraction for representing and customizing sparse formats. Unlike the existing attribute-based frameworks, UniSparse decouples the logical representation of the sparse tensor (i.e., the data structure) from its low-level memory layout, enabling the customization of both. As a result, a rich set of format customizations can be succinctly expressed in a small set of well-defined query, mutation, and layout primitives. We also develop a compiler leveraging the MLIR infrastructure, which supports adaptive customization of formats, and automatic code generation of format conversion and compute operations for heterogeneous architectures. We demonstrate the efficacy of our approach through experiments running commonly-used sparse linear algebra operations with specialized formats on multiple different hardware targets, including an Intel CPU, an NVIDIA GPU, an AMD Xilinx FPGA, and a simulated processing-in-memory (PIM) device.

Figures (14)

• Figure 1: Different formats of a sparse matrix ($A$) --- Yellow-shaded blocks refer to the data structure of the same tensor element. The size, ptr, idx, and val labels on the left indicate the size of a dimension, a pointer array, an index array, or a value array. The index expressions on the right denote the index map at each level. Horizontal dotted lines separate different tensor dimensions, while vertical dotted lines divide the matrix $A$ into sub-matrices.
• Figure 2: An overview of UniSparse.
• Figure 3: An illustration of UniSparse.
• Figure 4: The metadata trees of the matrix $A$ in Figure \ref{['matrix']} --- Yellow-shaded blocks refer to the data structure of the same tensor element. The root of a metadata tree is the symbol of the tensor. $A_0$ and $A_1$ are sub-matrices of the matrix $A$. The circular nodes are metadata nodes that contain indices, and the arrows represent pointers from a higher major dimension to a lower one. Index expressions are on the left of each coordinate level. Nodes with dashed lines are elements padded for memory alignment.
• Figure 5: The syntax for index maps and query primitives. denotes dimension identifiers, denotes tensor element values, and denotes constant numbers. Expressions enclosed in curly braces can be repeated zero or more times.