MLIR

Multi-Level IR Compiler Framework

'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 pointers, indices, 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 sparse compiler generate sparse code automatically was pioneered for dense 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).

The MLIR implementation 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-dimension level types dense/compressed together with a specification of the order on the dimensions (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 indices with respect to the dimensions of each tensor, is constructed to ensure that all tensors are visited in natural index 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 indices 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 indices 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.

  • [Bik96] Aart J.C. Bik. Compiler Support for Sparse Matrix Computations. PhD thesis, Leiden University, May 1996.
  • [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.
  • [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation. PhD thesis, MIT, February, 2020.
  • [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.

Attribute definition 

SparseTensorEncodingAttr 

An attribute to encode TACO-style information on sparsity properties of tensors. The encoding is eventually used by a sparse compiler pass to generate sparse code fully automatically for all tensor expressions that involve tensors with a sparse encoding. Compiler passes that run before this sparse compiler pass need to be aware of the semantics of tensor types with such an encoding.

Example:

#DCSC = #sparse_tensor.encoding<{
  dimLevelType = [ "compressed", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (j,i)>,
  pointerBitWidth = 32,
  indexBitWidth = 8
}>


... tensor<8x8xf64, #DCSC> ...

Parameters: 

ParameterC++ typeDescription
dimLevelType::llvm::ArrayRef<SparseTensorEncodingAttr::DimLevelType>Per-dimension level type
dimOrderingAffineMap
pointerBitWidthunsigned
indexBitWidthunsigned

Operation definition 

sparse_tensor.convert (::mlir::sparse_tensor::ConvertOp) 

Converts between different tensor types

Syntax:

operation ::= `sparse_tensor.convert` $source attr-dict `:` type($source) `to` type($dest)

Converts one sparse or dense tensor type to another tensor type. The rank of the source and destination types must match exactly, and the dimension sizes must either match exactly or relax from a static to a dynamic size. The sparse encoding of the two types can obviously be completely different. The name convert was preferred over cast, since the operation may incur a non-trivial cost.

When converting between two different sparse tensor types, only explicitly stored values are moved from one underlying sparse storage format to the other. When converting from an unannotated dense tensor type to a sparse tensor type, an explicit test for nonzero values is used. When converting to an unannotated dense tensor type, implicit zeroes in the sparse storage format are made explicit. Note that the conversions can have non-trivial costs associated with them, since they may involve elaborate data structure transformations. Also, conversions from sparse tensor types into dense tensor types may be infeasible in terms of storage requirements.

Examples:

%0 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor<32x32xf32, #CSR>
%1 = sparse_tensor.convert %a : tensor<32x32xf32> to tensor<?x?xf32, #CSR>
%2 = sparse_tensor.convert %b : tensor<8x8xi32, #CSC> to tensor<8x8xi32, #CSR>
%3 = sparse_tensor.convert %c : tensor<4x8xf64, #CSR> to tensor<4x?xf64, #CSC>

// The following conversion is not allowed (since it would require a
// runtime assertion that the source's dimension size is actually 100).
%4 = sparse_tensor.convert %d : tensor<?xf64> to tensor<100xf64, #SV>

Traits: SameOperandsAndResultElementType

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
sourcetensor of any type values

Results: 

ResultDescription
desttensor of any type values

sparse_tensor.init (::mlir::sparse_tensor::InitOp) 

Materializes an unitialized sparse tensor

Syntax:

operation ::= `sparse_tensor.init` `[` $sizes `]` attr-dict `:` type($result)

Materializes an uninitialized sparse tensor with given shape (either static or dynamic). The operation is provided as an anchor that materializes a properly typed but uninitialized sparse tensor into the output clause of a subsequent operation that yields a sparse tensor as the result.

Example:

%c = sparse_tensor.init_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
%0 = linalg.matmul
  ins(%a, %b: tensor<?x?xf32>, tensor<?x?xf32>)
  outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
sizesindex

Results: 

ResultDescription
resulttensor of any type values

sparse_tensor.lex_insert (::mlir::sparse_tensor::LexInsertOp) 

Inserts a value into given sparse tensor in lexicograph index order

Syntax:

operation ::= `sparse_tensor.lex_insert` $tensor `,` $indices `,` $value attr-dict `:` type($tensor) `,` type($indices) `,` type($value)

Inserts the given value at given indices into the underlying sparse storage format of the given tensor with the given indices. This operation can only be applied when a tensor materializes unintialized with an init operation, the insertions occur in strict lexicographic index order, and the final tensor is constructed with a tensor operation that has the hasInserts attribute set.

Note that this operation is “impure” in the sense that its behavior is solely defined by side-effects and not SSA values. The semantics may be refined over time as our sparse abstractions evolve.

sparse_tensor.lex_insert %tensor, %indices, %val
  : tensor<1024x1024xf64, #CSR>, memref<?xindex>, f64

Operands: 

OperandDescription
tensortensor of any type values
indices1D memref of index values
valueany type

sparse_tensor.load (::mlir::sparse_tensor::LoadOp) 

Rematerializes tensor from underlying sparse storage format

Syntax:

operation ::= `sparse_tensor.load` $tensor (`hasInserts` $hasInserts^)? attr-dict `:` type($tensor)

Rematerializes a tensor from the underlying sparse storage format of the given tensor. This is similar to the memref.load operation in the sense that it provides a bridge between a bufferized world view and a tensor world view. Unlike the memref.load operation, however, this sparse operation is used only temporarily to maintain a correctly typed intermediate representation during progressive bufferization.

The hasInserts attribute denote whether insertions to the underlying sparse storage format may have occurred, in which case the underlying sparse storage format needs to be finalized. Otherwise, the operation simply folds away.

Note that this operation is “impure” in the sense that its behavior is solely defined by side-effects and not SSA values. The semantics may be refined over time as our sparse abstractions evolve.

Example:

%1 = sparse_tensor.load %0 : tensor<8xf64, #SV>

Traits: SameOperandsAndResultType

Attributes: 

AttributeMLIR TypeDescription
hasInserts::mlir::UnitAttrunit attribute

Operands: 

OperandDescription
tensortensor of any type values

Results: 

ResultDescription
resulttensor of any type values

sparse_tensor.new (::mlir::sparse_tensor::NewOp) 

Materializes a new sparse tensor from given source

Syntax:

operation ::= `sparse_tensor.new` $source attr-dict `:` type($source) `to` type($result)

Materializes a sparse tensor with contents taken from an opaque pointer provided by source. For targets that have access to a file system, for example, this pointer may be a filename (or file) of a sparse tensor in a particular external storage format. The form of the operation is kept deliberately very general to allow for alternative implementations in the future, such as pointers to buffers or runnable initialization code. The operation is provided as an anchor that materializes a properly typed sparse tensor with inital contents into a computation.

Example:

sparse_tensor.new %source : !Source to tensor<1024x1024xf64, #CSR>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
sourceany type

Results: 

ResultDescription
resulttensor of any type values

sparse_tensor.release (::mlir::sparse_tensor::ReleaseOp) 

Releases underlying sparse storage format of given tensor

Syntax:

operation ::= `sparse_tensor.release` $tensor attr-dict `:` type($tensor)

Releases the underlying sparse storage format for a tensor that materialized earlier through a new operator, init operator, or a convert operator with an annotated tensor type as destination (unless that convert is folded away since the source and destination types were identical). This operation should only be called once for any materialized tensor. Also, after this operation, any subsequent memref querying operation on the tensor returns undefined results.

Note that this operation is “impure” in the sense that its behavior is solely defined by side-effects and not SSA values. The semantics may be refined over time as our sparse abstractions evolve.

Example:

sparse_tensor.release %tensor : tensor<1024x1024xf64, #CSR>

Operands: 

OperandDescription
tensortensor of any type values

sparse_tensor.indices (::mlir::sparse_tensor::ToIndicesOp) 

Extracts indices array at given dimension from a tensor

Syntax:

operation ::= `sparse_tensor.indices` $tensor `,` $dim attr-dict `:` type($tensor) `to` type($result)

Returns the indices array of the sparse storage format at the given dimension for the given sparse tensor. This is similar to the memref.buffer_cast operation in the sense that it provides a bridge between a tensor world view and a bufferized world view. Unlike the memref.buffer_cast operation, however, this sparse operation actually lowers into a call into a support library to obtain access to the indices array.

Example:

%1 = sparse_tensor.indices %0, %c1
   : tensor<64x64xf64, #CSR> to memref<?xindex>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
tensortensor of any type values
dimindex

Results: 

ResultDescription
resultstrided memref of any type values of rank 1

sparse_tensor.pointers (::mlir::sparse_tensor::ToPointersOp) 

Extracts pointers array at given dimension from a tensor

Syntax:

operation ::= `sparse_tensor.pointers` $tensor `,` $dim attr-dict `:` type($tensor) `to` type($result)

Returns the pointers array of the sparse storage format at the given dimension for the given sparse tensor. This is similar to the memref.buffer_cast operation in the sense that it provides a bridge between a tensor world view and a bufferized world view. Unlike the memref.buffer_cast operation, however, this sparse operation actually lowers into a call into a support library to obtain access to the pointers array.

Example:

%1 = sparse_tensor.pointers %0, %c1
   : tensor<64x64xf64, #CSR> to memref<?xindex>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
tensortensor of any type values
dimindex

Results: 

ResultDescription
resultstrided memref of any type values of rank 1

sparse_tensor.values (::mlir::sparse_tensor::ToValuesOp) 

Extracts numerical values array from a tensor

Syntax:

operation ::= `sparse_tensor.values` $tensor attr-dict `:` type($tensor) `to` type($result)

Returns the values array of the sparse storage format for the given sparse tensor, independent of the actual dimension. This is similar to the memref.buffer_cast operation in the sense that it provides a bridge between a tensor world view and a bufferized world view. Unlike the memref.buffer_cast operation, however, this sparse operation actually lowers into a call into a support library to obtain access to the values array.

Example:

%1 = sparse_tensor.values %0 : tensor<64x64xf64, #CSR> to memref<?xf64>

Interfaces: NoSideEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
tensortensor of any type values

Results: 

ResultDescription
resultstrided memref of any type values of rank 1