UniSparse: An Intermediate Language for General Sparse Format Customization
This work addresses the problem of inefficient sparse workload processing due to rigid format support in compilers, which is crucial for developers and researchers in high-performance computing and machine learning, though it is incremental as it builds on existing MLIR infrastructure.
The paper tackles the lack of support for customizing sparse tensor formats in existing programming models and compilers by proposing UniSparse, an intermediate language that decouples logical representation from memory layout, enabling flexible format customization and automatic code generation for heterogeneous hardware, resulting in demonstrated efficacy across multiple hardware targets including CPU, GPU, FPGA, and PIM devices.
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.