[all-commits] [llvm/llvm-project] 652b39: [mlir][sparse][linalg] add linalg rewriting specif...

[Aart Bik via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 652b39b46f85ad826a20d3e0cec5d0db91b43daf
      https://github.com/llvm/llvm-project/commit/652b39b46f85ad826a20d3e0cec5d0db91b43daf
  Author: Aart Bik <ajcbik at google.com>
  Date:   2022-02-23 (Wed, 23 Feb 2022)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
    M mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
    M mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp
    A mlir/lib/Dialect/Linalg/Transforms/SparseTensorRewriting.cpp
    M mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
    M mlir/test/Integration/Dialect/SparseTensor/taco/test_SDDMM.py

  Log Message:
  -----------
  [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