[mlir][sparse] fix bug in reduction chain
authorAart Bik <ajcbik@google.com>
Wed, 3 Mar 2021 20:37:19 +0000 (12:37 -0800)
committerAart Bik <ajcbik@google.com>
Thu, 4 Mar 2021 01:38:22 +0000 (17:38 -0800)
commit553cb6d473c55556146e5c9cc9cd31af99ce99cb
treee5a642d0900283e47afc051e437bd4afa7afd243
parentdbf41ddaa3200fc5b66a5694d84b322a26b13a5a
[mlir][sparse] fix bug in reduction chain

Found with exhaustive testing, it is possible that a while loop
appears in between chainable for loops. As long as we don't
scalarize reductions in while loops, this means we need to
terminate the chain at the while. This also refactors the
reduction code into more readable helper methods.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D97886
mlir/lib/Dialect/Linalg/Transforms/Sparsification.cpp
mlir/test/Dialect/Linalg/sparse_1d.mlir