[Mlir-commits] [mlir] 2060155 - [mlir] NFC Replace some code snippets with equivalent method calls

[Tres Popp]

Author: Tres Popp
Date: 2021-08-10T08:22:08+02:00
New Revision: 20601554807a227383ba4ffa700f193e8abf9aeb

URL: https://github.com/llvm/llvm-project/commit/20601554807a227383ba4ffa700f193e8abf9aeb
DIFF: https://github.com/llvm/llvm-project/commit/20601554807a227383ba4ffa700f193e8abf9aeb.diff

LOG: [mlir] NFC Replace some code snippets with equivalent method calls

Replace some code snippets With scf::ForOp methods. Additionally,
share a listener at one more point (although this pattern is still
not safe to roll back currently)

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
index f4045d087d1c1..cbdc26a644dad 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp
@@ -132,8 +132,7 @@ LogicalResult mlir::scf::peelForLoop(RewriterBase &b, ForOp forOp,
       b.create<CmpIOp>(loc, CmpIPredicate::slt, splitBound, previousUb);
 
   // Create IfOp for last iteration.
-  auto resultTypes = llvm::to_vector<4>(
-      llvm::map_range(forOp.initArgs(), [](Value v) { return v.getType(); }));
+  auto resultTypes = forOp.getResultTypes();
   ifOp = b.create<scf::IfOp>(loc, resultTypes, hasMoreIter,
                              /*withElseRegion=*/!resultTypes.empty());
   forOp.replaceAllUsesWith(ifOp->getResults());
@@ -141,15 +140,15 @@ LogicalResult mlir::scf::peelForLoop(RewriterBase &b, ForOp forOp,
   // Build then case.
   BlockAndValueMapping bvm;
   bvm.map(forOp.region().getArgument(0), splitBound);
-  for (auto it : llvm::zip(forOp.region().getArguments().drop_front(),
-                           forOp->getResults())) {
+  for (auto it : llvm::zip(forOp.getRegionIterArgs(), forOp->getResults())) {
     bvm.map(std::get<0>(it), std::get<1>(it));
   }
   b.cloneRegionBefore(forOp.region(), ifOp.thenRegion(),
                       ifOp.thenRegion().begin(), bvm);
   // Build else case.
   if (!resultTypes.empty())
-    ifOp.getElseBodyBuilder().create<scf::YieldOp>(loc, forOp->getResults());
+    ifOp.getElseBodyBuilder(b.getListener())
+        .create<scf::YieldOp>(loc, forOp->getResults());
 
   return success();
 }