RFC: Introduce a PyTACO-MLIR Bridge for integration testing

As announced at the end of the Sparse Tensors in MLIR thread, we have many exciting new sparse tensor features and improvements planned for 2022, and this is the first topical post on such a feature: adding a PyTACO-MLIR bridge for integration testing.

A lot of the new sparse tensor support in MLIR borrows heavily from the great foundation laid by the TACO project. One nifty feature of TACO is PyTACO, which is an extension to Python that provides easy access to sparsity annotations by means of tensor index notation. For example, a sparse matrix in CSR format times dense vector kernel can be expressed in PyTACO as follows.

import pytaco as pt
...
csr = pt.format([dense, compressed])
A = pt.read("pwtk.mtx", csr)
...
i, j = pt.get_index_vars(2)
...  
y[i] = A[i, j] * x[j] + z[i]

Thanks to excellent work by Bixia Zheng, we now can run PyTACO programs using MLIR by essentially just changing the import line into

import mlir_pytaco_api as pt

After that, a new bridge lowers PyTACO expressions into Linalg expressions and subsequently runnable code using the sparse compiler in MLIR. A JIT compilation using LLVM IR takes care of evaluating the tensor index expression for the Python environment after that.

Since the bridge is currently merely intended to simplify integration testing (writing Python is so much easier than writing direct MLIR code), and is not “production ready” yet, all code related to the new bridge will reside in the following directory.

mlir/test/Integration/Dialect/SparseTensor/taco

Please let us know your thoughts on this new functionality. We will have a revision for review up on Phabricator real soon.

2 Likes

Here is Bixia’s revision: D117260 Upstream MLIR PyTACO implementation

It’s pretty nice to see what we can already achieve with MLIR without writing a single line of C++! Well done @bixia1 :slight_smile:

I like that this patch is not intrusive and can be living in just the “integration testing” directory right now, which makes it really not load-bearing: this LGTM right now.

1 Like

I think everyone is silently excited and waiting for this to land, seems like we can move forward with it at this point!

1 Like