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

`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`.

Author: vporpo

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()`

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);

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())

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;
 }

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

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);
+}