[Mlir-commits] [mlir] 38c407b - [mlir][SCF] Add single iteration scf.for promotion to the FuncOp level helper.

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2020-06-05T11:28:21-04:00
New Revision: 38c407bf00b5a9867f512e4beb3955b0cb387d94

URL: https://github.com/llvm/llvm-project/commit/38c407bf00b5a9867f512e4beb3955b0cb387d94
DIFF: https://github.com/llvm/llvm-project/commit/38c407bf00b5a9867f512e4beb3955b0cb387d94.diff

LOG: [mlir][SCF] Add single iteration scf.for promotion to the FuncOp level helper.

Previously only the Affine version would be folded.

Differential Revision: https://reviews.llvm.org/D81261

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/LoopUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 0d1966fcaea5..2d68f283e381 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -220,7 +220,12 @@ LogicalResult mlir::promoteIfSingleIteration(scf::ForOp forOp) {
 /// their body into the containing Block.
 void mlir::promoteSingleIterationLoops(FuncOp f) {
   // Gathers all innermost loops through a post order pruned walk.
-  f.walk([](AffineForOp forOp) { promoteIfSingleIteration(forOp); });
+  f.walk([](Operation *op) {
+    if (auto forOp = dyn_cast<AffineForOp>(op))
+      promoteIfSingleIteration(forOp);
+    else if (auto forOp = dyn_cast<scf::ForOp>(op))
+      promoteIfSingleIteration(forOp);
+  });
 }
 
 /// Generates an affine.for op with the specified lower and upper bounds