[Mlir-commits] [mlir] 7885bf8 - [mlir][DialectConversion] Fix recursive `clone` calls.

[Sean Silva]

Author: Sean Silva
Date: 2020-10-19T15:51:46-07:00
New Revision: 7885bf8b78e21a29e27c598e0aed602de8f15260

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

LOG: [mlir][DialectConversion] Fix recursive `clone` calls.

The framework was not tracking ops created in any regions of the cloned
op.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/Builders.h
    mlir/lib/IR/Builders.cpp
    mlir/test/Dialect/Standard/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index ccf11489add0..5e12e477dd86 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -458,10 +458,8 @@ class OpBuilder : public Builder {
   /// ( leaving them alone if no entry is present).  Replaces references to
   /// cloned sub-operations to the corresponding operation that is copied,
   /// and adds those mappings to the map.
-  Operation *clone(Operation &op, BlockAndValueMapping &mapper) {
-    return insert(op.clone(mapper));
-  }
-  Operation *clone(Operation &op) { return insert(op.clone()); }
+  Operation *clone(Operation &op, BlockAndValueMapping &mapper);
+  Operation *clone(Operation &op);
 
   /// Creates a deep copy of this operation but keep the operation regions
   /// empty. Operands are remapped using `mapper` (if present), and `mapper` is

diff  --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index dfd7b10f4f0b..6833221f1672 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Matchers.h"
@@ -459,3 +460,23 @@ LogicalResult OpBuilder::tryFold(Operation *op,
 
   return success();
 }
+
+Operation *OpBuilder::clone(Operation &op, BlockAndValueMapping &mapper) {
+  Operation *newOp = op.clone(mapper);
+  // The `insert` call below handles the notification for inserting `newOp`
+  // itself. But if `newOp` has any regions, we need to notify the listener
+  // about any ops that got inserted inside those regions as part of cloning.
+  if (listener) {
+    auto walkFn = [&](Operation *walkedOp) {
+      listener->notifyOperationInserted(walkedOp);
+    };
+    for (Region &region : newOp->getRegions())
+      region.walk(walkFn);
+  }
+  return insert(newOp);
+}
+
+Operation *OpBuilder::clone(Operation &op) {
+  BlockAndValueMapping mapper;
+  return clone(op, mapper);
+}

diff  --git a/mlir/test/Dialect/Standard/bufferize.mlir b/mlir/test/Dialect/Standard/bufferize.mlir
index 1ba092aa127c..61259985c286 100644
--- a/mlir/test/Dialect/Standard/bufferize.mlir
+++ b/mlir/test/Dialect/Standard/bufferize.mlir
@@ -86,3 +86,21 @@ func @tensor_from_elements(%arg0: index, %arg1: index) -> tensor<2xindex> {
   %0 = tensor_from_elements %arg0, %arg1 : tensor<2xindex>
   return %0 : tensor<2xindex>
 }
+
+// The dynamic_tensor_from_elements op clones each op in its body.
+// Make sure that regions nested within such ops are recursively converted.
+// CHECK-LABEL: func @recursively_convert_cloned_regions
+func @recursively_convert_cloned_regions(%arg0: tensor<?xindex>, %arg1: index, %arg2: i1) -> tensor<?xindex> {
+  %tensor = dynamic_tensor_from_elements %arg1 {
+  ^bb0(%iv: index):
+    %48 = scf.if %arg2 -> (index) {
+      scf.yield %iv : index
+    } else {
+      // CHECK-NOT: extract_element
+      %50 = extract_element %arg0[%iv] : tensor<?xindex>
+      scf.yield %50 : index
+    }
+    yield %48 : index
+  } : tensor<?xindex>
+  return %tensor : tensor<?xindex>
+}