A High Performance Sparse Tensor Algebra Compiler in MLIR

Tian, Ruiqin; Guo, Luanzheng; Li, Jiajia; Ren, Bin; Kestor, Gokcen

doi:10.1109/llvmhpc54804.2021.00009

Cited by 24 publications

(8 citation statements)

References 37 publications

Supporting

Mentioning

Contrasting

Order By: Relevance

“…COMET [53,77] is a MLIR-based compiler infrastructure for dense and sparse tensor algebra computations. It uses a dimension-wise sparse storage scheme and a code generation algorithm similar to TACO, but has more performance portability building upon MLIR.…”

Section: Sparse Compilersmentioning

confidence: 99%

Compiler Support for Sparse Tensor Computations in MLIR

Bik¹,

Koanantakool²,

Shpeisman³

et al. 2022

Preprint

View full text Add to dashboard Cite

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.

show abstract

Section: Sparse Compilersmentioning

confidence: 99%

Compiler Support for Sparse Tensor Computations in MLIR

Bik¹,

Koanantakool²,

Shpeisman³

et al. 2022

Preprint

View full text Add to dashboard Cite

show abstract

“…TACO does not generate multi-threaded code when the output tensor is sparse. Prior work has evaluated against single-threaded code in such situations [16,34]. Following the strategy of Senanayeke et al [30], we modified the TACO generated code manually to add multithreading…”

Section: Benchmarksmentioning

confidence: 99%

“…Automated sparse code generation [5,6,17,34] is a heavilyresearched topic. Even though these methods are highly effective, they lack fine-grained optimizations like ours that applies across kernels to reduce the time-complexity of the computation.…”

Section: Sparse Tensor Algebramentioning

confidence: 99%

“…The problem of compiler optimizations for sparse codes is well known [5,17,18,34,35], and there are several challenges that compilers face: (1) tensor computations have to deal with specific data formats; (2) load imbalance can arise due to irregular structure; and (3) data locality issues arise due to the sparsity of the tensors. TACO provides a compiler for automatically generating kernels for dense and sparse tensor algebra operations [17].…”

Section: Introductionmentioning

confidence: 99%

See 1 more Smart Citation

SparseLNR: Accelerating Sparse Tensor Computations Using Loop Nest Restructuring

Dias,

Sundararajah,

Saumya

et al. 2022

Preprint

View full text Add to dashboard Cite

Sparse tensor algebra computations have become important in many real-world applications like machine learning, scientific simulations, and data mining. Hence, automated code generation and performance optimizations for tensor algebra kernels are paramount. Recent advancements such as the Tensor Algebra Compiler (TACO) greatly generalize and automate the code generation for tensor algebra expressions. However, the code generated by TACO for many important tensor computations remains suboptimal due to the absence of a scheduling directive to support transformations such as distribution/fusion.This paper extends TACO's scheduling space to support kernel distribution/loop fusion in order to reduce asymptotic time complexity and improve locality of complex tensor algebra computations. We develop an intermediate representation (IR) for tensor operations called branched iteration graph which specifies breakdown of the computation into smaller ones (kernel distribution) and then fuse (loop fusion) outermost dimensions of the loop nests, while the innermost dimensions are distributed, to increase data locality. We describe exchanges of intermediate results between space iteration spaces, transformation in the IR, and its programmatic invocation. Finally, we show that the transformation can be used to optimize sparse tensor kernels. Our results show that this new transformation significantly improves the performance of several real-world tensor algebra computations compared to TACO-generated code.

show abstract

“…The rise of Multi-Level Intermediate Representation (MLIR) [15] as a new compiler infrastructure has enabled the expansion of code transformations to higher abstraction levels. Though MLIR has already been applied in a number of fields, from linear algebra [4,32] to quantum technology [20], it has not yet been leveraged as a viable replacement to traditional code generators in frameworks that model electrophysiology. In electrophysiology, ionic models describe the interactions between cells composing tissues, as for example the human heart.…”

Section: Introductionmentioning

confidence: 99%

Lifting Code Generation of Cardiac Physiology Simulation to Novel Compiler Technology

Thangamani

Jost

Loechner

et al. 2023

Proceedings of the 21st ACM/IEEE International Symposium on Code Generation and Optimization

View full text Add to dashboard Cite

The study of numerical models for the human body has become a major focus of the research community in biology and medicine. For instance, numerical ionic models of a complex organ, such as the heart, must be able to represent individual cells and their interconnections through ionic channels, forming a system with billions of cells, and requiring efficient code to handle such a large system. The modeling of the electrical system of the heart combines a compute-intensive kernel that calculates the intensity of current flowing through cell membranes, and feeds a linear solver for computing the electrical potential of each cell.Considering this context, we propose limpetMLIR, a code generator and compiler transformer to accelerate the kernel phase of ionic models and bridge the gap between compiler technology and electrophysiology simulation. LimpetMLIR makes use of the MLIR infrastructure, its dialects, and transformations to drive forward the study of ionic models, and accelerate the execution of multi-cell systems. Experiments conducted in 43 ionic models show that our limpetMLIR based code generation greatly outperforms current state-ofthe-art simulation systems by an average of 2.9×, reaching peak speedups of more than 15× in some cases. To our knowledge, this is the first work that deeply connects an optimizing compiler infrastructure to electrophysiology models of the human body, showing the potential benefits of using compiler technology in the simulation of human cell interactions.

show abstract

A High Performance Sparse Tensor Algebra Compiler in MLIR

Cited by 24 publications

References 37 publications

Compiler Support for Sparse Tensor Computations in MLIR

Compiler Support for Sparse Tensor Computations in MLIR

SparseLNR: Accelerating Sparse Tensor Computations Using Loop Nest Restructuring

Lifting Code Generation of Cardiac Physiology Simulation to Novel Compiler Technology

Contact Info

Product

Resources

About