SparseConflicts: Handling Conflicting Data Layouts in Sparse Tensor Contractions

Abstract

Optimizing sparse tensor computations is challenging due to the use of compressed storage formats, which leads to non-affine loop nests and a vast, complex schedule space. The performance of a given schedule is sensitive to the sparsity pattern of the input tensors, making it difficult to find a single optimal solution. When input tensors in the same tensor contraction have conflicting data layouts in relation to the iteration order, it requires costly-both in time and memory-layout transformation, such as transposition. A promising but under-explored alternative is to generate a schedule that avoids explicit transposition, but this has not been systematically supported in existing compilers. This paper presents a new code generation strategy that generalizes the intermediate representation of the TACO sparse tensor compiler to generate a single loop nest, circumventing explicit transposition of tensors when the tensors have conflicting data layout iteration orders. We extend TACOś iteration graph to express a search-based strategy for locating elements in tensors with conflicting layouts, and we introduce new intermediate representation nodes to lower these schedules to efficient code. This enables the systematic generation of loops that do not require explicit transposition, thus avoiding the overhead of materializing temporary tensors. We evaluate our approach on a set of sparse tensor contractions using both real-world and synthetic datasets. Our results demonstrate that for computations with misaligned data layouts, our fused approach achieves up to 2x speedup for some sparsity patterns over the traditional approach of explicitly creating a transposed temporary. We also provide guidelines for when this new scheduling strategy is likely to be beneficial.