[flang-commits] [flang] [mlir] [mlir][AliasAnalysis] Add precise listener-driven cache invalidation (PR #207032)
via flang-commits
flang-commits at lists.llvm.org
Wed Jul 1 09:47:51 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-flang-fir-hlfir
Author: Slava Zakharin (vzakhari)
<details>
<summary>Changes</summary>
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
---
Patch is 30.16 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/207032.diff
7 Files Affected:
- (modified) flang/include/flang/Optimizer/Analysis/AliasAnalysis.h (+119)
- (modified) flang/lib/Optimizer/Analysis/AliasAnalysis.cpp (+165)
- (modified) flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp (+10)
- (added) flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp (+151)
- (modified) flang/unittests/Optimizer/CMakeLists.txt (+1)
- (modified) mlir/include/mlir/Analysis/AliasAnalysis.h (+75)
- (modified) mlir/lib/Analysis/AliasAnalysis.cpp (+10)
``````````diff
diff --git a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
index 832634a708dba..913270d4586ac 100644
--- a/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
+++ b/flang/include/flang/Optimizer/Analysis/AliasAnalysis.h
@@ -20,13 +20,29 @@
#include "llvm/ADT/PointerUnion.h"
#include "llvm/ADT/SmallVector.h"
#include <memory>
+#include <utility>
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.
@@ -351,10 +367,39 @@ 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.
+ ///
+ /// 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);
@@ -370,6 +415,48 @@ 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.
+ fir::AliasAnalysis::Source getSourceImpl(mlir::Value v,
+ bool getLastInstantiationPoint,
+ bool collectScopedOrigins);
+
+ /// 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
/// classification (with collectScopedOrigins=false), then overrides
@@ -446,6 +533,38 @@ 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 (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 02a628968ddc4..3fb73fcc78280 100644
--- a/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
+++ b/flang/lib/Optimizer/Analysis/AliasAnalysis.cpp
@@ -22,14 +22,17 @@
#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"
#include "llvm/Support/Debug.h"
#include <optional>
+#include <utility>
using namespace mlir;
@@ -1108,9 +1111,164 @@ static mlir::Value walkBlockArgPassThroughs(mlir::Value v) {
return v;
}
+//===----------------------------------------------------------------------===//
+// 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) {
+ 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.
+ SourceCacheKey key{v, (getLastInstantiationPoint ? 1u : 0u) |
+ (collectScopedOrigins ? 2u : 0u)};
+ auto it = getSourceCache.find(key);
+ 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);
+ 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;
+}
+
+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
@@ -1147,6 +1305,9 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
// 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);
@@ -1414,6 +1575,10 @@ AliasAnalysis::Source AliasAnalysis::getSource(mlir::Value v,
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/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/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...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/207032
More information about the flang-commits
mailing list