一句话:提出 UniSparse 中间语言,解耦稀疏张量的逻辑结构与内存布局,支持通用格式定制及异构代码生成。
稀疏负载常使用定制数据格式以提升内存性能,但现有稀疏张量编程模型与编译器对格式定制的支持有限,且基于逐维属性的表示难以覆盖多样的自定义结构与布局。本文提出 UniSparse,一种为通用稀疏格式定制设计的中间语言。它将稀疏张量的逻辑表示(数据结构)与底层内存布局解耦,使用一组查询、修改和布局原语简洁表达丰富的格式定制。作者基于 MLIR 实现编译器,支持自适应格式定制,并面向异构架构自动生成格式转换与计算代码。在 Intel CPU、NVIDIA GPU、AMD Xilinx FPGA 及模拟 PIM 设备上,对常用稀疏线性代数运算使用定制格式进行实验,验证了方法的有效性。
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.
@article{paperbot2542,
title = {UniSparse: An Intermediate Language for General Sparse Format Customization},
author = {Jie Liu and Zhongyuan Zhao and Zijian Ding and Benjamin Brock and Hongbo Rong and Zhiru Zhang},
journal = {Proceedings of the ACM on Programming Languages},
volume = {8},
number = {OOPSLA1},
year = {2024},
doi = {10.1145/3649816}
}