Add cloning functionality to Block and Function, this also adds support for remapping successor block operands of terminator operations. We define a new BlockAndValueMapping class to simplify mapping between cloned values.

PiperOrigin-RevId: 230768759

mlir/include/mlir/IR/Block.h

@@ -28,6 +28,7 @@
 namespace mlir {
 class IfInst;
 class BlockList;
+class BlockAndValueMapping;
 
 template <typename BlockType> class PredecessorIterator;
 template <typename BlockType> class SuccessorIterator;
@@ -358,6 +359,13 @@ public:
     return const_cast<BlockList *>(this)->getContainingFunction();
   }
 
+  /// Clone the internal blocks from this block list into dest. Any
+  /// 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(BlockList *dest, BlockAndValueMapping &mapper,
+                 MLIRContext *context) const;
+
 private:
   BlockListType blocks;
 

mlir/include/mlir/IR/BlockAndValueMapping.h

@@ -0,0 +1,85 @@
+//===- BlockAndValueMapping.h -----------------------------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines a utility class for maintaining a mapping for multiple
+// value types.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_BLOCKANDVALUEMAPPING_H
+#define MLIR_IR_BLOCKANDVALUEMAPPING_H
+
+#include "mlir/IR/Block.h"
+
+namespace mlir {
+// This is a utility class for mapping one set of values to another. New
+// mappings can be inserted via 'map'. Existing mappings can be
+// found via the 'lookup*' functions. There are two variants that differ only in
+// return value when an existing is not found for the provided key.
+// 'lookupOrNull' returns nullptr where as 'lookupOrDefault' will return the
+// lookup key.
+class BlockAndValueMapping {
+public:
+  /// Inserts a new mapping for 'from' to 'to'. If there is an existing mapping,
+  /// it is overwritten.
+  void map(const Block *from, Block *to) { valueMap[from] = to; }
+  void map(const Value *from, Value *to) { valueMap[from] = to; }
+
+  /// Erases a mapping for 'from'.
+  void erase(const IRObjectWithUseList *from) { valueMap.erase(from); }
+
+  /// Checks to see if a mapping for 'from' exists.
+  bool contains(const IRObjectWithUseList *from) const {
+    return valueMap.count(from);
+  }
+
+  /// Lookup a mapped value within the map. If a mapping for the provided value
+  /// does not exist then return nullptr.
+  Block *lookupOrNull(const Block *from) const {
+    return lookupOrValue(from, (Block *)nullptr);
+  }
+  Value *lookupOrNull(const Value *from) const {
+    return lookupOrValue(from, (Value *)nullptr);
+  }
+
+  /// Lookup a mapped value within the map. If a mapping for the provided value
+  /// does not exist then return the provided value.
+  Block *lookupOrDefault(Block *from) const {
+    return lookupOrValue(from, from);
+  }
+  Value *lookupOrDefault(Value *from) const {
+    return lookupOrValue(from, from);
+  }
+
+  /// Clears all mappings held by the mapper.
+  void clear() { valueMap.clear(); }
+
+private:
+  /// Utility lookupOrValue that looks up an existing key or returns the
+  /// provided value. This function assumes that if a mapping does exist, then
+  /// it is of 'T' type.
+  template <typename T> T *lookupOrValue(const T *from, T *value) const {
+    auto it = valueMap.find(from);
+    return it != valueMap.end() ? static_cast<T *>(it->second) : value;
+  }
+
+  llvm::DenseMap<const IRObjectWithUseList *, IRObjectWithUseList *> valueMap;
+};
+
+} // end namespace mlir
+
+#endif // MLIR_IR_BLOCKANDVALUEMAPPING_H

mlir/include/mlir/IR/Builders.h

@@ -24,6 +24,7 @@
 namespace mlir {
 
 class AffineExpr;
+class BlockAndValueMapping;
 class Module;
 class UnknownLoc;
 class UniquedFilename;
@@ -265,9 +266,13 @@ public:
   /// ( leaving them alone if no entry is present).  Replaces references to
   /// cloned sub-instructions to the corresponding instruction that is copied,
   /// and adds those mappings to the map.
-  Instruction *clone(const Instruction &inst,
-                     OperationInst::OperandMapTy &operandMapping) {
-    Instruction *cloneInst = inst.clone(operandMapping, getContext());
+  Instruction *clone(const Instruction &inst, BlockAndValueMapping &mapper) {
+    Instruction *cloneInst = inst.clone(mapper, getContext());
+    block->getInstructions().insert(insertPoint, cloneInst);
+    return cloneInst;
+  }
+  Instruction *clone(const Instruction &inst) {
+    Instruction *cloneInst = inst.clone(getContext());
     block->getInstructions().insert(insertPoint, cloneInst);
     return cloneInst;
   }

mlir/include/mlir/IR/Function.h

@@ -34,6 +34,7 @@
 
 namespace mlir {
 class AttributeListStorage;
+class BlockAndValueMapping;
 class FunctionType;
 class MLIRContext;
 class Module;
@@ -185,7 +186,24 @@ public:
   /// target linked.
   void viewGraph() const;
 
+  /// Create a deep copy of this function and all of its blocks, remapping
+  /// any operands that use values outside of the function using the map that is
+  /// provided (leaving them alone if no entry is present). If the mapper
+  /// contains entries for function arguments, these arguments are not included
+  /// in the new function. Replaces references to cloned sub-values with the
+  /// corresponding value that is copied, and adds those mappings to the mapper.
+  Function *clone(BlockAndValueMapping &mapper) const;
+  Function *clone() const;
+
+  /// Clone the internal blocks and attributes from this function into dest. Any
+  /// cloned blocks are appended to the back of dest. This function asserts that
+  /// the attributes of the current function and dest are compatible.
+  void cloneInto(Function *dest, BlockAndValueMapping &mapper) const;
+
 private:
+  /// Set the attributes held by this function.
+  void setAttributes(ArrayRef<NamedAttribute> attrs = {});
+
   /// The name of the function.
   Identifier name;
 

mlir/include/mlir/IR/Instruction.h

@@ -30,6 +30,7 @@
 
 namespace mlir {
 class Block;
+class BlockAndValueMapping;
 class Location;
 class ForInst;
 class MLIRContext;
@@ -82,17 +83,12 @@ public:
   /// Remove this instruction from its parent block and delete it.
   void erase();
 
-  // This is a verbose type used by the clone method below.
-  using OperandMapTy =
-      DenseMap<const Value *, Value *, llvm::DenseMapInfo<const Value *>,
-               llvm::detail::DenseMapPair<const Value *, Value *>>;
-
   /// Create a deep copy of this instruction, remapping any operands that use
   /// values outside of the instruction using the map that is provided (leaving
   /// them alone if no entry is present).  Replaces references to cloned
   /// sub-instructions to the corresponding instruction that is copied, and adds
   /// those mappings to the map.
-  Instruction *clone(OperandMapTy &operandMap, MLIRContext *context) const;
+  Instruction *clone(BlockAndValueMapping &mapper, MLIRContext *context) const;
   Instruction *clone(MLIRContext *context) const;
 
   /// Returns the instruction block that contains this instruction.

mlir/lib/Analysis/Utils.cpp

@@ -457,8 +457,7 @@ ForInst *mlir::insertBackwardComputationSlice(
   // of the loop at 'dstLoopDepth' in 'dstLoopIVs'.
   auto *dstForInst = dstLoopIVs[dstLoopDepth - 1];
   FuncBuilder b(dstForInst->getBody(), dstForInst->getBody()->begin());
-  DenseMap<const Value *, Value *> operandMap;
-  auto *sliceLoopNest = cast<ForInst>(b.clone(*srcLoopIVs[0], operandMap));
+  auto *sliceLoopNest = cast<ForInst>(b.clone(*srcLoopIVs[0]));
 
   Instruction *sliceInst =
       getInstAtPosition(positions, /*level=*/0, sliceLoopNest->getBody());

mlir/lib/IR/Block.cpp

@@ -16,7 +16,9 @@
 // =============================================================================
 
 #include "mlir/IR/Block.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/InstVisitor.h"
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
@@ -217,6 +219,63 @@ Function *BlockList::getContainingFunction() {
   return container.dyn_cast<Function *>();
 }
 
+/// Clone the internal blocks from this block list into dest. Any
+/// cloned blocks are appended to the back of dest.
+void BlockList::cloneInto(BlockList *dest, BlockAndValueMapping &mapper,
+                          MLIRContext *context) const {
+  assert(dest && "expected valid block list to clone into");
+
+  // If the list is empty there is nothing to clone.
+  if (empty())
+    return;
+
+  Block *lastOldBlock = &dest->back();
+  for (const Block &block : *this) {
+    Block *newBlock = new Block();
+    mapper.map(&block, newBlock);
+
+    // Clone the block arguments. The user might be deleting arguments to the
+    // block by specifying them in the mapper. If so, we don't add the
+    // argument to the cloned block.
+    for (const auto *arg : block.getArguments())
+      if (!mapper.contains(arg))
+        mapper.map(arg, newBlock->addArgument(arg->getType()));
+
+    // Clone and remap the instructions within this block.
+    for (const auto &inst : block)
+      newBlock->push_back(inst.clone(mapper, context));
+
+    dest->push_back(newBlock);
+  }
+
+  // Now that each of the blocks have been cloned, go through and remap the
+  // operands of each of the instructions.
+  struct Walker : public InstWalker<Walker> {
+    BlockAndValueMapping &mapper;
+    Walker(BlockAndValueMapping &mapper) : mapper(mapper) {}
+
+    /// Remap the instruction operands.
+    void visitInstruction(Instruction *inst) {
+      for (auto &instOp : inst->getInstOperands())
+        if (auto *mappedOp = mapper.lookupOrNull(instOp.get()))
+          instOp.set(mappedOp);
+    }
+    // Remap the successor block operands.
+    void visitOperationInst(OperationInst *opInst) {
+      if (!opInst->isTerminator())
+        return;
+      for (auto &succOp : opInst->getBlockOperands())
+        if (auto *mappedOp = mapper.lookupOrNull(succOp.get()))
+          succOp.set(mappedOp);
+    }
+  };
+
+  Walker v(mapper);
+  for (auto it = std::next(lastOldBlock->getIterator()), e = dest->end();
+       it != e; ++it)
+    v.walk(it->begin(), it->end());
+}
+
 BlockList *llvm::ilist_traits<::mlir::Block>::getContainingBlockList() {
   size_t Offset(
       size_t(&((BlockList *)nullptr->*BlockList::getSublistAccess(nullptr))));

mlir/lib/IR/Function.cpp

@@ -18,10 +18,12 @@
 #include "mlir/IR/Function.h"
 #include "AttributeListStorage.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/InstVisitor.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Types.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringRef.h"
 using namespace mlir;
@@ -30,7 +32,7 @@ Function::Function(Location location, StringRef name, FunctionType type,
                    ArrayRef<NamedAttribute> attrs)
     : name(Identifier::get(name, type.getContext())), location(location),
       type(type), blocks(this) {
-  this->attrs = AttributeListStorage::get(attrs, getContext());
+  setAttributes(attrs);
 }
 
 Function::~Function() {
@@ -143,6 +145,62 @@ bool Function::emitError(const Twine &message) const {
   return getContext()->emitError(getLoc(), message);
 }
 
+/// Clone the internal blocks from this function into dest and all attributes
+/// from this function to dest.
+void Function::cloneInto(Function *dest, BlockAndValueMapping &mapper) const {
+  // Add the attributes of this function to dest.
+  llvm::MapVector<Identifier, Attribute> newAttrs;
+  for (auto &attr : dest->getAttrs())
+    newAttrs.insert(attr);
+  for (auto &attr : getAttrs()) {
+    auto insertPair = newAttrs.insert(attr);
+
+    // TODO(riverriddle) Verify that the two functions have compatible
+    // attributes.
+    (void)insertPair;
+    assert((insertPair.second || insertPair.first->second == attr.second) &&
+           "the two functions have incompatible attributes");
+  }
+  dest->setAttributes(newAttrs.takeVector());
+
+  // Clone the block list.
+  blocks.cloneInto(&dest->blocks, mapper, dest->getContext());
+}
+
+/// Create a deep copy of this function and all of its blocks, remapping
+/// any operands that use values outside of the function using the map that is
+/// provided (leaving them alone if no entry is present). Replaces references
+/// to cloned sub-values with the corresponding value that is copied, and adds
+/// those mappings to the mapper.
+Function *Function::clone(BlockAndValueMapping &mapper) const {
+  FunctionType newType = type;
+
+  // If the function has a body, then the user might be deleting arguments to
+  // the function by specifying them in the mapper. If so, we don't add the
+  // argument to the input type vector.
+  if (!empty()) {
+    SmallVector<Type, 4> inputTypes;
+    for (unsigned i = 0, e = getNumArguments(); i != e; ++i)
+      if (!mapper.contains(getArgument(i)))
+        inputTypes.push_back(type.getInput(i));
+    newType = FunctionType::get(inputTypes, type.getResults(), getContext());
+  }
+
+  // Create a new function and clone the current function into it.
+  Function *newFunc = new Function(getLoc(), getName(), newType);
+  cloneInto(newFunc, mapper);
+  return newFunc;
+}
+Function *Function::clone() const {
+  BlockAndValueMapping mapper;
+  return clone(mapper);
+}
+
+/// Set the attributes held by this function.
+void Function::setAttributes(ArrayRef<NamedAttribute> attrs) {
+  this->attrs = AttributeListStorage::get(attrs, getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // Function implementation.
 //===----------------------------------------------------------------------===//

mlir/lib/IR/Instruction.cpp

@@ -1,5 +1,4 @@
-//===- Instruction.cpp - MLIR Instruction Classes
-//----------------------------===//
+//===- Instruction.cpp - MLIR Instruction Classes -------------------------===//
 //
 // Copyright 2019 The MLIR Authors.
 //
@@ -19,6 +18,7 @@
 #include "AttributeListStorage.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/InstVisitor.h"
@@ -786,15 +786,8 @@ MLIRContext *IfInst::getContext() const {
 /// them alone if no entry is present).  Replaces references to cloned
 /// sub-instructions to the corresponding instruction that is copied, and adds
 /// those mappings to the map.
-Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
+Instruction *Instruction::clone(BlockAndValueMapping &mapper,
                                 MLIRContext *context) const {
-  // If the specified value is in operandMap, return the remapped value.
-  // Otherwise return the value itself.
-  auto remapOperand = [&](const Value *value) -> Value * {
-    auto it = operandMap.find(value);
-    return it != operandMap.end() ? it->second : const_cast<Value *>(value);
-  };
-
   SmallVector<Value *, 8> operands;
   SmallVector<Block *, 2> successors;
   if (auto *opInst = dyn_cast<OperationInst>(this)) {
@@ -803,7 +796,8 @@ Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
     if (!opInst->isTerminator()) {
       // Non-terminators just add all the operands.
       for (auto *opValue : getOperands())
-        operands.push_back(remapOperand(opValue));
+        operands.push_back(
+            mapper.lookupOrDefault(const_cast<Value *>(opValue)));
     } else {
       // We add the operands separated by nullptr's for each successor.
       unsigned firstSuccOperand = opInst->getNumSuccessors()
@@ -813,19 +807,22 @@ Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
 
       unsigned i = 0;
       for (; i != firstSuccOperand; ++i)
-        operands.push_back(remapOperand(InstOperands[i].get()));
+        operands.push_back(
+            mapper.lookupOrDefault(const_cast<Value *>(InstOperands[i].get())));
 
       successors.reserve(opInst->getNumSuccessors());
       for (unsigned succ = 0, e = opInst->getNumSuccessors(); succ != e;
            ++succ) {
-        successors.push_back(const_cast<Block *>(opInst->getSuccessor(succ)));
+        successors.push_back(mapper.lookupOrDefault(
+            const_cast<Block *>(opInst->getSuccessor(succ))));
 
         // Add sentinel to delineate successor operands.
         operands.push_back(nullptr);
 
         // Remap the successors operands.
         for (auto *operand : opInst->getSuccessorOperands(succ))
-          operands.push_back(remapOperand(operand));
+          operands.push_back(
+              mapper.lookupOrDefault(const_cast<Value *>(operand)));
       }
     }
 
@@ -838,13 +835,13 @@ Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
                                         successors, context);
     // Remember the mapping of any results.
     for (unsigned i = 0, e = opInst->getNumResults(); i != e; ++i)
-      operandMap[opInst->getResult(i)] = newOp->getResult(i);
+      mapper.map(opInst->getResult(i), newOp->getResult(i));
     return newOp;
   }
 
   operands.reserve(getNumOperands());
   for (auto *opValue : getOperands())
-    operands.push_back(remapOperand(opValue));
+    operands.push_back(mapper.lookupOrDefault(const_cast<Value *>(opValue)));
 
   if (auto *forInst = dyn_cast<ForInst>(this)) {
     auto lbMap = forInst->getLowerBoundMap();
@@ -856,11 +853,11 @@ Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
         ubMap, forInst->getStep());
 
     // Remember the induction variable mapping.
-    operandMap[forInst] = newFor;
+    mapper.map(forInst, newFor);
 
     // Recursively clone the body of the for loop.
     for (auto &subInst : *forInst->getBody())
-      newFor->getBody()->push_back(subInst.clone(operandMap, context));
+      newFor->getBody()->push_back(subInst.clone(mapper, context));
 
     return newFor;
   }
@@ -871,18 +868,18 @@ Instruction *Instruction::clone(DenseMap<const Value *, Value *> &operandMap,
 
   auto *resultThen = newIf->getThen();
   for (auto &childInst : *ifInst->getThen())
-    resultThen->push_back(childInst.clone(operandMap, context));
+    resultThen->push_back(childInst.clone(mapper, context));
 
   if (ifInst->hasElse()) {
     auto *resultElse = newIf->createElse();
     for (auto &childInst : *ifInst->getElse())
-      resultElse->push_back(childInst.clone(operandMap, context));
+      resultElse->push_back(childInst.clone(mapper, context));
   }
 
   return newIf;
 }
 
 Instruction *Instruction::clone(MLIRContext *context) const {
-  DenseMap<const Value *, Value *> operandMap;
-  return clone(operandMap, context);
+  BlockAndValueMapping mapper;
+  return clone(mapper, context);
 }

mlir/lib/Transforms/LoopUnrollAndJam.cpp

@@ -46,6 +46,7 @@
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/InstVisitor.h"
@@ -190,11 +191,10 @@ bool mlir::loopUnrollJamByFactor(ForInst *forInst, uint64_t unrollJamFactor) {
   // unrollJamFactor.
   if (mayBeConstantTripCount.hasValue() &&
       mayBeConstantTripCount.getValue() % unrollJamFactor != 0) {
-    DenseMap<const Value *, Value *> operandMap;
     // Insert the cleanup loop right after 'forInst'.
     FuncBuilder builder(forInst->getBlock(),
                         std::next(Block::iterator(forInst)));
-    auto *cleanupForInst = cast<ForInst>(builder.clone(*forInst, operandMap));
+    auto *cleanupForInst = cast<ForInst>(builder.clone(*forInst));
     cleanupForInst->setLowerBoundMap(
         getCleanupLoopLowerBound(*forInst, unrollJamFactor, &builder));
 
@@ -217,7 +217,7 @@ bool mlir::loopUnrollJamByFactor(ForInst *forInst, uint64_t unrollJamFactor) {
 
     // Unroll and jam (appends unrollJamFactor-1 additional copies).
     for (unsigned i = 1; i < unrollJamFactor; i++) {
-      DenseMap<const Value *, Value *> operandMapping;
+      BlockAndValueMapping operandMapping;
 
       // If the induction variable is used, create a remapping to the value for
       // this unrolled instance.
@@ -228,7 +228,7 @@ bool mlir::loopUnrollJamByFactor(ForInst *forInst, uint64_t unrollJamFactor) {
         auto *ivUnroll =
             builder.create<AffineApplyOp>(forInst->getLoc(), bumpMap, forInst)
                 ->getResult(0);
-        operandMapping[forInst] = ivUnroll;
+        operandMapping.map(forInst, ivUnroll);
       }
       // Clone the sub-block being unroll-jammed.
       for (auto it = subBlock.first; it != std::next(subBlock.second); ++it) {

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -24,6 +24,7 @@
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/InstVisitor.h"
@@ -161,7 +162,7 @@ generateLoop(AffineMap lbMap, AffineMap ubMap,
   auto *loopChunk = b->createFor(srcForInst->getLoc(), lbOperands, lbMap,
                                  ubOperands, ubMap, srcForInst->getStep());
 
-  OperationInst::OperandMapTy operandMap;
+  BlockAndValueMapping operandMap;
 
   for (auto it = instGroupQueue.begin() + offset, e = instGroupQueue.end();
        it != e; ++it) {
@@ -179,9 +180,9 @@ generateLoop(AffineMap lbMap, AffineMap ubMap,
                                srcForInst->getStep() * shift)),
                            loopChunk)
                           ->getResult(0);
-      operandMap[srcForInst] = ivRemap;
+      operandMap.map(srcForInst, ivRemap);
     } else {
-      operandMap[srcForInst] = loopChunk;
+      operandMap.map(srcForInst, loopChunk);
     }
     for (auto *inst : insts) {
       loopChunk->getBody()->push_back(inst->clone(operandMap, b->getContext()));
@@ -386,9 +387,8 @@ bool mlir::loopUnrollByFactor(ForInst *forInst, uint64_t unrollFactor) {
 
   // Generate the cleanup loop if trip count isn't a multiple of unrollFactor.
   if (getLargestDivisorOfTripCount(*forInst) % unrollFactor != 0) {
-    DenseMap<const Value *, Value *> operandMap;
     FuncBuilder builder(forInst->getBlock(), ++Block::iterator(forInst));
-    auto *cleanupForInst = cast<ForInst>(builder.clone(*forInst, operandMap));
+    auto *cleanupForInst = cast<ForInst>(builder.clone(*forInst));
     auto clLbMap = getCleanupLoopLowerBound(*forInst, unrollFactor, &builder);
     assert(clLbMap &&
            "cleanup loop lower bound map for single result bound maps can "
@@ -420,7 +420,7 @@ bool mlir::loopUnrollByFactor(ForInst *forInst, uint64_t unrollFactor) {
 
   // Unroll the contents of 'forInst' (append unrollFactor-1 additional copies).
   for (unsigned i = 1; i < unrollFactor; i++) {
-    DenseMap<const Value *, Value *> operandMap;
+    BlockAndValueMapping operandMap;
 
     // If the induction variable is used, create a remapping to the value for
     // this unrolled instance.
@@ -431,7 +431,7 @@ bool mlir::loopUnrollByFactor(ForInst *forInst, uint64_t unrollFactor) {
       auto *ivUnroll =
           builder.create<AffineApplyOp>(forInst->getLoc(), bumpMap, forInst)
               ->getResult(0);
-      operandMap[forInst] = ivUnroll;
+      operandMap.map(forInst, ivUnroll);
     }
 
     // Clone the original body of 'forInst'.

mlir/lib/Transforms/Vectorize.cpp

@@ -1201,8 +1201,7 @@ static bool vectorizeRootMatches(MLFunctionMatches matches,
       continue;
     }
     FuncBuilder builder(loop); // builder to insert in place of loop
-    DenseMap<const Value *, Value *> nomap;
-    ForInst *clonedLoop = cast<ForInst>(builder.clone(*loop, nomap));
+    ForInst *clonedLoop = cast<ForInst>(builder.clone(*loop));
     auto fail = doVectorize(m, &state);
     /// Sets up error handling for this root loop. This is how the root match
     /// maintains a clone for handling failure and restores the proper state via