[Mlir-commits] [mlir] [mlir][affine] Use alias analysis to redetermine intervening memory effects in affine-scalrep (PR #90859)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu May 2 07:48:06 PDT 2024
https://github.com/asraa created https://github.com/llvm/llvm-project/pull/90859
This fixes a TODO to use alias analysis to determine whether a read op intervenes between two write operations to the same memref.
>From 100a4cb9c8a76ac3cb43833d2a230faecba7d3a6 Mon Sep 17 00:00:00 2001
From: Asra <asraa at google.com>
Date: Thu, 2 May 2024 14:45:17 +0000
Subject: [PATCH] [mlir][affine] Use alias analysis to redetermine intervening
memory effects in scalrep
Signed-off-by: Asra <asraa at google.com>
---
mlir/include/mlir/Dialect/Affine/Utils.h | 7 ++-
.../Transforms/AffineScalarReplacement.cpp | 4 +-
mlir/lib/Dialect/Affine/Utils/Utils.cpp | 52 ++++++++++---------
mlir/test/Dialect/Affine/scalrep.mlir | 13 +++++
4 files changed, 49 insertions(+), 27 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h
index 67c7a964feefd7..7f25db029781c8 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 ed94fb690af2cd..707bba2f1e6f9c 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 8b8ed2578ca5cc..f46381403bc522 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 22d394bfcf0979..baafedbc307d3e 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -682,6 +682,19 @@ 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-LABEL: func @redundant_store_elim_fail
func.func @redundant_store_elim_fail(%out : memref<512xf32>) {
More information about the Mlir-commits
mailing list