Cleanup the 'clone' methods and remove the need to explicitly pass in the context.

This also adds a new 'Region::cloneInto' method that accepts an insertion position.

PiperOrigin-RevId: 255504640

diff --git a/mlir/include/mlir/IR/BlockAndValueMapping.h b/mlir/include/mlir/IR/BlockAndValueMapping.h
index e370d8b..bd69aa2 100644 (file)
--- a/mlir/include/mlir/IR/BlockAndValueMapping.h
+++ b/mlir/include/mlir/IR/BlockAndValueMapping.h
@@ -65,6 +65,14 @@ public:
     return lookupOrValue(from, from);
   }
 
+  /// Lookup a mapped value within the map. This asserts the provided value
+  /// exists within the map.
+  template <typename T> T *lookup(T *from) const {
+    auto *result = lookupOrNull(from);
+    assert(result && "expected 'from' to be contained within the map");
+    return result;
+  }
+
   /// Clears all mappings held by the mapper.
   void clear() { valueMap.clear(); }
 

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 27d0c28..6ce5c22 100644 (file)
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -346,12 +346,12 @@ public:
   /// cloned sub-operations to the corresponding operation that is copied,
   /// and adds those mappings to the map.
   Operation *clone(Operation &op, BlockAndValueMapping &mapper) {
-    Operation *cloneOp = op.clone(mapper, getContext());
+    Operation *cloneOp = op.clone(mapper);
     block->getOperations().insert(insertPoint, cloneOp);
     return cloneOp;
   }
   Operation *clone(Operation &op) {
-    Operation *cloneOp = op.clone(getContext());
+    Operation *cloneOp = op.clone();
     block->getOperations().insert(insertPoint, cloneOp);
     return cloneOp;
   }
@@ -360,12 +360,12 @@ public:
   /// empty. Operands are remapped using `mapper` (if present), and `mapper` is
   /// updated to contain the results.
   Operation *cloneWithoutRegions(Operation &op, BlockAndValueMapping &mapper) {
-    Operation *cloneOp = op.cloneWithoutRegions(mapper, getContext());
+    Operation *cloneOp = op.cloneWithoutRegions(mapper);
     block->getOperations().insert(insertPoint, cloneOp);
     return cloneOp;
   }
   Operation *cloneWithoutRegions(Operation &op) {
-    Operation *cloneOp = op.cloneWithoutRegions(getContext());
+    Operation *cloneOp = op.cloneWithoutRegions();
     block->getOperations().insert(insertPoint, cloneOp);
     return cloneOp;
   }

diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 9b310aa..e532399 100644 (file)
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -91,15 +91,14 @@ public:
   /// 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(BlockAndValueMapping &mapper, MLIRContext *context);
-  Operation *clone(MLIRContext *context);
+  Operation *clone(BlockAndValueMapping &mapper);
+  Operation *clone();
 
   /// Create 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(BlockAndValueMapping &mapper,
-                                 MLIRContext *context);
-  Operation *cloneWithoutRegions(MLIRContext *context);
+  Operation *cloneWithoutRegions(BlockAndValueMapping &mapper);
+  Operation *cloneWithoutRegions();
 
   /// Returns the operation block that contains this operation.
   Block *getBlock() { return block; }

diff --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index cf366af..2189ad4 100644 (file)
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -92,8 +92,10 @@ public:
   /// cloned blocks are appended to the back of dest. If the mapper
   /// contains entries for block arguments, these arguments are not included
   /// in the respective cloned block.
-  void cloneInto(Region *dest, BlockAndValueMapping &mapper,
-                 MLIRContext *context);
+  void cloneInto(Region *dest, BlockAndValueMapping &mapper);
+  /// Clone this region into 'dest' before the given position in 'dest'.
+  void cloneInto(Region *dest, Region::iterator destPos,
+                 BlockAndValueMapping &mapper);
 
   /// Takes body of another region (that region will have no body after this
   /// operation completes).  The current body of this region is cleared.

diff --git a/mlir/lib/IR/Function.cpp b/mlir/lib/IR/Function.cpp
index 424548a..6f68397 100644 (file)
--- a/mlir/lib/IR/Function.cpp
+++ b/mlir/lib/IR/Function.cpp
@@ -135,7 +135,7 @@ void Function::cloneInto(Function *dest, BlockAndValueMapping &mapper) {
   dest->setAttrs(newAttrs.takeVector());
 
   // Clone the body.
-  body.cloneInto(&dest->body, mapper, getContext());
+  body.cloneInto(&dest->body, mapper);
 }
 
 /// Create a deep copy of this function and all of its blocks, remapping

diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index d26ddeb..83171f1 100644 (file)
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -544,8 +544,7 @@ InFlightDiagnostic Operation::emitOpError(const Twine &message) {
 /// Create 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 *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper,
-                                          MLIRContext *context) {
+Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper) {
   SmallVector<Value *, 8> operands;
   SmallVector<Block *, 2> successors;
 
@@ -582,7 +581,7 @@ Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper,
   unsigned numRegions = getNumRegions();
   auto *newOp = Operation::create(getLoc(), getName(), operands, resultTypes,
                                   attrs, successors, numRegions,
-                                  hasResizableOperandsList(), context);
+                                  hasResizableOperandsList(), getContext());
 
   // Remember the mapping of any results.
   for (unsigned i = 0, e = getNumResults(); i != e; ++i)
@@ -591,9 +590,9 @@ Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper,
   return newOp;
 }
 
-Operation *Operation::cloneWithoutRegions(MLIRContext *context) {
+Operation *Operation::cloneWithoutRegions() {
   BlockAndValueMapping mapper;
-  return cloneWithoutRegions(mapper, context);
+  return cloneWithoutRegions(mapper);
 }
 
 /// Create a deep copy of this operation, remapping any operands that use
@@ -601,20 +600,19 @@ Operation *Operation::cloneWithoutRegions(MLIRContext *context) {
 /// 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 *Operation::clone(BlockAndValueMapping &mapper,
-                            MLIRContext *context) {
-  auto *newOp = cloneWithoutRegions(mapper, context);
+Operation *Operation::clone(BlockAndValueMapping &mapper) {
+  auto *newOp = cloneWithoutRegions(mapper);
 
   // Clone the regions.
   for (unsigned i = 0; i != numRegions; ++i)
-    getRegion(i).cloneInto(&newOp->getRegion(i), mapper, context);
+    getRegion(i).cloneInto(&newOp->getRegion(i), mapper);
 
   return newOp;
 }
 
-Operation *Operation::clone(MLIRContext *context) {
+Operation *Operation::clone() {
   BlockAndValueMapping mapper;
-  return clone(mapper, context);
+  return clone(mapper);
 }
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index baf401f..992d911 100644 (file)
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -77,15 +77,20 @@ bool Region::isProperAncestor(Region *other) {
 
 /// Clone the internal blocks from this region into `dest`. Any
 /// cloned blocks are appended to the back of dest.
-void Region::cloneInto(Region *dest, BlockAndValueMapping &mapper,
-                       MLIRContext *context) {
+void Region::cloneInto(Region *dest, BlockAndValueMapping &mapper) {
+  assert(dest && "expected valid region to clone into");
+  cloneInto(dest, dest->end(), mapper);
+}
+
+/// Clone this region into 'dest' before the given position in 'dest'.
+void Region::cloneInto(Region *dest, Region::iterator destPos,
+                       BlockAndValueMapping &mapper) {
   assert(dest && "expected valid region to clone into");
 
   // If the list is empty there is nothing to clone.
   if (empty())
     return;
 
-  iterator lastOldBlock = --dest->end();
   for (Block &block : *this) {
     Block *newBlock = new Block();
     mapper.map(&block, newBlock);
@@ -99,9 +104,9 @@ void Region::cloneInto(Region *dest, BlockAndValueMapping &mapper,
 
     // Clone and remap the operations within this block.
     for (auto &op : block)
-      newBlock->push_back(op.clone(mapper, context));
+      newBlock->push_back(op.clone(mapper));
 
-    dest->push_back(newBlock);
+    dest->getBlocks().insert(destPos, newBlock);
   }
 
   // Now that each of the blocks have been cloned, go through and remap the
@@ -115,7 +120,7 @@ void Region::cloneInto(Region *dest, BlockAndValueMapping &mapper,
         succOp.set(mappedOp);
   };
 
-  for (auto it = std::next(lastOldBlock), e = dest->end(); it != e; ++it)
+  for (iterator it(mapper.lookup(&front())); it != destPos; ++it)
     it->walk(remapOperands);
 }