[MLIR][SideEffects] Added 'Init' Memory Effect which defines an Idempotent MemWrite effect and modified LICM pass #153281
+158
−2
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…t and modified LICM pass. Allows speculatable ops with 'Init' Memory Effects to be moved out of loops if op does not have other, non-Init, Memory Effects and no other operations within it's nested region(s) have Memory Effects that apply to the same resources as the original op.
|
Thank you for submitting a Pull Request (PR) to the LLVM Project! This PR will be automatically labeled and the relevant teams will be notified. If you wish to, you can add reviewers by using the "Reviewers" section on this page. If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers. If you have further questions, they may be answered by the LLVM GitHub User Guide. You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums. |
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-core Author: Mo Bagherbeik (mbagherbeikTT) ChangesAllows speculatable ops with 'Init' Memory Effects to be moved out of loops if op does not have other, non-Init, Memory Effects and no other operations within it's nested region(s) have Memory Effects that apply to the same resources as the original op. Full diff: https://github.com/llvm/llvm-project/pull/153281.diff 4 Files Affected:
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index aef7ec622fe4f..f534a35d4bd8b 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -377,6 +377,13 @@ struct Read : public Effect::Base<Read> {};
/// 'write' effect implies only mutating a resource, and not any visible
/// dereference or read.
struct Write : public Effect::Base<Write> {};
+
+// The following effect indicates that the operation initializes some
+// memory resource to a known value i.e., an idempotent MemWrite.
+// An 'init' effect implies only mutating a resource in a way that's
+// identical across calls if inputs are the same, and not any visible
+// dereference or read.
+struct Init : public Effect::Base<Init> {};
} // namespace MemoryEffects
//===----------------------------------------------------------------------===//
@@ -421,6 +428,15 @@ bool isOpTriviallyDead(Operation *op);
/// Note: Terminators and symbols are never considered to be trivially dead.
bool wouldOpBeTriviallyDead(Operation *op);
+/// Returns true if the given operation is movable under memory effects.
+///
+/// An operation is movable if any of the following are true:
+/// (1) isMemoryEffectFree(op) --> true
+/// (2) isMemoryInitMovable(op) --> true
+///
+/// If the operation meets either criteria, then it is movable
+bool isMemoryEffectMovable(Operation *op);
+
/// Returns true if the given operation is free of memory effects.
///
/// An operation is free of memory effects if its implementation of
@@ -433,6 +449,33 @@ bool wouldOpBeTriviallyDead(Operation *op);
/// conditions are satisfied.
bool isMemoryEffectFree(Operation *op);
+/// Returns true if the given operation has a collision-free 'Init' memory
+/// effect.
+///
+/// An operation is movable if:
+/// (1) it has memory effects AND all of its memory effects are of type 'Init'
+/// (2) there are no other ops with memory effects on any ofthose same resources
+/// within the operation's region(s)
+///
+/// If the operation meets both criteria, then it is movable
+bool isMemoryInitMovable(Operation *op);
+
+/// Returns true if op and all operations within its nested regions
+/// have >1 Memory Effects on ANY of the input resources.
+///
+/// The first call to this function is by an op with >=1 MemInit effect on
+/// >=1 unique resources. To check that none of these resources are in conflict
+/// with other Memory Effects, we scan the entire parent region and maintain
+/// a count of Memory Effects that apply to the resources of the original op.
+/// If any resource has more than 1 Memory Effect in that region, the resource
+/// is in conflict and the op can't be moved by LICM.
+///
+/// Function mutates resources map
+///
+/// If no resources are in conflict, the op is movable.
+bool hasMemoryEffectInitConflict(
+ Operation *op, std::unordered_map<std::string, int> &resources);
+
/// Returns the side effects of an operation. If the operation has
/// RecursiveMemoryEffects, include all side effects of child operations.
///
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
index b292174fccb36..37083690bae52 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
@@ -87,6 +87,18 @@ def MemWrite : MemWrite<DefaultResource, 0, PartialEffect>;
class MemWriteAt<int stage, EffectRange range = PartialEffect>
: MemWrite<DefaultResource, stage, range>;
+// The following effect indicates that the operation initializes some
+// memory resource to a known value i.e., an idempotent MemWrite.
+// An 'init' effect implies only mutating a resource in a way that's
+// identical across calls if inputs are the same, and not any visible
+// dereference or read.
+class MemInit<Resource resource, int stage = 0,
+ EffectRange range = PartialEffect>
+ : MemoryEffect<"::mlir::MemoryEffects::Init", resource, stage, range>;
+def MemInit : MemInit<DefaultResource, 0, PartialEffect>;
+class MemInitAt<int stage, EffectRange range = PartialEffect>
+ : MemInit<DefaultResource, stage, range>;
+
//===----------------------------------------------------------------------===//
// Effect Traits
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 266f6dbacce89..de0eb58c8fedb 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -9,6 +9,8 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/IR/SymbolTable.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include <unordered_set>
#include <utility>
using namespace mlir;
@@ -25,7 +27,7 @@ using namespace mlir;
//===----------------------------------------------------------------------===//
bool MemoryEffects::Effect::classof(const SideEffects::Effect *effect) {
- return isa<Allocate, Free, Read, Write>(effect);
+ return isa<Allocate, Free, Read, Write, Init>(effect);
}
//===----------------------------------------------------------------------===//
@@ -130,6 +132,7 @@ template bool mlir::hasSingleEffect<MemoryEffects::Allocate>(Operation *);
template bool mlir::hasSingleEffect<MemoryEffects::Free>(Operation *);
template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *);
template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *);
+template bool mlir::hasSingleEffect<MemoryEffects::Init>(Operation *);
template <typename EffectTy>
bool mlir::hasSingleEffect(Operation *op, Value value) {
@@ -159,6 +162,8 @@ template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *,
Value value);
template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *,
Value value);
+template bool mlir::hasSingleEffect<MemoryEffects::Init>(Operation *,
+ Value value);
template <typename ValueTy, typename EffectTy>
bool mlir::hasSingleEffect(Operation *op, ValueTy value) {
@@ -193,6 +198,9 @@ template bool
mlir::hasSingleEffect<OpOperand *, MemoryEffects::Write>(Operation *,
OpOperand *);
template bool
+mlir::hasSingleEffect<OpOperand *, MemoryEffects::Init>(Operation *,
+ OpOperand *);
+template bool
mlir::hasSingleEffect<OpResult, MemoryEffects::Allocate>(Operation *, OpResult);
template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Free>(Operation *,
OpResult);
@@ -200,6 +208,8 @@ template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Read>(Operation *,
OpResult);
template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Write>(Operation *,
OpResult);
+template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Init>(Operation *,
+ OpResult);
template bool
mlir::hasSingleEffect<BlockArgument, MemoryEffects::Allocate>(Operation *,
BlockArgument);
@@ -212,6 +222,9 @@ mlir::hasSingleEffect<BlockArgument, MemoryEffects::Read>(Operation *,
template bool
mlir::hasSingleEffect<BlockArgument, MemoryEffects::Write>(Operation *,
BlockArgument);
+template bool
+mlir::hasSingleEffect<BlockArgument, MemoryEffects::Init>(Operation *,
+ BlockArgument);
template <typename... EffectTys>
bool mlir::hasEffect(Operation *op) {
@@ -228,6 +241,7 @@ template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *);
template bool mlir::hasEffect<MemoryEffects::Free>(Operation *);
template bool mlir::hasEffect<MemoryEffects::Read>(Operation *);
template bool mlir::hasEffect<MemoryEffects::Write>(Operation *);
+template bool mlir::hasEffect<MemoryEffects::Init>(Operation *);
template bool
mlir::hasEffect<MemoryEffects::Write, MemoryEffects::Free>(Operation *);
@@ -249,6 +263,7 @@ template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *,
template bool mlir::hasEffect<MemoryEffects::Free>(Operation *, Value value);
template bool mlir::hasEffect<MemoryEffects::Read>(Operation *, Value value);
template bool mlir::hasEffect<MemoryEffects::Write>(Operation *, Value value);
+template bool mlir::hasEffect<MemoryEffects::Init>(Operation *, Value value);
template bool
mlir::hasEffect<MemoryEffects::Write, MemoryEffects::Free>(Operation *,
Value value);
@@ -274,6 +289,8 @@ template bool mlir::hasEffect<OpOperand *, MemoryEffects::Read>(Operation *,
OpOperand *);
template bool mlir::hasEffect<OpOperand *, MemoryEffects::Write>(Operation *,
OpOperand *);
+template bool mlir::hasEffect<OpOperand *, MemoryEffects::Init>(Operation *,
+ OpOperand *);
template bool
mlir::hasEffect<OpOperand *, MemoryEffects::Write, MemoryEffects::Free>(
Operation *, OpOperand *);
@@ -286,6 +303,8 @@ template bool mlir::hasEffect<OpResult, MemoryEffects::Read>(Operation *,
OpResult);
template bool mlir::hasEffect<OpResult, MemoryEffects::Write>(Operation *,
OpResult);
+template bool mlir::hasEffect<OpResult, MemoryEffects::Init>(Operation *,
+ OpResult);
template bool
mlir::hasEffect<OpResult, MemoryEffects::Write, MemoryEffects::Free>(
Operation *, OpResult);
@@ -301,6 +320,8 @@ template bool
mlir::hasEffect<BlockArgument, MemoryEffects::Write>(Operation *,
BlockArgument);
template bool
+mlir::hasEffect<BlockArgument, MemoryEffects::Init>(Operation *, BlockArgument);
+template bool
mlir::hasEffect<BlockArgument, MemoryEffects::Write, MemoryEffects::Free>(
Operation *, BlockArgument);
@@ -312,14 +333,20 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
return wouldOpBeTriviallyDeadImpl(op);
}
+bool mlir::isMemoryEffectMovable(Operation *op) {
+ return (isMemoryEffectFree(op) || isMemoryInitMovable(op));
+}
+
bool mlir::isMemoryEffectFree(Operation *op) {
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
- if (!memInterface.hasNoEffect())
+ if (!memInterface.hasNoEffect()) {
return false;
+ }
// If the op does not have recursive side effects, then it is memory effect
// free.
- if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
+ if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
return true;
+ }
} else if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
// Otherwise, if the op does not implement the memory effect interface and
// it does not have recursive side effects, then it cannot be known that the
@@ -329,13 +356,99 @@ bool mlir::isMemoryEffectFree(Operation *op) {
// Recurse into the regions and ensure that all nested ops are memory effect
// free.
- for (Region ®ion : op->getRegions())
- for (Operation &op : region.getOps())
- if (!isMemoryEffectFree(&op))
+ for (Region ®ion : op->getRegions()) {
+ for (Operation &op : region.getOps()) {
+ if (!isMemoryEffectFree(&op)) {
return false;
+ }
+ }
+ }
return true;
}
+bool mlir::isMemoryInitMovable(Operation *op) {
+ if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ // gather all effects on op
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+ memInterface.getEffects(effects);
+
+ // op has interface but no effects, be conservative
+ if (effects.empty()) {
+ return false;
+ }
+
+ std::unordered_map<std::string, int> resources;
+
+ // ensure op only has Init effects and gather unique
+ // resource names
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ if (!isa<MemoryEffects::Init>(effect.getEffect())) {
+ return false;
+ }
+
+ std::string name = effect.getResource()->getName().str();
+ resources.try_emplace(name, 0);
+ }
+
+ // op itself is good, need to check rest of its parent region
+ Operation *parent = op->getParentOp();
+
+ for (Region ®ion : parent->getRegions()) {
+ for (Operation &op_i : region.getOps()) {
+ if (hasMemoryEffectInitConflict(&op_i, resources)) {
+ return false;
+ }
+ }
+ }
+ return true;
+ }
+
+ // op does not implement the memory effect op interface
+ // meaning it doesn't have any memory init effects and
+ // shouldn't be flagged as movable to be conservative
+ return false;
+}
+
+bool mlir::hasMemoryEffectInitConflict(
+ Operation *op, std::unordered_map<std::string, int> &resources) {
+ if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ if (!memInterface.hasNoEffect()) {
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+ memInterface.getEffects(effects);
+
+ // ensure op only has Init effects and gather unique
+ // resource names
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ if (!isa<MemoryEffects::Init>(effect.getEffect())) {
+ return true;
+ }
+
+ // only care about resources of the op that called
+ // this recursive function for the first time
+ std::string name = effect.getResource()->getName().str();
+
+ if (resources.find(name) != resources.end()) {
+ if (++resources[name] > 1) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+ }
+
+ // Recurse into the regions and ensure that nested ops don't
+ // conflict with each others MemInits
+ for (Region ®ion : op->getRegions()) {
+ for (Operation &op : region.getOps()) {
+ if (hasMemoryEffectInitConflict(&op, resources)) {
+ return true;
+ }
+ }
+ }
+ return false;
+}
+
// the returned vector may contain duplicate effects
std::optional<llvm::SmallVector<MemoryEffects::EffectInstance>>
mlir::getEffectsRecursively(Operation *rootOp) {
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index cb3f2c52e2116..06fabcf8a2ad5 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -110,7 +110,7 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
return loopLike.isDefinedOutsideOfLoop(value);
},
[&](Operation *op, Region *) {
- return isMemoryEffectFree(op) && isSpeculatable(op);
+ return isMemoryEffectMovable(op) && isSpeculatable(op);
},
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
|
joker-eph
reviewed
Aug 12, 2025
joker-eph
reviewed
Aug 12, 2025
joker-eph
reviewed
Aug 12, 2025
joker-eph
reviewed
Aug 12, 2025
|
|
||
| bool mlir::hasMemoryEffectInitConflict( | ||
| Operation *op, std::unordered_map<std::string, int> &resources) { | ||
| if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) { |
There was a problem hiding this comment.
Same as above: easy-return
There was a problem hiding this comment.
I'm not sure about this one since there's a loop that comes after that if statement
| @@ -110,7 +110,7 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) { | |||
| return loopLike.isDefinedOutsideOfLoop(value); | |||
| }, | |||
| [&](Operation *op, Region *) { | |||
| return isMemoryEffectFree(op) && isSpeculatable(op); | |||
| return isMemoryEffectMovable(op) && isSpeculatable(op); | |||
There was a problem hiding this comment.
This seems pretty expensive actually? Have you tried thinking about the complexity of the algorithm here?
There was a problem hiding this comment.
It is relatively expensive. It's similar to isMemoryEffectFree() when checking an op with HasRecursiveMemoryEffects. The difference is, when checking each op, isMemoryEffectMovable() goes up to the op's parent region before recursing down. There are ways to reduce redundant recursions down to already checked regions at the cost of maintaining some extra data during the pass.
| // memory resource to a known value i.e., an idempotent MemWrite. | ||
| // An 'init' effect implies only mutating a resource in a way that's | ||
| // identical across calls if inputs are the same, and not any visible | ||
| // dereference or read. |
There was a problem hiding this comment.
How it is different than a write-only effect?
An op with write-effect on a resource should also be idempotent under the conditions that it takes "identical inputs" right?
There was a problem hiding this comment.
Hey. Thanks for replying and going through the PR.I was trying to see if I can go through the CI tests with skip-precommit-approval (and failed) but your feedback is very helpful. I just finished writing the RFC for this and didn't mean to bother anyone with the code yet.
To answer you question: based on how the LICM pass uses memory effects, a 'MemWrite' effect on an op, even if it also has 'Idempotent', does not allow the op to be moved. Conceptually, MemWrite<someResource> + Idempotent implies the entire op is idempotent, whereas MemInit<someResource> only implies that the op has an idempotent write effect on someResource.
There was a problem hiding this comment.
Conceptually, MemWrite + Idempotent implies the entire op is idempotent, whereas MemInit only implies that the op has an idempotent write effect on someResource.
I don't get it, I need examples for see the nuance. Maybe we can keep the discussion on the RFC.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Allows speculatable ops with 'Init' Memory Effects to be moved out of loops if op does not have other, non-Init, Memory Effects and no other operations within it's nested region(s) have Memory Effects that apply to the same resources as the original op.