[Mlir-commits] [mlir] [MLIR][SideEffects][MemoryEffects] Modified LICM to be more aggressive when checking movability of ops with MemWrite effects (PR #155344)
Mo Bagherbeik
llvmlistbot at llvm.org
Sun Sep 21 17:11:13 PDT 2025
https://github.com/mbagherbeikTT updated https://github.com/llvm/llvm-project/pull/155344
>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 01/22] 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 02/22] 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 03/22] 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 04/22] 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); });
}
>From 8fb4b974fbac5b1176ded32ee8968c36f445d354 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 26 Aug 2025 00:47:51 +0000
Subject: [PATCH 05/22] removed memInit and refactored LICM
---
.../mlir/Interfaces/SideEffectInterfaces.h | 57 +++--
.../mlir/Interfaces/SideEffectInterfaces.td | 11 -
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 95 ++++----
.../loop-invariant-code-motion.mlir | 204 ++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOps.h | 23 ++
mlir/test/lib/Dialect/Test/TestOps.td | 54 +++++
6 files changed, 353 insertions(+), 91 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 8f49812016ac6..ca48d23574feb 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -15,6 +15,7 @@
#define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Dominance.h"
namespace mlir {
namespace SideEffects {
@@ -378,12 +379,6 @@ struct Read : public Effect::Base<Read> {};
/// 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
//===----------------------------------------------------------------------===//
@@ -428,13 +423,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.
+/// Returns true if the given operation is allowed to be moved under the
+/// memory effects interface.
///
-/// An operation is movable if any of the following are true:
-/// (1) isMemoryEffectFree(op) --> true
-/// (2) isMemoryInitMovable(op) --> true
+/// An operation is movable if either case is true:
+/// (a) free of memory effects as defined in isMemoryEffectFree()
+/// (b) if the operation does have memory effects, it must be conflict-free
+/// as defined in isMemoryEffectConflictFree()
///
-/// If the operation meets either criteria, then it is movable
+/// If the operation meets either criteria, then it is movable under memory effects
bool isMemoryEffectMovable(Operation *op);
/// Returns true if the given operation is free of memory effects.
@@ -449,32 +446,32 @@ bool isMemoryEffectMovable(Operation *op);
/// conditions are satisfied.
bool isMemoryEffectFree(Operation *op);
-/// Returns true if the given operation has a collision-free 'Init' memory
-/// effect.
+/// Returns true if the given operation has conflict-free write effects
///
-/// 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)
+/// An operation is conflict free:
+/// (1) all of its memory effects are of type Write
+/// (2) there are no other ops with Alloc/Free/Write effects on the same
+/// resources within the ops parent region
+/// (3) all ops in the parent region with Read effects on the same resources
+/// are dominated by the operation
///
-/// If the operation meets both criteria, then it is movable
-bool isMemoryInitMovable(Operation *op);
+/// If the operation meets all 3 criteria, then it is conflict free
+bool isMemoryEffectConflictFree(Operation *op);
-/// Returns true if op and all operations within its nested regions
-/// have >1 Memory Effects on ANY of the input resources.
+/// Returns true if op and/or any operations within its nested regions
+/// have a memory effect conflict with mainOp as defined below:
///
-/// 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.
+/// op has a memory effect conflict with mainOp if op and/or any of
+/// the operations in its nested regions meet any of these criteria:
+/// (a) they have any Alloc/Free/Write effects on the resources used by mainOp
+/// (b) they dominate mainOp and have any read effect on the resources used by mainOp
///
/// Function mutates resources map
///
-/// If no resources are in conflict, the op is movable.
-bool hasMemoryEffectInitConflict(
- Operation *op, DenseMap<TypeID, int> &resourceCounts);
+/// If none of the critera above are met, mainOp and op are conflict free
+bool hasMemoryEffectConflict(
+ Operation *mainOp, Operation *op,
+ mlir::DominanceInfo &dom, 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/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
index 37083690bae52..841da64a509bd 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
@@ -87,17 +87,6 @@ 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 2b88d286fcce2..245850014f507 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -8,6 +8,7 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/IR/Dominance.h"
#include "mlir/IR/SymbolTable.h"
#include "llvm/ADT/SmallPtrSet.h"
#include <unordered_set>
@@ -27,7 +28,7 @@ using namespace mlir;
//===----------------------------------------------------------------------===//
bool MemoryEffects::Effect::classof(const SideEffects::Effect *effect) {
- return isa<Allocate, Free, Read, Write, Init>(effect);
+ return isa<Allocate, Free, Read, Write>(effect);
}
//===----------------------------------------------------------------------===//
@@ -132,7 +133,6 @@ 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) {
@@ -162,8 +162,6 @@ 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) {
@@ -198,9 +196,6 @@ 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);
@@ -208,8 +203,6 @@ 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);
@@ -222,9 +215,6 @@ 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) {
@@ -241,7 +231,6 @@ 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 *);
@@ -263,7 +252,6 @@ 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);
@@ -289,8 +277,6 @@ 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 *);
@@ -303,8 +289,6 @@ 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);
@@ -320,8 +304,6 @@ 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);
@@ -334,7 +316,13 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
}
bool mlir::isMemoryEffectMovable(Operation *op) {
- return (isMemoryEffectFree(op) || isMemoryInitMovable(op));
+ if (isMemoryEffectFree(op))
+ return true;
+
+ if (isMemoryEffectConflictFree(op))
+ return true;
+
+ return false;
}
bool mlir::isMemoryEffectFree(Operation *op) {
@@ -358,16 +346,15 @@ bool mlir::isMemoryEffectFree(Operation *op) {
// free.
for (Region ®ion : op->getRegions())
for (Operation &op : region.getOps())
- if (!isMemoryEffectFree(&op))
+ if (!isMemoryEffectMovable(&op))
return false;
return true;
}
-bool mlir::isMemoryInitMovable(Operation *op) {
+bool mlir::isMemoryEffectConflictFree(Operation *op) {
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;
@@ -378,13 +365,12 @@ bool mlir::isMemoryInitMovable(Operation *op) {
// op has interface but no effects, be conservative
if (effects.empty()) return false;
-
DenseMap<TypeID, int> resourceCounts;
- // ensure op only has Init effects and gather unique
+ // ensure op only has Write effects and gather unique
// resource names
for (const MemoryEffects::EffectInstance &effect : effects) {
- if (!isa<MemoryEffects::Init>(effect.getEffect()))
+ if (!isa<MemoryEffects::Write>(effect.getEffect()))
return false;
resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
@@ -392,46 +378,55 @@ bool mlir::isMemoryInitMovable(Operation *op) {
// op itself is good, need to check rest of its parent region
Operation *parent = op->getParentOp();
+ mlir::DominanceInfo dom(parent);
for (Region ®ion : parent->getRegions())
- for (Operation &op_i : region.getOps())
- if (hasMemoryEffectInitConflict(&op_i, resourceCounts))
+ for (Operation &opI : region.getOps())
+ if (hasMemoryEffectConflict(op, &opI, dom, resourceCounts))
return false;
return true;
}
-bool mlir::hasMemoryEffectInitConflict(
- Operation *op, DenseMap<TypeID, int> &resourceCounts) {
+bool mlir::hasMemoryEffectConflict(
+ Operation *mainOp, Operation *op,
+ mlir::DominanceInfo &dom, DenseMap<TypeID, int> &resourceCounts) {
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;
+
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+ memInterface.getEffects(effects);
+
+ auto isDominated = dom.properlyDominates(mainOp, op);
- // only care about resources of the op that called
- // this recursive function for the first time
- auto resourceID = effect.getResource()->getResourceID();
+ // ensure op only has Write or dominated Read effects
+ // check used resources
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ auto resourceID = effect.getResource()->getResourceID();
- if (resourceCounts.contains(resourceID))
- if (++resourceCounts[resourceID] > 1)
- return true;
+ if (resourceCounts.contains(resourceID)) {
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ if (isDominated) {
+ continue; // skip dominated reads
+ }
+ }
+ else if (!isa<MemoryEffects::Write>(effect.getEffect())) {
+ return true; // count alloc/free in same region as conflict, be conservative
+ }
+
+ // update write counts, should always be <=1 per resource in region
+ if (++resourceCounts[resourceID] > 1) {
+ return true;
+ }
}
- return false;
}
}
// Recurse into the regions and ensure that nested ops don't
- // conflict with each others MemInits
+ // conflict with each other's MemWrites
for (Region ®ion : op->getRegions())
- for (Operation &op : region.getOps())
- if (hasMemoryEffectInitConflict(&op, resourceCounts))
+ for (Operation &opI : region.getOps())
+ if (hasMemoryEffectConflict(mainOp, &opI, dom, resourceCounts))
return true;
return false;
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index c1604e226a334..50882f91ea86e 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1437,3 +1437,207 @@ func.func @do_not_hoist_vector_transfer_ops_memref(
}
func.return %final : vector<4x4xf32>
}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_basic
+func.func @move_single_resource_basic() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c0_i32_0 = arith.constant 0 : i32
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_A"() : () -> ()
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_write_dominant
+func.func @move_single_resource_write_dominant() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c0_i32_0 = arith.constant 0 : i32
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_read_dominant_negative
+func.func @move_single_resource_read_dominant_negative() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c0_i32_0 = arith.constant 0 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ "test.test_effects_write_A"() : () -> ()
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_basic_conflict
+func.func @move_single_resource_basic_conflict() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c0_i32_0 = arith.constant 0 : i32
+ %c0_i32_1 = arith.constant 0 : i32
+ %c0_i32_2 = arith.constant 0 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ // CHECK: "test.test_effects_write_AC"() : () -> ()
+ "test.test_effects_write_AC"() : () -> ()
+ // CHECK: "test.test_effects_read_AC"() : () -> ()
+ "test.test_effects_read_AC"() : () -> ()
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_if_region_negative
+func.func @move_single_resource_if_region_negative() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c3_i32 = arith.constant 5 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ %1 = arith.cmpi slt, %arg0, %c3_i32 : i32
+
+ scf.if %1 {
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ }
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_for_inside_if_region
+func.func @move_single_resource_for_inside_if_region() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c3_i32 = arith.constant 5 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ %1 = arith.cmpi slt, %arg0, %c3_i32 : i32
+
+ scf.if %1 {
+ %c0_i32_0 = arith.constant 0 : i32
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_A"() : () -> ()
+ }
+ }
+ }
+ }
+ return
+}
+
+// -----
+
+// CHECK-LABEL func.func @move_multi_resource_comprehensive
+func.func @move_multi_resource_comprehensive() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c0_i32_2 = arith.constant 0 : i32
+ %c3_i32 = arith.constant 5 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_write_CD"() : () -> ()
+ // CHECK: "test.test_effects_write_EF"() : () -> ()
+ scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_CD"() : () -> ()
+ // CHECK: "test.test_effects_read_CD"() : () -> ()
+ "test.test_effects_read_CD"() : () -> ()
+
+ "test.test_effects_write_EF"() : () -> ()
+ // CHECK: "test.test_effects_read_EF"() : () -> ()
+ "test.test_effects_read_EF"() : () -> ()
+ }
+ }
+
+ %1 = arith.cmpi slt, %arg0, %c3_i32 : i32
+ scf.if %1 {
+ %c0_i32_0 = arith.constant 0 : i32
+ %c0_i32_1 = arith.constant 0 : i32
+
+ // CHECK: "test.test_effects_write_B"() : () -> ()
+ scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ scf.for %arg4 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ }
+
+ scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+ "test.test_effects_read_B"() : () -> ()
+ }
+
+ // CHECK: "test.test_effects_write_AC"() : () -> ()
+ scf.for %arg6 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_AC"() : () -> ()
+ // CHECK: "test.test_effects_read_AC"() : () -> ()
+ "test.test_effects_read_AC"() : () -> ()
+ }
+ }
+ }
+ else {
+ // CHECK: "test.test_effects_write_F"() : () -> ()
+ "test.test_effects_write_F"() : () -> ()
+ // CHECK: "test.test_effects_read_F"() : () -> ()
+ "test.test_effects_read_F"() : () -> ()
+ }
+ }
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.h b/mlir/test/lib/Dialect/Test/TestOps.h
index b414b47c87425..ab803663d7847 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.h
+++ b/mlir/test/lib/Dialect/Test/TestOps.h
@@ -56,6 +56,29 @@ struct TestResource : public mlir::SideEffects::Resource::Base<TestResource> {
llvm::StringRef getName() final { return "<Test>"; }
};
+struct TestResourceA : public mlir::SideEffects::Resource::Base<TestResourceA> {
+ llvm::StringRef getName() final { return "<TestA>"; }
+};
+
+struct TestResourceB : public mlir::SideEffects::Resource::Base<TestResourceB> {
+ llvm::StringRef getName() final { return "<TestB>"; }
+};
+
+struct TestResourceC : public mlir::SideEffects::Resource::Base<TestResourceC> {
+ llvm::StringRef getName() final { return "<TestC>"; }
+};
+
+struct TestResourceD : public mlir::SideEffects::Resource::Base<TestResourceD> {
+ llvm::StringRef getName() final { return "<TestD>"; }
+};
+
+struct TestResourceE : public mlir::SideEffects::Resource::Base<TestResourceE> {
+ llvm::StringRef getName() final { return "<TestE>"; }
+};
+
+struct TestResourceF : public mlir::SideEffects::Resource::Base<TestResourceF> {
+ llvm::StringRef getName() final { return "<TestF>"; }
+};
//===----------------------------------------------------------------------===//
// PropertiesWithCustomPrint
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 231400ec9cd29..550678e9e91d1 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2977,6 +2977,60 @@ def TestEffectsResult : TEST_Op<"test_effects_result"> {
let results = (outs Res<I32, "", [MemAlloc, MemWrite]>);
}
+//===----------------------------------------------------------------------===//
+// Test Ops with multiple effects for Loop Invariant Code Motion
+//===----------------------------------------------------------------------===//
+
+def TestResourceA : Resource<"TestResourceA">;
+def TestResourceB : Resource<"TestResourceB">;
+def TestResourceC : Resource<"TestResourceC">;
+def TestResourceD : Resource<"TestResourceD">;
+def TestResourceE : Resource<"TestResourceE">;
+def TestResourceF : Resource<"TestResourceF">;
+
+def TestEffectsWriteA : TEST_Op<"test_effects_write_A",
+ [MemoryEffects<[MemWrite<TestResourceA>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadA : TEST_Op<"test_effects_read_A",
+ [MemoryEffects<[MemRead<TestResourceA>]>]>;
+
+def TestEffectsWriteB : TEST_Op<"test_effects_write_B",
+ [MemoryEffects<[MemWrite<TestResourceB>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadB : TEST_Op<"test_effects_read_B",
+ [MemoryEffects<[MemRead<TestResourceB>]>]>;
+
+def TestEffectsWriteF : TEST_Op<"test_effects_write_F",
+ [MemoryEffects<[MemWrite<TestResourceF>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadF : TEST_Op<"test_effects_read_F",
+ [MemoryEffects<[MemRead<TestResourceF>]>]>;
+
+def TestEffectsWriteAC : TEST_Op<"test_effects_write_AC",
+ [MemoryEffects<[MemWrite<TestResourceA>, MemWrite<TestResourceC>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadAC : TEST_Op<"test_effects_read_AC",
+ [MemoryEffects<[MemRead<TestResourceA>, MemRead<TestResourceC>]>]>;
+
+def TestEffectsWriteCD : TEST_Op<"test_effects_write_CD",
+ [MemoryEffects<[MemWrite<TestResourceC>, MemWrite<TestResourceD>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadCD : TEST_Op<"test_effects_read_CD",
+ [MemoryEffects<[MemRead<TestResourceC>, MemRead<TestResourceD>]>]>;
+
+def TestEffectsWriteEF : TEST_Op<"test_effects_write_EF",
+ [MemoryEffects<[MemWrite<TestResourceE>, MemWrite<TestResourceF>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadEF : TEST_Op<"test_effects_read_EF",
+ [MemoryEffects<[MemRead<TestResourceE>, MemRead<TestResourceF>]>]>;
+
+
//===----------------------------------------------------------------------===//
// Test Ops with verifiers
//===----------------------------------------------------------------------===//
>From 8bad9d4a6bee2d61b06d0810da4d2b9bfb1454f6 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 26 Aug 2025 01:44:03 +0000
Subject: [PATCH 06/22] typo fix
---
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 245850014f507..e66c6480872dd 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -346,7 +346,7 @@ bool mlir::isMemoryEffectFree(Operation *op) {
// free.
for (Region ®ion : op->getRegions())
for (Operation &op : region.getOps())
- if (!isMemoryEffectMovable(&op))
+ if (!isMemoryEffectFree(&op))
return false;
return true;
>From ccb7f415a5b83bc4b74111e5e3d483bb462e3e3e Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 26 Aug 2025 20:39:38 +0000
Subject: [PATCH 07/22] LICM now checks if parent loop has constant
bounds/steps and isn't zero trip
---
.../mlir/Interfaces/SideEffectInterfaces.h | 26 +++++++---
mlir/lib/Interfaces/CMakeLists.txt | 16 +++++-
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 52 +++++++++++++++++--
.../loop-invariant-code-motion.mlir | 41 +++++++++++++--
4 files changed, 119 insertions(+), 16 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index ca48d23574feb..7c6ce6c8a51ee 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -14,8 +14,9 @@
#ifndef MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
#define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
-#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
namespace mlir {
namespace SideEffects {
@@ -423,6 +424,15 @@ bool isOpTriviallyDead(Operation *op);
/// Note: Terminators and symbols are never considered to be trivially dead.
bool wouldOpBeTriviallyDead(Operation *op);
+/// Returns TRUE if the loop is dead/zero-trip,
+/// FALSE if loop has constant bounds/steps and has at least 1 iteration
+/// on every dimension, returns nullopt otherwise
+///
+/// Can only infer if loop is dead if it has constant loop bounds/steps.
+/// Otherwise we assume that it's dead to be conservative.
+///
+std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface loop);
+
/// Returns true if the given operation is allowed to be moved under the
/// memory effects interface.
///
@@ -449,13 +459,15 @@ bool isMemoryEffectFree(Operation *op);
/// Returns true if the given operation has conflict-free write effects
///
/// An operation is conflict free:
-/// (1) all of its memory effects are of type Write
-/// (2) there are no other ops with Alloc/Free/Write effects on the same
-/// resources within the ops parent region
-/// (3) all ops in the parent region with Read effects on the same resources
-/// are dominated by the operation
+/// (1) Parent is a loop with the LoopLikeOpInterface
+/// (2) Parent loop is not a zero trip loop and has constant bounds/steps
+/// (3) all of the op's memory effects are of type Write
+/// (4) there are no other ops with Alloc/Free/Write effects on the same
+/// resources within the op's parent loop region
+/// (5) all ops in the parent loop region with Read effects on the same
+/// resources are dominated by the operation
///
-/// If the operation meets all 3 criteria, then it is conflict free
+/// If the operation meets all criteria, then it is conflict free
bool isMemoryEffectConflictFree(Operation *op);
/// Returns true if op and/or any operations within its nested regions
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index af923d98c76ff..ea699922dc9ec 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -85,7 +85,21 @@ add_mlir_interface_library(MemorySlotInterfaces)
add_mlir_interface_library(ParallelCombiningOpInterface)
add_mlir_interface_library(RuntimeVerifiableOpInterface)
add_mlir_interface_library(ShapedOpInterfaces)
-add_mlir_interface_library(SideEffectInterfaces)
+
+add_mlir_library(MLIRSideEffectInterfaces
+ SideEffectInterfaces.cpp
+
+ ADDITIONAL_HEADER_DIRS
+ ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
+
+ DEPENDS
+ MLIRSideEffectInterfacesIncGen
+
+ LINK_LIBS PUBLIC
+ MLIRIR
+ MLIRDialectUtils
+ MLIRLoopLikeInterface
+)
add_mlir_library(MLIRSubsetOpInterface
SubsetOpInterface.cpp
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index e66c6480872dd..da926b66adf5f 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -8,10 +8,9 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/SymbolTable.h"
-#include "llvm/ADT/SmallPtrSet.h"
-#include <unordered_set>
#include <utility>
using namespace mlir;
@@ -23,6 +22,13 @@ using namespace mlir;
/// Include the definitions of the side effect interfaces.
#include "mlir/Interfaces/SideEffectInterfaces.cpp.inc"
+// //===----------------------------------------------------------------------===//
+// // LoopLike Interfaces
+// //===----------------------------------------------------------------------===//
+
+// /// Include the definitions of the loop like interfaces.
+// #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
+
//===----------------------------------------------------------------------===//
// MemoryEffects
//===----------------------------------------------------------------------===//
@@ -315,6 +321,33 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
return wouldOpBeTriviallyDeadImpl(op);
}
+std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface loop) {
+ auto lbs = loop.getLoopLowerBounds();
+ auto ubs = loop.getLoopUpperBounds();
+ auto steps = loop.getLoopSteps();
+
+ if (!lbs || !ubs || !steps)
+ return std::nullopt;
+
+ if (lbs->size() != ubs->size() || ubs->size() != steps->size())
+ return std::nullopt;
+
+ for (size_t i = 0; i < steps->size(); ++i) {
+ auto lb = getConstantIntValue((*lbs)[i]);
+ auto ub = getConstantIntValue((*ubs)[i]);
+ auto st = getConstantIntValue((*steps)[i]);
+
+ if (!lb || !ub || !st)
+ return std::nullopt; // non-constant -> unknown
+
+ if (*st >= 0 && *lb >= *ub)
+ return true;
+ if (*st < 0 && *lb <= *ub)
+ return true;
+ }
+ return false;
+}
+
bool mlir::isMemoryEffectMovable(Operation *op) {
if (isMemoryEffectFree(op))
return true;
@@ -358,6 +391,19 @@ bool mlir::isMemoryEffectConflictFree(Operation *op) {
// shouldn't be flagged as movable to be conservative
if (!memInterface) return false;
+ // check parent loop to make sure it's not dead
+ Operation *parent = op->getParentOp();
+ if (!parent)
+ return false;
+
+ auto loopInterface = dyn_cast<LoopLikeOpInterface>(parent);
+ if (!loopInterface)
+ return false;
+
+ auto isDead = isZeroTrip(loopInterface);
+ if (!isDead.has_value() || isDead.value())
+ return false;
+
// gather all effects on op
llvm::SmallVector<MemoryEffects::EffectInstance> effects;
memInterface.getEffects(effects);
@@ -376,8 +422,6 @@ bool mlir::isMemoryEffectConflictFree(Operation *op) {
resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
}
- // op itself is good, need to check rest of its parent region
- Operation *parent = op->getParentOp();
mlir::DominanceInfo dom(parent);
for (Region ®ion : parent->getRegions())
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 50882f91ea86e..769e8c3a1b9ec 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1480,8 +1480,8 @@ func.func @move_single_resource_write_dominant() attributes {} {
// -----
-// CHECK-LABEL func.func @move_single_resource_read_dominant_negative
-func.func @move_single_resource_read_dominant_negative() attributes {} {
+// CHECK-LABEL func.func @move_single_resource_read_dominant
+func.func @move_single_resource_read_dominant() attributes {} {
%c0_i32 = arith.constant 0 : i32
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
@@ -1530,8 +1530,8 @@ func.func @move_single_resource_basic_conflict() attributes {} {
// -----
-// CHECK-LABEL func.func @move_single_resource_if_region_negative
-func.func @move_single_resource_if_region_negative() attributes {} {
+// CHECK-LABEL func.func @move_single_resource_if_region
+func.func @move_single_resource_if_region() attributes {} {
%c0_i32 = arith.constant 0 : i32
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
@@ -1641,3 +1641,36 @@ func.func @move_multi_resource_comprehensive() attributes {} {
}
return
}
+
+// -----
+
+// CHECK-LABEL func.func @move_single_resource_dead_loops
+func.func @move_single_resource_dead_loops() attributes {} {
+ %c0_i32 = arith.constant 0 : i32
+ %c1_i32 = arith.constant 10 : i32
+ %c2_i32 = arith.constant 1 : i32
+ %c3_i32 = arith.constant -1 : i32
+
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c3_i32 : i32 {
+ // CHECK: "test.test_effects_write_C"() : () -> ()
+ "test.test_effects_write_C"() : () -> ()
+
+ scf.for %arg1 = %c1_i32 to %c0_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_write_B"() : () -> ()
+ "test.test_effects_write_B"() : () -> ()
+
+ scf.for %arg2 = %c0_i32 to %c0_i32 step %c0_i32 : i32 {
+ // CHECK: "test.test_effects_write_A"() : () -> ()
+ "test.test_effects_write_A"() : () -> ()
+
+ // CHECK: "test.test_effects_write_EF"() : () -> ()
+ scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_EF"() : () -> ()
+ // CHECK: "test.test_effects_read_EF"() : () -> ()
+ "test.test_effects_read_EF"() : () -> ()
+ }
+ }
+ }
+ }
+ return
+}
>From f74d2f02c8e2218d9e7a610f2343563552a4c6ce Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 26 Aug 2025 20:58:35 +0000
Subject: [PATCH 08/22] some comments/blanks cleanup
---
mlir/include/mlir/Interfaces/SideEffectInterfaces.h | 1 -
mlir/include/mlir/Interfaces/SideEffectInterfaces.td | 1 -
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 7 -------
3 files changed, 9 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 5f8d1d3e5ac05..5363185f13d38 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -379,7 +379,6 @@ 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> {};
-
} // namespace MemoryEffects
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
index 841da64a509bd..b292174fccb36 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.td
@@ -87,7 +87,6 @@ def MemWrite : MemWrite<DefaultResource, 0, PartialEffect>;
class MemWriteAt<int stage, EffectRange range = PartialEffect>
: MemWrite<DefaultResource, stage, range>;
-
//===----------------------------------------------------------------------===//
// Effect Traits
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 0e5f5e87883ad..d9c2fd44fe50f 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -22,13 +22,6 @@ using namespace mlir;
/// Include the definitions of the side effect interfaces.
#include "mlir/Interfaces/SideEffectInterfaces.cpp.inc"
-// //===----------------------------------------------------------------------===//
-// // LoopLike Interfaces
-// //===----------------------------------------------------------------------===//
-
-// /// Include the definitions of the loop like interfaces.
-// #include "mlir/Interfaces/LoopLikeInterface.cpp.inc"
-
//===----------------------------------------------------------------------===//
// MemoryEffects
//===----------------------------------------------------------------------===//
>From 2f3c151d12e817fb8b141912d516789172497077 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 26 Aug 2025 21:30:22 +0000
Subject: [PATCH 09/22] made isZeroDrop pass-by-reference
---
mlir/include/mlir/Interfaces/SideEffectInterfaces.h | 2 +-
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 2 +-
2 files changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 5363185f13d38..3f19fbcdf5d13 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -468,7 +468,7 @@ bool wouldOpBeTriviallyDead(Operation *op);
/// Can only infer if loop is dead if it has constant loop bounds/steps.
/// Otherwise we assume that it's dead to be conservative.
///
-std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface loop);
+std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface &loop);
/// Returns true if the given operation is allowed to be moved under the
/// memory effects interface.
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index d9c2fd44fe50f..09c993667859b 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -319,7 +319,7 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
return wouldOpBeTriviallyDeadImpl(op);
}
-std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface loop) {
+std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface &loop) {
auto lbs = loop.getLoopLowerBounds();
auto ubs = loop.getLoopUpperBounds();
auto steps = loop.getLoopSteps();
>From efa955057967241cb9b713f950e805fe8fe284c3 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Wed, 27 Aug 2025 13:50:45 -0400
Subject: [PATCH 10/22] applying minor fixes from code review
Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
---
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 10 +++++++---
1 file changed, 7 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 09c993667859b..3d23225c57927 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -439,7 +439,7 @@ bool mlir::hasMemoryEffectConflict(
llvm::SmallVector<MemoryEffects::EffectInstance> effects;
memInterface.getEffects(effects);
- auto isDominated = dom.properlyDominates(mainOp, op);
+ bool isDominated = dom.properlyDominates(mainOp, op);
// ensure op only has Write or dominated Read effects
// check used resources
@@ -448,9 +448,8 @@ bool mlir::hasMemoryEffectConflict(
if (resourceCounts.contains(resourceID)) {
if (isa<MemoryEffects::Read>(effect.getEffect())) {
- if (isDominated) {
+ if (isDominated)
continue; // skip dominated reads
- }
}
else if (!isa<MemoryEffects::Write>(effect.getEffect())) {
return true; // count alloc/free in same region as conflict, be conservative
@@ -462,6 +461,11 @@ bool mlir::hasMemoryEffectConflict(
}
}
}
+ } 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
+ // op is conflicting or not.
+ return true;
}
// Recurse into the regions and ensure that nested ops don't
>From 9463b0c4434888b814091f2bf2fcb810a80d5a41 Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Tue, 2 Sep 2025 00:34:20 +0000
Subject: [PATCH 11/22] wip
---
.../mlir/Interfaces/SideEffectInterfaces.h | 76 ++++-----
.../Transforms/LoopInvariantCodeMotionUtils.h | 16 +-
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 123 +++-----------
.../Utils/LoopInvariantCodeMotionUtils.cpp | 153 +++++++++++++++++-
4 files changed, 216 insertions(+), 152 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 3f19fbcdf5d13..3e61ecc2df0e6 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -348,6 +348,14 @@ struct AlwaysSpeculatableImplTrait
//===----------------------------------------------------------------------===//
namespace MemoryEffects {
+enum Priority {
+ kDefault = 0,
+ kAllocPriority = 1,
+ kFreePriority = 2,
+ kReadPriority = 3,
+ kWritePriority = 4
+};
+
/// This class represents the base class used for memory effects.
struct Effect : public SideEffects::Effect {
using SideEffects::Effect::Effect;
@@ -357,28 +365,48 @@ struct Effect : public SideEffects::Effect {
using Base = SideEffects::Effect::Base<DerivedEffect, Effect>;
static bool classof(const SideEffects::Effect *effect);
+
+ /// Return the priority associated with this memory effect.
+ Priority getPriority() const { return priority; }
+
+protected:
+ /// Priority value for this effect. Lower numbers indicate higher precedence.
+ Priority priority = kDefault;
};
using EffectInstance = SideEffects::EffectInstance<Effect>;
+/// Returns vector of effects sorted by effect stage then priority
+/// priority order: allocate -> free -> read -> write
+llvm::SmallVector<MemoryEffects::EffectInstance>
+getMemoryEffectsSorted(Operation *op);
+
/// The following effect indicates that the operation allocates from some
/// resource. An 'allocate' effect implies only allocation of the resource, and
/// not any visible mutation or dereference.
-struct Allocate : public Effect::Base<Allocate> {};
+struct Allocate : public Effect::Base<Allocate> {
+ Allocate() : Effect::Base<Allocate>() { this->priority = kAllocPriority; }
+};
/// The following effect indicates that the operation frees some resource that
/// has been allocated. An 'allocate' effect implies only de-allocation of the
/// resource, and not any visible allocation, mutation or dereference.
-struct Free : public Effect::Base<Free> {};
+struct Free : public Effect::Base<Free> {
+ Free() : Effect::Base<Free>() { this->priority = kFreePriority; }
+};
/// The following effect indicates that the operation reads from some resource.
/// A 'read' effect implies only dereferencing of the resource, and not any
/// visible mutation.
-struct Read : public Effect::Base<Read> {};
+struct Read : public Effect::Base<Read> {
+ Read() : Effect::Base<Read>() { this->priority = kReadPriority; }
+};
/// The following effect indicates that the operation writes to some resource. A
/// 'write' effect implies only mutating a resource, and not any visible
/// dereference or read.
-struct Write : public Effect::Base<Write> {};
+struct Write : public Effect::Base<Write> {
+ Write() : Effect::Base<Write>() { this->priority = kWritePriority; }
+};
} // namespace MemoryEffects
//===----------------------------------------------------------------------===//
@@ -470,17 +498,6 @@ bool wouldOpBeTriviallyDead(Operation *op);
///
std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface &loop);
-/// Returns true if the given operation is allowed to be moved under the
-/// memory effects interface.
-///
-/// An operation is movable if either case is true:
-/// (a) free of memory effects as defined in isMemoryEffectFree()
-/// (b) if the operation does have memory effects, it must be conflict-free
-/// as defined in isMemoryEffectConflictFree()
-///
-/// If the operation meets either criteria, then it is movable under memory effects
-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
@@ -493,35 +510,6 @@ bool isMemoryEffectMovable(Operation *op);
/// conditions are satisfied.
bool isMemoryEffectFree(Operation *op);
-/// Returns true if the given operation has conflict-free write effects
-///
-/// An operation is conflict free:
-/// (1) Parent is a loop with the LoopLikeOpInterface
-/// (2) Parent loop is not a zero trip loop and has constant bounds/steps
-/// (3) all of the op's memory effects are of type Write
-/// (4) there are no other ops with Alloc/Free/Write effects on the same
-/// resources within the op's parent loop region
-/// (5) all ops in the parent loop region with Read effects on the same
-/// resources are dominated by the operation
-///
-/// If the operation meets all criteria, then it is conflict free
-bool isMemoryEffectConflictFree(Operation *op);
-
-/// Returns true if op and/or any operations within its nested regions
-/// have a memory effect conflict with mainOp as defined below:
-///
-/// op has a memory effect conflict with mainOp if op and/or any of
-/// the operations in its nested regions meet any of these criteria:
-/// (a) they have any Alloc/Free/Write effects on the resources used by mainOp
-/// (b) they dominate mainOp and have any read effect on the resources used by mainOp
-///
-/// Function mutates resources map
-///
-/// If none of the critera above are met, mainOp and op are conflict free
-bool hasMemoryEffectConflict(
- Operation *mainOp, Operation *op,
- mlir::DominanceInfo &dom, 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/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index 3ceef44d799e8..ec483f60e3732 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -9,9 +9,12 @@
#ifndef MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
#define MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
+#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
#include "llvm/ADT/SmallVector.h"
+#include <utility>
namespace mlir {
@@ -21,6 +24,17 @@ class Region;
class RewriterBase;
class Value;
+/// Gathers potential conflicts on all memory resources used within loop
+///
+/// Given a target loop and an op within it (or the loop op itself),
+/// gathers op's memory effects and flags potential resource conflicts
+/// in a map and then recurses into the op's regions to gather nested
+/// resource conflicts
+///
+/// First call should use loop = someLoop and op = someLoop.getOperation()
+void gatherResourceConflicts(LoopLikeOpInterface loop, Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts);
+
/// Given a list of regions, perform loop-invariant code motion. An operation is
/// loop-invariant if it depends only of values defined outside of the loop.
/// LICM moves these operations out of the loop body so that they are not
@@ -63,7 +77,7 @@ class Value;
///
/// Returns the number of operations moved.
size_t moveLoopInvariantCode(
- ArrayRef<Region *> regions,
+ LoopLikeOpInterface loopLike,
function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
function_ref<void(Operation *, Region *)> moveOutOfRegion);
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 3d23225c57927..7987bb86f4a7f 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -11,6 +11,7 @@
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/SymbolTable.h"
+#include <optional>
#include <utility>
using namespace mlir;
@@ -346,14 +347,29 @@ std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface &loop) {
return false;
}
-bool mlir::isMemoryEffectMovable(Operation *op) {
- if (isMemoryEffectFree(op))
- return true;
+llvm::SmallVector<MemoryEffects::EffectInstance>
+mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) {
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
- if (isMemoryEffectConflictFree(op))
- return true;
+ llvm::SmallVector<MemoryEffects::EffectInstance> effectsSorted;
+ memInterface.getEffects(effectsSorted);
- return false;
+ auto sortEffects =
+ [](llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+ llvm::stable_sort(effects, [](const MemoryEffects::EffectInstance &a,
+ const MemoryEffects::EffectInstance &b) {
+ if (a.getStage() < b.getStage())
+ return true;
+
+ if (a.getStage() == b.getStage())
+ return a.getEffect()->getPriority() < b.getEffect()->getPriority();
+
+ return false; // b before a
+ });
+ };
+ sortEffects(effectsSorted);
+
+ return effectsSorted;
}
bool mlir::isMemoryEffectFree(Operation *op) {
@@ -383,101 +399,6 @@ bool mlir::isMemoryEffectFree(Operation *op) {
return true;
}
-bool mlir::isMemoryEffectConflictFree(Operation *op) {
- auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
- // op does not implement the memory effect op interface
- // shouldn't be flagged as movable to be conservative
- if (!memInterface) return false;
-
- // check parent loop to make sure it's not dead
- Operation *parent = op->getParentOp();
- if (!parent)
- return false;
-
- auto loopInterface = dyn_cast<LoopLikeOpInterface>(parent);
- if (!loopInterface)
- return false;
-
- auto isDead = isZeroTrip(loopInterface);
- if (!isDead.has_value() || isDead.value())
- return false;
-
- // 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;
-
- DenseMap<TypeID, int> resourceCounts;
-
- // ensure op only has Write effects and gather unique
- // resource names
- for (const MemoryEffects::EffectInstance &effect : effects) {
- if (!isa<MemoryEffects::Write>(effect.getEffect()))
- return false;
-
- resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
- }
-
- mlir::DominanceInfo dom(parent);
-
- for (Region ®ion : parent->getRegions())
- for (Operation &opI : region.getOps())
- if (hasMemoryEffectConflict(op, &opI, dom, resourceCounts))
- return false;
-
- return true;
-}
-
-bool mlir::hasMemoryEffectConflict(
- Operation *mainOp, Operation *op,
- mlir::DominanceInfo &dom, DenseMap<TypeID, int> &resourceCounts) {
-
- if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
-
- llvm::SmallVector<MemoryEffects::EffectInstance> effects;
- memInterface.getEffects(effects);
-
- bool isDominated = dom.properlyDominates(mainOp, op);
-
- // ensure op only has Write or dominated Read effects
- // check used resources
- for (const MemoryEffects::EffectInstance &effect : effects) {
- auto resourceID = effect.getResource()->getResourceID();
-
- if (resourceCounts.contains(resourceID)) {
- if (isa<MemoryEffects::Read>(effect.getEffect())) {
- if (isDominated)
- continue; // skip dominated reads
- }
- else if (!isa<MemoryEffects::Write>(effect.getEffect())) {
- return true; // count alloc/free in same region as conflict, be conservative
- }
-
- // update write counts, should always be <=1 per resource in region
- if (++resourceCounts[resourceID] > 1) {
- 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
- // op is conflicting or not.
- return true;
- }
-
- // Recurse into the regions and ensure that nested ops don't
- // conflict with each other's MemWrites
- for (Region ®ion : op->getRegions())
- for (Operation &opI : region.getOps())
- if (hasMemoryEffectConflict(mainOp, &opI, dom, resourceCounts))
- 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 4968707ebf21f..0cfa8331f3c70 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -21,6 +21,7 @@
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include <queue>
+#include <utility>
#define DEBUG_TYPE "licm"
@@ -58,14 +59,149 @@ static bool canBeHoisted(Operation *op,
op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
}
+static void mergeResource(
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &dstMap,
+ const MemoryEffects::EffectInstance &srcEffect,
+ bool srcHasConflict) {
+
+ TypeID srcResourceID = srcEffect.getEffect()->getEffectID();
+
+ bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect())
+ || isa<MemoryEffects::Free>(srcEffect.getEffect());
+
+ bool conflict = srcHasConflict || srcIsAllocOrFree;
+
+ DenseMap<mlir::Value, std::pair<int, bool>> myMap;
+
+ auto dstIt = dstMap.find(srcResourceID);
+
+ // if it doesnt already exist, create entry for resource in map
+ if (dstIt == dstMap.end()) {
+ dstMap.try_emplace(srcResourceID, std::make_pair(conflict, srcEffect));
+ return;
+ }
+
+ // resource already in use
+ bool dstHasConflict = dstIt->second.first;
+ auto dstEffect = dstIt->second.second;
+
+ if (dstHasConflict)
+ return;
+
+ bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
+ bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
+ bool readBeforeWrite = dstRead && srcWrite;
+
+ conflict = conflict || readBeforeWrite;
+
+ dstMap.try_emplace(srcResourceID, std::make_pair(conflict, srcEffect));
+}
+
+static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
+ for (const auto &input : op->getOperands())
+ if (loopLike.isDefinedOutsideOfLoop(input))
+ return true;
+
+ return false;
+}
+
+static bool mayHaveMemoryEffectConflict(Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
+
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+
+ // op does not implement the memory effect op interface
+ // shouldn't be flagged as movable to be conservative
+ if (!memInterface)
+ return true;
+
+ // 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 true;
+
+ // RFC moving ops with HasRecursiveMemoryEffects that have nested ops
+ // with effects is not supported
+
+ // op has no conflicts IFF all resources are flagged as having no conflicts
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ auto resourceID = effect.getResource()->getResourceID();
+
+ auto resConIt = resourceConflicts.find(resourceID);
+ if (resConIt == resourceConflicts.end())
+ return true; // RFC realistically shouldn't reach here but just in case?
+
+ bool hasConflict = resConIt->second.first;
+ if (hasConflict)
+ return true;
+ }
+
+ return false;
+}
+
+void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
+
+ if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ // gather all effects on op
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects =
+ MemoryEffects::getMemoryEffectsSorted(op);
+
+ if (!effects.empty()) {
+ // any variant input to the op could be the data source
+ // for writes, set all writes to conflict
+ bool writesConflict = hasLoopVariantInput(loopLike, op);
+
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ bool conflict = false;
+ bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
+
+ // all writes to a resource that follow a read on any other resource
+ // have to be considered a conflict as guaranteeing that the read
+ // is invariant and won't affect the write requires more robust logic
+ if (isa<MemoryEffects::Read>(effect.getEffect()))
+ writesConflict = true;
+
+ if (isWrite)
+ if (writesConflict)
+ conflict = true;
+
+ mergeResource(resourceConflicts, effect, conflict);
+ }
+ }
+ }
+
+ for (Region ®ion : op->getRegions())
+ for (Operation &opInner : region.getOps())
+ gatherResourceConflicts(loopLike, &opInner, resourceConflicts);
+}
+
size_t mlir::moveLoopInvariantCode(
- ArrayRef<Region *> regions,
+ LoopLikeOpInterface loopLike,
function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
function_ref<void(Operation *, Region *)> moveOutOfRegion) {
size_t numMoved = 0;
- for (Region *region : regions) {
+ // check that the loop isn't dead
+ auto isDead = isZeroTrip(loopLike);
+ if (!isDead.has_value() || isDead.value())
+ return numMoved;
+
+ // go through loop body and map out resource usages
+ // op->regions are essentially merged sequentially
+ // e.g. an if's "then" and "else" regions are treated like one
+ // continuous region --> need to add fork checking
+ //
+ // loop "do" and "then" regions also merged
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> resourceConflicts;
+ mlir::gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
+
+
+ for (Region *region : loopLike.getLoopRegions()) {
LDBG() << "Original loop:\n" << *region->getParentOp();
std::queue<Operation *> worklist;
@@ -86,8 +222,13 @@ size_t mlir::moveLoopInvariantCode(
LDBG() << "Checking op: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
- if (!shouldMoveOutOfRegion(op, region) ||
- !canBeHoisted(op, definedOutside))
+
+ bool noMemoryConflicts = isMemoryEffectFree(op)
+ || !mayHaveMemoryEffectConflict(op, resourceConflicts);
+
+ if (!noMemoryConflicts
+ || !shouldMoveOutOfRegion(op, region)
+ || !canBeHoisted(op, definedOutside))
continue;
LDBG() << "Moving loop-invariant op: " << *op;
@@ -107,12 +248,12 @@ size_t mlir::moveLoopInvariantCode(
size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
return moveLoopInvariantCode(
- loopLike.getLoopRegions(),
+ loopLike,
[&](Value value, Region *) {
return loopLike.isDefinedOutsideOfLoop(value);
},
[&](Operation *op, Region *) {
- return isSpeculatable(op) && isMemoryEffectMovable(op);
+ return isSpeculatable(op);
},
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
>From a87023b62231e313ffe9fe5c32e169460aad749b Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Thu, 4 Sep 2025 23:30:32 +0000
Subject: [PATCH 12/22] working version without isZeroTrip
---
.../mlir/Interfaces/SideEffectInterfaces.h | 22 +---
.../Transforms/LoopInvariantCodeMotionUtils.h | 1 -
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 34 +-----
.../Utils/LoopInvariantCodeMotionUtils.cpp | 25 ++--
.../loop-invariant-code-motion.mlir | 115 ++++++++++++------
mlir/test/lib/Dialect/Test/TestOps.td | 40 +++++-
6 files changed, 136 insertions(+), 101 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 3e61ecc2df0e6..e2b3ba10e40cd 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -16,7 +16,6 @@
#include "mlir/IR/Dominance.h"
#include "mlir/IR/OpDefinition.h"
-#include "mlir/Interfaces/LoopLikeInterface.h"
namespace mlir {
namespace SideEffects {
@@ -349,7 +348,7 @@ struct AlwaysSpeculatableImplTrait
namespace MemoryEffects {
enum Priority {
- kDefault = 0,
+ kDefaultPriority = 0,
kAllocPriority = 1,
kFreePriority = 2,
kReadPriority = 3,
@@ -371,7 +370,7 @@ struct Effect : public SideEffects::Effect {
protected:
/// Priority value for this effect. Lower numbers indicate higher precedence.
- Priority priority = kDefault;
+ Priority priority = Priority::kDefaultPriority;
};
using EffectInstance = SideEffects::EffectInstance<Effect>;
@@ -384,28 +383,28 @@ getMemoryEffectsSorted(Operation *op);
/// resource. An 'allocate' effect implies only allocation of the resource, and
/// not any visible mutation or dereference.
struct Allocate : public Effect::Base<Allocate> {
- Allocate() : Effect::Base<Allocate>() { this->priority = kAllocPriority; }
+ Allocate() : Effect::Base<Allocate>() { this->priority = Priority::kAllocPriority; }
};
/// The following effect indicates that the operation frees some resource that
/// has been allocated. An 'allocate' effect implies only de-allocation of the
/// resource, and not any visible allocation, mutation or dereference.
struct Free : public Effect::Base<Free> {
- Free() : Effect::Base<Free>() { this->priority = kFreePriority; }
+ Free() : Effect::Base<Free>() { this->priority = Priority::kFreePriority; }
};
/// The following effect indicates that the operation reads from some resource.
/// A 'read' effect implies only dereferencing of the resource, and not any
/// visible mutation.
struct Read : public Effect::Base<Read> {
- Read() : Effect::Base<Read>() { this->priority = kReadPriority; }
+ Read() : Effect::Base<Read>() { this->priority = Priority::kReadPriority; }
};
/// The following effect indicates that the operation writes to some resource. A
/// 'write' effect implies only mutating a resource, and not any visible
/// dereference or read.
struct Write : public Effect::Base<Write> {
- Write() : Effect::Base<Write>() { this->priority = kWritePriority; }
+ Write() : Effect::Base<Write>() { this->priority = Priority::kWritePriority; }
};
} // namespace MemoryEffects
@@ -489,15 +488,6 @@ bool isOpTriviallyDead(Operation *op);
/// Note: Terminators and symbols are never considered to be trivially dead.
bool wouldOpBeTriviallyDead(Operation *op);
-/// Returns TRUE if the loop is dead/zero-trip,
-/// FALSE if loop has constant bounds/steps and has at least 1 iteration
-/// on every dimension, returns nullopt otherwise
-///
-/// Can only infer if loop is dead if it has constant loop bounds/steps.
-/// Otherwise we assume that it's dead to be conservative.
-///
-std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface &loop);
-
/// Returns true if the given operation is free of memory effects.
///
/// An operation is free of memory effects if its implementation of
diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index ec483f60e3732..82581efdadbed 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -13,7 +13,6 @@
#include "mlir/Support/LLVM.h"
#include "mlir/Support/TypeID.h"
-#include "llvm/ADT/SmallVector.h"
#include <utility>
namespace mlir {
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 7987bb86f4a7f..ec8618e24e308 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -8,7 +8,6 @@
#include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/SymbolTable.h"
#include <optional>
@@ -320,38 +319,15 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
return wouldOpBeTriviallyDeadImpl(op);
}
-std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface &loop) {
- auto lbs = loop.getLoopLowerBounds();
- auto ubs = loop.getLoopUpperBounds();
- auto steps = loop.getLoopSteps();
-
- if (!lbs || !ubs || !steps)
- return std::nullopt;
-
- if (lbs->size() != ubs->size() || ubs->size() != steps->size())
- return std::nullopt;
-
- for (size_t i = 0; i < steps->size(); ++i) {
- auto lb = getConstantIntValue((*lbs)[i]);
- auto ub = getConstantIntValue((*ubs)[i]);
- auto st = getConstantIntValue((*steps)[i]);
-
- if (!lb || !ub || !st)
- return std::nullopt; // non-constant -> unknown
-
- if (*st >= 0 && *lb >= *ub)
- return true;
- if (*st < 0 && *lb <= *ub)
- return true;
- }
- return false;
-}
-
llvm::SmallVector<MemoryEffects::EffectInstance>
mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) {
+ llvm::SmallVector<MemoryEffects::EffectInstance> effectsSorted;
+
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
- llvm::SmallVector<MemoryEffects::EffectInstance> effectsSorted;
+ if (!memInterface)
+ return effectsSorted; // return empty vec
+
memInterface.getEffects(effectsSorted);
auto sortEffects =
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 0cfa8331f3c70..960bddd139ea8 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -15,6 +15,7 @@
#include "mlir/IR/Operation.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Interfaces/SubsetOpInterface.h"
@@ -64,20 +65,18 @@ static void mergeResource(
const MemoryEffects::EffectInstance &srcEffect,
bool srcHasConflict) {
- TypeID srcResourceID = srcEffect.getEffect()->getEffectID();
+ TypeID srcResourceID = srcEffect.getResource()->getResourceID();
bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect())
|| isa<MemoryEffects::Free>(srcEffect.getEffect());
bool conflict = srcHasConflict || srcIsAllocOrFree;
- DenseMap<mlir::Value, std::pair<int, bool>> myMap;
-
auto dstIt = dstMap.find(srcResourceID);
- // if it doesnt already exist, create entry for resource in map
+ // if it doesn't already exist, create entry for resource in map
if (dstIt == dstMap.end()) {
- dstMap.try_emplace(srcResourceID, std::make_pair(conflict, srcEffect));
+ dstMap.insert(std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
return;
}
@@ -94,12 +93,12 @@ static void mergeResource(
conflict = conflict || readBeforeWrite;
- dstMap.try_emplace(srcResourceID, std::make_pair(conflict, srcEffect));
+ dstIt->second =std::make_pair(conflict, srcEffect);
}
static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
- for (const auto &input : op->getOperands())
- if (loopLike.isDefinedOutsideOfLoop(input))
+ for (OpOperand &input : op->getOpOperands())
+ if (!loopLike.isDefinedOutsideOfLoop(input.get()))
return true;
return false;
@@ -187,9 +186,9 @@ size_t mlir::moveLoopInvariantCode(
size_t numMoved = 0;
// check that the loop isn't dead
- auto isDead = isZeroTrip(loopLike);
- if (!isDead.has_value() || isDead.value())
- return numMoved;
+ // auto isDead = loopLike.isZeroTrip();
+ // if (!isDead.has_value() || isDead.value())
+ // return numMoved;
// go through loop body and map out resource usages
// op->regions are essentially merged sequentially
@@ -200,8 +199,8 @@ size_t mlir::moveLoopInvariantCode(
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> resourceConflicts;
mlir::gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
-
- for (Region *region : loopLike.getLoopRegions()) {
+ auto regions = loopLike.getLoopRegions();
+ for (Region *region : regions) {
LDBG() << "Original loop:\n" << *region->getParentOp();
std::queue<Operation *> worklist;
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 769e8c3a1b9ec..e97351a4f57f2 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1446,7 +1446,9 @@ func.func @move_single_resource_basic() attributes {} {
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
%c0_i32_0 = arith.constant 0 : i32
+
// CHECK: "test.test_effects_write_A"() : () -> ()
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
@@ -1465,12 +1467,16 @@ func.func @move_single_resource_write_dominant() attributes {} {
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
%c0_i32_0 = arith.constant 0 : i32
+
// CHECK: "test.test_effects_write_A"() : () -> ()
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
- "test.test_effects_write_A"() : () -> ()
+
// CHECK: "test.test_effects_read_A"() : () -> ()
+
+ "test.test_effects_write_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
}
@@ -1490,9 +1496,11 @@ func.func @move_single_resource_read_dominant() attributes {} {
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+
// CHECK: "test.test_effects_read_A"() : () -> ()
- "test.test_effects_read_A"() : () -> ()
// CHECK: "test.test_effects_write_A"() : () -> ()
+
+ "test.test_effects_read_A"() : () -> ()
"test.test_effects_write_A"() : () -> ()
}
}
@@ -1514,13 +1522,15 @@ func.func @move_single_resource_basic_conflict() attributes {} {
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+
// CHECK: "test.test_effects_write_A"() : () -> ()
- "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
- "test.test_effects_read_A"() : () -> ()
// CHECK: "test.test_effects_write_AC"() : () -> ()
- "test.test_effects_write_AC"() : () -> ()
// CHECK: "test.test_effects_read_AC"() : () -> ()
+
+ "test.test_effects_write_A"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ "test.test_effects_write_AC"() : () -> ()
"test.test_effects_read_AC"() : () -> ()
}
}
@@ -1543,9 +1553,11 @@ func.func @move_single_resource_if_region() attributes {} {
%1 = arith.cmpi slt, %arg0, %c3_i32 : i32
scf.if %1 {
+
// CHECK: "test.test_effects_write_A"() : () -> ()
- "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
+
+ "test.test_effects_write_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
}
@@ -1568,7 +1580,9 @@ func.func @move_single_resource_for_inside_if_region() attributes {} {
scf.if %1 {
%c0_i32_0 = arith.constant 0 : i32
+
// CHECK: "test.test_effects_write_A"() : () -> ()
+
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_A"() : () -> ()
@@ -1590,16 +1604,17 @@ func.func @move_multi_resource_comprehensive() attributes {} {
%c3_i32 = arith.constant 5 : i32
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+
// CHECK: "test.test_effects_write_CD"() : () -> ()
+ // CHECK: "test.test_effects_read_CD"() : () -> ()
// CHECK: "test.test_effects_write_EF"() : () -> ()
+ // CHECK: "test.test_effects_read_EF"() : () -> ()
+
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_CD"() : () -> ()
- // CHECK: "test.test_effects_read_CD"() : () -> ()
"test.test_effects_read_CD"() : () -> ()
-
"test.test_effects_write_EF"() : () -> ()
- // CHECK: "test.test_effects_read_EF"() : () -> ()
"test.test_effects_read_EF"() : () -> ()
}
}
@@ -1610,32 +1625,37 @@ func.func @move_multi_resource_comprehensive() attributes {} {
%c0_i32_1 = arith.constant 0 : i32
// CHECK: "test.test_effects_write_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+
scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+
// CHECK: "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+
scf.for %arg4 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_A"() : () -> ()
- // CHECK: "test.test_effects_read_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
- scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_B"() : () -> ()
- // CHECK: "test.test_effects_read_B"() : () -> ()
"test.test_effects_read_B"() : () -> ()
}
// CHECK: "test.test_effects_write_AC"() : () -> ()
+ // CHECK: "test.test_effects_read_AC"() : () -> ()
+
scf.for %arg6 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_AC"() : () -> ()
- // CHECK: "test.test_effects_read_AC"() : () -> ()
"test.test_effects_read_AC"() : () -> ()
}
}
}
else {
// CHECK: "test.test_effects_write_F"() : () -> ()
- "test.test_effects_write_F"() : () -> ()
// CHECK: "test.test_effects_read_F"() : () -> ()
+
+ "test.test_effects_write_F"() : () -> ()
"test.test_effects_read_F"() : () -> ()
}
}
@@ -1644,33 +1664,56 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// -----
-// CHECK-LABEL func.func @move_single_resource_dead_loops
-func.func @move_single_resource_dead_loops() attributes {} {
+// CHECK-LABEL func.func @move_multi_resource_write_conflicts
+func.func @move_multi_resource_write_conflicts() attributes {} {
%c0_i32 = arith.constant 0 : i32
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
- %c3_i32 = arith.constant -1 : i32
-
- scf.for %arg0 = %c0_i32 to %c1_i32 step %c3_i32 : i32 {
- // CHECK: "test.test_effects_write_C"() : () -> ()
- "test.test_effects_write_C"() : () -> ()
- scf.for %arg1 = %c1_i32 to %c0_i32 step %c2_i32 : i32 {
- // CHECK: "test.test_effects_write_B"() : () -> ()
- "test.test_effects_write_B"() : () -> ()
+ // CHECK: "test.test_effects_write_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_write_A_with_input"(%c7) : (index) -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
- scf.for %arg2 = %c0_i32 to %c0_i32 step %c0_i32 : i32 {
- // CHECK: "test.test_effects_write_A"() : () -> ()
- "test.test_effects_write_A"() : () -> ()
+ %input = arith.constant 7 : index
+ "test.test_effects_write_A_with_input"(%input) : (index) -> ()
+ "test.test_effects_read_A"() : () -> ()
- // CHECK: "test.test_effects_write_EF"() : () -> ()
- scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
- "test.test_effects_write_EF"() : () -> ()
- // CHECK: "test.test_effects_read_EF"() : () -> ()
- "test.test_effects_read_EF"() : () -> ()
- }
- }
- }
+ "test.test_effects_write_B"() : () -> ()
+ "test.test_effects_read_B"() : () -> ()
}
+
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_read_A_write_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+
+ "test.test_effects_read_A_write_B"() : () -> ()
+ "test.test_effects_read_B"() : () -> ()
+
+ "test.test_effects_read_A"() : () -> ()
+ }
+
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // CHECK: "test.test_effects_read_A_then_write_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+
+ "test.test_effects_read_A_then_write_B"() : () -> ()
+ "test.test_effects_read_B"() : () -> ()
+
+ "test.test_effects_read_A"() : () -> ()
+ }
+
+ // CHECK: "test.test_effects_write_A_then_read_B"() : () -> ()
+ // CHECK: "test.test_effects_read_B"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ "test.test_effects_write_A_then_read_B"() : () -> ()
+ "test.test_effects_read_B"() : () -> ()
+ "test.test_effects_read_A"() : () -> ()
+ }
+
return
-}
+}
\ No newline at end of file
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 550678e9e91d1..9616560dd5fd5 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2992,43 +2992,71 @@ def TestEffectsWriteA : TEST_Op<"test_effects_write_A",
[MemoryEffects<[MemWrite<TestResourceA>]>,
AlwaysSpeculatable]>;
+def TestEffectsWriteAWithInput : TEST_Op<"test_effects_write_A_with_input",
+ [MemoryEffects<[MemWrite<TestResourceA>]>,
+ AlwaysSpeculatable]> {
+
+ let arguments = (ins AnyType:$arg);
+}
+
def TestEffectsReadA : TEST_Op<"test_effects_read_A",
- [MemoryEffects<[MemRead<TestResourceA>]>]>;
+ [MemoryEffects<[MemRead<TestResourceA>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadAWriteB : TEST_Op<"test_effects_read_A_write_B",
+ [MemoryEffects<[MemRead<TestResourceA>,
+ MemWrite<TestResourceB>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsReadAThenWriteB : TEST_Op<"test_effects_read_A_then_write_B",
+ [MemoryEffects<[MemRead<TestResourceA, 0>,
+ MemWrite<TestResourceB, 1>]>,
+ AlwaysSpeculatable]>;
+
+def TestEffectsWriteAThenReadB : TEST_Op<"test_effects_write_A_then_read_B",
+ [MemoryEffects<[MemWrite<TestResourceA, 0>,
+ MemRead<TestResourceB, 1>]>,
+ AlwaysSpeculatable]>;
def TestEffectsWriteB : TEST_Op<"test_effects_write_B",
[MemoryEffects<[MemWrite<TestResourceB>]>,
AlwaysSpeculatable]>;
def TestEffectsReadB : TEST_Op<"test_effects_read_B",
- [MemoryEffects<[MemRead<TestResourceB>]>]>;
+ [MemoryEffects<[MemRead<TestResourceB>]>,
+ AlwaysSpeculatable]>;
def TestEffectsWriteF : TEST_Op<"test_effects_write_F",
[MemoryEffects<[MemWrite<TestResourceF>]>,
AlwaysSpeculatable]>;
def TestEffectsReadF : TEST_Op<"test_effects_read_F",
- [MemoryEffects<[MemRead<TestResourceF>]>]>;
+ [MemoryEffects<[MemRead<TestResourceF>]>,
+ AlwaysSpeculatable]>;
def TestEffectsWriteAC : TEST_Op<"test_effects_write_AC",
[MemoryEffects<[MemWrite<TestResourceA>, MemWrite<TestResourceC>]>,
AlwaysSpeculatable]>;
def TestEffectsReadAC : TEST_Op<"test_effects_read_AC",
- [MemoryEffects<[MemRead<TestResourceA>, MemRead<TestResourceC>]>]>;
+ [MemoryEffects<[MemRead<TestResourceA>, MemRead<TestResourceC>]>,
+ AlwaysSpeculatable]>;
def TestEffectsWriteCD : TEST_Op<"test_effects_write_CD",
[MemoryEffects<[MemWrite<TestResourceC>, MemWrite<TestResourceD>]>,
AlwaysSpeculatable]>;
def TestEffectsReadCD : TEST_Op<"test_effects_read_CD",
- [MemoryEffects<[MemRead<TestResourceC>, MemRead<TestResourceD>]>]>;
+ [MemoryEffects<[MemRead<TestResourceC>, MemRead<TestResourceD>]>,
+ AlwaysSpeculatable]>;
def TestEffectsWriteEF : TEST_Op<"test_effects_write_EF",
[MemoryEffects<[MemWrite<TestResourceE>, MemWrite<TestResourceF>]>,
AlwaysSpeculatable]>;
def TestEffectsReadEF : TEST_Op<"test_effects_read_EF",
- [MemoryEffects<[MemRead<TestResourceE>, MemRead<TestResourceF>]>]>;
+ [MemoryEffects<[MemRead<TestResourceE>, MemRead<TestResourceF>]>,
+ AlwaysSpeculatable]>;
//===----------------------------------------------------------------------===//
>From 85d37f302e67e2aaa0549f091b8c03f397482502 Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Fri, 5 Sep 2025 15:31:00 +0000
Subject: [PATCH 13/22] docstrings
---
.../Utils/LoopInvariantCodeMotionUtils.cpp | 17 +++++++++++++----
1 file changed, 13 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 960bddd139ea8..d3dbda8b7c2a6 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -60,8 +60,11 @@ static bool canBeHoisted(Operation *op,
op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
}
+/// Merges srcEffect's Memory Effect on its resource into the
+/// resourceConflicts map, flagging resources if the srcEffect
+/// results in a conflict
static void mergeResource(
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &dstMap,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts,
const MemoryEffects::EffectInstance &srcEffect,
bool srcHasConflict) {
@@ -72,11 +75,11 @@ static void mergeResource(
bool conflict = srcHasConflict || srcIsAllocOrFree;
- auto dstIt = dstMap.find(srcResourceID);
+ auto dstIt = resourceConflicts.find(srcResourceID);
// if it doesn't already exist, create entry for resource in map
- if (dstIt == dstMap.end()) {
- dstMap.insert(std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+ if (dstIt == resourceConflicts.end()) {
+ resourceConflicts.insert(std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
return;
}
@@ -96,6 +99,7 @@ static void mergeResource(
dstIt->second =std::make_pair(conflict, srcEffect);
}
+/// Returns true if any of op's OpOperands are defined outside of loopLike
static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
for (OpOperand &input : op->getOpOperands())
if (!loopLike.isDefinedOutsideOfLoop(input.get()))
@@ -104,6 +108,11 @@ static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
return false;
}
+/// Returns true if:
+/// (a) any of the resources used by op's Memory Effects have been
+/// flagged as having a conflict within the resourceConflicts map
+/// (b) op doesn't have a MemoryEffectOpInterface or has one but
+/// without any specific effects
static bool mayHaveMemoryEffectConflict(Operation *op,
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
>From 871ecf41a731c6440e80b637fd02dc7ae9356402 Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Fri, 12 Sep 2025 21:23:29 +0000
Subject: [PATCH 14/22] PR cleanup
---
.../mlir/Interfaces/SideEffectInterfaces.h | 30 +++++++-----
.../Transforms/LoopInvariantCodeMotionUtils.h | 24 +++++++++-
mlir/lib/Interfaces/CMakeLists.txt | 16 +------
.../Utils/LoopInvariantCodeMotionUtils.cpp | 48 +++++++++++++------
mlir/test/lib/Dialect/Test/TestOps.td | 2 +-
5 files changed, 76 insertions(+), 44 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index e2b3ba10e40cd..14aa6fd8550c2 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -347,12 +347,18 @@ struct AlwaysSpeculatableImplTrait
//===----------------------------------------------------------------------===//
namespace MemoryEffects {
+/// Defines the priority of the different memory effects.
+///
+/// Sorting/ordering memory effects of an operation is done based on
+/// their defined stage and priority, in that order. If stage values for two effect instances are
+/// equal, they are then sorted by priority. Lower priority values indicate higher
+/// precedence.
enum Priority {
- kDefaultPriority = 0,
- kAllocPriority = 1,
- kFreePriority = 2,
- kReadPriority = 3,
- kWritePriority = 4
+ DefaultPriority = 0,
+ AllocPriority = 1,
+ FreePriority = 2,
+ ReadPriority = 3,
+ WritePriority = 4
};
/// This class represents the base class used for memory effects.
@@ -370,12 +376,12 @@ struct Effect : public SideEffects::Effect {
protected:
/// Priority value for this effect. Lower numbers indicate higher precedence.
- Priority priority = Priority::kDefaultPriority;
+ Priority priority = Priority::DefaultPriority;
};
using EffectInstance = SideEffects::EffectInstance<Effect>;
-/// Returns vector of effects sorted by effect stage then priority
-/// priority order: allocate -> free -> read -> write
+/// Returns vector of the op's memory effects sorted in increasing stage order
+/// and in increasing priority order within each stage.
llvm::SmallVector<MemoryEffects::EffectInstance>
getMemoryEffectsSorted(Operation *op);
@@ -383,28 +389,28 @@ getMemoryEffectsSorted(Operation *op);
/// resource. An 'allocate' effect implies only allocation of the resource, and
/// not any visible mutation or dereference.
struct Allocate : public Effect::Base<Allocate> {
- Allocate() : Effect::Base<Allocate>() { this->priority = Priority::kAllocPriority; }
+ Allocate() : Effect::Base<Allocate>() { this->priority = Priority::AllocPriority; }
};
/// The following effect indicates that the operation frees some resource that
/// has been allocated. An 'allocate' effect implies only de-allocation of the
/// resource, and not any visible allocation, mutation or dereference.
struct Free : public Effect::Base<Free> {
- Free() : Effect::Base<Free>() { this->priority = Priority::kFreePriority; }
+ Free() : Effect::Base<Free>() { this->priority = Priority::FreePriority; }
};
/// The following effect indicates that the operation reads from some resource.
/// A 'read' effect implies only dereferencing of the resource, and not any
/// visible mutation.
struct Read : public Effect::Base<Read> {
- Read() : Effect::Base<Read>() { this->priority = Priority::kReadPriority; }
+ Read() : Effect::Base<Read>() { this->priority = Priority::ReadPriority; }
};
/// The following effect indicates that the operation writes to some resource. A
/// 'write' effect implies only mutating a resource, and not any visible
/// dereference or read.
struct Write : public Effect::Base<Write> {
- Write() : Effect::Base<Write>() { this->priority = Priority::kWritePriority; }
+ Write() : Effect::Base<Write>() { this->priority = Priority::WritePriority; }
};
} // namespace MemoryEffects
diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index 82581efdadbed..a05dce51e4e3a 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -23,14 +23,34 @@ class Region;
class RewriterBase;
class Value;
-/// Gathers potential conflicts on all memory resources used within loop
+/// Gathers potential conflicts on all memory resources used within loop.
///
/// Given a target loop and an op within it (or the loop op itself),
/// gathers op's memory effects and flags potential resource conflicts
/// in a map and then recurses into the op's regions to gather nested
-/// resource conflicts
+/// resource conflicts.
///
+/// Typical usage:
+/// \code
+/// LoopLikeOpInterface myLoop = ...;
+/// DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> myConflicts;
+/// gatherResourceConflicts(myLoop, myLoop.getOperation(), resourceConflicts);
+/// \endcode
+///
+/// \param loop The loop to gather resource conflicts for.
+/// \param op The operation to gather resource conflicts for,
+/// typically the loop op itself via loop.getOperation().
+/// \param resourceConflicts Map to store potential resource conflicts.
+/// Key is the resource ID that effects are applied to. Value is a pair of
+/// a boolean, indicating if the resource has a conflict, and the last effect
+/// that was applied to the resource (if no conflicts exist) or the effect
+/// that caused the conflict (if conflicts exist). This was done so we
+/// check the effect that causes the conflict for debugging purposes.
/// First call should use loop = someLoop and op = someLoop.getOperation()
+///
+/// resourceConflicts is modified by the function and will be non-empty
+/// as long as there are memory effects within the loop, even if there are
+/// no conflicts.
void gatherResourceConflicts(LoopLikeOpInterface loop, Operation *op,
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts);
diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt
index ea699922dc9ec..af923d98c76ff 100644
--- a/mlir/lib/Interfaces/CMakeLists.txt
+++ b/mlir/lib/Interfaces/CMakeLists.txt
@@ -85,21 +85,7 @@ add_mlir_interface_library(MemorySlotInterfaces)
add_mlir_interface_library(ParallelCombiningOpInterface)
add_mlir_interface_library(RuntimeVerifiableOpInterface)
add_mlir_interface_library(ShapedOpInterfaces)
-
-add_mlir_library(MLIRSideEffectInterfaces
- SideEffectInterfaces.cpp
-
- ADDITIONAL_HEADER_DIRS
- ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
-
- DEPENDS
- MLIRSideEffectInterfacesIncGen
-
- LINK_LIBS PUBLIC
- MLIRIR
- MLIRDialectUtils
- MLIRLoopLikeInterface
-)
+add_mlir_interface_library(SideEffectInterfaces)
add_mlir_library(MLIRSubsetOpInterface
SubsetOpInterface.cpp
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index d3dbda8b7c2a6..b977ab318534c 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -61,8 +61,15 @@ static bool canBeHoisted(Operation *op,
}
/// Merges srcEffect's Memory Effect on its resource into the
-/// resourceConflicts map, flagging resources if the srcEffect
-/// results in a conflict
+/// resourceConflicts map, flagging the resource if the srcEffect
+/// results in a conflict.
+///
+/// \param resourceConflicts The map to store resources' conflicts status.
+/// \param srcEffect The effect to merge into the resourceConflicts map.
+/// \param srcHasConflict Whether the srcEffect results in a conflict based
+/// on higher level analysis.
+///
+/// resourceConflicts is modified by the function and will be non-empty
static void mergeResource(
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts,
const MemoryEffects::EffectInstance &srcEffect,
@@ -110,9 +117,9 @@ static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
/// Returns true if:
/// (a) any of the resources used by op's Memory Effects have been
-/// flagged as having a conflict within the resourceConflicts map
+/// flagged as having a conflict within the resourceConflicts map OR
/// (b) op doesn't have a MemoryEffectOpInterface or has one but
-/// without any specific effects
+/// without any specific effects.
static bool mayHaveMemoryEffectConflict(Operation *op,
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
@@ -155,7 +162,7 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
// gather all effects on op
- llvm::SmallVector<MemoryEffects::EffectInstance> effects =
+ SmallVector<MemoryEffects::EffectInstance> effects =
MemoryEffects::getMemoryEffectsSorted(op);
if (!effects.empty()) {
@@ -164,6 +171,8 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
bool writesConflict = hasLoopVariantInput(loopLike, op);
for (const MemoryEffects::EffectInstance &effect : effects) {
+ LDBG() << "Effect: " << effect.getEffect()->getPriority() << " with resource: " << effect.getResource()->getName() << " op " << op->getName();
+
bool conflict = false;
bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
@@ -173,8 +182,7 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
if (isa<MemoryEffects::Read>(effect.getEffect()))
writesConflict = true;
- if (isWrite)
- if (writesConflict)
+ if (isWrite && writesConflict)
conflict = true;
mergeResource(resourceConflicts, effect, conflict);
@@ -194,19 +202,25 @@ size_t mlir::moveLoopInvariantCode(
function_ref<void(Operation *, Region *)> moveOutOfRegion) {
size_t numMoved = 0;
+ // TODO: see if this can be spec'd in a meaningful way to add back later.
+ //
// check that the loop isn't dead
// auto isDead = loopLike.isZeroTrip();
// if (!isDead.has_value() || isDead.value())
// return numMoved;
- // go through loop body and map out resource usages
- // op->regions are essentially merged sequentially
- // e.g. an if's "then" and "else" regions are treated like one
- // continuous region --> need to add fork checking
+ // Go through loop body and map out resource usages.
+ // op->regions are essentially merged sequentially.
+ // E.g., an if's "then" and "else" regions are treated like one
+ // continuous region --> need to add fork checking.
//
- // loop "do" and "then" regions also merged
+ // loop "do" and "then" regions also merged.
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> resourceConflicts;
- mlir::gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
+ gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
+
+ for (auto &item : resourceConflicts) {
+ LDBG() << "Resource: " << item.second.second.getResource()->getName() << " has conflict: " << item.second.first << " with effect: " << item.second.second.getEffect()->getPriority();
+ }
auto regions = loopLike.getLoopRegions();
for (Region *region : regions) {
@@ -233,7 +247,13 @@ size_t mlir::moveLoopInvariantCode(
bool noMemoryConflicts = isMemoryEffectFree(op)
|| !mayHaveMemoryEffectConflict(op, resourceConflicts);
-
+
+ LDBG() << "isMemoryEffectFree: " << isMemoryEffectFree(op)
+ << " mayHaveMemoryEffectConflict: " << mayHaveMemoryEffectConflict(op, resourceConflicts)
+ << " noMemoryConflicts: " << noMemoryConflicts
+ << " shouldMoveOutOfRegion: " << shouldMoveOutOfRegion(op, region)
+ << " canBeHoisted: " << canBeHoisted(op, definedOutside);
+
if (!noMemoryConflicts
|| !shouldMoveOutOfRegion(op, region)
|| !canBeHoisted(op, definedOutside))
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 9616560dd5fd5..79d57de3729a3 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -2978,7 +2978,7 @@ def TestEffectsResult : TEST_Op<"test_effects_result"> {
}
//===----------------------------------------------------------------------===//
-// Test Ops with multiple effects for Loop Invariant Code Motion
+// Test Ops with multiple effects for Loop Invariant Code Motion.
//===----------------------------------------------------------------------===//
def TestResourceA : Resource<"TestResourceA">;
>From 5c0c1008b5ad7aa9a0a05ce02ab94aaebd6e8a6a Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Fri, 12 Sep 2025 21:24:31 +0000
Subject: [PATCH 15/22] removed LDBGs
---
.../Utils/LoopInvariantCodeMotionUtils.cpp | 12 ------------
1 file changed, 12 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index b977ab318534c..34e33bccf6cd1 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -171,8 +171,6 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
bool writesConflict = hasLoopVariantInput(loopLike, op);
for (const MemoryEffects::EffectInstance &effect : effects) {
- LDBG() << "Effect: " << effect.getEffect()->getPriority() << " with resource: " << effect.getResource()->getName() << " op " << op->getName();
-
bool conflict = false;
bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
@@ -218,10 +216,6 @@ size_t mlir::moveLoopInvariantCode(
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> resourceConflicts;
gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
- for (auto &item : resourceConflicts) {
- LDBG() << "Resource: " << item.second.second.getResource()->getName() << " has conflict: " << item.second.first << " with effect: " << item.second.second.getEffect()->getPriority();
- }
-
auto regions = loopLike.getLoopRegions();
for (Region *region : regions) {
LDBG() << "Original loop:\n" << *region->getParentOp();
@@ -248,12 +242,6 @@ size_t mlir::moveLoopInvariantCode(
bool noMemoryConflicts = isMemoryEffectFree(op)
|| !mayHaveMemoryEffectConflict(op, resourceConflicts);
- LDBG() << "isMemoryEffectFree: " << isMemoryEffectFree(op)
- << " mayHaveMemoryEffectConflict: " << mayHaveMemoryEffectConflict(op, resourceConflicts)
- << " noMemoryConflicts: " << noMemoryConflicts
- << " shouldMoveOutOfRegion: " << shouldMoveOutOfRegion(op, region)
- << " canBeHoisted: " << canBeHoisted(op, definedOutside);
-
if (!noMemoryConflicts
|| !shouldMoveOutOfRegion(op, region)
|| !canBeHoisted(op, definedOutside))
>From 30307b6c2b68e6861dd482f897ef26402fd2c337 Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Fri, 12 Sep 2025 22:14:03 +0000
Subject: [PATCH 16/22] fixed formatting
---
.../mlir/Interfaces/SideEffectInterfaces.h | 12 +--
.../Transforms/LoopInvariantCodeMotionUtils.h | 21 +++---
mlir/lib/Interfaces/SideEffectInterfaces.cpp | 30 ++++----
.../Utils/LoopInvariantCodeMotionUtils.cpp | 75 ++++++++++---------
4 files changed, 74 insertions(+), 64 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
index 14aa6fd8550c2..5442d652cc345 100644
--- a/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
+++ b/mlir/include/mlir/Interfaces/SideEffectInterfaces.h
@@ -348,11 +348,11 @@ struct AlwaysSpeculatableImplTrait
namespace MemoryEffects {
/// Defines the priority of the different memory effects.
-///
+///
/// Sorting/ordering memory effects of an operation is done based on
-/// their defined stage and priority, in that order. If stage values for two effect instances are
-/// equal, they are then sorted by priority. Lower priority values indicate higher
-/// precedence.
+/// their defined stage and priority, in that order. If stage values for two
+/// effect instances are equal, they are then sorted by priority. Lower priority
+/// values indicate higher precedence.
enum Priority {
DefaultPriority = 0,
AllocPriority = 1,
@@ -389,7 +389,9 @@ getMemoryEffectsSorted(Operation *op);
/// resource. An 'allocate' effect implies only allocation of the resource, and
/// not any visible mutation or dereference.
struct Allocate : public Effect::Base<Allocate> {
- Allocate() : Effect::Base<Allocate>() { this->priority = Priority::AllocPriority; }
+ Allocate() : Effect::Base<Allocate>() {
+ this->priority = Priority::AllocPriority;
+ }
};
/// The following effect indicates that the operation frees some resource that
diff --git a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
index a05dce51e4e3a..3a97e137fe933 100644
--- a/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
+++ b/mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h
@@ -25,16 +25,17 @@ class Value;
/// Gathers potential conflicts on all memory resources used within loop.
///
-/// Given a target loop and an op within it (or the loop op itself),
-/// gathers op's memory effects and flags potential resource conflicts
-/// in a map and then recurses into the op's regions to gather nested
-/// resource conflicts.
+/// Given a target loop and an op within it (or the loop op itself),
+/// gathers op's memory effects and flags potential resource conflicts
+/// in a map and then recurses into the op's regions to gather nested
+/// resource conflicts.
///
/// Typical usage:
/// \code
/// LoopLikeOpInterface myLoop = ...;
-/// DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> myConflicts;
-/// gatherResourceConflicts(myLoop, myLoop.getOperation(), resourceConflicts);
+/// DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+/// myConflicts; gatherResourceConflicts(myLoop, myLoop.getOperation(),
+/// resourceConflicts);
/// \endcode
///
/// \param loop The loop to gather resource conflicts for.
@@ -44,15 +45,17 @@ class Value;
/// Key is the resource ID that effects are applied to. Value is a pair of
/// a boolean, indicating if the resource has a conflict, and the last effect
/// that was applied to the resource (if no conflicts exist) or the effect
-/// that caused the conflict (if conflicts exist). This was done so we
+/// that caused the conflict (if conflicts exist). This was done so we
/// check the effect that causes the conflict for debugging purposes.
/// First call should use loop = someLoop and op = someLoop.getOperation()
///
/// resourceConflicts is modified by the function and will be non-empty
/// as long as there are memory effects within the loop, even if there are
/// no conflicts.
-void gatherResourceConflicts(LoopLikeOpInterface loop, Operation *op,
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts);
+void gatherResourceConflicts(
+ LoopLikeOpInterface loop, Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts);
/// Given a list of regions, perform loop-invariant code motion. An operation is
/// loop-invariant if it depends only of values defined outside of the loop.
diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index ec8618e24e308..6603f91936bc3 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -330,19 +330,19 @@ mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) {
memInterface.getEffects(effectsSorted);
- auto sortEffects =
- [](llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
- llvm::stable_sort(effects, [](const MemoryEffects::EffectInstance &a,
- const MemoryEffects::EffectInstance &b) {
- if (a.getStage() < b.getStage())
- return true;
-
- if (a.getStage() == b.getStage())
- return a.getEffect()->getPriority() < b.getEffect()->getPriority();
-
- return false; // b before a
- });
- };
+ auto sortEffects =
+ [](llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+ llvm::stable_sort(effects, [](const MemoryEffects::EffectInstance &a,
+ const MemoryEffects::EffectInstance &b) {
+ if (a.getStage() < b.getStage())
+ return true;
+
+ if (a.getStage() == b.getStage())
+ return a.getEffect()->getPriority() < b.getEffect()->getPriority();
+
+ return false; // b before a
+ });
+ };
sortEffects(effectsSorted);
return effectsSorted;
@@ -352,12 +352,12 @@ bool mlir::isMemoryEffectFree(Operation *op) {
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
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>())
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
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 34e33bccf6cd1..11edc85448f6c 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -60,25 +60,26 @@ static bool canBeHoisted(Operation *op,
op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
}
-/// Merges srcEffect's Memory Effect on its resource into the
+/// Merges srcEffect's Memory Effect on its resource into the
/// resourceConflicts map, flagging the resource if the srcEffect
/// results in a conflict.
///
/// \param resourceConflicts The map to store resources' conflicts status.
/// \param srcEffect The effect to merge into the resourceConflicts map.
-/// \param srcHasConflict Whether the srcEffect results in a conflict based
+/// \param srcHasConflict Whether the srcEffect results in a conflict based
/// on higher level analysis.
///
/// resourceConflicts is modified by the function and will be non-empty
-static void mergeResource(
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts,
- const MemoryEffects::EffectInstance &srcEffect,
- bool srcHasConflict) {
+static void
+mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts,
+ const MemoryEffects::EffectInstance &srcEffect,
+ bool srcHasConflict) {
TypeID srcResourceID = srcEffect.getResource()->getResourceID();
- bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect())
- || isa<MemoryEffects::Free>(srcEffect.getEffect());
+ bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect()) ||
+ isa<MemoryEffects::Free>(srcEffect.getEffect());
bool conflict = srcHasConflict || srcIsAllocOrFree;
@@ -86,7 +87,8 @@ static void mergeResource(
// if it doesn't already exist, create entry for resource in map
if (dstIt == resourceConflicts.end()) {
- resourceConflicts.insert(std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+ resourceConflicts.insert(
+ std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
return;
}
@@ -100,10 +102,10 @@ static void mergeResource(
bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
bool readBeforeWrite = dstRead && srcWrite;
-
+
conflict = conflict || readBeforeWrite;
- dstIt->second =std::make_pair(conflict, srcEffect);
+ dstIt->second = std::make_pair(conflict, srcEffect);
}
/// Returns true if any of op's OpOperands are defined outside of loopLike
@@ -120,14 +122,16 @@ static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
/// flagged as having a conflict within the resourceConflicts map OR
/// (b) op doesn't have a MemoryEffectOpInterface or has one but
/// without any specific effects.
-static bool mayHaveMemoryEffectConflict(Operation *op,
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
+static bool mayHaveMemoryEffectConflict(
+ Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts) {
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
-
+
// op does not implement the memory effect op interface
// shouldn't be flagged as movable to be conservative
- if (!memInterface)
+ if (!memInterface)
return true;
// gather all effects on op
@@ -135,7 +139,7 @@ static bool mayHaveMemoryEffectConflict(Operation *op,
memInterface.getEffects(effects);
// op has interface but no effects, be conservative
- if (effects.empty())
+ if (effects.empty())
return true;
// RFC moving ops with HasRecursiveMemoryEffects that have nested ops
@@ -145,10 +149,10 @@ static bool mayHaveMemoryEffectConflict(Operation *op,
for (const MemoryEffects::EffectInstance &effect : effects) {
auto resourceID = effect.getResource()->getResourceID();
- auto resConIt = resourceConflicts.find(resourceID);
+ auto resConIt = resourceConflicts.find(resourceID);
if (resConIt == resourceConflicts.end())
return true; // RFC realistically shouldn't reach here but just in case?
-
+
bool hasConflict = resConIt->second.first;
if (hasConflict)
return true;
@@ -157,13 +161,15 @@ static bool mayHaveMemoryEffectConflict(Operation *op,
return false;
}
-void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts) {
+void mlir::gatherResourceConflicts(
+ LoopLikeOpInterface loopLike, Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts) {
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
// gather all effects on op
SmallVector<MemoryEffects::EffectInstance> effects =
- MemoryEffects::getMemoryEffectsSorted(op);
+ MemoryEffects::getMemoryEffectsSorted(op);
if (!effects.empty()) {
// any variant input to the op could be the data source
@@ -173,7 +179,7 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
for (const MemoryEffects::EffectInstance &effect : effects) {
bool conflict = false;
bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
-
+
// all writes to a resource that follow a read on any other resource
// have to be considered a conflict as guaranteeing that the read
// is invariant and won't affect the write requires more robust logic
@@ -181,7 +187,7 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
writesConflict = true;
if (isWrite && writesConflict)
- conflict = true;
+ conflict = true;
mergeResource(resourceConflicts, effect, conflict);
}
@@ -189,8 +195,8 @@ void mlir::gatherResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
}
for (Region ®ion : op->getRegions())
- for (Operation &opInner : region.getOps())
- gatherResourceConflicts(loopLike, &opInner, resourceConflicts);
+ for (Operation &opInner : region.getOps())
+ gatherResourceConflicts(loopLike, &opInner, resourceConflicts);
}
size_t mlir::moveLoopInvariantCode(
@@ -213,7 +219,8 @@ size_t mlir::moveLoopInvariantCode(
// continuous region --> need to add fork checking.
//
// loop "do" and "then" regions also merged.
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> resourceConflicts;
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ resourceConflicts;
gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
auto regions = loopLike.getLoopRegions();
@@ -239,12 +246,12 @@ size_t mlir::moveLoopInvariantCode(
LDBG() << "Checking op: "
<< OpWithFlags(op, OpPrintingFlags().skipRegions());
- bool noMemoryConflicts = isMemoryEffectFree(op)
- || !mayHaveMemoryEffectConflict(op, resourceConflicts);
-
- if (!noMemoryConflicts
- || !shouldMoveOutOfRegion(op, region)
- || !canBeHoisted(op, definedOutside))
+ bool noMemoryConflicts =
+ isMemoryEffectFree(op) ||
+ !mayHaveMemoryEffectConflict(op, resourceConflicts);
+
+ if (!noMemoryConflicts || !shouldMoveOutOfRegion(op, region) ||
+ !canBeHoisted(op, definedOutside))
continue;
LDBG() << "Moving loop-invariant op: " << *op;
@@ -268,9 +275,7 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
[&](Value value, Region *) {
return loopLike.isDefinedOutsideOfLoop(value);
},
- [&](Operation *op, Region *) {
- return isSpeculatable(op);
- },
+ [&](Operation *op, Region *) { return isSpeculatable(op); },
[&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
}
>From 1dbe5db73f32e29303dd7e891dc759498b289574 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Sat, 13 Sep 2025 17:54:54 -0400
Subject: [PATCH 17/22] Apply suggestions from code review
Co-authored-by: Mehdi Amini <joker.eph at gmail.com>
---
.../Utils/LoopInvariantCodeMotionUtils.cpp | 18 ++++++------------
1 file changed, 6 insertions(+), 12 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 11edc85448f6c..bfbd0ee1c6cae 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -83,14 +83,10 @@ mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
bool conflict = srcHasConflict || srcIsAllocOrFree;
- auto dstIt = resourceConflicts.find(srcResourceID);
-
- // if it doesn't already exist, create entry for resource in map
- if (dstIt == resourceConflicts.end()) {
- resourceConflicts.insert(
+ auto [dstIt, inserted] = resourceConflicts.insert(
std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+ if (inserted)
return;
- }
// resource already in use
bool dstHasConflict = dstIt->second.first;
@@ -108,13 +104,11 @@ mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
dstIt->second = std::make_pair(conflict, srcEffect);
}
-/// Returns true if any of op's OpOperands are defined outside of loopLike
+/// Returns true if any of op's operands is defined inside the loop.
static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
- for (OpOperand &input : op->getOpOperands())
- if (!loopLike.isDefinedOutsideOfLoop(input.get()))
- return true;
-
- return false;
+ return llvm::any_of(op->getOperands(), [] (Value v) {
+ return !loopLike.isDefinedOutsideOfLoop(v);
+ });
}
/// Returns true if:
>From 6333c0059b2dfaa6c814a3e745b182d44291ce9f Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Sun, 14 Sep 2025 03:11:15 +0000
Subject: [PATCH 18/22] cleanup + LDBG
---
.../Utils/LoopInvariantCodeMotionUtils.cpp | 65 +++++++++++++------
1 file changed, 46 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index bfbd0ee1c6cae..00dad89d4f2f3 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -81,19 +81,25 @@ mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect()) ||
isa<MemoryEffects::Free>(srcEffect.getEffect());
+ LDBG() << " Merging Effect \"" << srcEffect.getEffect() << "<" << srcEffect.getResource()->getName() << ">\"";
+
bool conflict = srcHasConflict || srcIsAllocOrFree;
auto [dstIt, inserted] = resourceConflicts.insert(
std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
- if (inserted)
+ if (inserted) {
+ LDBG() << " Effect inserted to map";
return;
+ }
// resource already in use
bool dstHasConflict = dstIt->second.first;
auto dstEffect = dstIt->second.second;
- if (dstHasConflict)
+ if (dstHasConflict) {
+ LDBG() << " Resource has existing conflict from Effect \"" << dstEffect.getEffect() << "<" << dstEffect.getResource()->getName() << ">\"";
return;
+ }
bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
@@ -101,12 +107,13 @@ mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
conflict = conflict || readBeforeWrite;
+ LDBG() << " Resource " << dstEffect.getResource()->getName() << " updated to conflict status = " << conflict;
dstIt->second = std::make_pair(conflict, srcEffect);
}
/// Returns true if any of op's operands is defined inside the loop.
static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
- return llvm::any_of(op->getOperands(), [] (Value v) {
+ return llvm::any_of(op->getOperands(), [&] (Value v) {
return !loopLike.isDefinedOutsideOfLoop(v);
});
}
@@ -120,6 +127,8 @@ static bool mayHaveMemoryEffectConflict(
Operation *op,
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
&resourceConflicts) {
+ LDBG() << "Checking for memory effect conflict on op: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
@@ -144,12 +153,14 @@ static bool mayHaveMemoryEffectConflict(
auto resourceID = effect.getResource()->getResourceID();
auto resConIt = resourceConflicts.find(resourceID);
- if (resConIt == resourceConflicts.end())
- return true; // RFC realistically shouldn't reach here but just in case?
+ assert(resConIt != resourceConflicts.end());
bool hasConflict = resConIt->second.first;
- if (hasConflict)
+ if (hasConflict) {
+ LDBG() << " has conflict on resource " << effect.getResource()->getName()
+ << " from Memory Effect Mem" << effect.getEffect();
return true;
+ }
}
return false;
@@ -160,30 +171,46 @@ void mlir::gatherResourceConflicts(
DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
&resourceConflicts) {
+ LDBG() << "Gather conflicts: on loop "
+ << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions())
+ << ", visiting op: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
// gather all effects on op
SmallVector<MemoryEffects::EffectInstance> effects =
MemoryEffects::getMemoryEffectsSorted(op);
if (!effects.empty()) {
- // any variant input to the op could be the data source
- // for writes, set all writes to conflict
+ LDBG() << " # of effects = " << effects.size();
+
+ // Any input to the op could be the input data source
+ // for write effects in the same op. E.g.,
+ // scf.for ... {
+ // %0 = foo.bar(...) : ...
+ // foo.baz(%0) // foo.baz has a MemWrite<SomeResource, 0> effect
+ // }
+ // The input %0 could be the data source for the write effect in
+ // foo.baz. Since %0 is loop-variant, this may cause a conflict on
+ // SomeResource as the MemWrite contents may change between loop iterations.
+ // A more complex analysis would be needed to determine
+ // if this is a true conflict or not.
bool writesConflict = hasLoopVariantInput(loopLike, op);
+ LDBG() << " has loop-variant input? "
+ << (writesConflict ? "true" : "false");
for (const MemoryEffects::EffectInstance &effect : effects) {
- bool conflict = false;
- bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
-
- // all writes to a resource that follow a read on any other resource
- // have to be considered a conflict as guaranteeing that the read
- // is invariant and won't affect the write requires more robust logic
- if (isa<MemoryEffects::Read>(effect.getEffect()))
- writesConflict = true;
-
- if (isWrite && writesConflict)
- conflict = true;
+ bool conflict = writesConflict && isa<MemoryEffects::Write>(effect.getEffect());
mergeResource(resourceConflicts, effect, conflict);
+
+ // All writes to a resource that follow a read on any other resource
+ // have to be considered as causing a conflict, similar
+ // to how we handle loop-variant inputs above; guaranteeing that the read
+ // is invariant and won't affect the write requires more robust logic.
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ LDBG() << " has MemRead; subsequent Writes will be treated as conflicting";
+ writesConflict = true;
+ }
}
}
}
>From aaa7cc5a39c14486192628b6b2bf247be1158e4b Mon Sep 17 00:00:00 2001
From: mbagherbeikTT <mbagherbeik at tenstorrent.com>
Date: Sun, 14 Sep 2025 03:11:30 +0000
Subject: [PATCH 19/22] clarifying test cases
---
.../loop-invariant-code-motion.mlir | 63 +++++++++++++++++--
1 file changed, 59 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index e97351a4f57f2..9b3f84f6afe20 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1447,7 +1447,12 @@ func.func @move_single_resource_basic() attributes {} {
%c2_i32 = arith.constant 1 : i32
%c0_i32_0 = arith.constant 0 : i32
+ // Single write effect on one resource in a triple-nested loop
+ // No loop-variant inputs to op and no read effects -> movable
// CHECK: "test.test_effects_write_A"() : () -> ()
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
@@ -1468,14 +1473,18 @@ func.func @move_single_resource_write_dominant() attributes {} {
%c2_i32 = arith.constant 1 : i32
%c0_i32_0 = arith.constant 0 : i32
+ // Write effect on one resource followed by a Read.
+ // No loop-variant inputs to Write op, no conflict on
+ // "A" --> both ops movable
// CHECK: "test.test_effects_write_A"() : () -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
-
- // CHECK: "test.test_effects_read_A"() : () -> ()
-
"test.test_effects_write_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
@@ -1493,10 +1502,16 @@ func.func @move_single_resource_read_dominant() attributes {} {
%c2_i32 = arith.constant 1 : i32
%c0_i32_0 = arith.constant 0 : i32
+ // CHECK: scf.for %arg0
+ // CHECK: scf.for %arg1
+ // CHECK: scf.for %arg2
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
+ // Read effect on "A" dominates write.
+ // Causes conflict "A" --> not movable.
// CHECK: "test.test_effects_read_A"() : () -> ()
// CHECK: "test.test_effects_write_A"() : () -> ()
@@ -1519,12 +1534,19 @@ func.func @move_single_resource_basic_conflict() attributes {} {
%c0_i32_1 = arith.constant 0 : i32
%c0_i32_2 = arith.constant 0 : i32
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
+ // Read of "A" dominates Write "A" --> "A" is in conflict.
+ // "C" is not in conflict but, since all resources used
+ // by op aren't conflict free, they're not movable.
// CHECK: "test.test_effects_write_AC"() : () -> ()
// CHECK: "test.test_effects_read_AC"() : () -> ()
@@ -1547,13 +1569,18 @@ func.func @move_single_resource_if_region() attributes {} {
%c2_i32 = arith.constant 1 : i32
%c3_i32 = arith.constant 5 : i32
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
%1 = arith.cmpi slt, %arg0, %c3_i32 : i32
scf.if %1 {
-
+ // Checking that we're not moving ops out of
+ // non-loop regions.
// CHECK: "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
@@ -1575,13 +1602,21 @@ func.func @move_single_resource_for_inside_if_region() attributes {} {
%c2_i32 = arith.constant 1 : i32
%c3_i32 = arith.constant 5 : i32
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
%1 = arith.cmpi slt, %arg0, %c3_i32 : i32
+ // CHECK: scf.if
scf.if %1 {
%c0_i32_0 = arith.constant 0 : i32
+ // Checking that we can move ops out of loops nested
+ // within other regions, without moving ops out of
+ // the parent, non-loop region.
// CHECK: "test.test_effects_write_A"() : () -> ()
+ // CHECK: scf.for
+ // CHECK: scf.for
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
@@ -1603,6 +1638,10 @@ func.func @move_multi_resource_comprehensive() attributes {} {
%c0_i32_2 = arith.constant 0 : i32
%c3_i32 = arith.constant 5 : i32
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_write_CD"() : () -> ()
@@ -1610,6 +1649,7 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// CHECK: "test.test_effects_write_EF"() : () -> ()
// CHECK: "test.test_effects_read_EF"() : () -> ()
+ // Both of these should be moved out of their parent
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_CD"() : () -> ()
@@ -1624,19 +1664,25 @@ func.func @move_multi_resource_comprehensive() attributes {} {
%c0_i32_0 = arith.constant 0 : i32
%c0_i32_1 = arith.constant 0 : i32
+ // CHECK: scf.for
// CHECK: "test.test_effects_write_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
+ // CHECK: scf.for
+ // CHECK: scf.for
+ // CHECK: scf.for
scf.for %arg3 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
+ // Loop should be moved out of parent
scf.for %arg4 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
+ // Loop should be moved out of parent
scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_B"() : () -> ()
"test.test_effects_read_B"() : () -> ()
@@ -1645,6 +1691,7 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// CHECK: "test.test_effects_write_AC"() : () -> ()
// CHECK: "test.test_effects_read_AC"() : () -> ()
+ // Loop should be moved out of parent
scf.for %arg6 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_AC"() : () -> ()
"test.test_effects_read_AC"() : () -> ()
@@ -1672,6 +1719,8 @@ func.func @move_multi_resource_write_conflicts() attributes {} {
// CHECK: "test.test_effects_write_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_write_A_with_input"(%c7) : (index) -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
@@ -1685,6 +1734,8 @@ func.func @move_multi_resource_write_conflicts() attributes {} {
}
// CHECK: "test.test_effects_read_A"() : () -> ()
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_read_A_write_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
@@ -1696,6 +1747,8 @@ func.func @move_multi_resource_write_conflicts() attributes {} {
}
// CHECK: "test.test_effects_read_A"() : () -> ()
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
// CHECK: "test.test_effects_read_A_then_write_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
@@ -1709,6 +1762,8 @@ func.func @move_multi_resource_write_conflicts() attributes {} {
// CHECK: "test.test_effects_write_A_then_read_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
+ // CHECK: scf.for
+
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_A_then_read_B"() : () -> ()
"test.test_effects_read_B"() : () -> ()
>From d56ccc3231ca6c9bb361a0c50c5d7af3c0a403ad Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Mon, 15 Sep 2025 22:14:33 +0000
Subject: [PATCH 20/22] addin isZeroTrip back
---
.../mlir/Interfaces/LoopLikeInterface.h | 1 +
.../mlir/Interfaces/LoopLikeInterface.td | 29 +++++++++++++++++++
.../Utils/LoopInvariantCodeMotionUtils.cpp | 10 +++++--
3 files changed, 37 insertions(+), 3 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.h b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
index 9925fc6ce6ca9..d6c23131cf242 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.h
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.h
@@ -14,6 +14,7 @@
#define MLIR_INTERFACES_LOOPLIKEINTERFACE_H_
#include "mlir/IR/OpDefinition.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
namespace mlir {
class RewriterBase;
diff --git a/mlir/include/mlir/Interfaces/LoopLikeInterface.td b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
index 6c95b4802837b..72ba34ba77da9 100644
--- a/mlir/include/mlir/Interfaces/LoopLikeInterface.td
+++ b/mlir/include/mlir/Interfaces/LoopLikeInterface.td
@@ -318,6 +318,35 @@ def LoopLikeOpInterface : OpInterface<"LoopLikeOpInterface"> {
$_op->operand_begin() + firstOperandIndex + initsMutable.size());
}
+ /// Return whether the loop is known to have zero iterations.
+ /// Returns std::nullopt if not enough constant information is available.
+ ::std::optional<bool> isZeroTrip() {
+ auto lbs = $_op.getLoopLowerBounds();
+ auto ubs = $_op.getLoopUpperBounds();
+ auto steps = $_op.getLoopSteps();
+
+ if (!lbs || !ubs || !steps)
+ return ::std::nullopt;
+
+ if (lbs->size() != ubs->size() || ubs->size() != steps->size())
+ return ::std::nullopt;
+
+ for (size_t i = 0; i < steps->size(); ++i) {
+ auto lb = ::mlir::getConstantIntValue((*lbs)[i]);
+ auto ub = ::mlir::getConstantIntValue((*ubs)[i]);
+ auto st = ::mlir::getConstantIntValue((*steps)[i]);
+
+ if (!lb || !ub || !st)
+ return ::std::nullopt; // non-constant -> unknown
+
+ if (*st >= 0 && *lb >= *ub)
+ return true;
+ if (*st < 0 && *lb <= *ub)
+ return true;
+ }
+ return false;
+ }
+
/// Return the region iter_arg that corresponds to the given init operand.
/// Return an "empty" block argument if the given operand is not an init
/// operand of this loop op.
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 00dad89d4f2f3..064df1a9f200f 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -230,9 +230,13 @@ size_t mlir::moveLoopInvariantCode(
// TODO: see if this can be spec'd in a meaningful way to add back later.
//
// check that the loop isn't dead
- // auto isDead = loopLike.isZeroTrip();
- // if (!isDead.has_value() || isDead.value())
- // return numMoved;
+ auto isDead = loopLike.isZeroTrip();
+
+ LDBG() << "Loop " << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions())
+ << " has constant bounds and steps? isZeroTrip()? " << (isDead.has_value() ? (isDead.value() ? "YES, YES" : "YES, NO") : "NO, NULL");
+
+ if (!isDead.has_value() || isDead.value())
+ return numMoved;
// Go through loop body and map out resource usages.
// op->regions are essentially merged sequentially.
>From a396617902f0c1fe6dc23558a7a6b964ae5c3766 Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 16 Sep 2025 15:37:38 +0000
Subject: [PATCH 21/22] isZeroTrip() check will only skip loops that are
verifiably dead
---
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 064df1a9f200f..903804e5f6a5d 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -235,7 +235,7 @@ size_t mlir::moveLoopInvariantCode(
LDBG() << "Loop " << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions())
<< " has constant bounds and steps? isZeroTrip()? " << (isDead.has_value() ? (isDead.value() ? "YES, YES" : "YES, NO") : "NO, NULL");
- if (!isDead.has_value() || isDead.value())
+ if (isDead.has_value() && isDead.value())
return numMoved;
// Go through loop body and map out resource usages.
>From 4618db97d019247e6bf520a9d7869b6e0f1e46bf Mon Sep 17 00:00:00 2001
From: Mo Bagherbeik <mbagherbeik at tenstorrent.com>
Date: Tue, 16 Sep 2025 17:49:36 +0000
Subject: [PATCH 22/22] added outer loop to LICM pass
---
.../Utils/LoopInvariantCodeMotionUtils.cpp | 107 ++++++++++--------
.../loop-invariant-code-motion.mlir | 17 +--
2 files changed, 68 insertions(+), 56 deletions(-)
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 903804e5f6a5d..a946fc3b88e48 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -225,7 +225,7 @@ size_t mlir::moveLoopInvariantCode(
function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
function_ref<void(Operation *, Region *)> moveOutOfRegion) {
- size_t numMoved = 0;
+ size_t numMovedTotal = 0;
// TODO: see if this can be spec'd in a meaningful way to add back later.
//
@@ -236,62 +236,71 @@ size_t mlir::moveLoopInvariantCode(
<< " has constant bounds and steps? isZeroTrip()? " << (isDead.has_value() ? (isDead.value() ? "YES, YES" : "YES, NO") : "NO, NULL");
if (isDead.has_value() && isDead.value())
- return numMoved;
+ return numMovedTotal;
- // Go through loop body and map out resource usages.
- // op->regions are essentially merged sequentially.
- // E.g., an if's "then" and "else" regions are treated like one
- // continuous region --> need to add fork checking.
- //
- // loop "do" and "then" regions also merged.
- DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
- resourceConflicts;
- gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
-
- auto regions = loopLike.getLoopRegions();
- for (Region *region : regions) {
- LDBG() << "Original loop:\n" << *region->getParentOp();
-
- std::queue<Operation *> worklist;
- // Add top-level operations in the loop body to the worklist.
- for (Operation &op : region->getOps())
- worklist.push(&op);
-
- auto definedOutside = [&](Value value) {
- return isDefinedOutsideRegion(value, region);
- };
-
- while (!worklist.empty()) {
- Operation *op = worklist.front();
- worklist.pop();
- // Skip ops that have already been moved. Check if the op can be hoisted.
- if (op->getParentRegion() != region)
- continue;
+ int numMoved = 0;
- LDBG() << "Checking op: "
- << OpWithFlags(op, OpPrintingFlags().skipRegions());
+ do {
+ // reset value for iteration
+ numMoved = 0;
- bool noMemoryConflicts =
- isMemoryEffectFree(op) ||
- !mayHaveMemoryEffectConflict(op, resourceConflicts);
+ // Go through loop body and map out resource usages.
+ // op->regions are essentially merged sequentially.
+ // E.g., an if's "then" and "else" regions are treated like one
+ // continuous region --> need to add fork checking.
+ //
+ // loop "do" and "then" regions also merged.
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ resourceConflicts;
+ gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
- if (!noMemoryConflicts || !shouldMoveOutOfRegion(op, region) ||
- !canBeHoisted(op, definedOutside))
- continue;
+ auto regions = loopLike.getLoopRegions();
+ for (Region *region : regions) {
+ LDBG() << "Original loop:\n" << *region->getParentOp();
- LDBG() << "Moving loop-invariant op: " << *op;
- moveOutOfRegion(op, region);
- ++numMoved;
+ std::queue<Operation *> worklist;
+ // Add top-level operations in the loop body to the worklist.
+ for (Operation &op : region->getOps())
+ worklist.push(&op);
- // Since the op has been moved, we need to check its users within the
- // top-level of the loop body.
- for (Operation *user : op->getUsers())
- if (user->getParentRegion() == region)
- worklist.push(user);
+ auto definedOutside = [&](Value value) {
+ return isDefinedOutsideRegion(value, region);
+ };
+
+ while (!worklist.empty()) {
+ Operation *op = worklist.front();
+ worklist.pop();
+ // Skip ops that have already been moved. Check if the op can be hoisted.
+ if (op->getParentRegion() != region)
+ continue;
+
+ LDBG() << "Checking op: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ bool noMemoryConflicts =
+ isMemoryEffectFree(op) ||
+ !mayHaveMemoryEffectConflict(op, resourceConflicts);
+
+ if (!noMemoryConflicts || !shouldMoveOutOfRegion(op, region) ||
+ !canBeHoisted(op, definedOutside))
+ continue;
+
+ LDBG() << "Moving loop-invariant op: " << *op;
+ moveOutOfRegion(op, region);
+ ++numMoved;
+
+ // Since the op has been moved, we need to check its users within the
+ // top-level of the loop body.
+ for (Operation *user : op->getUsers())
+ if (user->getParentRegion() == region)
+ worklist.push(user);
+ }
}
- }
- return numMoved;
+ numMovedTotal += numMoved;
+ } while (numMoved > 0);
+
+ return numMovedTotal;
}
size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 9b3f84f6afe20..02bd3faca3f58 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1649,7 +1649,7 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// CHECK: "test.test_effects_write_EF"() : () -> ()
// CHECK: "test.test_effects_read_EF"() : () -> ()
- // Both of these should be moved out of their parent
+ // Both of these should be empty and moved out of their parent
scf.for %arg1 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
scf.for %arg2 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_CD"() : () -> ()
@@ -1676,13 +1676,13 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// CHECK: "test.test_effects_write_A"() : () -> ()
// CHECK: "test.test_effects_read_A"() : () -> ()
- // Loop should be moved out of parent
+ // emptyoop should be empty and moved out of parent
scf.for %arg4 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_A"() : () -> ()
"test.test_effects_read_A"() : () -> ()
}
- // Loop should be moved out of parent
+ // Loop should be empty and moved out of parent
scf.for %arg5 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_B"() : () -> ()
"test.test_effects_read_B"() : () -> ()
@@ -1691,7 +1691,7 @@ func.func @move_multi_resource_comprehensive() attributes {} {
// CHECK: "test.test_effects_write_AC"() : () -> ()
// CHECK: "test.test_effects_read_AC"() : () -> ()
- // Loop should be moved out of parent
+ // Loop should be empty and moved out of parent
scf.for %arg6 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
"test.test_effects_write_AC"() : () -> ()
"test.test_effects_read_AC"() : () -> ()
@@ -1717,14 +1717,17 @@ func.func @move_multi_resource_write_conflicts() attributes {} {
%c1_i32 = arith.constant 10 : i32
%c2_i32 = arith.constant 1 : i32
+ // CHECK: %{{.*}} = arith.constant 7 : index
+
// CHECK: "test.test_effects_write_B"() : () -> ()
// CHECK: "test.test_effects_read_B"() : () -> ()
+
+ // CHECK: "test.test_effects_write_A_with_input"(%c7) : (index) -> ()
+ // CHECK: "test.test_effects_read_A"() : () -> ()
+
// CHECK: scf.for
scf.for %arg0 = %c0_i32 to %c1_i32 step %c2_i32 : i32 {
- // CHECK: "test.test_effects_write_A_with_input"(%c7) : (index) -> ()
- // CHECK: "test.test_effects_read_A"() : () -> ()
-
%input = arith.constant 7 : index
"test.test_effects_write_A_with_input"(%input) : (index) -> ()
"test.test_effects_read_A"() : () -> ()
More information about the Mlir-commits
mailing list