[Mlir-commits] [mlir] [MLIR][SideEffects] Added 'Init' Memory Effect which defines an Idempotent MemWrite effect and modified LICM pass (PR #153281)
Mo Bagherbeik
llvmlistbot at llvm.org
Wed Aug 13 13:47:35 PDT 2025
https://github.com/mbagherbeikTT updated https://github.com/llvm/llvm-project/pull/153281
>From 34913fe68d21cccbe91862f9a9145acaeb895d55 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 12 Aug 2025 20:28:54 +0000
Subject: [PATCH 1/4] Added 'Init' Memory Effect which defines an Idempotent
MemWrite effect 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.
---
.../mlir/Interfaces/SideEffectInterfaces.h | 43 ++++++
.../mlir/Interfaces/SideEffectInterfaces.td | 12 ++
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 124 +++++++++++++++++-
.../Utils/LoopInvariantCodeMotionUtils.cpp | 2 +-
4 files changed, 174 insertions(+), 7 deletions(-)
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 59fd19310cea5..de0eb58c8fedb 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -10,6 +10,7 @@
#include "mlir/IR/SymbolTable.h"
#include "llvm/ADT/SmallPtrSet.h"
+#include <unordered_set>
#include <utility>
using namespace mlir;
@@ -26,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);
}
//===----------------------------------------------------------------------===//
@@ -131,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) {
@@ -160,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) {
@@ -194,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);
@@ -201,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);
@@ -213,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) {
@@ -229,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 *);
@@ -250,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);
@@ -275,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 *);
@@ -287,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);
@@ -302,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);
@@ -313,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
@@ -330,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); });
}
>From 0438627801b7250333eb3254faba8c97b6dbe90a Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Wed, 13 Aug 2025 19:35:18 +0000
Subject: [PATCH 2/4] fixed braces and early returns
---
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 97 +++++++++-----------
1 file changed, 43 insertions(+), 54 deletions(-)
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index de0eb58c8fedb..eea8b418ef5d8 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -339,14 +339,14 @@ bool mlir::isMemoryEffectMovable(Operation *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
@@ -356,61 +356,55 @@ 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);
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+ // 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
+ if (!memInterface) return false;
- // op has interface but no effects, be conservative
- if (effects.empty()) {
- return false;
- }
+ // gather all effects on op
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+ memInterface.getEffects(effects);
- std::unordered_map<std::string, int> resources;
+ // op has interface but no effects, be conservative
+ if (effects.empty()) return false;
- // 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);
- }
+ std::unordered_map<std::string, int> resources;
- // op itself is good, need to check rest of its parent region
- Operation *parent = op->getParentOp();
+ // 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;
- for (Region ®ion : parent->getRegions()) {
- for (Operation &op_i : region.getOps()) {
- if (hasMemoryEffectInitConflict(&op_i, resources)) {
- return false;
- }
- }
- }
- return true;
+ std::string name = effect.getResource()->getName().str();
+ resources.try_emplace(name, 0);
}
- // 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;
+ // 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;
}
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;
@@ -419,19 +413,16 @@ bool mlir::hasMemoryEffectInitConflict(
// ensure op only has Init effects and gather unique
// resource names
for (const MemoryEffects::EffectInstance &effect : effects) {
- if (!isa<MemoryEffects::Init>(effect.getEffect())) {
+ 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) {
+ if (resources.find(name) != resources.end())
+ if (++resources[name] > 1)
return true;
- }
- }
}
return false;
}
@@ -439,13 +430,11 @@ bool mlir::hasMemoryEffectInitConflict(
// 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)) {
+ for (Region ®ion : op->getRegions())
+ for (Operation &op : region.getOps())
+ if (hasMemoryEffectInitConflict(&op, resources))
return true;
- }
- }
- }
+
return false;
}
>From 872d60e786dc35615cd5e0182bde37fa85e4e7fa Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Wed, 13 Aug 2025 19:38:18 +0000
Subject: [PATCH 3/4] switched to DenseMap
---
.../mlir/Interfaces/SideEffectInterfaces.h | 2 +-
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 17 ++++++++---------
2 files changed, 9 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index f534a35d4bd8b..8f49812016ac6 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -474,7 +474,7 @@ bool isMemoryInitMovable(Operation *op);
///
/// If no resources are in conflict, the op is movable.
bool hasMemoryEffectInitConflict(
- Operation *op, std::unordered_map<std::string, int> &resources);
+ Operation *op, DenseMap<TypeID, int> &resourceCounts);
/// Returns the side effects of an operation. If the operation has
/// RecursiveMemoryEffects, include all side effects of child operations.
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index eea8b418ef5d8..2b88d286fcce2 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -379,7 +379,7 @@ bool mlir::isMemoryInitMovable(Operation *op) {
if (effects.empty()) return false;
- std::unordered_map<std::string, int> resources;
+ DenseMap<TypeID, int> resourceCounts;
// ensure op only has Init effects and gather unique
// resource names
@@ -387,8 +387,7 @@ bool mlir::isMemoryInitMovable(Operation *op) {
if (!isa<MemoryEffects::Init>(effect.getEffect()))
return false;
- std::string name = effect.getResource()->getName().str();
- resources.try_emplace(name, 0);
+ resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
}
// op itself is good, need to check rest of its parent region
@@ -396,14 +395,14 @@ bool mlir::isMemoryInitMovable(Operation *op) {
for (Region ®ion : parent->getRegions())
for (Operation &op_i : region.getOps())
- if (hasMemoryEffectInitConflict(&op_i, resources))
+ if (hasMemoryEffectInitConflict(&op_i, resourceCounts))
return false;
return true;
}
bool mlir::hasMemoryEffectInitConflict(
- Operation *op, std::unordered_map<std::string, int> &resources) {
+ Operation *op, DenseMap<TypeID, int> &resourceCounts) {
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
if (!memInterface.hasNoEffect()) {
@@ -418,10 +417,10 @@ bool mlir::hasMemoryEffectInitConflict(
// only care about resources of the op that called
// this recursive function for the first time
- std::string name = effect.getResource()->getName().str();
+ auto resourceID = effect.getResource()->getResourceID();
- if (resources.find(name) != resources.end())
- if (++resources[name] > 1)
+ if (resourceCounts.contains(resourceID))
+ if (++resourceCounts[resourceID] > 1)
return true;
}
return false;
@@ -432,7 +431,7 @@ bool mlir::hasMemoryEffectInitConflict(
// conflict with each others MemInits
for (Region ®ion : op->getRegions())
for (Operation &op : region.getOps())
- if (hasMemoryEffectInitConflict(&op, resources))
+ if (hasMemoryEffectInitConflict(&op, resourceCounts))
return true;
return false;
>From 0c23f0ccb2e4af55db74d12559c837afb86a767a Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Wed, 13 Aug 2025 20:10:49 +0000
Subject: [PATCH 4/4] reordered shouldMoveOutofRegion condition checks for LICM
---
mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 06fabcf8a2ad5..dc3baba865bf1 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 isMemoryEffectMovable(op) && isSpeculatable(op);
+ return isSpeculatable(op) && isMemoryEffectMovable(op);
},
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
More information about the Mlir-commits
mailing list