Partitioning Unstructured Sparse Tensor Algebra for Load-Balanced Parallel Execution

2026-04-18

AI Summary

Sparse tensor algebra poses significant challenges for efficient load balancing on parallel systems due to its irregular structure, strong data dependencies, and skewed distributions. This work proposes the first provably load-balanced partitioning algorithm capable of handling arbitrary numbers of operands and multi-dimensional hierarchical sparsity patterns. The algorithm is integrated into a sparse tensor algebra compilation framework that automatically generates high-performance parallel kernels for both multicore CPUs and GPUs. By overcoming the limitations of conventional parallel merge strategies, the approach produces code that matches or exceeds the performance of hand-optimized libraries such as Intel MKL and cuSPARSE, achieving geometric mean speedups of 0.73–3.4× across diverse hardware platforms. Moreover, it substantially outperforms existing general-purpose parallelization strategies, with geometric mean speedups ranging from 2.0× to 6.4×.

Technology Category

• Machine Learning: Matrix & Tensor Methods
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Systems and Infrastructure for Web, Mobile and WoT: Experiences and lessons learnt from Web-based algorithms and system deployments
• Semantics and Knowledge: Methods, algorithms and applications for the development of semantic models, knowledge graphs and other forms of structured data models with machine-interpretable semantics

Abstract

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$).

Problem

• sparse tensor algebra
• load balancing
• parallel execution
• irregular computation
• partitioning

Innovation

• sparse tensor algebra
• load balancing
• parallel partitioning
• hierarchical sparse data structures
• automatic code generation