[flang-commits] [flang] [mlir] [mlir][AliasAnalysis] Add precise listener-driven cache invalidation (PR #207032)
Slava Zakharin via flang-commits
flang-commits at lists.llvm.org
Wed Jul 1 09:47:25 PDT 2026
https://github.com/vzakhari created https://github.com/llvm/llvm-project/pull/207032
Extend the opt-in query cache so it can be used safely by passes that mutate
the IR through a rewriter. The query-caching hooks gain a RewriterBase*
parameter (defaulted, so QueryCacheScope callers are unaffected): in frozen
mode (null rewriter) the cache is a snapshot as before; with a rewriter the
FIR AliasAnalysis installs a ForwardingListener on it and evicts cached
getSource() results precisely as operations are erased, modified, or replaced.
Each cached entry now records the operations it depends on (the def-use chain
and the dominating dummy_scope op visited while computing it), maintained via a
reverse-dependency index. Nested queries thread a dependency sink and union a
cache hit's dependencies into the enclosing entry, so transitive dependencies
are tracked. Erasure notifications also walk nested operations and fire before
the storage is freed, closing the window where a freed-and-reused pointer could
yield a stale hit.
Adds AliasAnalysisCacheTest covering precise eviction on modify, eviction on
erase, and frozen-mode memoization.
Assisted by Cursor
>From 585f5854516df94c0ee343486b184e7868df4d8f Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Tue, 30 Jun 2026 12:04:20 -0700
Subject: [PATCH 1/2] [mlir][AliasAnalysis] Add opt-in query caching; use it in
flang LICM
Add a generic, opt-in query-caching capability to the mlir::AliasAnalysis
aggregator. The aggregator gains a public RAII QueryCacheScope and private
enable/disableQueryCaching() fan-out; the Concept gains enableQueryCaching()/
disableQueryCaching() and the Model forwards them to implementations that
provide enableSourceCache()/disableSourceCache() (detected via
llvm::is_detected) and no-ops otherwise, so existing implementations need no
changes and transformation passes never name a concrete implementation.
The FIR AliasAnalysis implements the hook by memoizing getSource(): getSource()
becomes a thin wrapper over getSourceImpl() that consults a frozen-snapshot
cache keyed on (value, flags) when caching is enabled.
Flang's LICM enables caching for the duration of the pass via QueryCacheScope.
This is sound because LICM only moves operations and getSource()'s
position-dependent inputs are already frozen for the pass run by the
AliasAnalysis's own dominance/scope caches, making caching observationally
equivalent to no caching across the hoists.
Assisted by Cursor
---
.../flang/Optimizer/Analysis/AliasAnalysis.h | 38 ++++++++++++
.../lib/Optimizer/Analysis/AliasAnalysis.cpp | 29 +++++++++
.../Transforms/LoopInvariantCodeMotion.cpp | 10 ++++
mlir/include/mlir/Analysis/AliasAnalysis.h | 60 +++++++++++++++++++
mlir/lib/Analysis/AliasAnalysis.cpp | 10 ++++
5 files changed, 147 insertions(+)
diff --git a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
index 832634a708dba..a40632c23efce 100644
--- a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
+++ b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
@@ -20,6 +20,7 @@
#include "llvm/ADT/PointerUnion.h"
#include "llvm/ADT/SmallVector.h"
#include <memory>
+#include <utility>
namespace fir {
@@ -351,10 +352,25 @@ struct AliasAnalysis {
/// snapshots are not collected, and getSource performs only the
/// SourceKind/origin classification without that bookkeeping side
/// effect.
+ ///
+ /// When source caching is enabled (see enableSourceCache()), the result is
+ /// memoized keyed on (value, flags), and recursive sub-queries share the
+ /// same cache. The cache is a frozen snapshot: it is only valid while the IR
+ /// reachable from the queried values is not mutated in a way that would
+ /// change the source. Callers must scope caching no wider than such a region
+ /// (see mlir::AliasAnalysis::QueryCacheScope).
fir::AliasAnalysis::Source getSource(mlir::Value,
bool getLastInstantiationPoint = false,
bool collectScopedOrigins = true);
+ /// Enable/disable memoization of getSource() results. These are invoked by
+ /// the mlir::AliasAnalysis aggregator (via mlir::AliasAnalysis::
+ /// QueryCacheScope) and are not meant to be called directly by passes.
+ /// enableSourceCache() turns on the frozen snapshot cache; disableSourceCache
+ /// clears it and turns it off.
+ void enableSourceCache();
+ void disableSourceCache();
+
/// Return true, if `ty` is a reference type to a boxed
/// POINTER object or a raw fir::PointerType.
static bool isPointerReference(mlir::Type ty);
@@ -370,6 +386,16 @@ struct AliasAnalysis {
bool functionHasMultipleScopes(mlir::Value v);
private:
+ /// Compute the memory source of a value. This is the uncached
+ /// implementation of getSource(); getSource() is a thin wrapper that
+ /// memoizes the result when source caching is enabled.
+ fir::AliasAnalysis::Source getSourceImpl(mlir::Value v,
+ bool getLastInstantiationPoint,
+ bool collectScopedOrigins);
+
+ /// Clear the getSource() memoization cache.
+ void clearSourceCache() { getSourceCache.clear(); }
+
/// Build an intermediate Source rooted at the declare captured by the
/// snapshot. Reuses getSource(declValue) for the SourceKind / origin
/// classification (with collectScopedOrigins=false), then overrides
@@ -446,6 +472,18 @@ struct AliasAnalysis {
/// functionHasMultipleScopes(); both true and false are cached so that
/// repeated queries are O(1) without re-walking the function body.
llvm::DenseMap<mlir::Operation *, bool> multiScopeCache;
+
+ /// Opt-in memoization of getSource() results, keyed on the queried value
+ /// and the two boolean flags (getLastInstantiationPoint,
+ /// collectScopedOrigins) packed into the unsigned. Only consulted/populated
+ /// while \c sourceCacheEnabled is set. This is a frozen snapshot cache: it
+ /// has no automatic invalidation, so it is only enabled (via
+ /// enableSourceCache(), driven by mlir::AliasAnalysis::QueryCacheScope) for
+ /// the duration of a region in which the relevant IR is not mutated.
+ llvm::DenseMap<std::pair<mlir::Value, unsigned>, Source> getSourceCache;
+
+ /// Whether getSource() should consult/populate getSourceCache.
+ bool sourceCacheEnabled = false;
};
inline bool operator==(const AliasAnalysis::Source::SourceOrigin &lhs,
diff --git a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
index 02a628968ddc4..f3e77a4fc0826 100644
--- a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
+++ b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -30,6 +30,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include <optional>
+#include <utility>
using namespace mlir;
@@ -1108,9 +1109,37 @@ static mlir::Value walkBlockArgPassThroughs(mlir::Value v) {
return v;
}
+void AliasAnalysis::enableSourceCache() { sourceCacheEnabled = true; }
+
+void AliasAnalysis::disableSourceCache() {
+ sourceCacheEnabled = false;
+ clearSourceCache();
+}
+
AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
bool getLastInstantiationPoint,
bool collectScopedOrigins) {
+ if (!sourceCacheEnabled)
+ return getSourceImpl(v, getLastInstantiationPoint, collectScopedOrigins);
+
+ // Key on the queried value and the two boolean flags. Recursive sub-queries
+ // go through this same wrapper, so the whole walk is memoized.
+ std::pair<mlir::Value, unsigned> key{v,
+ (getLastInstantiationPoint ? 1u : 0u) |
+ (collectScopedOrigins ? 2u : 0u)};
+ auto it = getSourceCache.find(key);
+ if (it != getSourceCache.end())
+ return it->second;
+
+ Source source =
+ getSourceImpl(v, getLastInstantiationPoint, collectScopedOrigins);
+ getSourceCache.try_emplace(key, source);
+ return source;
+}
+
+AliasAnalysis::Source
+AliasAnalysis::getSourceImpl(mlir::Value v, bool getLastInstantiationPoint,
+ bool collectScopedOrigins) {
// If v is a pass-through block argument (see walkBlockArgPassThroughs),
// continue from the underlying operand so the tracking loop below has a
// defining op to chew on. Without this, a recursive query like the one in
diff --git a/flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp b/flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp
index 14fd5c6dd7fbc..2f640199252e4 100644
--- a/flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp
+++ b/flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp
@@ -268,6 +268,16 @@ void LoopInvariantCodeMotion::runOnOperation() {
auto &aliasAnalysis = getAnalysis<AliasAnalysis>();
aliasAnalysis.addAnalysisImplementation(fir::AliasAnalysis{});
+ // Enable alias-query caching for the duration of this pass. The scope
+ // enables caching on every registered implementation that supports it
+ // (the FIR AliasAnalysis memoizes getSource()) and no-ops on the rest, and
+ // reliably disables it on every exit. This is a frozen-snapshot cache with
+ // no listener: it is sound here because LICM only moves operations, and
+ // getSource()'s position-dependent inputs are already frozen for the pass
+ // run by the AliasAnalysis's own dominance/scope caches, so caching is
+ // observationally equivalent to no caching across the hoists.
+ mlir::AliasAnalysis::QueryCacheScope cacheScope(aliasAnalysis);
+
std::function<bool(Operation *, LoopLikeOpInterface, bool)>
shouldMoveOutOfLoop = [&](Operation *op, LoopLikeOpInterface loopLike,
bool maybeConditionallyExecuted) {
diff --git a/mlir/include/mlir/Analysis/AliasAnalysis.h b/mlir/include/mlir/Analysis/AliasAnalysis.h
index 0fdcf01d9b6a3..71a08780c48a6 100644
--- a/mlir/include/mlir/Analysis/AliasAnalysis.h
+++ b/mlir/include/mlir/Analysis/AliasAnalysis.h
@@ -15,6 +15,7 @@
#define MLIR_ANALYSIS_ALIASANALYSIS_H_
#include "mlir/IR/Operation.h"
+#include "llvm/ADT/STLForwardCompat.h"
namespace mlir {
@@ -187,8 +188,22 @@ struct AliasAnalysisTraits {
/// Return the modify-reference behavior of `op` on `location`.
virtual ModRefResult getModRef(Operation *op, Value location) = 0;
+
+ /// Enable opt-in caching of alias query results on this implementation.
+ /// Implementations that do not support caching ignore this request.
+ virtual void enableQueryCaching() = 0;
+
+ /// Disable and clear any query caching on this implementation.
+ /// Implementations that do not support caching ignore this request.
+ virtual void disableQueryCaching() = 0;
};
+ /// Detection trait: true if `ImplT` provides `enableSourceCache()` /
+ /// `disableSourceCache()`. Mirrors the `has_is_invalidated` idiom used by
+ /// the pass analysis manager.
+ template <typename T>
+ using has_query_caching_t = decltype(std::declval<T &>().enableSourceCache());
+
/// This class represents the `Model` of an alias analysis implementation
/// `ImplT`. A model is instantiated for each alias analysis implementation
/// to implement the `Concept` without the need for the derived
@@ -209,6 +224,17 @@ struct AliasAnalysisTraits {
return impl.getModRef(op, location);
}
+ /// Forward query-caching control to the implementation when it provides
+ /// the cache hooks; for implementations without them these are no-ops.
+ void enableQueryCaching() final {
+ if constexpr (llvm::is_detected<has_query_caching_t, ImplT>::value)
+ impl.enableSourceCache();
+ }
+ void disableQueryCaching() final {
+ if constexpr (llvm::is_detected<has_query_caching_t, ImplT>::value)
+ impl.disableSourceCache();
+ }
+
private:
ImplT impl;
};
@@ -279,7 +305,41 @@ class AliasAnalysis {
/// Return the modify-reference behavior of `op` on `location`.
ModRefResult getModRef(Operation *op, Value location);
+ //===--------------------------------------------------------------------===//
+ // Query Caching
+ //===--------------------------------------------------------------------===//
+
+ /// RAII scope that enables opt-in query caching on all registered
+ /// implementations that support it for the duration of the scope, and
+ /// disables (and clears) it on destruction.
+ ///
+ /// This is the only entry point to query caching: enabling and disabling
+ /// are private so that caching is always paired with a guaranteed teardown.
+ /// A frozen (snapshot) cache is only valid while the IR is not mutated in a
+ /// way that affects the cached queries, so callers must keep the scope no
+ /// wider than such a no-mutation (or move-only) region.
+ class QueryCacheScope {
+ public:
+ explicit QueryCacheScope(AliasAnalysis &aa) : aa(aa) {
+ aa.enableQueryCaching();
+ }
+ ~QueryCacheScope() { aa.disableQueryCaching(); }
+
+ QueryCacheScope(const QueryCacheScope &) = delete;
+ QueryCacheScope &operator=(const QueryCacheScope &) = delete;
+ QueryCacheScope(QueryCacheScope &&) = delete;
+ QueryCacheScope &operator=(QueryCacheScope &&) = delete;
+
+ private:
+ AliasAnalysis &aa;
+ };
+
private:
+ /// Enable/disable query caching on every registered implementation that
+ /// supports it. Private: use QueryCacheScope to guarantee teardown.
+ void enableQueryCaching();
+ void disableQueryCaching();
+
/// A set of internal alias analysis implementations.
SmallVector<std::unique_ptr<Concept>, 4> aliasImpls;
};
diff --git a/mlir/lib/Analysis/AliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis.cpp
index 1382060f4dd90..9db48e18042cc 100644
--- a/mlir/lib/Analysis/AliasAnalysis.cpp
+++ b/mlir/lib/Analysis/AliasAnalysis.cpp
@@ -99,3 +99,13 @@ ModRefResult AliasAnalysis::getModRef(Operation *op, Value location) {
}
return result;
}
+
+void AliasAnalysis::enableQueryCaching() {
+ for (const std::unique_ptr<Concept> &aliasImpl : aliasImpls)
+ aliasImpl->enableQueryCaching();
+}
+
+void AliasAnalysis::disableQueryCaching() {
+ for (const std::unique_ptr<Concept> &aliasImpl : aliasImpls)
+ aliasImpl->disableQueryCaching();
+}
>From 13632fdf7ed39af3e0c6403d51d7c7ee61c0aba3 Mon Sep 17 00:00:00 2001
From: Slava Zakharin <szakharin at nvidia.com>
Date: Tue, 30 Jun 2026 12:23:03 -0700
Subject: [PATCH 2/2] [mlir][AliasAnalysis] Add precise listener-driven cache
invalidation
Extend the opt-in query cache so it can be used safely by passes that mutate
the IR through a rewriter. The query-caching hooks gain a RewriterBase*
parameter (defaulted, so QueryCacheScope callers are unaffected): in frozen
mode (null rewriter) the cache is a snapshot as before; with a rewriter the
FIR AliasAnalysis installs a ForwardingListener on it and evicts cached
getSource() results precisely as operations are erased, modified, or replaced.
Each cached entry now records the operations it depends on (the def-use chain
and the dominating dummy_scope op visited while computing it), maintained via a
reverse-dependency index. Nested queries thread a dependency sink and union a
cache hit's dependencies into the enclosing entry, so transitive dependencies
are tracked. Erasure notifications also walk nested operations and fire before
the storage is freed, closing the window where a freed-and-reused pointer could
yield a stale hit.
Adds AliasAnalysisCacheTest covering precise eviction on modify, eviction on
erase, and frozen-mode memoization.
Assisted by Cursor
---
.../flang/Optimizer/Analysis/AliasAnalysis.h | 107 +++++++++++--
.../lib/Optimizer/Analysis/AliasAnalysis.cpp | 150 ++++++++++++++++-
.../Optimizer/AliasAnalysisCacheTest.cpp | 151 ++++++++++++++++++
flang/unittests/Optimizer/CMakeLists.txt | 1 +
mlir/include/mlir/Analysis/AliasAnalysis.h | 35 ++--
mlir/lib/Analysis/AliasAnalysis.cpp | 4 +-
6 files changed, 416 insertions(+), 32 deletions(-)
create mode 100644 flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp
diff --git a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
index a40632c23efce..913270d4586ac 100644
--- a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
+++ b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
@@ -24,10 +24,25 @@
namespace fir {
+/// Listener used to evict getSource() cache entries as the IR is mutated
+/// through a rewriter. Defined in AliasAnalysis.cpp.
+class AliasAnalysisCacheListener;
+
//===----------------------------------------------------------------------===//
// AliasAnalysis
//===----------------------------------------------------------------------===//
struct AliasAnalysis {
+ // Special members are user-declared (and mostly defined out of line) because
+ // the source cache owns a std::unique_ptr to the incomplete listener type
+ // AliasAnalysisCacheListener. A move constructor is required so instances can
+ // be registered via mlir::AliasAnalysis::addAnalysisImplementation().
+ AliasAnalysis();
+ ~AliasAnalysis();
+ AliasAnalysis(AliasAnalysis &&);
+ AliasAnalysis &operator=(AliasAnalysis &&);
+ AliasAnalysis(const AliasAnalysis &) = delete;
+ AliasAnalysis &operator=(const AliasAnalysis &) = delete;
+
// Structures to describe the memory source of a value.
/// Kind of the memory source referenced by a value.
@@ -366,11 +381,25 @@ struct AliasAnalysis {
/// Enable/disable memoization of getSource() results. These are invoked by
/// the mlir::AliasAnalysis aggregator (via mlir::AliasAnalysis::
/// QueryCacheScope) and are not meant to be called directly by passes.
- /// enableSourceCache() turns on the frozen snapshot cache; disableSourceCache
- /// clears it and turns it off.
- void enableSourceCache();
+ ///
+ /// If \p rewriter is null, the cache is a frozen snapshot with no automatic
+ /// invalidation: it is only valid while the IR reachable from the queried
+ /// values is not mutated. If \p rewriter is non-null, a listener is installed
+ /// on it (chained ahead of any existing listener) so that cached entries are
+ /// evicted precisely as operations are erased / modified / replaced through
+ /// that rewriter; the cache then stays valid as long as all such mutations
+ /// flow through the rewriter.
+ ///
+ /// disableSourceCache() clears the cache (and reverse index), uninstalls the
+ /// listener if any, and turns caching off.
+ void enableSourceCache(mlir::RewriterBase *rewriter = nullptr);
void disableSourceCache();
+ /// Testing only: number of entries currently held in the getSource() cache.
+ std::size_t getSourceCacheSizeForTesting() const {
+ return getSourceCache.size();
+ }
+
/// Return true, if `ty` is a reference type to a boxed
/// POINTER object or a raw fir::PointerType.
static bool isPointerReference(mlir::Type ty);
@@ -386,6 +415,22 @@ struct AliasAnalysis {
bool functionHasMultipleScopes(mlir::Value v);
private:
+ friend class fir::AliasAnalysisCacheListener;
+
+ /// Cache key for getSource(): the queried value together with the two
+ /// boolean flags (getLastInstantiationPoint, collectScopedOrigins) packed
+ /// into the unsigned.
+ using SourceCacheKey = std::pair<mlir::Value, unsigned>;
+
+ /// A memoized getSource() result together with the operations it depends on
+ /// (the def-use chain and scope ops visited while computing it). The deps
+ /// drive precise eviction: if any of them is erased / modified / replaced,
+ /// this entry must be dropped.
+ struct CachedSource {
+ Source result;
+ llvm::SmallVector<mlir::Operation *, 4> deps;
+ };
+
/// Compute the memory source of a value. This is the uncached
/// implementation of getSource(); getSource() is a thin wrapper that
/// memoizes the result when source caching is enabled.
@@ -393,8 +438,24 @@ struct AliasAnalysis {
bool getLastInstantiationPoint,
bool collectScopedOrigins);
- /// Clear the getSource() memoization cache.
- void clearSourceCache() { getSourceCache.clear(); }
+ /// Record \p op as a dependency of the getSource() computation currently in
+ /// progress (no-op when not collecting dependencies). Used to build the
+ /// reverse index that drives precise cache eviction.
+ void recordSourceDep(mlir::Operation *op) {
+ if (currentDepSink && op)
+ currentDepSink->push_back(op);
+ }
+
+ /// Evict every cached getSource() entry that depends on \p op (and, for
+ /// erasure, on any op nested within its regions), dropping the corresponding
+ /// reverse-index entries.
+ void evictSourceDependents(mlir::Operation *op);
+
+ /// Clear the getSource() memoization cache and its reverse index.
+ void clearSourceCache() {
+ getSourceCache.clear();
+ opDependents.clear();
+ }
/// Build an intermediate Source rooted at the declare captured by the
/// snapshot. Reuses getSource(declValue) for the SourceKind / origin
@@ -473,17 +534,37 @@ struct AliasAnalysis {
/// repeated queries are O(1) without re-walking the function body.
llvm::DenseMap<mlir::Operation *, bool> multiScopeCache;
- /// Opt-in memoization of getSource() results, keyed on the queried value
- /// and the two boolean flags (getLastInstantiationPoint,
- /// collectScopedOrigins) packed into the unsigned. Only consulted/populated
- /// while \c sourceCacheEnabled is set. This is a frozen snapshot cache: it
- /// has no automatic invalidation, so it is only enabled (via
- /// enableSourceCache(), driven by mlir::AliasAnalysis::QueryCacheScope) for
- /// the duration of a region in which the relevant IR is not mutated.
- llvm::DenseMap<std::pair<mlir::Value, unsigned>, Source> getSourceCache;
+ /// Opt-in memoization of getSource() results, keyed on the queried value and
+ /// the two boolean flags (see SourceCacheKey). Only consulted/populated while
+ /// \c sourceCacheEnabled is set. Each entry also records the operations it
+ /// depends on so it can be evicted precisely. Without a listener (frozen
+ /// mode) this is a snapshot with no automatic invalidation, valid only while
+ /// the relevant IR is not mutated; with a listener installed (see
+ /// enableSourceCache(rewriter)) entries are evicted as the IR is mutated
+ /// through that rewriter. Driven by mlir::AliasAnalysis::QueryCacheScope.
+ llvm::DenseMap<SourceCacheKey, CachedSource> getSourceCache;
+
+ /// Reverse dependency index: for each operation, the set of getSourceCache
+ /// keys whose result depends on it. Maintained alongside getSourceCache and
+ /// used to evict only the affected entries when an op is mutated/erased.
+ llvm::DenseMap<mlir::Operation *, llvm::SmallVector<SourceCacheKey, 2>>
+ opDependents;
+
+ /// When non-null, getSourceImpl() and nested getSource() hits append the
+ /// operations they depend on here, so the in-progress getSource() result can
+ /// be stored with its dependencies. Saved/restored around nested queries.
+ llvm::SmallVectorImpl<mlir::Operation *> *currentDepSink = nullptr;
/// Whether getSource() should consult/populate getSourceCache.
bool sourceCacheEnabled = false;
+
+ /// The rewriter on which \c cacheListener is installed (listener mode), or
+ /// null in frozen mode. Used to uninstall the listener on disable.
+ mlir::RewriterBase *cacheRewriter = nullptr;
+
+ /// Listener installed on \c cacheRewriter that evicts cache entries as the
+ /// IR is mutated. Null in frozen mode. Owns a back-pointer to this analysis.
+ std::unique_ptr<AliasAnalysisCacheListener> cacheListener;
};
inline bool operator==(const AliasAnalysis::Source::SourceOrigin &lhs,
diff --git a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
index f3e77a4fc0826..3fb73fcc78280 100644
--- a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
+++ b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -22,9 +22,11 @@
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/Dialect/OpenMP/OpenMPInterfaces.h"
#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
@@ -1109,13 +1111,116 @@ static mlir::Value walkBlockArgPassThroughs(mlir::Value v) {
return v;
}
-void AliasAnalysis::enableSourceCache() { sourceCacheEnabled = true; }
+//===----------------------------------------------------------------------===//
+// Source cache invalidation listener
+//===----------------------------------------------------------------------===//
+
+/// Listener installed on a rewriter while source caching is enabled in listener
+/// mode. It evicts the affected getSource() cache entries on each mutation and
+/// then forwards the notification to any previously installed listener so the
+/// cache composes with greedy/conversion driver listeners.
+class AliasAnalysisCacheListener
+ : public mlir::RewriterBase::ForwardingListener {
+public:
+ AliasAnalysisCacheListener(AliasAnalysis &aa,
+ mlir::OpBuilder::Listener *previous)
+ : mlir::RewriterBase::ForwardingListener(previous), aa(aa),
+ previous(previous) {}
+
+ /// The listener that was installed on the rewriter before this one, so it
+ /// can be restored when caching is disabled.
+ mlir::OpBuilder::Listener *getPreviousListener() const { return previous; }
+
+ void notifyOperationErased(mlir::Operation *op) override {
+ // Runs before the operation's storage is freed, so evicting here closes
+ // the window where a freed-and-reused Operation* could produce a stale hit.
+ aa.evictSourceDependents(op);
+ ForwardingListener::notifyOperationErased(op);
+ }
+
+ void notifyOperationModified(mlir::Operation *op) override {
+ aa.evictSourceDependents(op);
+ ForwardingListener::notifyOperationModified(op);
+ }
+
+ void notifyOperationReplaced(mlir::Operation *op,
+ mlir::Operation *newOp) override {
+ aa.evictSourceDependents(op);
+ ForwardingListener::notifyOperationReplaced(op, newOp);
+ }
+
+ void notifyOperationReplaced(mlir::Operation *op,
+ mlir::ValueRange replacement) override {
+ aa.evictSourceDependents(op);
+ ForwardingListener::notifyOperationReplaced(op, replacement);
+ }
+
+private:
+ AliasAnalysis &aa;
+ mlir::OpBuilder::Listener *previous;
+};
+
+//===----------------------------------------------------------------------===//
+// AliasAnalysis special members
+//===----------------------------------------------------------------------===//
+
+// Defined out of line because cacheListener is a std::unique_ptr to the
+// (here-complete) AliasAnalysisCacheListener.
+AliasAnalysis::AliasAnalysis() = default;
+AliasAnalysis::~AliasAnalysis() = default;
+AliasAnalysis::AliasAnalysis(AliasAnalysis &&) = default;
+AliasAnalysis &AliasAnalysis::operator=(AliasAnalysis &&) = default;
+
+//===----------------------------------------------------------------------===//
+// Source cache
+//===----------------------------------------------------------------------===//
+
+void AliasAnalysis::enableSourceCache(mlir::RewriterBase *rewriter) {
+ sourceCacheEnabled = true;
+ if (rewriter && !cacheListener) {
+ cacheRewriter = rewriter;
+ cacheListener = std::make_unique<AliasAnalysisCacheListener>(
+ *this, rewriter->getListener());
+ rewriter->setListener(cacheListener.get());
+ }
+}
void AliasAnalysis::disableSourceCache() {
sourceCacheEnabled = false;
+ if (cacheListener) {
+ // Restore the listener that was installed before ours.
+ cacheRewriter->setListener(cacheListener->getPreviousListener());
+ cacheListener.reset();
+ cacheRewriter = nullptr;
+ }
clearSourceCache();
}
+void AliasAnalysis::evictSourceDependents(mlir::Operation *op) {
+ if (getSourceCache.empty())
+ return;
+
+ // Collect the ops whose dependents must be evicted. For erased ops, this
+ // includes ops nested in their regions, since cached entries may depend on
+ // those nested ops too and their storage is freed along with the parent.
+ llvm::SmallVector<mlir::Operation *, 8> worklist{op};
+ while (!worklist.empty()) {
+ mlir::Operation *cur = worklist.pop_back_val();
+ auto it = opDependents.find(cur);
+ if (it != opDependents.end()) {
+ // Copy the keys out before erasing, then drop the cache entries.
+ llvm::SmallVector<SourceCacheKey, 2> keys = std::move(it->second);
+ opDependents.erase(it);
+ for (const SourceCacheKey &key : keys)
+ getSourceCache.erase(key);
+ }
+ for (mlir::Region ®ion : cur->getRegions())
+ for (mlir::Block &block : region)
+ for (mlir::Operation &nested : block)
+ worklist.push_back(&nested);
+ }
+}
+
AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
bool getLastInstantiationPoint,
bool collectScopedOrigins) {
@@ -1124,16 +1229,40 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
// Key on the queried value and the two boolean flags. Recursive sub-queries
// go through this same wrapper, so the whole walk is memoized.
- std::pair<mlir::Value, unsigned> key{v,
- (getLastInstantiationPoint ? 1u : 0u) |
- (collectScopedOrigins ? 2u : 0u)};
+ SourceCacheKey key{v, (getLastInstantiationPoint ? 1u : 0u) |
+ (collectScopedOrigins ? 2u : 0u)};
auto it = getSourceCache.find(key);
- if (it != getSourceCache.end())
- return it->second;
+ if (it != getSourceCache.end()) {
+ // Propagate this entry's dependencies into the enclosing query (if any) so
+ // transitive dependencies are tracked: if a dep of this entry is later
+ // erased, the enclosing entry that incorporated it is evicted too.
+ if (currentDepSink)
+ currentDepSink->append(it->second.deps.begin(), it->second.deps.end());
+ return it->second.result;
+ }
+ // Miss: compute with a fresh dependency sink, then propagate to the parent.
+ llvm::SmallVector<mlir::Operation *, 4> deps;
+ llvm::SmallVectorImpl<mlir::Operation *> *parentSink = currentDepSink;
+ currentDepSink = &deps;
Source source =
getSourceImpl(v, getLastInstantiationPoint, collectScopedOrigins);
- getSourceCache.try_emplace(key, source);
+ currentDepSink = parentSink;
+
+ // Deduplicate before storing / indexing (the walk may visit an op multiple
+ // times). Sorting by pointer only affects internal bookkeeping, not emitted
+ // IR, so dump determinism is preserved.
+ llvm::sort(deps);
+ deps.erase(llvm::unique(deps), deps.end());
+
+ if (parentSink)
+ parentSink->append(deps.begin(), deps.end());
+
+ bool inserted =
+ getSourceCache.try_emplace(key, CachedSource{source, deps}).second;
+ if (inserted)
+ for (mlir::Operation *depOp : deps)
+ opDependents[depOp].push_back(key);
return source;
}
@@ -1176,6 +1305,9 @@ AliasAnalysis::getSourceImpl(mlir::Value v, bool getLastInstantiationPoint,
// buildSourceAtDeclare reuses getSource purely for declare classification).
llvm::SmallVector<Source::ScopedOrigin, 4> scopedOrigins;
while (defOp && !breakFromLoop) {
+ // Record each operation visited along the def-use chain as a dependency
+ // of this query, so the cached result is evicted if any of them changes.
+ recordSourceDep(defOp);
// Operations may have multiple results, so we need to analyze
// the result for which the source is queried.
auto opResult = mlir::cast<OpResult>(v);
@@ -1443,6 +1575,10 @@ AliasAnalysis::getSourceImpl(mlir::Value v, bool getLastInstantiationPoint,
if (collectScopedOrigins) {
Source::ScopedOrigin scopedOrigin;
scopedOrigin.scope = getDeclarationScope(op);
+ // The chosen scope depends on the dominating fir.dummy_scope op;
+ // record it so the entry is evicted if that scope op changes.
+ if (scopedOrigin.scope)
+ recordSourceDep(scopedOrigin.scope.getDefiningOp());
scopedOrigin.declValue = opResult;
scopedOrigin.accessPath.steps.assign(pathSteps.rbegin(),
pathSteps.rend());
diff --git a/flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp b/flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp
new file mode 100644
index 0000000000000..b199adc4191b6
--- /dev/null
+++ b/flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp
@@ -0,0 +1,151 @@
+//===- AliasAnalysisCacheTest.cpp -----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Tests for the listener-driven invalidation of fir::AliasAnalysis's getSource
+// cache: mutating the IR through a rewriter must evict precisely the cached
+// entries that depend on the mutated operation, and an erase must drop the
+// affected entries before the operation's storage is freed (so a freed,
+// possibly reused, pointer can never produce a stale cache hit).
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "flang/Optimizer/Analysis/AliasAnalysis.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Support/InitFIR.h"
+
+struct AliasAnalysisCacheTest : public testing::Test {
+public:
+ void SetUp() override {
+ fir::support::loadDialects(context);
+ builder = std::make_unique<mlir::OpBuilder>(&context);
+ mlir::Location loc = builder->getUnknownLoc();
+
+ moduleOp = mlir::ModuleOp::create(*builder, loc);
+ builder->setInsertionPointToStart(moduleOp->getBody());
+ mlir::func::FuncOp func = mlir::func::FuncOp::create(
+ *builder, loc, "test", builder->getFunctionType({}, {}));
+ builder->setInsertionPointToStart(func.addEntryBlock());
+ }
+
+ mlir::Location getLoc() { return builder->getUnknownLoc(); }
+
+ // Build an `alloca -> declare` chain for a scalar i32 variable and return the
+ // declared variable address.
+ mlir::Value createScalarVariable(llvm::StringRef name) {
+ mlir::Location loc = getLoc();
+ mlir::Type eleType = mlir::IntegerType::get(&context, 32);
+ mlir::Value addr = fir::AllocaOp::create(*builder, loc, eleType);
+ auto declare = fir::DeclareOp::create(*builder, loc, addr.getType(), addr,
+ /*shape=*/mlir::Value{}, /*typeParams=*/mlir::ValueRange{},
+ /*dummy_scope=*/nullptr, /*storage=*/nullptr, /*storage_offset=*/0,
+ mlir::StringAttr::get(&context, name),
+ /*fortran_attrs=*/fir::FortranVariableFlagsAttr{},
+ /*data_attr=*/cuf::DataAttributeAttr{},
+ /*dummy_arg_no=*/mlir::IntegerAttr{});
+ return declare.getResult();
+ }
+
+ mlir::MLIRContext context;
+ std::unique_ptr<mlir::OpBuilder> builder;
+ mlir::OwningOpRef<mlir::ModuleOp> moduleOp;
+};
+
+// Modifying an operation that a cached query transitively depends on (here the
+// alloca underlying a declared variable) must evict that query's entry while
+// leaving unrelated entries intact.
+TEST_F(AliasAnalysisCacheTest, PreciseEvictionOnModify) {
+ mlir::Value var1 = createScalarVariable("x1");
+ mlir::Value var2 = createScalarVariable("x2");
+
+ fir::AliasAnalysis aa;
+ mlir::IRRewriter rewriter(&context);
+ aa.enableSourceCache(&rewriter);
+
+ fir::AliasAnalysis::Source s1 = aa.getSource(var1);
+ fir::AliasAnalysis::Source s2 = aa.getSource(var2);
+ std::size_t before = aa.getSourceCacheSizeForTesting();
+ EXPECT_GE(before, 2u);
+
+ // var1's source depends on its alloca; modifying the alloca must evict the
+ // var1 entry (a transitive dependency, not the queried op itself).
+ mlir::Operation *alloca1 =
+ var1.getDefiningOp()->getOperand(0).getDefiningOp();
+ ASSERT_TRUE(alloca1);
+ rewriter.modifyOpInPlace(alloca1, [] {});
+
+ std::size_t after = aa.getSourceCacheSizeForTesting();
+ EXPECT_LT(after, before); // some entry was evicted
+ EXPECT_GT(after, 0u); // but not everything (var2 is unaffected)
+
+ // Re-querying var2 must still return the same source (it was not evicted).
+ fir::AliasAnalysis::Source s2again = aa.getSource(var2);
+ EXPECT_EQ(s2again.origin, s2.origin);
+ EXPECT_EQ(s2again.kind, s2.kind);
+
+ // Sanity: var1 and var2 resolve to distinct origins.
+ EXPECT_NE(s1.origin, s2.origin);
+
+ aa.disableSourceCache();
+ EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+}
+
+// Erasing a cached query's defining operation must drop the affected entries
+// (the notification fires before the storage is freed), so no stale entry can
+// survive to alias with a later, pointer-reused operation.
+TEST_F(AliasAnalysisCacheTest, EvictionOnEraseClosesReuseHole) {
+ mlir::Value var1 = createScalarVariable("x1");
+ mlir::Value var2 = createScalarVariable("x2");
+
+ // A load whose result we query; it has no uses, so it can be erased.
+ mlir::Value load1 = fir::LoadOp::create(*builder, getLoc(), var1);
+
+ fir::AliasAnalysis aa;
+ mlir::IRRewriter rewriter(&context);
+ aa.enableSourceCache(&rewriter);
+
+ aa.getSource(load1);
+ aa.getSource(var2);
+ std::size_t before = aa.getSourceCacheSizeForTesting();
+ EXPECT_GE(before, 2u);
+
+ // Erase the load: its cache entry must be gone afterwards, while var2's
+ // entry survives.
+ mlir::Operation *loadOp = load1.getDefiningOp();
+ rewriter.eraseOp(loadOp);
+
+ std::size_t after = aa.getSourceCacheSizeForTesting();
+ EXPECT_LT(after, before);
+ EXPECT_GT(after, 0u);
+
+ // var2 is still cached and valid.
+ fir::AliasAnalysis::Source s2 = aa.getSource(var2);
+ EXPECT_EQ(s2.kind, fir::AliasAnalysis::SourceKind::Allocate);
+
+ aa.disableSourceCache();
+}
+
+// Without a rewriter the cache is a frozen snapshot: it still memoizes results,
+// but installs no listener (so the caller is responsible for scoping).
+TEST_F(AliasAnalysisCacheTest, FrozenModeMemoizes) {
+ mlir::Value var1 = createScalarVariable("x1");
+
+ fir::AliasAnalysis aa;
+ aa.enableSourceCache(/*rewriter=*/nullptr);
+ EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+
+ fir::AliasAnalysis::Source first = aa.getSource(var1);
+ EXPECT_GT(aa.getSourceCacheSizeForTesting(), 0u);
+ fir::AliasAnalysis::Source second = aa.getSource(var1);
+ EXPECT_EQ(first.origin, second.origin);
+
+ aa.disableSourceCache();
+ EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+}
diff --git a/flang/unittests/Optimizer/CMakeLists.txt b/flang/unittests/Optimizer/CMakeLists.txt
index 3af6f9d2a3724..6991d67191dd6 100644
--- a/flang/unittests/Optimizer/CMakeLists.txt
+++ b/flang/unittests/Optimizer/CMakeLists.txt
@@ -37,6 +37,7 @@ add_flang_unittest(FlangOptimizerTests
Builder/Runtime/StopTest.cpp
Builder/Runtime/TransformationalTest.cpp
OpenACC/FIROpenACCPointerLikeTypeInterfaceTest.cpp
+ AliasAnalysisCacheTest.cpp
FIRCallInterfaceTest.cpp
FIRContextTest.cpp
FIRTypesTest.cpp
diff --git a/mlir/include/mlir/Analysis/AliasAnalysis.h b/mlir/include/mlir/Analysis/AliasAnalysis.h
index 71a08780c48a6..b52d5a9a6f146 100644
--- a/mlir/include/mlir/Analysis/AliasAnalysis.h
+++ b/mlir/include/mlir/Analysis/AliasAnalysis.h
@@ -19,6 +19,8 @@
namespace mlir {
+class RewriterBase;
+
//===----------------------------------------------------------------------===//
// AliasResult
//===----------------------------------------------------------------------===//
@@ -191,7 +193,11 @@ struct AliasAnalysisTraits {
/// Enable opt-in caching of alias query results on this implementation.
/// Implementations that do not support caching ignore this request.
- virtual void enableQueryCaching() = 0;
+ /// If `rewriter` is non-null, the implementation may install a listener on
+ /// it to invalidate cached results precisely as the IR is mutated; if null
+ /// the cache is a frozen snapshot valid only while the IR is not mutated in
+ /// a way that affects the cached queries.
+ virtual void enableQueryCaching(RewriterBase *rewriter) = 0;
/// Disable and clear any query caching on this implementation.
/// Implementations that do not support caching ignore this request.
@@ -202,7 +208,8 @@ struct AliasAnalysisTraits {
/// `disableSourceCache()`. Mirrors the `has_is_invalidated` idiom used by
/// the pass analysis manager.
template <typename T>
- using has_query_caching_t = decltype(std::declval<T &>().enableSourceCache());
+ using has_query_caching_t = decltype(std::declval<T &>().enableSourceCache(
+ std::declval<RewriterBase *>()));
/// This class represents the `Model` of an alias analysis implementation
/// `ImplT`. A model is instantiated for each alias analysis implementation
@@ -226,9 +233,9 @@ struct AliasAnalysisTraits {
/// Forward query-caching control to the implementation when it provides
/// the cache hooks; for implementations without them these are no-ops.
- void enableQueryCaching() final {
+ void enableQueryCaching(RewriterBase *rewriter) final {
if constexpr (llvm::is_detected<has_query_caching_t, ImplT>::value)
- impl.enableSourceCache();
+ impl.enableSourceCache(rewriter);
}
void disableQueryCaching() final {
if constexpr (llvm::is_detected<has_query_caching_t, ImplT>::value)
@@ -315,13 +322,21 @@ class AliasAnalysis {
///
/// This is the only entry point to query caching: enabling and disabling
/// are private so that caching is always paired with a guaranteed teardown.
- /// A frozen (snapshot) cache is only valid while the IR is not mutated in a
- /// way that affects the cached queries, so callers must keep the scope no
- /// wider than such a no-mutation (or move-only) region.
+ ///
+ /// If `rewriter` is null (the default), the cache is a frozen snapshot that
+ /// is only valid while the IR is not mutated in a way that affects the
+ /// cached queries, so callers must keep the scope no wider than such a
+ /// no-mutation (or move-only) region. If `rewriter` is non-null,
+ /// implementations that support it install a listener on the rewriter and
+ /// invalidate cached results precisely as the IR is mutated through it; the
+ /// cache then stays valid as long as all mutations flow through that
+ /// rewriter.
class QueryCacheScope {
public:
- explicit QueryCacheScope(AliasAnalysis &aa) : aa(aa) {
- aa.enableQueryCaching();
+ explicit QueryCacheScope(AliasAnalysis &aa,
+ RewriterBase *rewriter = nullptr)
+ : aa(aa) {
+ aa.enableQueryCaching(rewriter);
}
~QueryCacheScope() { aa.disableQueryCaching(); }
@@ -337,7 +352,7 @@ class AliasAnalysis {
private:
/// Enable/disable query caching on every registered implementation that
/// supports it. Private: use QueryCacheScope to guarantee teardown.
- void enableQueryCaching();
+ void enableQueryCaching(RewriterBase *rewriter = nullptr);
void disableQueryCaching();
/// A set of internal alias analysis implementations.
diff --git a/mlir/lib/Analysis/AliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis.cpp
index 9db48e18042cc..79d7ced564a33 100644
--- a/mlir/lib/Analysis/AliasAnalysis.cpp
+++ b/mlir/lib/Analysis/AliasAnalysis.cpp
@@ -100,9 +100,9 @@ ModRefResult AliasAnalysis::getModRef(Operation *op, Value location) {
return result;
}
-void AliasAnalysis::enableQueryCaching() {
+void AliasAnalysis::enableQueryCaching(RewriterBase *rewriter) {
for (const std::unique_ptr<Concept> &aliasImpl : aliasImpls)
- aliasImpl->enableQueryCaching();
+ aliasImpl->enableQueryCaching(rewriter);
}
void AliasAnalysis::disableQueryCaching() {
More information about the flang-commits
mailing list