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

[silvas]

Summary

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

[rriddle]

mlir/include/mlir/IR/Builders.h

471

This looks really awkward, can you just walk the cloned operation after you've created it?

[rriddle]

Thanks for catching this!

[silvas]

update

[silvas]

mlir/include/mlir/IR/Builders.h

471

Good idea! I was worried about retraversing all the cloned IR, but I suppose it's not a big deal, especially since it will be retraversed later as part of the dialect conversion.

[rriddle]

mlir/include/mlir/IR/Builders.h

471

If it shows up on a profile it should be easy to adopt a more complicated scheme.

mlir/lib/IR/Builders.cpp

470–473

[silvas]

adopt fix

mlir/include/mlir/IR/Builders.h

public:
   /// Note: This function does not erase the operation on a successful fold.
   LogicalResult tryFold(Operation *op, SmallVectorImpl<Value> &results);
 
   /// Creates a deep copy of the specified operation, remapping any operands
   /// that use values outside of the operation using the map that is provided
   /// ( 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
   /// updated to contain the results.
   Operation *cloneWithoutRegions(Operation &op, BlockAndValueMapping &mapper) {
     return insert(op.cloneWithoutRegions(mapper));
   }
   Operation *cloneWithoutRegions(Operation &op) {
     return insert(op.cloneWithoutRegions());
   }
   template <typename OpT> OpT cloneWithoutRegions(OpT op) {
     return cast<OpT>(cloneWithoutRegions(*op.getOperation()));
   }
 
 private:
   /// The current block this builder is inserting into.
   Block *block = nullptr;

mlir/lib/IR/Builders.cpp

 //===- Builders.cpp - Helpers for constructing MLIR Classes ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/Support/raw_ostream.h"
 using namespace mlir;
 
LogicalResult OpBuilder::tryFold(Operation *op,
   }
 
   // If we were successful, insert any generated constants.
   for (Operation *cst : generatedConstants)
     insert(cst);
 
   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);
+}

mlir/test/Dialect/Standard/bufferize.mlir

 // CHECK:           %[[C1:.*]] = constant 1 : index
 // CHECK:           store %[[ELEM1]], %[[MEMREF]][%[[C1]]]
 // CHECK:           %[[RET:.*]] = tensor_load %[[MEMREF]]
 // CHECK:           return %[[RET]] : tensor<2xindex>
 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>
+}