To wrap-up the discussion in [RFC] Sparse tensor support in torch-mlir, this is the current status, spread over three different repositories (PyTorch, Torch-MLIR, and MLIR), and now even a fourth (MPACT).
(1) Sparsity in PyTorch is now propagated through sparse metadata in exported FX graphs (see this feature request, which also contains links to the PyTorch discussion forums and actual pull-requests).
(2) This sparse metadata is translated to actual sparse tensor types in Torch-MLIR, by means of the SparseTensor dialect and e.g. lowering the Torch dialect to StableHLO or Linalg.
(3) Sparse tensor types were already supported in MLIR, although various improvements have been made since to clean up the incoming IR as well as simplify dealing with ABI issues for sparse tensor types when interacting with runtimes like PyTorch or JAX.
and now also
(4) A new repo The MPACT Compiler, which puts all this together in a ML compiler that can run PyTorch models using MLIR and LLVM as backend, including sparsity support.
More about this can be found at the MPACT Project Page, which also includes links to the MLIR Sparsifier page and other related projects.
Note that the current compiler is functional but not necessarily performant. Also, the current approach focuses on inference, leaving some of the autograd issues for training for later inspection. But with a lot of the endless “engineering plumbing” to connect torch sparse with MLIR out of the way, hopefully the more interesting research can finally start!
