Partitioning Unstructured Sparse Tensor Algebra for Load-Balanced Parallel Execution

Fuente: arXiv
Gespeichert in:
Bibliographische Detailangaben
Hauptverfasser: Chougule, Atharva, Root, Alexander J, Lacouture, Rubens, Yan, Bobby, Yadav, Rohan, Kjolstad, Fredrik
Format: Preprint
Veröffentlicht: 2026
Schlagworte:
Online-Zugang:
Tags: Tag hinzufügen
Keine Tags, Fügen Sie den ersten Tag hinzu!
_version_ 1866908986649870336
author Chougule, Atharva
Root, Alexander J
Lacouture, Rubens
Yan, Bobby
Yadav, Rohan
Kjolstad, Fredrik
author_facet Chougule, Atharva
Root, Alexander J
Lacouture, Rubens
Yan, Bobby
Yadav, Rohan
Kjolstad, Fredrik
contents Sparse tensor algebra is challenging to efficiently parallelize due to the irregular, data-dependent, and potentially skewed structure of sparse computation. We propose the first partitioning algorithm that provably load balances the computation of any sparse tensor algebra expression across parallel execution units. Our algorithm generalizes parallel merging algorithms to any number of operands, and to multi-dimensional, hierarchical sparse data structures. We implement our algorithm within an existing sparse tensor algebra compilation framework to automatically generate parallel sparse tensor algebra kernels that target multi-core CPUs and GPUs. We show that our generated code is competitive with hand-implemented parallelization strategies used by vendor libraries like Intel MKL and NVIDIA cuSPARSE (geo-means of $0.73$--$3.4\times$) and \textsc{Taco} (geo-means of $1.0$--$2.4\times$), and significantly outperforms general-purpose strategies for sparse tensor expressions where specialized algorithms have not been developed (geo-means of $2.0$--$6.4\times$).
format Preprint
id arxiv_https___arxiv_org_abs_2604_17198
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Partitioning Unstructured Sparse Tensor Algebra for Load-Balanced Parallel Execution
Chougule, Atharva
Root, Alexander J
Lacouture, Rubens
Yan, Bobby
Yadav, Rohan
Kjolstad, Fredrik
Programming Languages
D.3
Sparse tensor algebra is challenging to efficiently parallelize due to the irregular, data-dependent, and potentially skewed structure of sparse computation. We propose the first partitioning algorithm that provably load balances the computation of any sparse tensor algebra expression across parallel execution units. Our algorithm generalizes parallel merging algorithms to any number of operands, and to multi-dimensional, hierarchical sparse data structures. We implement our algorithm within an existing sparse tensor algebra compilation framework to automatically generate parallel sparse tensor algebra kernels that target multi-core CPUs and GPUs. We show that our generated code is competitive with hand-implemented parallelization strategies used by vendor libraries like Intel MKL and NVIDIA cuSPARSE (geo-means of $0.73$--$3.4\times$) and \textsc{Taco} (geo-means of $1.0$--$2.4\times$), and significantly outperforms general-purpose strategies for sparse tensor expressions where specialized algorithms have not been developed (geo-means of $2.0$--$6.4\times$).
title Partitioning Unstructured Sparse Tensor Algebra for Load-Balanced Parallel Execution
topic Programming Languages
D.3
url https://arxiv.org/abs/2604.17198