[llvm] [SandboxVec] Clear Context's state within runOnFunction() (PR #124842)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-vectorizers

Author: vporpo (vporpo)

<details>
<summary>Changes</summary>

`sandboxir::Context` is defined at a pass-level scope with the `SandboxVectorizerPass` class because the function pass manager `FPM` object depends on it, and that is in pass-level scope to avoid recreating the pass pipeline every single time `runOnFunction()` is called.

This means that the Context's state lives on across function passes. The problem is twofold:
(i) the LLVM IR to Sandbox IR map can grow very large including objects from different functions, which is of no use to the vectorizer, as it's a function-level pass.
(ii) this can result in stale data in the LLVM IR to Sandbox IR object map, as other passes may delete LLVM IR objects.

To fix both issues this patch introduces a `Context::clear()` function that clears the `LLVMValueToValueMap`.

---
Full diff: https://github.com/llvm/llvm-project/pull/124842.diff


6 Files Affected:

- (modified) llvm/include/llvm/SandboxIR/Context.h (+2) 
- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h (+9-1) 
- (modified) llvm/lib/SandboxIR/Context.cpp (+6) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp (+5-2) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt (+1) 
- (added) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerTest.cpp (+63) 


``````````diff
diff --git a/llvm/include/llvm/SandboxIR/Context.h b/llvm/include/llvm/SandboxIR/Context.h
index 7fe97d984b9589..a88b0003f55bdb 100644
--- a/llvm/include/llvm/SandboxIR/Context.h
+++ b/llvm/include/llvm/SandboxIR/Context.h
@@ -218,6 +218,8 @@ class Context {
 public:
   Context(LLVMContext &LLVMCtx);
   ~Context();
+  /// Clears function-level state.
+  void clear();
 
   Tracker &getTracker() { return IRTracker; }
   /// Convenience function for `getTracker().save()`
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
index 09369dbb496fce..7ea9386f08beef 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
@@ -24,10 +24,18 @@ class SandboxVectorizerPass : public PassInfoMixin<SandboxVectorizerPass> {
   TargetTransformInfo *TTI = nullptr;
   AAResults *AA = nullptr;
   ScalarEvolution *SE = nullptr;
-
+  // NOTE: We define the Context as a pass-scope object instead of local object
+  // in runOnFunction() because the passes defined in the pass-manager need
+  // access to it for registering/deregistering callbacks during construction
+  // and destruction.
   std::unique_ptr<sandboxir::Context> Ctx;
 
   // A pipeline of SandboxIR function passes run by the vectorizer.
+  // NOTE: We define this as a pass-scope object to avoid recreating the
+  // pass-pipeline every time in runOnFunction(). The downside is that the
+  // Context also needs to be defined as a pass-scope object because the passes
+  // within FPM may register/unregister callbacks, so they need access to
+  // Context.
   sandboxir::FunctionPassManager FPM;
 
   bool runImpl(Function &F);
diff --git a/llvm/lib/SandboxIR/Context.cpp b/llvm/lib/SandboxIR/Context.cpp
index 440210f5a1bf7b..830f2832853fe5 100644
--- a/llvm/lib/SandboxIR/Context.cpp
+++ b/llvm/lib/SandboxIR/Context.cpp
@@ -611,6 +611,12 @@ Context::Context(LLVMContext &LLVMCtx)
 
 Context::~Context() {}
 
+void Context::clear() {
+  // TODO: Ideally we should clear only function-scope objects, and keep global
+  // objects, like Constants to avoid recreating them.
+  LLVMValueToValueMap.clear();
+}
+
 Module *Context::getModule(llvm::Module *LLVMM) const {
   auto It = LLVMModuleToModuleMap.find(LLVMM);
   if (It != LLVMModuleToModuleMap.end())
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
index 542fcde71e83c6..798a0ad915375b 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
@@ -88,7 +88,10 @@ bool SandboxVectorizerPass::runImpl(Function &LLVMF) {
   sandboxir::Function &F = *Ctx->createFunction(&LLVMF);
   sandboxir::Analyses A(*AA, *SE, *TTI);
   bool Change = FPM.runOnFunction(F, A);
-  // TODO: This is a function pass, so we won't be needing the function-level
-  // Sandbox IR objects in the future. So we should clear them.
+  // Given that sandboxir::Context `Ctx` is defined at a pass-level scope, the
+  // maps from LLVM IR to Sandbox IR may go stale as later passes remove LLVM IR
+  // objects. To avoid issues caused by this clear the context's state.
+  // NOTE: The alternative would be to define Ctx and FPM within runOnFunction()
+  Ctx->clear();
   return Change;
 }
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
index bbfbcc730a4cbe..104a27975cfc0c 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/CMakeLists.txt
@@ -12,6 +12,7 @@ add_llvm_unittest(SandboxVectorizerTests
   InstrMapsTest.cpp
   IntervalTest.cpp
   LegalityTest.cpp
+  SandboxVectorizerTest.cpp
   SchedulerTest.cpp
   SeedCollectorTest.cpp	
   VecUtilsTest.cpp
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerTest.cpp
new file mode 100644
index 00000000000000..e9b304618a06ca
--- /dev/null
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerTest.cpp
@@ -0,0 +1,63 @@
+//===- SandboxVectorizerTest.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
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/ScalarEvolution.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Dominators.h"
+#include "llvm/IR/PassInstrumentation.h"
+#include "llvm/SandboxIR/Function.h"
+#include "llvm/SandboxIR/Instruction.h"
+#include "llvm/Support/SourceMgr.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+struct SandboxVectorizerTest : public testing::Test {
+  LLVMContext C;
+  std::unique_ptr<Module> M;
+
+  void parseIR(LLVMContext &C, const char *IR) {
+    SMDiagnostic Err;
+    M = parseAssemblyString(IR, Err, C);
+    if (!M)
+      Err.print("SandboxVectorizerTest", errs());
+  }
+};
+
+// Check that we can run the pass on the same function more than once without
+// issues. This basically checks that Sandbox IR Context gets cleared after we
+// run the function pass.
+TEST_F(SandboxVectorizerTest, ContextCleared) {
+  parseIR(C, R"IR(
+define void @foo() {
+  ret void
+}
+)IR");
+  auto &LLVMF = *M->getFunction("foo");
+  SandboxVectorizerPass SVecPass;
+  FunctionAnalysisManager AM;
+  AM.registerPass([] { return TargetIRAnalysis(); });
+  AM.registerPass([] { return AAManager(); });
+  AM.registerPass([] { return ScalarEvolutionAnalysis(); });
+  AM.registerPass([] { return PassInstrumentationAnalysis(); });
+  AM.registerPass([] { return TargetLibraryAnalysis(); });
+  AM.registerPass([] { return AssumptionAnalysis(); });
+  AM.registerPass([] { return DominatorTreeAnalysis(); });
+  AM.registerPass([] { return LoopAnalysis(); });
+  SVecPass.run(LLVMF, AM);
+  // This shouldn't crash.
+  SVecPass.run(LLVMF, AM);
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/124842