Index: mlir/include/mlir/IR/Operation.h =================================================================== --- mlir/include/mlir/IR/Operation.h +++ mlir/include/mlir/IR/Operation.h @@ -72,23 +72,54 @@ /// Remove the operation from its parent block, but don't delete it. void remove(); + /// Various options controlling the behaviour of clone. + class CloneOptions { + public: + /// Default constructs an option with all flags set to false. + CloneOptions(); + + /// Returns an instance with all flags set to true. + static CloneOptions all(); + + /// Configures whether cloning should traverse into any of the regions of + /// the operation. The resulting clone will have the same amount of regions + /// but they will be empty instead. + CloneOptions &cloneRegions(bool enable = true); + + /// Returns whether regions of the operation should be cloned as well. + bool shouldCloneRegions() const { return cloneRegionsFlag; } + + /// Configures whether operands should be cloned as well. Otherwise the + /// resulting clone will simply have no operands. + CloneOptions &cloneOperands(bool enable = true); + + /// Returns whether operands should be cloned as well. + bool shouldCloneOperands() const { return cloneOperandsFlag; } + + private: + /// Whether regions should be cloned. + bool cloneRegionsFlag : 1; + /// Whether operands should be cloned. + bool cloneOperandsFlag : 1; + }; + /// 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. - Operation *clone(BlockAndValueMapping &mapper); - Operation *clone(); + /// Optionally, one may configure to not clone parts of the operation using + /// the options parameter. + Operation *clone(BlockAndValueMapping &mapper, + CloneOptions options = CloneOptions::all()); + Operation *clone(CloneOptions options = CloneOptions::all()); /// Create a partial copy of this operation without traversing into attached /// regions. The new operation will have the same number of regions as the /// original one, but they will be left empty. /// Operands are remapped using `mapper` (if present), and `mapper` is updated /// to contain the results. - /// The `mapResults` argument specifies whether the results of the operation - /// should also be mapped. - Operation *cloneWithoutRegions(BlockAndValueMapping &mapper, - bool mapResults = true); + Operation *cloneWithoutRegions(BlockAndValueMapping &mapper); /// Create a partial copy of this operation without traversing into attached /// regions. The new operation will have the same number of regions as the Index: mlir/lib/IR/Operation.cpp =================================================================== --- mlir/lib/IR/Operation.cpp +++ mlir/lib/IR/Operation.cpp @@ -523,19 +523,59 @@ // Operation Cloning //===----------------------------------------------------------------------===// +/// Default constructs an option with all flags set to false. +Operation::CloneOptions::CloneOptions() + : cloneRegionsFlag(false), cloneOperandsFlag(false) {} + +/// Returns an instance with all flags set to true. +Operation::CloneOptions Operation::CloneOptions::all() { + return CloneOptions().cloneRegions().cloneOperands(); +} + +/// Configures whether cloning should traverse into any of the regions of +/// the operation. The resulting clone will have the same amount of regions +/// but they will be empty instead. +Operation::CloneOptions &Operation::CloneOptions::cloneRegions(bool enable) { + cloneRegionsFlag = enable; + return *this; +} + +/// Configures whether operands should be cloned as well. Otherwise the +/// resulting clone will simply have no operands. +Operation::CloneOptions &Operation::CloneOptions::cloneOperands(bool enable) { + cloneOperandsFlag = enable; + return *this; +} + /// 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. The `mapResults` flag specifies whether the results /// of the cloned operation should be added to the map. -Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper, - bool mapResults) { +Operation *Operation::cloneWithoutRegions(BlockAndValueMapping &mapper) { + return clone(mapper, CloneOptions::all().cloneRegions(false)); +} + +Operation *Operation::cloneWithoutRegions() { + BlockAndValueMapping mapper; + return cloneWithoutRegions(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. +Operation *Operation::clone(BlockAndValueMapping &mapper, + CloneOptions options) { SmallVector operands; SmallVector successors; // Remap the operands. - operands.reserve(getNumOperands()); - for (auto opValue : getOperands()) - operands.push_back(mapper.lookupOrDefault(opValue)); + if (options.shouldCloneOperands()) { + operands.reserve(getNumOperands()); + for (auto opValue : getOperands()) + operands.push_back(mapper.lookupOrDefault(opValue)); + } // Remap the successors. successors.reserve(getNumSuccessors()); @@ -546,41 +586,22 @@ auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs, successors, getNumRegions()); - // Remember the mapping of any results. - if (mapResults) { - for (unsigned i = 0, e = getNumResults(); i != e; ++i) - mapper.map(getResult(i), newOp->getResult(i)); - } - - return newOp; -} - -Operation *Operation::cloneWithoutRegions() { - BlockAndValueMapping mapper; - return cloneWithoutRegions(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. -Operation *Operation::clone(BlockAndValueMapping &mapper) { - auto *newOp = cloneWithoutRegions(mapper, /*mapResults=*/false); - // Clone the regions. - for (unsigned i = 0; i != numRegions; ++i) - getRegion(i).cloneInto(&newOp->getRegion(i), mapper); + if (options.shouldCloneRegions()) { + for (unsigned i = 0; i != numRegions; ++i) + getRegion(i).cloneInto(&newOp->getRegion(i), mapper); + } + // Remember the mapping of any results. for (unsigned i = 0, e = getNumResults(); i != e; ++i) mapper.map(getResult(i), newOp->getResult(i)); return newOp; } -Operation *Operation::clone() { +Operation *Operation::clone(CloneOptions options) { BlockAndValueMapping mapper; - return clone(mapper); + return clone(mapper, options); } //===----------------------------------------------------------------------===// Index: mlir/lib/IR/Region.cpp =================================================================== --- mlir/lib/IR/Region.cpp +++ mlir/lib/IR/Region.cpp @@ -9,6 +9,10 @@ #include "mlir/IR/Region.h" #include "mlir/IR/BlockAndValueMapping.h" #include "mlir/IR/Operation.h" +#include "mlir/IR/RegionGraphTraits.h" + +#include "llvm/ADT/DepthFirstIterator.h" + using namespace mlir; Region::Region(Operation *container) : container(container) {} @@ -82,6 +86,14 @@ if (empty()) return; + // 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. + + // 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 : *this) { Block *newBlock = new Block(); mapper.map(&block, newBlock); @@ -93,26 +105,42 @@ if (!mapper.contains(arg)) mapper.map(arg, newBlock->addArgument(arg.getType(), arg.getLoc())); - // Clone and remap the operations within this block. - for (auto &op : block) - newBlock->push_back(op.clone(mapper)); - dest->getBlocks().insert(destPos, newBlock); } - // Now that each of the blocks have been cloned, go through and remap the - // operands of each of the operations. - auto remapOperands = [&](Operation *op) { - for (auto &operand : op->getOpOperands()) - if (auto mappedOp = mapper.lookupOrNull(operand.get())) - operand.set(mappedOp); - for (auto &succOp : op->getBlockOperands()) - if (auto *mappedOp = mapper.lookupOrNull(succOp.get())) - succOp.set(mappedOp); - }; - - for (iterator it(mapper.lookup(&front())); it != destPos; ++it) - it->walk(remapOperands); + // 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. + SmallVector> insertedOps; + for (Block &block : *this) { + Block *newBlock = mapper.lookup(&block); + // Clone and remap the operations within this block. + for (auto &op : block) { + Operation *clone = op.clone( + mapper, + Operation::CloneOptions::all().cloneRegions(false).cloneOperands( + false)); + newBlock->push_back(clone); + insertedOps.emplace_back(&op, clone); + } + } + + // Finally now that all operation results have been mapped, set the operands + // and clone the regions. + for (auto &it : insertedOps) { + Operation *source = it.first; + Operation *clone = it.second; + clone->setOperands(llvm::to_vector( + llvm::map_range(source->getOperands(), [&](Value operand) { + return mapper.lookupOrDefault(operand); + }))); + + for (auto regions : llvm::zip(source->getRegions(), clone->getRegions())) + std::get<0>(regions).cloneInto(&std::get<1>(regions), mapper); + } } /// Returns 'block' if 'block' lies in this region, or otherwise finds the