The inevitable trend of hardware specialization drives an increasing use of custom data formats in processing sparse workloads, which are typically memory-bound. These formats facilitate the automated generation of target-aware data layouts to improve 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. Furthermore, these frameworks adopt an attribute-based approach for format abstraction, which makes it difficult to identify disparate memory layout schemes and support general format customization.

To overcome this deficiency, we propose UniSparse, an intermediate language that incorporates a unified abstraction to represent general sparse formats, and a compiler that automates format customization and compute kernel generation. UniSparse expresses a sparse format as a map from dense coordinates to a coordinate hierarchy tree using a small set of well-defined mutation primitives. By traversing the abstract tree representation in a specific order, the UniSparse compiler can generate a precise data layout and the corresponding compute kernels downstream. Our approach enables adaptive customization of formats and automatic code generation of format conversion and compute operations for heterogeneous architectures using the MLIR infrastructure.

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, and a simulated processing-in-memory (PIM) device.