From 7885bf8b78e21a29e27c598e0aed602de8f15260 Mon Sep 17 00:00:00 2001
From: Sean Silva <silvasean@google.com>
Date: Sun, 18 Oct 2020 21:10:55 -0700
Subject: [PATCH] [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
---
 mlir/include/mlir/IR/Builders.h           |  6 ++----
 mlir/lib/IR/Builders.cpp                  | 21 +++++++++++++++++++++
 mlir/test/Dialect/Standard/bufferize.mlir | 18 ++++++++++++++++++
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index ccf1148..5e12e47 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -458,10 +458,8 @@ public:
   /// ( 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 dfd7b10f..6833221 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 1ba092a..6125998 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>
+}
-- 
2.7.4