Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Won, Jaeyeon, Ahrens, Willow, Emer, Joel S., Amarasinghe, Saman
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917027949576192
author Won, Jaeyeon
Ahrens, Willow
Emer, Joel S.
Amarasinghe, Saman
author_facet Won, Jaeyeon
Ahrens, Willow
Emer, Joel S.
Amarasinghe, Saman
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2510_17505
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums
Won, Jaeyeon
Ahrens, Willow
Emer, Joel S.
Amarasinghe, Saman
Programming Languages
Performance
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.
title Insum: Sparse GPU Kernels Simplified and Optimized with Indirect Einsums
topic Programming Languages
Performance
url https://arxiv.org/abs/2510.17505