[mlir][sparse] implement full reduction "scalarization" across loop nests
authorAart Bik <ajcbik@google.com>
Wed, 3 Nov 2021 02:58:33 +0000 (19:58 -0700)
committerAart Bik <ajcbik@google.com>
Fri, 5 Nov 2021 00:38:47 +0000 (17:38 -0700)
commit7373cabcda8f5c0ed83cf40034ff69bc47a4a8c9
tree41b55064fcc5d14affe0d557a8b8f8186f44c219
parent3f1aca58df8f646c32737b9fc01c51325fba88c2
[mlir][sparse] implement full reduction "scalarization" across loop nests

The earlier reduction "scalarization" was only applied to a chain of
*innermost* and *for* loops. This revision generalizes this to any
nesting of for- and while-loops. This implies that reductions can be
implemented with a lot less load and store operations. The chaining
is implemented with a forest of yield statements (but not as bad as
when we would also include the while-induction).

Fixes https://bugs.llvm.org/show_bug.cgi?id=52311

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D113078
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
mlir/test/Dialect/SparseTensor/sparse_1d.mlir
mlir/test/Dialect/SparseTensor/sparse_2d.mlir
mlir/test/Dialect/SparseTensor/sparse_3d.mlir
mlir/test/Dialect/SparseTensor/sparse_lower.mlir
mlir/test/Dialect/SparseTensor/sparse_lower_inplace.mlir
mlir/test/Dialect/SparseTensor/sparse_perm_lower.mlir