Multidirectional Propagation of Sparsity Information Across Tensor Slices

A computational graph in a tensor compiler represents a program as a set of kernels connected by edges that describe how tensors flow between them. Tensors may be dense or sparse, the latter storing only nonzero values. Recognizing sparsity allows the compiler to avoid redundant computation and reduce memory usage. This paper introduces Sparsity Propagation Analysis (SPA), a static analysis that conservatively infers sparsity in tensor slices: sets of tensor elements that share a fixed subset of indices. SPA estimates which slices can be treated as zero, via forward, backward, and lateral propagation of sparsity information. SPA applies to tensors of arbitrary dimension and to operators defined over general semirings (e.g., addition and multiplication, MAX and MIN, logical AND and OR). Its runtime grows with the number of slices rather than the number of elements, making it asymptotically faster than executing even a single iteration of the graph. We implement SPA in the Tensor Algebra Compiler (TACO) and evaluate it on benchmarks from the Einsum Collection. Results indicate that multidirectional propagation increases the precision of the sparsity analysis, reduces memory consumption, and enables optimizations that improve runtime performance.