一句话:以格式无关 Einsum 重写为间接 Einsum,配合 Insum 编译器生成高效稀疏 GPU 内核。
编写高性能稀疏 GPU 内核难度大、需要深厚专业知识,而现有稀疏编译器主要面向 CPU,难以生成高性能 GPU 代码,且在稀疏与稠密混合计算中往往无法有效优化稠密部分。本文提出新的稀疏计算表达方式:从稀疏张量上的格式无关 Einsum 出发,将其重写为格式感知的间接 Einsum,通过间接索引把稀疏数据与元数据映射为稠密张量操作,从而显式编码格式信息。为执行间接 Einsum,作者实现了 Insum 编译器,将其下沉到 PyTorch 编译器并扩展以更好支持 Tensor Core 上的间接 Einsum;同时提出两种定长稀疏格式 GroupCOO 与 BlockGroupCOO,与间接 Einsum 自然契合。实验表明,在多种稀疏 GPU 应用上取得 1.14–3.81 倍加速,相比手写实现代码量减少 202–4491 倍,源码已公开。
Programming high-performance sparse GPU kernels is notoriously difficult, requiring both substantial effort and deep expertise. Sparse compilers aim to simplify this process, but existing systems fall short in two key ways. First, they are primarily designed for CPUs and rarely produce high-performance GPU code. Second, when computations involve both sparse and dense regions, these compilers often fail to optimize the dense portions effectively. In this paper, we propose a new approach for expressing sparse computations. We start from format-agnostic Einsums over sparse tensors and rewrite them into format-conscious indirect Einsums, which explicitly encode format information by mapping sparse data and metadata onto dense tensor operations through indirect indexing. To execute indirect Einsums, we introduce the Insum compiler, which generates efficient GPU code for these Einsums by lowering to the PyTorch compiler, extended to better support Tensor Core–enabled indirect Einsums. We also present two fixed-length sparse formats, GroupCOO and BlockGroupCOO, designed to fit naturally with indirect Einsums. Our approach achieves 1.14×–3.81× speedups across a range of sparse GPU applications while reducing lines of code by 202×–4491× compared to hand-written implementations. The source code for Insum is publicly available at https://github.com/nullplay/IndirectEinsum.
DOI 原文 · arXiv · PDF(开放获取) · DBLP
@inproceedings{WonAAE26,
title = {Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums},
author = {Jaeyeon Won and Willow Ahrens and Saman Amarasinghe and Joel S. Emer},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790176}
}