Title: Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums Show Chinese title

Title: Insum：通过间接Einsum简化和优化的稀疏GPU内核

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.14x to 3.81x speedups across a range of sparse GPU applications while reducing lines of code by 202x to 4491x compared to hand-written implementations. Show Chinese abstract

Abstract: 编程高性能的稀疏GPU内核以著称的困难，需要大量的努力和深厚的专长。 稀疏编译器旨在简化这一过程，但现有的系统在两个关键方面存在不足。 首先，它们主要为CPU设计，很少能生成高性能的GPU代码。 其次，当计算涉及稀疏和密集区域时，这些编译器往往无法有效地优化密集部分。 在本文中，我们提出了一种表达稀疏计算的新方法。 我们从与格式无关的稀疏张量上的Einsum开始，并将其重写为与格式相关的间接Einsum，通过间接索引将稀疏数据和元数据映射到密集张量操作，从而显式编码格式信息。 为了执行间接Einsum，我们引入了Insum编译器，它通过降低到PyTorch编译器来为这些Einsum生成高效的GPU代码，该编译器被扩展以更好地支持Tensor Core启用的间接Einsum。 我们还介绍了两种固定长度的稀疏格式，GroupCOO和BlockGroupCOO，它们被设计为自然地与间接Einsum配合使用。 我们的方法在一系列稀疏GPU应用中实现了1.14倍到3.81倍的速度提升，同时相比手工编写的实现，代码行数减少了202倍到4491倍。