[flang-commits] [flang] [mlir] [mlir][AliasAnalysis] Add opt-in query caching; use it in flang LICM (PR #207030)
via flang-commits
flang-commits at lists.llvm.org
Wed Jul 1 09:47:05 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Slava Zakharin (vzakhari)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/207030.diff
5 Files Affected:
- (modified) flang/include/flang/Optimizer/Analysis/AliasAnalysis.h (+38)
- (modified) flang/lib/Optimizer/Analysis/AliasAnalysis.cpp (+29)
- (modified) flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp (+10)
- (modified) mlir/include/mlir/Analysis/AliasAnalysis.h (+60)
- (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..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();
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/207030
More information about the flang-commits
mailing list