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 +#include 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; + + /// 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 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 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 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> + 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 *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 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 +#include 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( + *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 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 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 deps; + llvm::SmallVectorImpl *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 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(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.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 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(&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 builder; + mlir::OwningOpRef 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 0fdcf01d9b6a3..b52d5a9a6f146 100644 --- a/mlir/include/mlir/Analysis/AliasAnalysis.h +++ b/mlir/include/mlir/Analysis/AliasAnalysis.h @@ -15,9 +15,12 @@ #define MLIR_ANALYSIS_ALIASANALYSIS_H_ #include "mlir/IR/Operation.h" +#include "llvm/ADT/STLForwardCompat.h" namespace mlir { +class RewriterBase; + //===----------------------------------------------------------------------===// // AliasResult //===----------------------------------------------------------------------===// @@ -187,8 +190,27 @@ 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. + /// 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. + 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 + using has_query_caching_t = decltype(std::declval().enableSourceCache( + std::declval())); + /// 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 +231,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(RewriterBase *rewriter) final { + if constexpr (llvm::is_detected::value) + impl.enableSourceCache(rewriter); + } + void disableQueryCaching() final { + if constexpr (llvm::is_detected::value) + impl.disableSourceCache(); + } + private: ImplT impl; }; @@ -279,7 +312,49 @@ 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. + /// + /// 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, + RewriterBase *rewriter = nullptr) + : aa(aa) { + aa.enableQueryCaching(rewriter); + } + ~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(RewriterBase *rewriter = nullptr); + void disableQueryCaching(); + /// A set of internal alias analysis implementations. SmallVector, 4> aliasImpls; }; diff --git a/mlir/lib/Analysis/AliasAnalysis.cpp b/mlir/lib/Analysis/AliasAnalysis.cpp index 1382060f4dd90..79d7ced564a33 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(RewriterBase *rewriter) { + for (const std::unique_ptr &aliasImpl : aliasImpls) + aliasImpl->enableQueryCaching(rewriter); +} + +void AliasAnalysis::disableQueryCaching() { + for (const std::unique_ptr &aliasImpl : aliasImpls) + aliasImpl->disableQueryCaching(); +}