[Mlir-commits] [mlir] 5d51db7 - [mlir][affine] Use alias analysis to redetermine intervening memory effects in affine-scalrep (#90859)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 9 14:58:02 PDT 2024
Author: asraa
Date: 2024-05-09T14:57:59-07:00
New Revision: 5d51db75e46155d0f3d70cf3253d3b075cfcf93a
URL: https://github.com/llvm/llvm-project/commit/5d51db75e46155d0f3d70cf3253d3b075cfcf93a
DIFF: https://github.com/llvm/llvm-project/commit/5d51db75e46155d0f3d70cf3253d3b075cfcf93a.diff
LOG: [mlir][affine] Use alias analysis to redetermine intervening memory effects in affine-scalrep (#90859)
This fixes a TODO to use alias analysis to determine whether a read op
intervenes between two write operations to the same memref.
Signed-off-by: Asra <asraa at google.com>
Added:
Modified:
mlir/include/mlir/Dialect/Affine/Utils.h
mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
mlir/lib/Dialect/Affine/Utils/Utils.cpp
mlir/test/Dialect/Affine/scalrep.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index 67c7a964feefd..7f25db029781c 100644
--- a/mlir/include/mlir/Dialect/Affine/Utils.h
+++ b/mlir/include/mlir/Dialect/Affine/Utils.h
@@ -13,6 +13,7 @@
#ifndef MLIR_DIALECT_AFFINE_UTILS_H
#define MLIR_DIALECT_AFFINE_UTILS_H
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/IR/OpDefinition.h"
@@ -106,7 +107,8 @@ struct VectorizationStrategy {
/// loads and eliminate invariant affine loads; consequently, eliminate dead
/// allocs.
void affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
- PostDominanceInfo &postDomInfo);
+ PostDominanceInfo &postDomInfo,
+ AliasAnalysis &analysis);
/// Vectorizes affine loops in 'loops' using the n-D vectorization factors in
/// 'vectorSizes'. By default, each vectorization factor is applied
@@ -325,7 +327,8 @@ OpFoldResult linearizeIndex(ArrayRef<OpFoldResult> multiIndex,
/// will check if there is no write to the memory between `start` and `memOp`
/// that would change the read within `memOp`.
template <typename EffectType, typename T>
-bool hasNoInterveningEffect(Operation *start, T memOp);
+bool hasNoInterveningEffect(Operation *start, T memOp,
+ llvm::function_ref<bool(Value, Value)> mayAlias);
struct AffineValueExpr {
explicit AffineValueExpr(AffineExpr e) : e(e) {}
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
index ed94fb690af2c..707bba2f1e6f9 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
@@ -13,6 +13,7 @@
#include "mlir/Dialect/Affine/Passes.h"
+#include "mlir/Analysis/AliasAnalysis.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Dominance.h"
@@ -47,5 +48,6 @@ mlir::affine::createAffineScalarReplacementPass() {
void AffineScalarReplacement::runOnOperation() {
affineScalarReplace(getOperation(), getAnalysis<DominanceInfo>(),
- getAnalysis<PostDominanceInfo>());
+ getAnalysis<PostDominanceInfo>(),
+ getAnalysis<AliasAnalysis>());
}
diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
index 8b8ed2578ca5c..f46381403bc52 100644
--- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp
@@ -678,12 +678,9 @@ static bool mayHaveEffect(Operation *srcMemOp, Operation *destMemOp,
}
template <typename EffectType, typename T>
-bool mlir::affine::hasNoInterveningEffect(Operation *start, T memOp) {
- auto isLocallyAllocated = [](Value memref) {
- auto *defOp = memref.getDefiningOp();
- return defOp && hasSingleEffect<MemoryEffects::Allocate>(defOp, memref);
- };
-
+bool mlir::affine::hasNoInterveningEffect(
+ Operation *start, T memOp,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
// A boolean representing whether an intervening operation could have impacted
// memOp.
bool hasSideEffect = false;
@@ -704,11 +701,8 @@ bool mlir::affine::hasNoInterveningEffect(Operation *start, T memOp) {
// If op causes EffectType on a potentially aliasing location for
// memOp, mark as having the effect.
if (isa<EffectType>(effect.getEffect())) {
- // TODO: This should be replaced with a check for no aliasing.
- // Aliasing information should be passed to this method.
if (effect.getValue() && effect.getValue() != memref &&
- isLocallyAllocated(memref) &&
- isLocallyAllocated(effect.getValue()))
+ !mayAlias(effect.getValue(), memref))
continue;
opMayHaveEffect = true;
break;
@@ -832,10 +826,10 @@ bool mlir::affine::hasNoInterveningEffect(Operation *start, T memOp) {
/// other operations will overwrite the memory loaded between the given load
/// and store. If such a value exists, the replaced `loadOp` will be added to
/// `loadOpsToErase` and its memref will be added to `memrefsToErase`.
-static void forwardStoreToLoad(AffineReadOpInterface loadOp,
- SmallVectorImpl<Operation *> &loadOpsToErase,
- SmallPtrSetImpl<Value> &memrefsToErase,
- DominanceInfo &domInfo) {
+static void forwardStoreToLoad(
+ AffineReadOpInterface loadOp, SmallVectorImpl<Operation *> &loadOpsToErase,
+ SmallPtrSetImpl<Value> &memrefsToErase, DominanceInfo &domInfo,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
// The store op candidate for forwarding that satisfies all conditions
// to replace the load, if any.
@@ -872,7 +866,8 @@ static void forwardStoreToLoad(AffineReadOpInterface loadOp,
// 4. Ensure there is no intermediate operation which could replace the
// value in memory.
- if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp))
+ if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(storeOp, loadOp,
+ mayAlias))
continue;
// We now have a candidate for forwarding.
@@ -901,7 +896,8 @@ static void forwardStoreToLoad(AffineReadOpInterface loadOp,
template bool
mlir::affine::hasNoInterveningEffect<mlir::MemoryEffects::Read,
affine::AffineReadOpInterface>(
- mlir::Operation *, affine::AffineReadOpInterface);
+ mlir::Operation *, affine::AffineReadOpInterface,
+ llvm::function_ref<bool(Value, Value)>);
// This attempts to find stores which have no impact on the final result.
// A writing op writeA will be eliminated if there exists an op writeB if
@@ -910,7 +906,8 @@ mlir::affine::hasNoInterveningEffect<mlir::MemoryEffects::Read,
// 3) There is no potential read between writeA and writeB.
static void findUnusedStore(AffineWriteOpInterface writeA,
SmallVectorImpl<Operation *> &opsToErase,
- PostDominanceInfo &postDominanceInfo) {
+ PostDominanceInfo &postDominanceInfo,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
for (Operation *user : writeA.getMemRef().getUsers()) {
// Only consider writing operations.
@@ -939,7 +936,8 @@ static void findUnusedStore(AffineWriteOpInterface writeA,
// There cannot be an operation which reads from memory between
// the two writes.
- if (!affine::hasNoInterveningEffect<MemoryEffects::Read>(writeA, writeB))
+ if (!affine::hasNoInterveningEffect<MemoryEffects::Read>(writeA, writeB,
+ mayAlias))
continue;
opsToErase.push_back(writeA);
@@ -955,7 +953,8 @@ static void findUnusedStore(AffineWriteOpInterface writeA,
// 3) There is no write between loadA and loadB.
static void loadCSE(AffineReadOpInterface loadA,
SmallVectorImpl<Operation *> &loadOpsToErase,
- DominanceInfo &domInfo) {
+ DominanceInfo &domInfo,
+ llvm::function_ref<bool(Value, Value)> mayAlias) {
SmallVector<AffineReadOpInterface, 4> loadCandidates;
for (auto *user : loadA.getMemRef().getUsers()) {
auto loadB = dyn_cast<AffineReadOpInterface>(user);
@@ -976,7 +975,7 @@ static void loadCSE(AffineReadOpInterface loadA,
// 3. There should not be a write between loadA and loadB.
if (!affine::hasNoInterveningEffect<MemoryEffects::Write>(
- loadB.getOperation(), loadA))
+ loadB.getOperation(), loadA, mayAlias))
continue;
// Check if two values have the same shape. This is needed for affine vector
@@ -1034,16 +1033,21 @@ static void loadCSE(AffineReadOpInterface loadA,
// than dealloc) remain.
//
void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
- PostDominanceInfo &postDomInfo) {
+ PostDominanceInfo &postDomInfo,
+ AliasAnalysis &aliasAnalysis) {
// Load op's whose results were replaced by those forwarded from stores.
SmallVector<Operation *, 8> opsToErase;
// A list of memref's that are potentially dead / could be eliminated.
SmallPtrSet<Value, 4> memrefsToErase;
+ auto mayAlias = [&](Value val1, Value val2) -> bool {
+ return !aliasAnalysis.alias(val1, val2).isNo();
+ };
+
// Walk all load's and perform store to load forwarding.
f.walk([&](AffineReadOpInterface loadOp) {
- forwardStoreToLoad(loadOp, opsToErase, memrefsToErase, domInfo);
+ forwardStoreToLoad(loadOp, opsToErase, memrefsToErase, domInfo, mayAlias);
});
for (auto *op : opsToErase)
op->erase();
@@ -1051,7 +1055,7 @@ void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
// Walk all store's and perform unused store elimination
f.walk([&](AffineWriteOpInterface storeOp) {
- findUnusedStore(storeOp, opsToErase, postDomInfo);
+ findUnusedStore(storeOp, opsToErase, postDomInfo, mayAlias);
});
for (auto *op : opsToErase)
op->erase();
@@ -1084,7 +1088,7 @@ void mlir::affine::affineScalarReplace(func::FuncOp f, DominanceInfo &domInfo,
// stores. Otherwise, some stores are wrongly seen as having an intervening
// effect.
f.walk([&](AffineReadOpInterface loadOp) {
- loadCSE(loadOp, opsToErase, domInfo);
+ loadCSE(loadOp, opsToErase, domInfo, mayAlias);
});
for (auto *op : opsToErase)
op->erase();
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 22d394bfcf097..4a99dee50a280 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -682,6 +682,24 @@ func.func @redundant_store_elim(%out : memref<512xf32>) {
// CHECK-NEXT: affine.store
// CHECK-NEXT: }
+// CHECK-LABEL: func @redundant_store_elim_nonintervening
+
+func.func @redundant_store_elim_nonintervening(%in : memref<512xf32>) {
+ %cf1 = arith.constant 1.0 : f32
+ %out = memref.alloc() : memref<512xf32>
+ affine.for %i = 0 to 16 {
+ affine.store %cf1, %out[32*%i] : memref<512xf32>
+ %0 = affine.load %in[32*%i] : memref<512xf32>
+ affine.store %0, %out[32*%i] : memref<512xf32>
+ }
+ return
+}
+
+// CHECK: affine.for
+// CHECK-NEXT: affine.load
+// CHECK-NEXT: affine.store
+// CHECK-NEXT: }
+
// CHECK-LABEL: func @redundant_store_elim_fail
func.func @redundant_store_elim_fail(%out : memref<512xf32>) {
More information about the Mlir-commits
mailing list