[Mlir-commits] [mlir] c430c21 - [MLIR][Shape] Use callback builder again

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-07-16T13:58:38Z
New Revision: c430c21202c377cfb9fce0e7272f7208d1e8a531

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

LOG: [MLIR][Shape] Use callback builder again

The issue that callback builders caused during rollback of conversion patterns
has been resolved. We can use them again.
See https://bugs.llvm.org/show_bug.cgi?id=46731

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

Added: 
    

Modified: 
    mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
index 0caaacd75bed..f050c4eb2278 100644
--- a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
+++ b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
@@ -104,15 +104,14 @@ ShapeOfOpConverter::matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
   // Copy shape extents to stack-allocated memory.
   auto zeroVal = rewriter.create<ConstantIndexOp>(loc, 0);
   auto oneVal = rewriter.create<ConstantIndexOp>(loc, 1);
-  auto loop = rewriter.create<scf::ForOp>(loc, zeroVal, rankVal, oneVal);
-  {
-    OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToStart(loop.getBody());
-    auto iVal = loop.getInductionVar();
-    auto dimVal = rewriter.create<DimOp>(loc, tensorVal, iVal);
-    auto dimIntVal = rewriter.create<IndexCastOp>(loc, dimVal, i64Ty);
-    rewriter.create<StoreOp>(loc, dimIntVal, memVal, ValueRange{iVal});
-  }
+  rewriter.create<scf::ForOp>(
+      loc, zeroVal, rankVal, oneVal, llvm::None,
+      [&](OpBuilder &b, Location loc, Value iVal, ValueRange args) {
+        auto dimVal = rewriter.create<DimOp>(loc, tensorVal, iVal);
+        auto dimIntVal = rewriter.create<IndexCastOp>(loc, dimVal, i64Ty);
+        rewriter.create<StoreOp>(loc, dimIntVal, memVal, ValueRange{iVal});
+        rewriter.create<scf::YieldOp>(loc);
+      });
 
   // Load extents to tensor value.
   auto shapeIntVal = rewriter.create<TensorLoadOp>(loc, memVal);