SySTeC: A Symmetric Sparse Tensor Compiler

2024-06-13

arXiv.org

AI Summary

Symmetric sparse tensor computations suffer from the difficulty of jointly optimizing symmetry and sparsity, while manual implementation is error-prone and combinatorially explosive. To address this, we present the first compiler that automatically generates functionally complete, symmetry-aware sparse tensor kernels. We introduce a novel taxonomy of tensor symmetries, unifying symmetry constraints with sparse iteration logic. We design a domain-specific intermediate representation (DSIR) and a symmetry-aware scheduler that integrate triangular loop clipping, transposition-equivalence class enumeration, and adaptive traversal across sparse formats. Evaluated on representative kernels—including SSYMV and 5D MTTKRP—our approach achieves speedups of 1.36×–30.4× over state-of-the-art asymmetric methods, demonstrating substantial performance gains and correctness guarantees through automated, symmetry-preserving code generation.

Technology Category

• Machine Learning: Matrix & Tensor Methods
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization
• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Security and Privacy: Large-scale security measurements

Abstract

Symmetric and sparse tensors arise naturally in many domains including linear algebra, statistics, physics, chemistry, and graph theory. Symmetric tensors are equal to their transposes, so in the $n$-dimensional case we can save up to a factor of $n!$ by avoiding redundant operations. Sparse tensors, on the other hand, are mostly zero, and we can save asymptotically by processing only nonzeros. Unfortunately, specializing for both symmetry and sparsity at the same time is uniquely challenging. Optimizing for symmetry requires consideration of $n!$ transpositions of a triangular kernel, which can be complex and error prone. Considering multiple transposed iteration orders and triangular loop bounds also complicates iteration through intricate sparse tensor formats. Additionally, since each combination of symmetry and sparse tensor formats requires a specialized implementation, this leads to a combinatorial number of cases. A compiler is needed, but existing compilers cannot take advantage of both symmetry and sparsity within the same kernel. In this paper, we describe the first compiler which can automatically generate symmetry-aware code for sparse or structured tensor kernels. We introduce a taxonomy for symmetry in tensor kernels, and show how to target each kind of symmetry. Our implementation demonstrates significant speedups ranging from 1.36x for SSYMV to 30.4x for a 5-dimensional MTTKRP over the non-symmetric state of the art.

Problem

• Compiling symmetric sparse tensor kernels automatically
• Handling combinatorial symmetry and sparsity cases efficiently
• Generating optimized code for multiple tensor transpositions

Innovation

• Compiler for symmetric sparse tensor code generation
• Taxonomy for symmetry types in tensor kernels
• Automatic optimization for both symmetry and sparsity