paperbot · PL 论文追踪

全部论文 待读 收藏

P2 Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums

Insum:用间接 Einsum 简化并优化稀疏 GPU 内核
ASPLOS Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 22026匹配度 2/10 · 该论文属于编译与代码生成方向,面向稀疏GPU内核的Einsum改写与优化,与用户关注的静态分析、LLM/Agent、浮点/贝叶斯程序分析没有直接交集;仅在‘机器学习系统’这一宽泛层面相关,且不涉及程序分析或Agentic程序。引用 2 inbox
Jaeyeon Won, Willow Ahrens, Saman Amarasinghe, Joel S. Emer
crossref: 2026-09-12dblp: 2026-09-12arxiv: 2026-09-12

AI 速览

一句话:以格式无关 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 倍,源码已公开。

问题
现有稀疏编译器主要面向 CPU,难以生成高性能 GPU 代码,且在稀疏-稠密混合计算中无法有效优化稠密部分,手写稀疏 GPU 内核代价极高。
方法
将格式无关的稀疏张量 Einsum 重写为格式感知的间接 Einsum,把稀疏数据与元数据映射为间接索引下的稠密张量操作;实现 Insum 编译器并下沉到扩展后的 PyTorch 编译器以支持 Tensor Core,同时提出定长格式 GroupCOO 与 BlockGroupCOO。
评估
在多种稀疏 GPU 应用上与手写实现对比,测得多 1.14–3.81 倍加速,代码行数减少 202–4491 倍。
原文摘要(Abstract)

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

BibTeX
@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}
}