[mlir][sparse][linalg] add linalg rewriting specific to sparse tensors
authorAart Bik <ajcbik@google.com>
Thu, 24 Feb 2022 00:22:42 +0000 (16:22 -0800)
committerAart Bik <ajcbik@google.com>
Thu, 24 Feb 2022 01:29:41 +0000 (17:29 -0800)
commit652b39b46f85ad826a20d3e0cec5d0db91b43daf
tree04ba4a6a508262207c2c028d29fdc1d831cb504c
parent3e3e79a9e4c378b59f5f393f556e6a84edcd8898
[mlir][sparse][linalg] add linalg rewriting specific to sparse tensors

Now that sparse tensor types are first-class citizens and the sparse compiler
is taking shape, it is time to make sure other compiler optimizations compose
well with sparse tensors. Mostly, this should be completely transparent (i.e.,
dense and sparse take the same path). However, in some cases, optimizations
only make sense in the context of sparse tensors. This is a first example of
such an optimization, where fusing a sampled elt-wise multiplication only makes
sense when the resulting kernel has a potential lower asymptotic complexity due
to the sparsity.

As an extreme example, running SDDMM with 1024x1024 matrices and a sparse
sampling matrix with only two elements runs in 463.55ms in the unfused
case but just 0.032ms in the fused case, with a speedup of 14485x that
is only possible in the exciting world of sparse computations!

Reviewed By: mravishankar

Differential Revision: https://reviews.llvm.org/D120429
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp [new file with mode: 0644]
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
mlir/test/Integration/Dialect/SparseTensor/taco/test_SDDMM.py