Skip to content

[MLIR][SideEffects] Added 'Init' Memory Effect which defines an Idempotent MemWrite effect and modified LICM pass #153281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 6 commits into
base: main
Choose a base branch
from
Draft

[MLIR][SideEffects] Added 'Init' Memory Effect which defines an Idempotent MemWrite effect and modified LICM pass #153281

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
34913fe
Added 'Init' Memory Effect which defines an Idempotent MemWrite effec…
mbagherbeikTT Aug 12, 2025
8097e75
Merge branch 'main' into mbagherbeikTT/mem_init
mbagherbeikTT Aug 12, 2025
78a55f1
Merge branch 'main' into mbagherbeikTT/mem_init
mbagherbeikTT Aug 12, 2025
0438627
fixed braces and early returns
mbagherbeikTT Aug 13, 2025
872d60e
switched to DenseMap
mbagherbeikTT Aug 13, 2025
0c23f0c
reordered shouldMoveOutofRegion condition checks for LICM
mbagherbeikTT Aug 13, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions mlir/include/mlir/Interfaces/SideEffectInterfaces.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -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
Expand All @@ -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.
///
Expand Down
12 changes: 12 additions & 0 deletions mlir/include/mlir/Interfaces/SideEffectInterfaces.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Author

@mbagherbeikTT mbagherbeikTT Aug 12, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

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
//===----------------------------------------------------------------------===//
Expand Down
125 changes: 119 additions & 6 deletions mlir/lib/Interfaces/SideEffectInterfaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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);
}

//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -193,13 +198,18 @@ 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);
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);
Expand All @@ -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) {
Expand All @@ -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 *);

Expand All @@ -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);
Expand All @@ -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 *);
Expand All @@ -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);
Expand All @@ -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);

Expand All @@ -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
Expand All @@ -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 &region : op->getRegions())
for (Operation &op : region.getOps())
if (!isMemoryEffectFree(&op))
for (Region &region : 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 &region : 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)) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above: easy-return

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure about this one since there's a loop that comes after that if statement

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 &region : 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) {
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems pretty expensive actually? Have you tried thinking about the complexity of the algorithm here?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

},
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
Expand Down