[Mlir-commits] [mlir] ad49330 - [MLIR][Shape] Fix `shape_of` lowering to `scf`

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-07-15T15:38:09Z
New Revision: ad493300322099787cab5f3a9f7310af0f9b5e6c

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

LOG: [MLIR][Shape] Fix `shape_of` lowering to `scf`

The use of the `scf.for` callback builder does not allow for a rollback of the
emitted conversions. Instead, we populate the loop body through the conversion
rewriter directly.

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

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 55ebae99af53..1f1134757b3a 100644
--- a/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
+++ b/mlir/lib/Conversion/ShapeToSCF/ShapeToSCF.cpp
@@ -103,14 +103,15 @@ 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);
-  rewriter.create<scf::ForOp>(
-      loc, zeroVal, rankVal, oneVal, ValueRange(),
-      [&](OpBuilder &b, Location loc, Value iVal, ValueRange args) {
-        auto dimVal = b.create<DimOp>(loc, tensorVal, iVal);
-        auto dimIntVal = b.create<IndexCastOp>(loc, dimVal, i64Ty);
-        b.create<StoreOp>(loc, dimIntVal, memVal, ValueRange({iVal}));
-        b.create<scf::YieldOp>(loc);
-      });
+  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});
+  }
 
   // Load extents to tensor value.
   auto shapeIntVal = rewriter.create<TensorLoadOp>(loc, memVal);