[Mlir-commits] [mlir] a41aaf1 - [mlir] Make `Regions`s `cloneInto` multithread-readable
Markus Böck
llvmlistbot at llvm.org
Thu Apr 21 04:43:41 PDT 2022
Author: Markus Böck
Date: 2022-04-21T13:43:00+02:00
New Revision: a41aaf166fed03e18021885d0951f1dec63b25b9
URL: https://github.com/llvm/llvm-project/commit/a41aaf166fed03e18021885d0951f1dec63b25b9
DIFF: https://github.com/llvm/llvm-project/commit/a41aaf166fed03e18021885d0951f1dec63b25b9.diff
LOG: [mlir] Make `Regions`s `cloneInto` multithread-readable
Prior to this patch, `cloneInto` would do a simple walk over the blocks and contained operations and clone and map them as it encounters them. As finishing touch it then remaps any successor and operands it has remapped during that process.
This is generally fine, but sadly leads to a lot of uses of both operations and blocks from the source region, in the cloned operations in the target region. Those uses lead to writes in the use-def list of the operations, making `cloneInto` never thread safe.
This patch reimplements `cloneInto` in three steps to avoid ever creating any extra uses on elements in the source region:
* It first creates the mapping of all blocks and block operands
* It then clones all operations to create the mapping of all operation results, but does not yet clone any regions or set the operands
* After all operation results have been mapped, it now sets the operations operands and clones their regions.
That way it is now possible to call `cloneInto` from multiple threads if the Region or Operation is isolated-from-above. This allows creating copies of functions or to use `mlir::inlineCall` with the same source region from multiple threads. In the general case, the method is thread-safe if through cloning, no new uses of `Value`s from outside the cloned Operation/Region are created. This can be ensured by mapping any outside operands via the `BlockAndValueMapping` to `Value`s owned by the caller thread.
While I was at it, I also reworked the `clone` method of `Operation` a little bit and added a proper options class to avoid having a `cloneWithoutRegionsAndOperands` method, and be more extensible in the future. `cloneWithoutRegions` is now also a simple wrapper that calls `clone` with the proper options set. That way all the operation cloning code is now contained solely within `clone`.
Differential Revision: https://reviews.llvm.org/D123917
Added:
Modified:
mlir/include/mlir/IR/Operation.h
mlir/include/mlir/IR/Region.h
mlir/lib/IR/Operation.cpp
mlir/lib/IR/Region.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/IR/Operation.h b/mlir/include/mlir/IR/Operation.h
index 316c52f086bbc..4692104c119db 100644
--- a/mlir/include/mlir/IR/Operation.h
+++ b/mlir/include/mlir/IR/Operation.h
@@ -72,23 +72,75 @@ class alignas(8) Operation final
/// 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, the operation's 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 operation's 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
diff --git a/mlir/include/mlir/IR/Region.h b/mlir/include/mlir/IR/Region.h
index c9b419a89004c..671a981a54502 100644
--- a/mlir/include/mlir/IR/Region.h
+++ b/mlir/include/mlir/IR/Region.h
@@ -227,6 +227,11 @@ class Region {
/// 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,
diff --git a/mlir/lib/IR/Operation.cpp b/mlir/lib/IR/Operation.cpp
index b6c64ec773c05..262974914d13b 100644
--- a/mlir/lib/IR/Operation.cpp
+++ b/mlir/lib/IR/Operation.cpp
@@ -523,36 +523,32 @@ InFlightDiagnostic Operation::emitOpError(const Twine &message) {
// 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<Value, 8> operands;
- SmallVector<Block *, 2> 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 @@ Operation *Operation::cloneWithoutRegions() {
/// 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<Value, 8> operands;
+ SmallVector<Block *, 2> 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);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/IR/Region.cpp b/mlir/lib/IR/Region.cpp
index a5bca205fae1d..1b98926ca0b0f 100644
--- a/mlir/lib/IR/Region.cpp
+++ b/mlir/lib/IR/Region.cpp
@@ -82,6 +82,18 @@ void Region::cloneInto(Region *dest, Region::iterator destPos,
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. 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 : *this) {
Block *newBlock = new Block();
mapper.map(&block, newBlock);
@@ -93,26 +105,47 @@ void Region::cloneInto(Region *dest, Region::iterator destPos,
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 (Operation &op : sourceBlock)
+ clonedBlock.push_back(op.clone(mapper, 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(*this, 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);
+
+ operands.resize(source.getNumOperands());
+ llvm::transform(
+ source.getOperands(), operands.begin(),
+ [&](Value operand) { return mapper.lookupOrDefault(operand); });
+ clone.setOperands(operands);
+
+ 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
More information about the Mlir-commits
mailing list