Module melior::dialect::ods::sparse_tensor
source · Expand description
sparse_tensor
dialect.
The SparseTensor
dialect supports all the attributes, types,
operations, and passes that are required to make sparse tensor
types first class citizens within the MLIR compiler infrastructure.
The dialect forms a bridge between high-level operations on sparse
tensors types and lower-level operations on the actual sparse storage
schemes consisting of positions, coordinates, and values. Lower-level
support may consist of fully generated code or may be provided by
means of a small sparse runtime support library.
The concept of treating sparsity as a property, not a tedious implementation detail, by letting a sparsifier generate sparse code automatically was pioneered for linear algebra by [Bik96] in MT1 (see https://www.aartbik.com/sparse.php) and formalized to tensor algebra by [Kjolstad17,Kjolstad20] in the Sparse Tensor Algebra Compiler (TACO) project (see http://tensor-compiler.org). Please note that we started to prefer the term “sparsifier” over the also commonly used “sparse compiler” terminology to refer to such a pass to make it clear that the sparsifier pass is not a separate compiler, but should be an integral part of any compiler pipeline that is built with the MLIR compiler infrastructure
The MLIR implementation [Biketal22] closely follows the “sparse iteration theory” that forms the foundation of TACO. A rewriting rule is applied to each tensor expression in the Linalg dialect (MLIR’s tensor index notation) where the sparsity of tensors is indicated using the per-level level-types (e.g., dense, compressed, singleton) together with a specification of the order on the levels (see [Chou18] for an in-depth discussions and possible extensions to these level-types). Subsequently, a topologically sorted iteration graph, reflecting the required order on coordinates with respect to the levels of each tensor, is constructed to ensure that all tensors are visited in natural level-coordinate order. Next, iteration lattices are constructed for the tensor expression for every index in topological order. Each iteration lattice point consists of a conjunction of tensor coordinates together with a tensor (sub)expression that needs to be evaluated for that conjunction. Within the lattice, iteration points are ordered according to the way coordinates are exhausted. As such these iteration lattices drive actual sparse code generation, which consists of a relatively straightforward one-to-one mapping from iteration lattices to combinations of for-loops, while-loops, and if-statements. Sparse tensor outputs that materialize uninitialized are handled with direct insertions if all parallel loops are outermost or insertions that indirectly go through a 1-dimensional access pattern expansion (a.k.a. workspace) where feasible [Gustavson72,Bik96,Kjolstad19].
- [Bik96] Aart J.C. Bik. Compiler Support for Sparse Matrix Computations. PhD thesis, Leiden University, May 1996.
- [Biketal22] Aart J.C. Bik, Penporn Koanantakool, Tatiana Shpeisman, Nicolas Vasilache, Bixia Zheng, and Fredrik Kjolstad. Compiler Support for Sparse Tensor Computations in MLIR. ACM Transactions on Architecture and Code Optimization, June, 2022. See: https://dl.acm.org/doi/10.1145/3544559
- [Chou18] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe. Format Abstraction for Sparse Tensor Algebra Compilers. Proceedings of the ACM on Programming Languages, October 2018.
- [Chou20] Stephen Chou, Fredrik Berg Kjolstad, and Saman Amarasinghe. Automatic Generation of Efficient Sparse Tensor Format Conversion Routines. Proceedings of the 41st ACM SIGPLAN Conference on Programming Language Design and Implementation, June, 2020.
- [Gustavson72] Fred G. Gustavson. Some basic techniques for solving sparse systems of linear equations. In Sparse Matrices and Their Applications, pages 41–52. Plenum Press, New York, 1972.
- [Kjolstad17] Fredrik Berg Kjolstad, Shoaib Ashraf Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. The Tensor Algebra Compiler. Proceedings of the ACM on Programming Languages, October 2017.
- [Kjolstad19] Fredrik Berg Kjolstad, Peter Ahrens, Shoaib Ashraf Kamil, and Saman Amarasinghe. Tensor Algebra Compilation with Workspaces, Proceedings of the IEEE/ACM International Symposium on Code Generation and Optimization, 2019.
- [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation. PhD thesis, MIT, February, 2020.
Structs§
- An
assemble
operation. Returns a sparse tensor assembled from the given values and levels. - A builder for an
assemble
operation. - A
binary
operation. Binary set operation utilized within linalg.generic. - A builder for a
binary
operation. - A
compress
operation. Compressed an access pattern for insertion. - A builder for a
compress
operation. - A
concatenate
operation. Concatenates a list of tensors into a single tensor.. - A builder for a
concatenate
operation. - A
convert
operation. Converts between different tensor types. - A builder for a
convert
operation. - A
crd_translate
operation. Performs coordinate translation between level and dimension coordinate space.. - A builder for a
crd_translate
operation. - A
disassemble
operation. Returns the (values, coordinates) pair disassembled from the input tensor. - A builder for a
disassemble
operation. - An
expand
operation. Expands an access pattern for insertion. - A builder for an
expand
operation. - A
foreach
operation. Iterates over elements in a tensor. - A builder for a
foreach
operation. - A
storage_specifier.get
operation. - A builder for a
storage_specifier.get
operation. - An
insert
operation. Inserts a value into the sparse tensor. - A builder for an
insert
operation. - A
load
operation. Rematerializes tensor from underlying sparse storage format. - A builder for a
load
operation. - A
lvl
operation. Level index operation. - A builder for a
lvl
operation. - A
new
operation. Materializes a new sparse tensor from given source. - A builder for a
new
operation. - A
number_of_entries
operation. Returns the number of entries that are stored in the tensor.. - A builder for a
number_of_entries
operation. - An
out
operation. Outputs a sparse tensor to the given destination. - A builder for an
out
operation. - A
push_back
operation. Pushes a value to the back of a given buffer. - A builder for a
push_back
operation. - A
reduce
operation. Custom reduction operation utilized within linalg.generic. - A builder for a
reduce
operation. - A
reinterpret_map
operation. Reinterprets the dimension/level maps of the source tensor. - A builder for a
reinterpret_map
operation. - A
reorder_coo
operation. Reorder the input COO such that it has the the same order as the output COO. - A builder for a
reorder_coo
operation. - A
select
operation. Select operation utilized within linalg.generic. - A builder for a
select
operation. - A
storage_specifier.set
operation. - A builder for a
storage_specifier.set
operation. - A
sort
operation. Sorts the arrays in xs and ys lexicographically on the integral values found in the xs list. - A builder for a
sort
operation. - A
storage_specifier.init
operation. - A builder for a
storage_specifier.init
operation. - A
coordinates_buffer
operation. Extracts the linear coordinates array from a tensor. - A builder for a
coordinates_buffer
operation. - A builder for a
coordinates
operation. - A builder for a
positions
operation. - A
slice.offset
operation. Extracts the offset of the sparse tensor slice at the given dimension. - A builder for a
slice.offset
operation. - A
slice.stride
operation. Extracts the stride of the sparse tensor slice at the given dimension. - A builder for a
slice.stride
operation. - A
values
operation. Extracts numerical values array from a tensor. - A builder for a
values
operation. - An
unary
operation. Unary set operation utilized within linalg.generic. - A builder for an
unary
operation. - A
yield
operation. Yield from sparse_tensor set-like operations. - A builder for a
yield
operation.
Functions§
- Creates an
assemble
operation. - Creates a
binary
operation. - Creates a
compress
operation. - Creates a
concatenate
operation. - Creates a
convert
operation. - Creates a
coordinates
operation. - Creates a
coordinates_buffer
operation. - Creates a
crd_translate
operation. - Creates a
disassemble
operation. - Creates an
expand
operation. - Creates a
foreach
operation. - Creates an
insert
operation. - Creates a
load
operation. - Creates a
lvl
operation. - Creates a
new
operation. - Creates a
number_of_entries
operation. - Creates an
out
operation. - Creates a
positions
operation. - Creates a
push_back
operation. - Creates a
reduce
operation. - Creates a
reinterpret_map
operation. - Creates a
reorder_coo
operation. - Creates a
select
operation. - Creates a
slice.offset
operation. - Creates a
slice.stride
operation. - Creates a
sort
operation. - Creates a
storage_specifier.get
operation. - Creates a
storage_specifier.init
operation. - Creates a
storage_specifier.set
operation. - Creates an
unary
operation. - Creates a
values
operation. - Creates a
yield
operation.