Compiler Support for Sparse Tensor Computations in MLIR

TL;DR

This paper introduces a compiler approach within MLIR that automatically generates efficient sparse tensor code from generic definitions, simplifying development and improving performance in scientific and machine learning applications.

Contribution

It presents a novel method to treat sparsity as a property and enables automatic sparse code generation in MLIR, reducing manual effort and errors.

Findings

Automated sparse code generation is feasible within MLIR.

Sparsity-aware compilation improves efficiency and reduces development complexity.

The approach integrates seamlessly with existing MLIR infrastructure.

Abstract

Sparse tensors arise in problems in science, engineering, machine learning, and data analytics. Programs that operate on such tensors can exploit sparsity to reduce storage requirements and computational time. Developing and maintaining sparse software by hand, however, is a complex and error-prone task. Therefore, we propose treating sparsity as a property of tensors, not a tedious implementation task, and letting a sparse compiler generate sparse code automatically from a sparsity-agnostic definition of the computation. This paper discusses integrating this idea into MLIR.