  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.
