[flang][AliasAnalysis] Add opt-in getSource memoization cache (#208319)

Add an opt-in cache to fir::AliasAnalysis that memoizes getSource()
results keyed on (value, flags). Caching is off by default and turned on
via enableSourceCache(); getSource() becomes a thin wrapper over the
uncached getSourceImpl(). The cache is a frozen snapshot with no
automatic invalidation and lives no longer than the AliasAnalysis
instance, so a client enables it only for a region in which it does not
mutate IR in a way that would change a source.

Flang's LICM enables the cache on its fir::AliasAnalysis before adding
it to the mlir::AliasAnalysis aggregate: LICM only moves operations, so
getSource()'s inputs are unchanged across the hoists, and the pass
manager drops the analysis (and its cache) after the pass.

Add a unit test covering the default-off behavior, cache fill/hit, and
that disableSourceCache() clears and bypasses the cache.

Assisted-by: Cursor

flang/include/flang/Optimizer/Analysis/AliasAnalysis.h

@@ -19,7 +19,9 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
+#include <cstddef>
 #include <memory>
+#include <utility>
 
 namespace fir {
 
@@ -351,10 +353,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 with no automatic
+  /// invalidation: it is only valid while the IR reachable from the queried
+  /// values is not mutated in a way that would change the source. It lives no
+  /// longer than this AliasAnalysis instance, so a client enables caching only
+  /// for a region in which it does not perform such mutations (e.g. a pass
+  /// that only moves operations).
   fir::AliasAnalysis::Source getSource(mlir::Value,
                                        bool getLastInstantiationPoint = false,
                                        bool collectScopedOrigins = true);
 
+  /// Enable memoization of getSource() results on this analysis instance.
+  /// Caching is opt-in and off by default; a client that does not mutate the
+  /// IR in a way that affects getSource() (see getSource()) may enable it to
+  /// avoid recomputing sources for repeated queries.
+  void enableSourceCache();
+
+  /// Disable getSource() memoization and drop any cached entries.
+  void disableSourceCache();
+
+  /// Testing only: number of entries currently held in the getSource() cache.
+  std::size_t getSourceCacheSizeForTesting() const {
+    return getSourceCache.size();
+  }
+
+  /// Testing only: cumulative getSource() cache hits / misses on this instance.
+  std::size_t getSourceCacheHitsForTesting() const { return sourceCacheHits; }
+  std::size_t getSourceCacheMissesForTesting() const {
+    return sourceCacheMisses;
+  }
+
   /// 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 +401,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 +487,21 @@ private:
   /// 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 with
+  /// no automatic invalidation (see getSource()).
+  llvm::DenseMap<std::pair<mlir::Value, unsigned>, Source> getSourceCache;
+
+  /// Whether getSource() should consult/populate getSourceCache.
+  bool sourceCacheEnabled = false;
+
+  /// Testing-only counters (see getSourceCacheHitsForTesting()). Incremented
+  /// on each getSource() cache hit / miss.
+  std::size_t sourceCacheHits = 0;
+  std::size_t sourceCacheMisses = 0;
 };
 
 inline bool operator==(const AliasAnalysis::Source::SourceOrigin &lhs,

flang/lib/Optimizer/Analysis/AliasAnalysis.cpp

@@ -29,6 +29,7 @@
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
+#include <utility>
 
 using namespace mlir;
 
@@ -1107,9 +1108,40 @@ 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()) {
+    ++sourceCacheHits;
+    return it->second;
+  }
+
+  ++sourceCacheMisses;
+  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

flang/lib/Optimizer/Transforms/LoopInvariantCodeMotion.cpp

@@ -24,6 +24,7 @@
 #include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/DebugLog.h"
+#include <utility>
 
 namespace fir {
 #define GEN_PASS_DEF_LOOPINVARIANTCODEMOTION
@@ -265,7 +266,15 @@ void LoopInvariantCodeMotion::runOnOperation() {
   LDBG() << "Enter [HL]FIR LoopInvariantCodeMotion()";
 
   auto &aliasAnalysis = getAnalysis<AliasAnalysis>();
-  aliasAnalysis.addAnalysisImplementation(fir::AliasAnalysis{});
+  // Enable getSource() memoization on the FIR AliasAnalysis for the duration
+  // of this pass. This is a frozen-snapshot cache with no automatic
+  // invalidation, but it is sound here because LICM only moves operations, so
+  // getSource()'s inputs are unchanged across the hoists. The cache lives no
+  // longer than this analysis instance, which the pass manager drops when the
+  // analysis is invalidated after the pass.
+  fir::AliasAnalysis firAliasAnalysis;
+  firAliasAnalysis.enableSourceCache();
+  aliasAnalysis.addAnalysisImplementation(std::move(firAliasAnalysis));
 
   std::function<bool(Operation *, LoopLikeOpInterface, bool)>
       shouldMoveOutOfLoop = [&](Operation *op, LoopLikeOpInterface loopLike,

flang/unittests/Optimizer/AliasAnalysisCacheTest.cpp

@@ -0,0 +1,107 @@
+//===- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Unit tests for the opt-in getSource() memoization cache on
+// fir::AliasAnalysis. The cache is off by default; when enabled it memoizes
+// getSource() results for the lifetime of the analysis instance and is a
+// frozen snapshot (no automatic invalidation).
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.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();
+
+    // Set up a module with a dummy function; insert into its entry block.
+    moduleOp = mlir::ModuleOp::create(*builder, loc);
+    builder->setInsertionPointToStart(moduleOp->getBody());
+    mlir::func::FuncOp func = mlir::func::FuncOp::create(*builder, loc,
+        "alias_analysis_cache_tests", builder->getFunctionType({}, {}));
+    auto *entryBlock = func.addEntryBlock();
+    builder->setInsertionPointToStart(entryBlock);
+  }
+
+  mlir::Location getLoc() { return builder->getUnknownLoc(); }
+
+  mlir::Value createAlloca() {
+    return fir::AllocaOp::create(
+        *builder, getLoc(), mlir::Float32Type::get(&context));
+  }
+
+  mlir::MLIRContext context;
+  std::unique_ptr<mlir::OpBuilder> builder;
+  mlir::OwningOpRef<mlir::ModuleOp> moduleOp;
+};
+
+// Caching is off by default: getSource() bypasses the cache, so nothing is
+// memoized and the hit/miss counters stay at zero.
+TEST_F(AliasAnalysisCacheTest, DisabledByDefault) {
+  mlir::Value a = createAlloca();
+  fir::AliasAnalysis aa;
+
+  (void)aa.getSource(a);
+  (void)aa.getSource(a);
+
+  EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+  EXPECT_EQ(aa.getSourceCacheHitsForTesting(), 0u);
+  EXPECT_EQ(aa.getSourceCacheMissesForTesting(), 0u);
+}
+
+// Once enabled, a query populates the cache (a miss) and a repeated query for
+// the same value is served from it (a hit) without growing the cache.
+TEST_F(AliasAnalysisCacheTest, FillAndHit) {
+  mlir::Value a = createAlloca();
+  mlir::Value b = createAlloca();
+  fir::AliasAnalysis aa;
+  aa.enableSourceCache();
+
+  (void)aa.getSource(a);
+  std::size_t sizeAfterFill = aa.getSourceCacheSizeForTesting();
+  std::size_t missesAfterFill = aa.getSourceCacheMissesForTesting();
+  EXPECT_GE(sizeAfterFill, 1u);
+  EXPECT_GE(missesAfterFill, 1u);
+  EXPECT_EQ(aa.getSourceCacheHitsForTesting(), 0u);
+
+  // Repeating the top-level query is a single cache hit: it returns the
+  // memoized result without recomputing (and thus without new misses or
+  // entries), regardless of any recursive sub-queries the first call made.
+  (void)aa.getSource(a);
+  EXPECT_EQ(aa.getSourceCacheSizeForTesting(), sizeAfterFill);
+  EXPECT_EQ(aa.getSourceCacheMissesForTesting(), missesAfterFill);
+  EXPECT_EQ(aa.getSourceCacheHitsForTesting(), 1u);
+
+  // A distinct value adds at least one new entry.
+  (void)aa.getSource(b);
+  EXPECT_GT(aa.getSourceCacheSizeForTesting(), sizeAfterFill);
+  EXPECT_GT(aa.getSourceCacheMissesForTesting(), missesAfterFill);
+}
+
+// disableSourceCache() clears the entries and turns memoization off, so later
+// queries bypass the cache and leave it empty.
+TEST_F(AliasAnalysisCacheTest, DisableClearsAndBypasses) {
+  mlir::Value a = createAlloca();
+  fir::AliasAnalysis aa;
+  aa.enableSourceCache();
+
+  (void)aa.getSource(a);
+  EXPECT_GE(aa.getSourceCacheSizeForTesting(), 1u);
+
+  aa.disableSourceCache();
+  EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+
+  (void)aa.getSource(a);
+  EXPECT_EQ(aa.getSourceCacheSizeForTesting(), 0u);
+}

flang/unittests/Optimizer/CMakeLists.txt

@@ -40,6 +40,7 @@ add_flang_unittest(FlangOptimizerTests
   Builder/Runtime/TransformationalTest.cpp
   OpenACC/FIROpenACCPointerLikeTypeInterfaceTest.cpp
   OpenACC/FIROpenACCSupportAnalysisTest.cpp
+  AliasAnalysisCacheTest.cpp
   FIRCallInterfaceTest.cpp
   FIRContextTest.cpp
   FIRTypesTest.cpp