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

[nicolasvasilache]

Summary

Previously only the Affine version would be folded.

[ftynse]

Have you considered having a single walk on Operation * that filters inside using both types? This will spare you a walk, I don't think it's optimized per-type somehow.

[nicolasvasilache]

Commit message

mlir/lib/Transforms/Utils/LoopUtils.cpp

LogicalResult mlir::promoteIfSingleIteration(scf::ForOp forOp) {
   forOp.erase();
   return success();
 }
 
 /// Promotes all single iteration 'for' ops in `f`, i.e., moves
 /// 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
 /// while generating the right IV remappings to realize shifts for operations in
 /// its body. The operations that go into the loop body are specified in
 /// opGroupQueue starting from the specified offset, and in that order. The
 /// first element of the pair specifies the shift applied to that group of
 /// operations; the shift is multiplied by the loop step before being applied.