Operations mappping when cloning moduleOp

Hi >_<,

I have a data structure maintaining a moduleOp and an iterator to some operations inside a function. However, if I want to clone a new temporary module and set a new iterator to the same operation as the one in original moduleOp. I found I can get a BlockAndValue Mapping. However, operations aren’t subclasses of value. How should I get the new iterator conveniently?

My current limited solution: get the operation by its result value. (silly and not work on no-result operations )

mlir::Operation *tmpCopy = module.getOperation()->clone(bavMap);
mlir::Value res = opit->getResult(0);
mlir::Value tmpres = bavMap.lookup(res);
opitInTmp = tmpres.getDefiningOp()->getIterator();

Thanks for your help!

This is indeed currently a missing feature of BlockAndValueMapping that I have run into a couple times myself. Until this is fixed upstream you can workaround it for the case of no result ops by mapping the block of the op and using it’s distance instead. Or in other words:

mlir::Operation* mapOperation(mlir::BlockAndValueMapping& mapping, mlir::Operation* operation)
    if (operation->getNumResult() != 0)
        return mapping.lookupOrDefault(operation->getResult(0)).getDefiningOp();
    mlir::Block* mappedBlock = mapping.lookupOrDefault(operation->getBlock());
    std::size_t distance = std::distance(operation->getBlock()->begin(), mlir::Block::iterator{operation});
    return &*std::next(mappedBlock->begin(), distance);
1 Like

woooow, this is good. Thanks for your explanation!

I agree this is a missing operation, I use code derived from the mainline MLIR implementation with the extra map:

//===- BetterCloning.cpp - Improved Operation cloning ---------------------===//
// This file implements logic to clone an operation and keep track of the
// operation mappings.

#include "Internals.h"
#include "mlir/IR/BlockAndValueMapping.h"
#include "mlir/IR/Operation.h"

using CloneOptions = mlir::Operation::CloneOptions;

static Operation *
cloneOperation(Operation *original, BlockAndValueMapping &mapper,
               DenseMap<Operation *, Operation *> &operationMap,
               CloneOptions options);

/// Clone this region into 'dest' before the given position in 'dest'.
static void cloneRegion(Region &original, Region &dest,
                        BlockAndValueMapping &mapper,
                        DenseMap<Operation *, Operation *> &operationMap) {
  assert(&original != &dest && "cannot clone region into itself");

  // If the list is empty there is nothing to clone.
  if (original.empty())

  Region::iterator destPos = dest.end();

  // The below clone implementation takes special care to be read only for the
  // sake of multi threading. That essentially means not adding any uses to any
  // of the blocks or operation results contained within this region as that
  // would lead to a write in their use-def list. This is unavoidable for
  // 'Value's from outside the region however, in which case it is not read
  // only. Using the BlockAndValueMapper it is possible to remap such 'Value's
  // to ones owned by the calling thread however, making it read only once
  // again.

  // First clone all the blocks and block arguments and map them, but don't yet
  // clone the operations, as they may otherwise add a use to a block that has
  // not yet been mapped
  for (Block &block : original) {
    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 (auto arg : block.getArguments())
      if (!mapper.contains(arg))
        mapper.map(arg, newBlock->addArgument(arg.getType(), arg.getLoc()));

    dest.getBlocks().insert(destPos, newBlock);

  auto newBlocksRange = llvm::make_range(
      Region::iterator(mapper.lookup(&original.front())), destPos);

  // Now follow up with creating the operations, but don't yet clone their
  // regions, nor set their operands. Setting the successors is safe as all have
  // already been mapped. We are essentially just creating the operation results
  // to be able to map them.
  // Cloning the operands and region as well would lead to uses of operations
  // not yet mapped.
  auto cloneOptions =
  for (auto zippedBlocks : llvm::zip(original, newBlocksRange)) {
    Block &sourceBlock = std::get<0>(zippedBlocks);
    Block &clonedBlock = std::get<1>(zippedBlocks);
    // Clone and remap the operations within this block.
    for (Operation &op : sourceBlock)
          cloneOperation(&op, mapper, operationMap, cloneOptions));

  // Finally now that all operation results have been mapped, set the operands
  // and clone the regions.
  SmallVector<Value> operands;
  for (auto zippedBlocks : llvm::zip(original, newBlocksRange)) {
    for (auto ops :
         llvm::zip(std::get<0>(zippedBlocks), std::get<1>(zippedBlocks))) {
      Operation &source = std::get<0>(ops);
      Operation &clone = std::get<1>(ops);

          source.getOperands(), operands.begin(),
          [&](Value operand) { return mapper.lookupOrDefault(operand); });

      for (auto regions : llvm::zip(source.getRegions(), clone.getRegions()))
        cloneRegion(std::get<0>(regions), std::get<1>(regions), mapper,

/// Create a deep copy of this 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.
static Operation *
cloneOperation(Operation *original, BlockAndValueMapping &mapper,
               DenseMap<Operation *, Operation *> &operationMap,
               CloneOptions options) {
  SmallVector<Value, 8> operands;
  SmallVector<Block *, 2> successors;

  // Remap the operands.
  if (options.shouldCloneOperands()) {
    for (auto opValue : original->getOperands())

  // Remap the successors.
  for (Block *successor : original->getSuccessors())

  // Create the new operation.
  auto *newOp = Operation::create(
      original->getLoc(), original->getName(), original->getResultTypes(),
      operands, original->getAttrs(), successors, original->getNumRegions());
  operationMap[original] = newOp;

  // Clone the regions.
  if (options.shouldCloneRegions()) {
    for (unsigned i = 0, e = original->getNumRegions(); i != e; ++i)
      cloneRegion(original->getRegion(i), newOp->getRegion(i), mapper,

  // Remember the mapping of any results.
  for (unsigned i = 0, e = original->getNumResults(); i != e; ++i)
    mapper.map(original->getResult(i), newOp->getResult(i));

  return newOp;

/// This is like `Operation::clone`, but instead of just keeping track of the
/// block and value mapping for the copy, it also keeps track of the
/// operation<->operation mapping.  This matters because not all operations have
/// results.
Operation *M::cloneOperation(Operation *original, BlockAndValueMapping &mapper,
                             DenseMap<Operation *, Operation *> &operationMap) {
  return ::cloneOperation(original, mapper, operationMap, CloneOptions::all());
1 Like

very exhaustive, thanks for your clear code with full comments! I learned a lot.