Index: mlir/include/mlir/IR/Operation.h =================================================================== --- mlir/include/mlir/IR/Operation.h +++ mlir/include/mlir/IR/Operation.h @@ -72,23 +72,75 @@ /// Remove the operation from its parent block, but don't delete it. void remove(); + /// Class encompassing various options related to cloning an operation. Users + /// of this class should pass it to Operation's 'clone' methods. + /// Current options include: + /// * Whether cloning should recursively traverse into the regions of the + /// operation or not + /// * Whether cloning should also clone the operands of the operation. + class CloneOptions { + public: + /// Default constructs an option with all flags set to false. That means all + /// parts of an operation that may optionally not be cloned, are not cloned. + CloneOptions(); + + /// Constructs an instance with the clone regions and clone operands flags + /// set accordingly. + CloneOptions(bool cloneRegions, bool cloneOperands); + + /// Returns an instance with all flags set to true. This is the default + /// when using the clone method and clones all parts of the operation. + static CloneOptions all(); + + /// Configures whether cloning should traverse into any of the regions of + /// the operation. If set to true, operations' regions are recursively + /// cloned. If set to false, cloned operations will have the same number of + /// regions, but they will be empty. + /// Cloning of nested operations in the operations' regions are currently + /// unaffected by other flags. + CloneOptions &cloneRegions(bool enable = true); + + /// Returns whether regions of the operation should be cloned as well. + bool shouldCloneRegions() const { return cloneRegionsFlag; } + + /// Configures whether operation' operands should be cloned. Otherwise the + /// resulting clones will simply have zero 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 what parts of the operation to clone using + /// the options parameter. + /// + /// Calling this method from multiple threads is generally safe if through the + /// process of cloning, no new uses of 'Value's from outside the operation are + /// created. Cloning an isolated-from-above operation with no operands, such + /// as top level function operations, is therefore always safe. Using the + /// mapper, it is possible to avoid adding uses to outside operands by + /// remapping them to 'Value's owned by the caller thread. + 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/include/mlir/IR/Region.h =================================================================== --- mlir/include/mlir/IR/Region.h +++ mlir/include/mlir/IR/Region.h @@ -227,6 +227,11 @@ /// 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. + /// + /// Calling this method from multiple threads is generally safe if through the + /// process of cloning, no new uses of 'Value's from outside the region are + /// created. Using the mapper, it is possible to avoid adding uses to outside + /// operands by remapping them to 'Value's owned by the caller thread. void cloneInto(Region *dest, BlockAndValueMapping &mapper); /// Clone this region into 'dest' before the given position in 'dest'. void cloneInto(Region *dest, Region::iterator destPos, Index: mlir/lib/IR/Operation.cpp =================================================================== --- mlir/lib/IR/Operation.cpp +++ mlir/lib/IR/Operation.cpp @@ -523,36 +523,32 @@ // Operation Cloning //===----------------------------------------------------------------------===// -/// 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) { - SmallVector operands; - SmallVector successors; +Operation::CloneOptions::CloneOptions() + : cloneRegionsFlag(false), cloneOperandsFlag(false) {} - // Remap the operands. - operands.reserve(getNumOperands()); - for (auto opValue : getOperands()) - operands.push_back(mapper.lookupOrDefault(opValue)); +Operation::CloneOptions::CloneOptions(bool cloneRegions, bool cloneOperands) + : cloneRegionsFlag(cloneRegions), cloneOperandsFlag(cloneOperands) {} - // Remap the successors. - successors.reserve(getNumSuccessors()); - for (Block *successor : getSuccessors()) - successors.push_back(mapper.lookupOrDefault(successor)); +Operation::CloneOptions Operation::CloneOptions::all() { + return CloneOptions().cloneRegions().cloneOperands(); +} - // Create the new operation. - auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs, - successors, getNumRegions()); +Operation::CloneOptions &Operation::CloneOptions::cloneRegions(bool enable) { + cloneRegionsFlag = enable; + return *this; +} - // Remember the mapping of any results. - if (mapResults) { - for (unsigned i = 0, e = getNumResults(); i != e; ++i) - mapper.map(getResult(i), newOp->getResult(i)); - } +Operation::CloneOptions &Operation::CloneOptions::cloneOperands(bool enable) { + cloneOperandsFlag = enable; + return *this; +} - return newOp; +/// 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) { + return clone(mapper, CloneOptions::all().cloneRegions(false)); } Operation *Operation::cloneWithoutRegions() { @@ -565,22 +561,43 @@ /// 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); +Operation *Operation::clone(BlockAndValueMapping &mapper, + CloneOptions options) { + SmallVector operands; + SmallVector successors; + + // Remap the operands. + if (options.shouldCloneOperands()) { + operands.reserve(getNumOperands()); + for (auto opValue : getOperands()) + operands.push_back(mapper.lookupOrDefault(opValue)); + } + + // Remap the successors. + successors.reserve(getNumSuccessors()); + for (Block *successor : getSuccessors()) + successors.push_back(mapper.lookupOrDefault(successor)); + + // Create the new operation. + auto *newOp = create(getLoc(), getName(), getResultTypes(), operands, attrs, + successors, getNumRegions()); // 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 @@ -82,6 +82,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 +101,49 @@ 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); + auto newBlocksRange = + llvm::make_range(Region::iterator{mapper.lookup(&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 = + Operation::CloneOptions::all().cloneRegions(false).cloneOperands(false); + for (auto zippedBlocks : llvm::zip(*this, newBlocksRange)) { + Block &sourceBlock = std::get<0>(zippedBlocks); + Block &clonedBlock = std::get<1>(zippedBlocks); + // Clone and remap the operations within this block. + for (auto &op : sourceBlock) { + Operation *clone = op.clone(mapper, cloneOptions); + clonedBlock.push_back(clone); + } + } + + // Finally now that all operation results have been mapped, set the operands + // and clone the regions. + for (auto zippedBlocks : llvm::zip(*this, newBlocksRange)) { + Block &sourceBlock = std::get<0>(zippedBlocks); + Block &clonedBlock = std::get<1>(zippedBlocks); + + for (auto ops : llvm::zip(sourceBlock, clonedBlock)) { + Operation &source = std::get<0>(ops); + Operation &clone = std::get<1>(ops); + + 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