[mlir][sparse] scalarize reductions in for-loops during sparse codegen
authorAart Bik <ajcbik@google.com>
Thu, 17 Dec 2020 23:42:23 +0000 (15:42 -0800)
committerAart Bik <ajcbik@google.com>
Fri, 18 Dec 2020 00:12:21 +0000 (16:12 -0800)
commit14da25b4b2eedf8a16aae34edfefd7bcaa5ceae5
tree571258ebd081a7aa283e483b1aa846b0d54377bb
parent9887097d802d0a585807e693727ab8836790da8d
[mlir][sparse] scalarize reductions in for-loops during sparse codegen

Reductions in innermost loops become harder for the backend to disambiguate
after bufferization into memrefs, resulting in less efficient load-update-store
cycles. By scalarizing innermost reductions, the backend is more likely to assign
a register to perform the reduction (also prepares vectorization). Even though
we could scalarize reductions for more outer loops and while-loops as well,
currently scalarization is only done for chains of innermost for-loops, where
it matters most, to avoid complicating codegen unnecessary (viz. adding lots
of yield instructions).

This CL also refactors condition simplification into the merger class,
where it belongs, so that conditions are simplified only once per loop
nest and not repeatedly as was currently done. This CL also fixes a few
minor bugs, some layout issues, and comments.

Reviewed By: penpornk

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