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.