Sparse Tensors in MLIR

Note that this really is a continuation of three previous, and by now rather lengthy threads:

After a lot of constructive debate, we made good progress. We now have a SparseTensorDialect, which will provide a home for anything related to sparse tensors (attributes, operations, rewriting, passes, flags, etc). Furthermore, we now also have a proper first-class citizen sparse tensor type in MLIR, implemented through an “encoding” field on the built-in tensor types.

There will be a few more consolidating revisions that move code into this new home (most notorious from Linalg; we thank @nicolasvasilache for providing us with a temporary home, realizing very well we overstayed our welcome). After that we expect the sparse compiler related code to become much more stable. I will use this new thread to keep you posted on this consolidation.

Also if you are interested in contributing to this effort (either with FE work that maps to the new sparse tensor type, or BE work in improving the actual sparse compiler), please let us know, either in this thread or through private channels.

2 Likes

Revision D101573 introduced the SparseTensor dialect (including a minor migration of some sparse primitives from Linalg).

The first in a series of consolidating revisions is D101669, which migrates the sparse tensor encoding from the Tensor dialect into the new SparseTensor dialect as its permanent home. This revision also includes some internal improvements of the attribute that were requested by @mehdi_amini, namely using struct-like storage, verification, and construction (even though the encoding still parses and prints as a dictionary).

As a result of this consolidation, the encoding of a sparse tensor type now looks as follows.

#CSR = #sparse_tensor.encoding<{
  dimLevelType = [ "dense", "compressed" ],
  dimOrdering = affine_map<(i,j) -> (i,j)>,
  pointerBitWidth = 64,
  indexBitWidth = 64
}>
...
  func @foo(%arga: tensor<1024x1024xf64, #CSR>, ...
...

Up next: more migration from Linalg into SparseTensor, and removal of glue and clutter…

Revision D101811 migrates all “sparse compiler” related files from Linalg dialect into their permanent new home in the SparseTensor dialect (this includes rewriting, options, and test files). In addition it introduces two proper compiler passes, which replace the previous and temporary test passes.

Besides the file migration, the user visible changes are:

(1) The test switch --test-sparsification has been replaced with the proper compiler switch --sparsification. This switch accepts options, as before, although the pointer/index bitwidth options are scheduled to be replaced by information in the proper sparse tensor type.

(2) The test switch --test-sparsification="lower" variant has now been extracted into the proper compiler switch --sparse-tensor-conversion, so that the conversion can be tested separately from the sparsification pass.

Up next: remove all last glue and clutter from linalg, and make SparseTensor dialect operate on properly typed sparse tensors only.

With the very elaborate, but also very fun revision D102095, my dream of migrating to a first-class citizen sparse tensor type has become a reality!

Some major difference are described below.

(1) The sparse annotations in Linalg ops have been replaced in favor of proper sparse tensor types. So, for example, rather than using an annotation on the operation trait

#trait_sum_reduce = {
  indexing_maps = [
    affine_map<(i,j) -> (i,j)>, // A
    affine_map<(i,j) -> ()>     // x (out)
  ],
  sparse = [                            ;
    [ "S", "S" ], // A                  ; THIS PART
    [          ]  // x                  ;   IS REMOVED
  ],                                    ;
  iterator_types = ["reduction", "reduction"],
  doc = "x += A(i,j)"
}

This information is now carried by the sparse tensor type itself:

#SparseMatrix = #sparse_tensor.encoding<{
  dimLevelType = [ "compressed", "compressed" ]
}>

(2) Removal of the secondary storage choices for pointer and index width (viz. --sparsification="ptr-type=2 ind-type=2"), which applied to all sparse tensors in favor of a per-tensor specification of these widths.

#SparseMatrix = #sparse_tensor.encoding<{
  dimLevelType = [ "compressed", "compressed" ],
  pointerBitWidth = 32,
  indexBitWidth = 32
}>

(3) Glue and clutter to materialize the sparse tensor storage as opaque pointer into proper tensors is completely gone! So no more:

func @kernel_sum_reduce(%argA: !SparseTensor ... ) {
    %arga = sparse_tensor.fromPtr %argA : !SparseTensor to tensor<?x?xf64>
    %0 = linalg.generic #trait_sum_reduce
      ins(%arga: tensor<?x?xf64>)
      ...

But simply:

func @kernel_sum_reduce(%arga: tensor<?x?xf64, #SparseMatrix> ...) {
    %0 = linalg.generic #trait_sum_reduce
      ins(%arga: tensor<?x?xf64, #SparseMatrix>)
     ...

Subsequent bufferization takes care of replacing the types with whatever underlying implementation is selected by the compiler.

Also, setting up a sparse tensor in an integration test sees the biggest improvement to connect with the sparse support library. No longer:

    %annotations = memref.alloc(%c2) : memref<?xi1>
    %sparse = constant true
    memref.store %sparse, %annotations[%c0] : memref<?xi1>
    memref.store %sparse, %annotations[%c1] : memref<?xi1>
    %i64 = constant 1 : index
    %f64 = constant 1 : index
    %a = call @newSparseTensor(%fileName, %annotations, %i64, %i64, %f64)
       : (!Filename, memref<?xi1>, index, index, index) -> (!SparseTensor)

But just the following (sic!):

%a = sparse_tensor.new %fileName : !Filename to tensor<?x?xf64, #SparseMatrix>

(4) More rigorous verification of type consistency of all ops, including the generated sparse primitives that connect the generated sparse code with the support library (note that the latter is simply there for convenience, in the long run even the library could be replaced with codegen).

Next: one last revision that removes all obsoleted sparse code form Linalg dialect.

1 Like

Finally, revision D102098 removes the now obsolete sparse annotations code in the Linalg dialect.

So long, and thanks for all the fish!

Hey Aart, this is fantastic! I have to admit I didn’t think we were this close.

I’d love to get some of this exposed through our Python API and give it a try. I think the things we need to wire up are:

  • C/Python API support for your sparse_tensor.encoding attribute.
  • Python Tablegen bindings for the sparse_tensor dialect
  • New APIs to construct a tensor type with an encoding attribute

Do you have any of that in-flight? If not, I could probably get it pounded out this weekend.

1 Like

Yes, please give it a look over the weekend! Getting proper API’s for setting up the sparse tensor encoding and sparse tensor types are indeed the proper next steps.

Most ops in the sparse tensor dialect are not needed externally (since they are generated by the sparse compiler), with the exception of the sparse_tensor.new operator, which is my “Swiss Army Knife” operator for materializing sparse tensor values into a computation (reading and initializing from an external file format like FROSTT being the most obvious implementation, but we would also generalize this to buffers or runnable initialize code).