SparseTIR: Composable Abstractions for Sparse Compilation in Deep Learning

1 INTRODUCTION

• Vendor libraries support only a few sparse operators, while manual optimization is tedious and format-specific; sparse tensor compilers also face diverse workloads and evolving heterogeneous backends.The paper identifies single-format support and single-shot compilation as challenges for applying sparse compilation to deep learning.Source: §1 · p. 1
• Composable formats store different parts of a sparse matrix in formats suited to their local patterns, with computations decomposed into routines for those patterns.This approach targets the varied sparsity patterns found in deep learning workloads.Source: §1 · pp. 1–2
• Composable transformations combine sparse-data-specific rewrites with hardware-specific optimizations such as tensorization and GPU mapping.This design is intended to incorporate hardware-backend advances more efficiently.Source: §1 · p. 2
• SparseTIR is an abstraction built from composable formats and transformations, represented in an IR and used to generate efficient sparse operators for deep learning.Its contributions include an IR for formats and schedules, a tuning system that searches their parameter space, and evaluation on sparse workloads.Source: §1 · p. 2
• On GPUs, SparseTIR achieves 1.20–2.34x speedups over vendor libraries for GNN operators and 1.05–2.98x for sparse transformer operators, and reports 0.56–7.45x for Sparse Convolution operators.It accelerates GraphSAGE training by 1.08–1.52x and RGCN inference by 4.20–40.18x.Source: §1 · p. 2

2 SYSTEM OVERVIEW

• Earlier sparse compilers accept one format per matrix and apply user schedules to high-level IR before lowering to device code, which does not expose loop-level details needed for some optimizations.The passage notes tensorization requires loop-level AST matching and replacement.Source: §2 · p. 2
• SparseTIR represents computations in coordinate space, then position space, then loop-level IR, using sparse iteration and sparse buffer lowering to move between stages.Position space indexes nonzero elements in compressed sparse data.Source: §2 · p. 2
• Schedules are composable transformations applied step by step without changing the IR stage; SparseTIR supports stage-I schedules and backend schedules at compatible later stages.Format decomposition can also be expressed as a stage-I transformation.Source: §2 · p. 2
• Composable formats and transformations are presented as enabling optimizations for a broader range of workloads and hardware-backend advances than single-shot compilers.The design uses composable components to extend the scope of sparse compilation.Source: Figure 2 · p. 3
• SparseTIR searches a joint space of composable formats and transformations; compile-time search can select formats and schedules for a fixed sparse structure, with overhead amortized across reuse.Users specify parameterized format and schedule templates, and the large search space may make compilation take time.Source: §2 · p. 3

3.1 Language Constructs

• Axes and sparse buffers separate sparse-layout metadata from values, while sparse iterations express computation over those axes, including nested iterations and non-index expressions.Axes encode dense/sparse and fixed/variable properties; buffers sharing a layout can reuse auxiliary data.Source: §3.1 · pp. 3–4

3.2 Stage I: Coordinate Space Computation

• Stage I expresses computation in coordinate space and supports format decomposition, which creates computation and data-copy iterations for new formats such as BSR and ELL.For stationary sparse matrices, conversion can occur during preprocessing; users provide the information needed to create the new sparse buffers.Source: §3.2.1 · p. 4, §3.2 · p. 4

3.2.2 Stage I Schedules.

• Stage I schedules reorder sparse axes or fuse iterators, changing generated loop order or combining iteration into one loop.Fusion is useful when one loop over non-zero elements is preferred, as in SDDMM.Source: §3.2.2 · p. 4

3.3 Stage II: Position Space Computation

• Stage II restructures sparse iterations as nested loops and accesses sparse buffers in position space, where positions index non-zero elements.Stage II extends TensorIR and supports loop, memory, parallelization, vectorization, and tensorization schedules.Source: §3.3 · p. 5, §3.3.2 · p. 6

3.3.1 Sparse Iteration Lowering.

• Stage II lowering materializes index buffers, emits loops over axes, and translates coordinate accesses to positions using compression and decompression, with binary search for sorted indices when needed.The lowering pass also gathers buffer read/write regions for block attributes.Source: §3.3.1 · pp. 5–6

3.4 Stage III: Loop-Level IR

• Stage III removes SparseTIR constructs and retains nested loops over flattened buffers in a loop-level IR compatible with tensorization.The implementation selects TensorIR in Apache TVM to use NVIDIA Tensor Cores.Source: §3.4 · p. 6

3.4.1 Sparse Buffer Lowering.

• Sparse buffer lowering flattens multidimensional buffers and rewrites accesses using computed offsets.For example, A[i, j] becomes A[J_indptr[i] + j].Source: §3.4.1 · p. 6

3.5 Target-Specific Code Generation

• SparseTIR reuses existing tensor-compiler backends for target-specific code generation and adds horizontal fusion to reduce overhead from multiple CUDA kernel launches.The fused kernels arise from composable formats.Source: §3.5 · p. 6

4 EVALUATION

• Against vendor libraries, SparseTIR achieves 1.20–2.34x speedup on GNN operators, 1.05–2.98x on sparse attention, and 0.56–7.44x on Sparse Convolution; end-to-end GraphSAGE training speeds up 1.08–1.52x and end-to-end RGCN inference 4.20–40.18x.Source: §4 · p. 7

4.1 Experiment Setup

• Experiments use RTX 3070 and Tesla V100 and compare against task-specific baselines including cuSPARSE, dgSPARSE, DGL, Triton, and TorchSparse.The reported SparseTIR kernels were checked against existing frameworks and libraries for numerical accuracy.Source: §4.1 · p. 7

4.2.1 SpMM.

• 1.22-2.34x on V100 and 1.20-1.91x on RTX 3070 versus cuSPARSE; hybrid decomposition helps power-law graphs, while padding can offset gains for centralized degrees.Column partitioning improves cache hit rates, but added result-matrix updates make its benefit saturate as partitions increase.Source: §4.2.1 · p. 8, Figure 12 · p. 8

4.2.3 End-to-end GraphSAGE training.

• Parameterized schedules generally outperform dgSPARSE on SDDMM; GraphSAGE training reaches 1.18-1.52x on V100 and 1.08-1.47x on RTX 3070 versus DGL.The SDDMM evaluation uses feature sizes 32, 64, 128, 256, and 512 and does not use composable formats.Source: §4.2.1 · p. 8, §4.2.3 · p. 9

4.3.1 Sparse Attention.

• 1.05-1.59x on multi-head SpMM and 1.50-2.98x on multi-head SDDMM versus Triton for the evaluated BSR sparse-attention operators.The evaluation fixes matrix size at 4096 × 4096, batch size at 12, band size at 256, and feature size per head at 64.Source: §4.3.1 · p. 9, Figure 16 · p. 10

4.3.2 Sparse Weight (Network Pruning).

• For block-pruned transformers, SparseTIR's doubly compressed BSR format skips all-zero rows and consistently outperforms its BSR-format kernel.The benchmark uses block size 32, average weight sparsity 93%, batch size 1, and sequence length 512.Source: §4.3.2 · p. 10, Figure 17 · p. 10
• For unstructured pruning, SR-BCRS beats BSR in most settings except density ≥2^-3; cuSPARSE CSRMM beats cuBLAS GeMM only at density ≤2^-6.The benchmark uses SR-BCRS(8, 32), average weight sparsity 94%, batch size 1, and sequence length 512.Source: §4.3.2 · p. 10, Figure 19 · p. 11

4.4.1 Relational Graph Convolution Network.

• SparseTIR's fused RGMS reaches 4.2-40.2x versus Graphiler for RGCN inference; sparse convolution beats TorchSparse for most operators but not at channel sizes > 128.The RGCN fusion consumes the intermediate in SRAM instead of explicitly storing it in HBM; for sparse convolution, large-channel matrix-multiplication overhead becomes dominant.Source: §4.4.1 · p. 11, §4.4.2 · p. 11, Figure 20 · p. 12

Tensor and deep learning compilers

• Dense tensor and deep-learning compilers separate kernel descriptions from schedules or optimize computation graphs, but provide limited support for sparse operators.TensorIR, Triton, and FreeTensor offer abstractions that could serve as SparseTIR's stage-III IR.Source: §5 · p. 12

Sparse compilers

• Sparse compilers cover format abstractions, iteration transformations, inspectors, and sparse IRs, while COMET and MLIR Sparse Dialect do not consider decomposable formats.SparseTIR also relates CoRA's operation splitting to a special case of format decomposition.Source: §5 · pp. 12–13

GNN systems and compilers

• GNN frameworks use vendor libraries and handwritten operators, while template-based GNN compilers have limited expressiveness across optimization choices.SparseTIR can describe FusedMM and could serve as a backend for Seastar and Graphiler; the cited GNN contributions are described as orthogonal.Source: §5 · p. 13

Sparse kernel optimizations

• GPU SpMM research explores diverse schedules that SparseTIR seeks to unify, while OSKI autotunes sparse operators but does not support customizing them.The cited GPU schedules include Merge-SpMM, ASpT, GE-SpMM, Sputnik, and DA-SpMM.Source: §5 · p. 13

Sparse format and hardware-efficient optimizations

• Sparse matrix reordering can improve block granularity or cache locality, and hardware-oriented sparsity methods target GPU acceleration units or load balancing.The paper describes these reordering algorithms as possible preprocessing steps for discovering composable formats.Source: §5 · p. 13

Automatic scheduling

• SparseTIR still requires users to specify schedule templates, and searching for an optimal schedule is time consuming.The authors identify automatic schedule-template generation and cost models as possible ways to help.Source: §6 · p. 13

Automatic format decomposition

• Format decomposition rules are manually designed; automatic format selection and decomposition are left for future work.Source: §6 · p. 13

Dynamic Sparsity

• When nonzero positions change over time, searching for a best schedule for each matrix becomes impractical.The paper points to shape-generic search spaces and lightweight runtime dispatch as potentially applicable approaches.Source: §6 · p. 13

Integration with graph-level IR

• SparseTIR models tensor-level sparsity only, with graph-level IR integration planned as an extension.Source: §6 · p. 13

7 CONCLUSION

• SparseTIR combines composable formats and transformations as a parameter search space for tuning sparse operators, with evaluations reporting improvements over existing libraries and frameworks.Source: §7 · p. 13

A PROGRAMMING INTERFACE FOR COMPOSABLE FORMATS

• Format-rewrite rules specify a new format, axis mappings, and affine index mappings, while decompose_format applies a list of rules to a SparseTIR program.Format conversion is a special case of decomposition using a single rewrite rule.Source: A PROGRAMMING INTERFACE FOR COMPOSABLE… · p. 14
• The interface example composes BSR(2) and ELL(2) rewrite rules to decompose an SpMM computation.Source: A PROGRAMMING INTERFACE FOR COMPOSABLE… · p. 14

B.1 Abstract

• The artifact provides experiment scripts and dependencies and requires an x86_64 host with an NVIDIA GPU of Turing architecture or later.A Docker image is provided for running benchmarks.Source: §B.1 · p. 14

B.2 Artifact check-list (meta-information) •

• The checklist covers OGB, SemanticKITTI, and DGL datasets; Turing-, Ampere-, or Hopper-architecture NVIDIA GPUs; and execution-time and GPU-memory metrics.Listed workloads include SpMM, SDDMM, GraphSAGE training, sparse transformers, 3D sparse convolution, and RGCN inference.Source: §B.2 · pp. 14–15

B.3 Description

• The artifact is available on GitHub and Zenodo, with installation and benchmark scripts, and the SparseTIR compiler is included as a submodule.Source: §B.3.1 · p. 15

B.3.2 Hardware dependencies.

• RTX 3070 and Tesla V100 systems were used; at least 16 GB of GPU memory is recommended to avoid possible OOM on large datasets such as Reddit.The authors expect other Turing, Ampere, or Hopper GPUs to work.Source: §B.3.2 · p. 15
• GNN experiments use Open Graph Benchmark and DGL datasets, sparse convolution uses SemanticKITTI, and PrunedBERT uses publicly available HuggingFace models.Source: §B.3.4 · p. 15

B.4 Installation

• Users can clone the artifact repository and build its Docker image themselves.Source: §B.4 · p. 15
• A prebuilt Docker image is also available, but is compatible only with Ampere NVIDIA GPUs.Source: §B.4 · p. 15

• Each benchmark folder provides a run.sh script that produces profiling logs and figures in the paper's style.The logs include average execution time and GPU memory usage.Source: §B.5 · p. 15
• SpMM and SDDMM experiments can be reproduced with scripts in their respective folders.Source: §B.5 · p. 15
• The e2e and sparse-attention folders contain scripts for GraphSAGE training and sparse-transformer operator experiments, respectively.Source: §B.5 · p. 15
• PrunedBERT, RGCN inference, and sparse-convolution experiments have scripts in the pruned-bert, rgcn, and sparse-conv folders.Source: §B.5 · p. 15
• CUDA-event profiling discards the first 10 runs as warm-up and repeats measurement for 100 cycles.Source: §B.5 · p. 15

• Runtime and speedup vary across platforms, but reproduced results are expected to roughly match those reported in the paper.The passage points to Figures 13, 14, 15, 16, 17, 19, 20, and 23.Source: §B.6 · p. 15

B.7 Experiment customization

• For custom GNN datasets, users can create data as DGLGraph objects.Source: §B.7 · p. 15
• Custom point clouds for sparse convolution must be converted to TorchSparse SparseTensor, while pruned weights can be supplied as SciPy sparse matrices.Source: §B.7 · p. 15

B.8 Notes

• Small-operator timing can be affected by data retained in L2 between runs; all reported experiments use FLUSH_L2=ON.The artifact lets users enable or disable L2 flushing with the FLUSH_L2 environment variable.Source: §B.8 · p. 15