From 1c1145af0b06bfe5e4b91610adc2b7c61ae8b17c Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Fri, 4 Oct 2024 16:26:31 -0700
Subject: [PATCH 1/8] [SandboxVectorizer] Use sbvec-passes flag to create a
 pipeline of Region passes after BottomUpVec.

The main change is that the main SandboxVectorizer pass no longer has a
pipeline of function passes. Now it consists of a BottomUpVec pass and
a pipeline of region passes after it.

BottomUpVec now takes a RegionPassManager as an argument to the
constructor. This argument is currently stored and not used, but the
idea is that after BottomUpVec vectorizes a region of a function, we'll
run successive optimization passes on that region.

I've also moved creation of the pass pipeline out of the `run` method to
the SandboxVectorizer constructor.
---
 .../SandboxVectorizer/Passes/BottomUpVec.h    |  7 +-
 .../SandboxVectorizer/Passes/NullPass.h       | 19 +++++
 .../SandboxVectorizer/SandboxVectorizer.h     | 19 ++++-
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  |  2 +
 .../SandboxVectorizer/SandboxVectorizer.cpp   | 80 +++++++++++++------
 .../default_pass_pipeline.ll                  |  3 +-
 .../SandboxVectorizer/user_pass_pipeline.ll   |  8 +-
 7 files changed, 106 insertions(+), 32 deletions(-)
 create mode 100644 llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index a2108f07c28e5..e3a35b09d0e29 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -19,14 +19,19 @@
 
 namespace llvm::sandboxir {
 
+class RegionPassManager;
+
 class BottomUpVec final : public FunctionPass {
   bool Change = false;
   LegalityAnalysis Legality;
   void vectorizeRec(ArrayRef<Value *> Bndl);
   void tryVectorize(ArrayRef<Value *> Seeds);
 
+  [[maybe_unused]] RegionPassManager *RPM;
+
 public:
-  BottomUpVec() : FunctionPass("bottom-up-vec") {}
+  BottomUpVec(RegionPassManager *RPM)
+      : FunctionPass("bottom-up-vec"), RPM(RPM) {}
   bool runOnFunction(Function &F) final;
 };
 
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h
new file mode 100644
index 0000000000000..75b9f42520156
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h
@@ -0,0 +1,19 @@
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_NULLPASS_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_NULLPASS_H
+
+#include "llvm/SandboxIR/Pass.h"
+
+namespace llvm::sandboxir {
+
+class Region;
+
+/// A Region pass that does nothing, for use as a placeholder in tests.
+class NullPass final : public RegionPass {
+public:
+  NullPass() : RegionPass("null") {}
+  bool runOnRegion(Region &R) final { return false; }
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_NULLPASS_H
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
index dd9f02d327264..4547217773cd0 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
@@ -8,7 +8,11 @@
 #ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SANDBOXVECTORIZER_H
 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SANDBOXVECTORIZER_H
 
+#include <memory>
+
 #include "llvm/IR/PassManager.h"
+#include "llvm/SandboxIR/PassManager.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
 
 namespace llvm {
 
@@ -17,10 +21,21 @@ class TargetTransformInfo;
 class SandboxVectorizerPass : public PassInfoMixin<SandboxVectorizerPass> {
   TargetTransformInfo *TTI = nullptr;
 
-public:
-  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+  // Used to build a RegionPass pipeline to be run on Regions created by the
+  // bottom-up vectorization pass.
+  sandboxir::PassRegistry PR;
+
+  // The main vectorizer pass.
+  std::unique_ptr<sandboxir::BottomUpVec> BottomUpVecPass;
+
+  // The PM containing the pipeline of region passes. It's owned by the pass
+  // registry.
+  sandboxir::RegionPassManager *RPM;
 
   bool runImpl(Function &F);
+public:
+  SandboxVectorizerPass();
+  PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
 } // namespace llvm
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index c59abd09d4362..69ddfaf7784ea 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -59,6 +59,8 @@ bool BottomUpVec::runOnFunction(Function &F) {
     // TODO: Replace with proper SeedCollector function.
     auto Seeds = collectSeeds(BB);
     // TODO: Slice Seeds into smaller chunks.
+    // TODO: If vectorization succeeds, run the RegionPassManager on the
+    // resulting region.
     if (Seeds.size() >= 2)
       tryVectorize(Seeds);
   }
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
index 80afcb499a2c2..1192cc5d9bb24 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
@@ -12,6 +12,7 @@
 #include "llvm/SandboxIR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
 
 using namespace llvm;
 
@@ -30,6 +31,57 @@ cl::opt<std::string> UserDefinedPassPipeline(
     cl::desc("Comma-separated list of vectorizer passes. If not set "
              "we run the predefined pipeline."));
 
+static void registerAllRegionPasses(sandboxir::PassRegistry &PR) {
+  PR.registerPass(std::make_unique<sandboxir::NullPass>());
+}
+
+static sandboxir::RegionPassManager &
+parseAndCreatePassPipeline(sandboxir::PassRegistry &PR, StringRef Pipeline) {
+  static constexpr const char EndToken = '\0';
+  // Add EndToken to the end to ease parsing.
+  std::string PipelineStr = std::string(Pipeline) + EndToken;
+  int FlagBeginIdx = 0;
+  auto &RPM = static_cast<sandboxir::RegionPassManager &>(
+      PR.registerPass(std::make_unique<sandboxir::RegionPassManager>("rpm")));
+
+  for (auto [Idx, C] : enumerate(PipelineStr)) {
+    // Keep moving Idx until we find the end of the pass name.
+    bool FoundDelim = C == EndToken || C == PR.PassDelimToken;
+    if (!FoundDelim)
+      continue;
+    unsigned Sz = Idx - FlagBeginIdx;
+    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
+    FlagBeginIdx = Idx + 1;
+
+    // Get the pass that corresponds to PassName and add it to the pass manager.
+    auto *Pass = PR.getPassByName(PassName);
+    if (Pass == nullptr) {
+      errs() << "Pass '" << PassName << "' not registered!\n";
+      exit(1);
+    }
+    // TODO: Add a type check here. The downcast is correct as long as
+    // registerAllRegionPasses only registers regions passes.
+    RPM.addPass(static_cast<sandboxir::RegionPass *>(Pass));
+  }
+  return RPM;
+}
+
+SandboxVectorizerPass::SandboxVectorizerPass() {
+  registerAllRegionPasses(PR);
+
+  // Create a pipeline to be run on each Region created by BottomUpVec.
+  if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
+    // Create the default pass pipeline.
+    RPM = &static_cast<sandboxir::RegionPassManager &>(PR.registerPass(
+        std::make_unique<sandboxir::FunctionPassManager>("rpm")));
+    // TODO: Add passes to the default pipeline.
+  } else {
+    // Create the user-defined pipeline.
+    RPM = &parseAndCreatePassPipeline(PR, UserDefinedPassPipeline);
+  }
+  BottomUpVecPass = std::make_unique<sandboxir::BottomUpVec>(RPM);
+}
+
 PreservedAnalyses SandboxVectorizerPass::run(Function &F,
                                              FunctionAnalysisManager &AM) {
   TTI = &AM.getResult<TargetIRAnalysis>(F);
@@ -56,31 +108,13 @@ bool SandboxVectorizerPass::runImpl(Function &LLVMF) {
     return false;
   }
 
-  sandboxir::Context Ctx(LLVMF.getContext());
-  // Create SandboxIR for `LLVMF`.
-  sandboxir::Function &F = *Ctx.createFunction(&LLVMF);
-  // Create the passes and register them with the PassRegistry.
-  sandboxir::PassRegistry PR;
-  auto &BottomUpVecPass = static_cast<sandboxir::FunctionPass &>(
-      PR.registerPass(std::make_unique<sandboxir::BottomUpVec>()));
-
-  sandboxir::FunctionPassManager *PM = nullptr;
-  if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
-    // Create the default pass pipeline.
-    PM = &static_cast<sandboxir::FunctionPassManager &>(PR.registerPass(
-        std::make_unique<sandboxir::FunctionPassManager>("pm")));
-    PM->addPass(&BottomUpVecPass);
-  } else {
-    // Create the user-defined pipeline.
-    PM = &PR.parseAndCreatePassPipeline(UserDefinedPassPipeline);
-  }
-
   if (PrintPassPipeline) {
-    PM->printPipeline(outs());
+    RPM->printPipeline(outs());
     return false;
   }
 
-  // Run the pass pipeline.
-  bool Change = PM->runOnFunction(F);
-  return Change;
+  // Create SandboxIR for LLVMF and run BottomUpVec on it.
+  sandboxir::Context Ctx(LLVMF.getContext());
+  sandboxir::Function &F = *Ctx.createFunction(&LLVMF);
+  return BottomUpVecPass->runOnFunction(F);
 }
diff --git a/llvm/test/Transforms/SandboxVectorizer/default_pass_pipeline.ll b/llvm/test/Transforms/SandboxVectorizer/default_pass_pipeline.ll
index 5ccd64d9f487a..86bfbee636478 100644
--- a/llvm/test/Transforms/SandboxVectorizer/default_pass_pipeline.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/default_pass_pipeline.ll
@@ -4,8 +4,7 @@
 
 ; This checks the default pass pipeline for the sandbox vectorizer.
 define void @pipeline() {
-; CHECK: pm
-; CHECK: bottom-up-vec
+; CHECK: rpm
 ; CHECK-EMPTY:
   ret void
 }
diff --git a/llvm/test/Transforms/SandboxVectorizer/user_pass_pipeline.ll b/llvm/test/Transforms/SandboxVectorizer/user_pass_pipeline.ll
index 2879fbba1b9c0..2e6dab0aa29c7 100644
--- a/llvm/test/Transforms/SandboxVectorizer/user_pass_pipeline.ll
+++ b/llvm/test/Transforms/SandboxVectorizer/user_pass_pipeline.ll
@@ -1,12 +1,12 @@
-; RUN: opt -passes=sandbox-vectorizer -sbvec-print-pass-pipeline -sbvec-passes=bottom-up-vec,bottom-up-vec %s -disable-output | FileCheck %s
+; RUN: opt -passes=sandbox-vectorizer -sbvec-print-pass-pipeline -sbvec-passes=null,null %s -disable-output | FileCheck %s
 
 ; !!!WARNING!!! This won't get updated by update_test_checks.py !
 
 ; This checks the user defined pass pipeline.
 define void @pipeline() {
-; CHECK: pm
-; CHECK: bottom-up-vec
-; CHECK: bottom-up-vec
+; CHECK: rpm
+; CHECK: null
+; CHECK: null
 ; CHECK-EMPTY:
   ret void
 }

From 674c91eded47b264299a877e10ea94b50887e3a5 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Fri, 4 Oct 2024 17:58:44 -0700
Subject: [PATCH 2/8] clang-format

---
 .../Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h   | 1 +
 1 file changed, 1 insertion(+)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
index 4547217773cd0..81c2a5e67cbbd 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
@@ -33,6 +33,7 @@ class SandboxVectorizerPass : public PassInfoMixin<SandboxVectorizerPass> {
   sandboxir::RegionPassManager *RPM;
 
   bool runImpl(Function &F);
+
 public:
   SandboxVectorizerPass();
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);

From ee3e3d3e7754e81977428dd6db87249f45b9a6cb Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Mon, 7 Oct 2024 12:25:18 -0700
Subject: [PATCH 3/8] Moved the region pass pipeline to BottomUpVec.

Everything related to the pass pipeline is now in BottomUpVec.cpp,
including pipeline creation and printing and the flags to control it.

I have also changed the `BottomUpVecPass` in `SandboxVectorizer` from
a `unique_ptr` to a plain old member, as I don't think we need that
level of indirection for anything.
---
 .../SandboxVectorizer/Passes/BottomUpVec.h    | 12 ++-
 .../SandboxVectorizer/SandboxVectorizer.h     | 12 +--
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 76 ++++++++++++++++++-
 .../SandboxVectorizer/SandboxVectorizer.cpp   | 73 +-----------------
 4 files changed, 83 insertions(+), 90 deletions(-)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index e3a35b09d0e29..13e85f8a03323 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -15,6 +15,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/SandboxIR/Constant.h"
 #include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/PassManager.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h"
 
 namespace llvm::sandboxir {
@@ -27,11 +28,16 @@ class BottomUpVec final : public FunctionPass {
   void vectorizeRec(ArrayRef<Value *> Bndl);
   void tryVectorize(ArrayRef<Value *> Seeds);
 
-  [[maybe_unused]] RegionPassManager *RPM;
+  // Used to build a RegionPass pipeline to be run on Regions created by the
+  // bottom-up vectorization pass.
+  PassRegistry PR;
+
+  // The PM containing the pipeline of region passes. It's owned by the pass
+  // registry.
+  RegionPassManager *RPM;
 
 public:
-  BottomUpVec(RegionPassManager *RPM)
-      : FunctionPass("bottom-up-vec"), RPM(RPM) {}
+  BottomUpVec();
   bool runOnFunction(Function &F) final;
 };
 
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
index 81c2a5e67cbbd..1dcd976bf751c 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h
@@ -11,7 +11,6 @@
 #include <memory>
 
 #include "llvm/IR/PassManager.h"
-#include "llvm/SandboxIR/PassManager.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
 
 namespace llvm {
@@ -21,21 +20,12 @@ class TargetTransformInfo;
 class SandboxVectorizerPass : public PassInfoMixin<SandboxVectorizerPass> {
   TargetTransformInfo *TTI = nullptr;
 
-  // Used to build a RegionPass pipeline to be run on Regions created by the
-  // bottom-up vectorization pass.
-  sandboxir::PassRegistry PR;
-
   // The main vectorizer pass.
-  std::unique_ptr<sandboxir::BottomUpVec> BottomUpVecPass;
-
-  // The PM containing the pipeline of region passes. It's owned by the pass
-  // registry.
-  sandboxir::RegionPassManager *RPM;
+  sandboxir::BottomUpVec BottomUpVecPass;
 
   bool runImpl(Function &F);
 
 public:
-  SandboxVectorizerPass();
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 69ddfaf7784ea..a87341c59e1f7 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -10,10 +10,73 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/SandboxIR/Function.h"
 #include "llvm/SandboxIR/Instruction.h"
-
-using namespace llvm::sandboxir;
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
 
 namespace llvm::sandboxir {
+
+cl::opt<bool>
+    PrintPassPipeline("sbvec-print-pass-pipeline", cl::init(false), cl::Hidden,
+                      cl::desc("Prints the pass pipeline and returns."));
+
+/// A magic string for the default pass pipeline.
+const char *DefaultPipelineMagicStr = "*";
+
+cl::opt<std::string> UserDefinedPassPipeline(
+    "sbvec-passes", cl::init(DefaultPipelineMagicStr), cl::Hidden,
+    cl::desc("Comma-separated list of vectorizer passes. If not set "
+             "we run the predefined pipeline."));
+
+static void registerAllRegionPasses(sandboxir::PassRegistry &PR) {
+  PR.registerPass(std::make_unique<sandboxir::NullPass>());
+}
+
+static sandboxir::RegionPassManager &
+parseAndCreatePassPipeline(sandboxir::PassRegistry &PR, StringRef Pipeline) {
+  static constexpr const char EndToken = '\0';
+  // Add EndToken to the end to ease parsing.
+  std::string PipelineStr = std::string(Pipeline) + EndToken;
+  int FlagBeginIdx = 0;
+  auto &RPM = static_cast<sandboxir::RegionPassManager &>(
+      PR.registerPass(std::make_unique<sandboxir::RegionPassManager>("rpm")));
+
+  for (auto [Idx, C] : enumerate(PipelineStr)) {
+    // Keep moving Idx until we find the end of the pass name.
+    bool FoundDelim = C == EndToken || C == PR.PassDelimToken;
+    if (!FoundDelim)
+      continue;
+    unsigned Sz = Idx - FlagBeginIdx;
+    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
+    FlagBeginIdx = Idx + 1;
+
+    // Get the pass that corresponds to PassName and add it to the pass manager.
+    auto *Pass = PR.getPassByName(PassName);
+    if (Pass == nullptr) {
+      errs() << "Pass '" << PassName << "' not registered!\n";
+      exit(1);
+    }
+    // TODO: Add a type check here. The downcast is correct as long as
+    // registerAllRegionPasses only registers regions passes.
+    RPM.addPass(static_cast<sandboxir::RegionPass *>(Pass));
+  }
+  return RPM;
+}
+
+BottomUpVec::BottomUpVec() : FunctionPass("bottom-up-vec") {
+  registerAllRegionPasses(PR);
+
+  // Create a pipeline to be run on each Region created by BottomUpVec.
+  if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
+    // Create the default pass pipeline.
+    RPM = &static_cast<sandboxir::RegionPassManager &>(PR.registerPass(
+        std::make_unique<sandboxir::FunctionPassManager>("rpm")));
+    // TODO: Add passes to the default pipeline.
+  } else {
+    // Create the user-defined pipeline.
+    RPM = &parseAndCreatePassPipeline(PR, UserDefinedPassPipeline);
+  }
+}
+
 // TODO: This is a temporary function that returns some seeds.
 //       Replace this with SeedCollector's function when it lands.
 static llvm::SmallVector<Value *, 4> collectSeeds(BasicBlock &BB) {
@@ -34,8 +97,6 @@ static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
   return Operands;
 }
 
-} // namespace llvm::sandboxir
-
 void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
   auto LegalityRes = Legality.canVectorize(Bndl);
   switch (LegalityRes.getSubclassID()) {
@@ -53,6 +114,11 @@ void BottomUpVec::vectorizeRec(ArrayRef<Value *> Bndl) {
 void BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) { vectorizeRec(Bndl); }
 
 bool BottomUpVec::runOnFunction(Function &F) {
+  if (PrintPassPipeline) {
+    RPM->printPipeline(outs());
+    return false;
+  }
+
   Change = false;
   // TODO: Start from innermost BBs first
   for (auto &BB : F) {
@@ -66,3 +132,5 @@ bool BottomUpVec::runOnFunction(Function &F) {
   }
   return Change;
 }
+
+} // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
index 1192cc5d9bb24..cbaf2b6288d92 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.cpp
@@ -9,79 +9,13 @@
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/SandboxIR/Constant.h"
-#include "llvm/SandboxIR/PassManager.h"
-#include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
-#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
 
 using namespace llvm;
 
 #define SV_NAME "sandbox-vectorizer"
 #define DEBUG_TYPE SV_NAME
 
-cl::opt<bool>
-    PrintPassPipeline("sbvec-print-pass-pipeline", cl::init(false), cl::Hidden,
-                      cl::desc("Prints the pass pipeline and returns."));
-
-/// A magic string for the default pass pipeline.
-const char *DefaultPipelineMagicStr = "*";
-
-cl::opt<std::string> UserDefinedPassPipeline(
-    "sbvec-passes", cl::init(DefaultPipelineMagicStr), cl::Hidden,
-    cl::desc("Comma-separated list of vectorizer passes. If not set "
-             "we run the predefined pipeline."));
-
-static void registerAllRegionPasses(sandboxir::PassRegistry &PR) {
-  PR.registerPass(std::make_unique<sandboxir::NullPass>());
-}
-
-static sandboxir::RegionPassManager &
-parseAndCreatePassPipeline(sandboxir::PassRegistry &PR, StringRef Pipeline) {
-  static constexpr const char EndToken = '\0';
-  // Add EndToken to the end to ease parsing.
-  std::string PipelineStr = std::string(Pipeline) + EndToken;
-  int FlagBeginIdx = 0;
-  auto &RPM = static_cast<sandboxir::RegionPassManager &>(
-      PR.registerPass(std::make_unique<sandboxir::RegionPassManager>("rpm")));
-
-  for (auto [Idx, C] : enumerate(PipelineStr)) {
-    // Keep moving Idx until we find the end of the pass name.
-    bool FoundDelim = C == EndToken || C == PR.PassDelimToken;
-    if (!FoundDelim)
-      continue;
-    unsigned Sz = Idx - FlagBeginIdx;
-    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
-    FlagBeginIdx = Idx + 1;
-
-    // Get the pass that corresponds to PassName and add it to the pass manager.
-    auto *Pass = PR.getPassByName(PassName);
-    if (Pass == nullptr) {
-      errs() << "Pass '" << PassName << "' not registered!\n";
-      exit(1);
-    }
-    // TODO: Add a type check here. The downcast is correct as long as
-    // registerAllRegionPasses only registers regions passes.
-    RPM.addPass(static_cast<sandboxir::RegionPass *>(Pass));
-  }
-  return RPM;
-}
-
-SandboxVectorizerPass::SandboxVectorizerPass() {
-  registerAllRegionPasses(PR);
-
-  // Create a pipeline to be run on each Region created by BottomUpVec.
-  if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
-    // Create the default pass pipeline.
-    RPM = &static_cast<sandboxir::RegionPassManager &>(PR.registerPass(
-        std::make_unique<sandboxir::FunctionPassManager>("rpm")));
-    // TODO: Add passes to the default pipeline.
-  } else {
-    // Create the user-defined pipeline.
-    RPM = &parseAndCreatePassPipeline(PR, UserDefinedPassPipeline);
-  }
-  BottomUpVecPass = std::make_unique<sandboxir::BottomUpVec>(RPM);
-}
-
 PreservedAnalyses SandboxVectorizerPass::run(Function &F,
                                              FunctionAnalysisManager &AM) {
   TTI = &AM.getResult<TargetIRAnalysis>(F);
@@ -108,13 +42,8 @@ bool SandboxVectorizerPass::runImpl(Function &LLVMF) {
     return false;
   }
 
-  if (PrintPassPipeline) {
-    RPM->printPipeline(outs());
-    return false;
-  }
-
   // Create SandboxIR for LLVMF and run BottomUpVec on it.
   sandboxir::Context Ctx(LLVMF.getContext());
   sandboxir::Function &F = *Ctx.createFunction(&LLVMF);
-  return BottomUpVecPass->runOnFunction(F);
+  return BottomUpVecPass.runOnFunction(F);
 }

From 4b4c0630581fb5e3be4bbe9d603f7825b4e48cc2 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Mon, 7 Oct 2024 14:35:42 -0700
Subject: [PATCH 4/8] Make BottomUpVec own its RegionPassManager.

Also made some flag-related globals static.
---
 llvm/include/llvm/SandboxIR/PassManager.h     |  1 +
 .../SandboxVectorizer/Passes/BottomUpVec.h    |  7 +++--
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 26 ++++++++-----------
 3 files changed, 15 insertions(+), 19 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index 54192c6bf1333..c14fc1c6d3dba 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/llvm/include/llvm/SandboxIR/PassManager.h
@@ -36,6 +36,7 @@ class PassManager : public ParentPass {
 
   PassManager(StringRef Name) : ParentPass(Name) {}
   PassManager(const PassManager &) = delete;
+  PassManager(PassManager &&) = default;
   virtual ~PassManager() = default;
   PassManager &operator=(const PassManager &) = delete;
 
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index 13e85f8a03323..ae06a85b2b819 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -29,12 +29,11 @@ class BottomUpVec final : public FunctionPass {
   void tryVectorize(ArrayRef<Value *> Seeds);
 
   // Used to build a RegionPass pipeline to be run on Regions created by the
-  // bottom-up vectorization pass.
+  // bottom-up vectorization pass. Owns the sub-passes.
   PassRegistry PR;
 
-  // The PM containing the pipeline of region passes. It's owned by the pass
-  // registry.
-  RegionPassManager *RPM;
+  // The PM containing the pipeline of region passes.
+  RegionPassManager RPM;
 
 public:
   BottomUpVec();
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index a87341c59e1f7..5c9b7639008b7 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -15,14 +15,14 @@
 
 namespace llvm::sandboxir {
 
-cl::opt<bool>
+static cl::opt<bool>
     PrintPassPipeline("sbvec-print-pass-pipeline", cl::init(false), cl::Hidden,
                       cl::desc("Prints the pass pipeline and returns."));
 
 /// A magic string for the default pass pipeline.
-const char *DefaultPipelineMagicStr = "*";
+static const char *DefaultPipelineMagicStr = "*";
 
-cl::opt<std::string> UserDefinedPassPipeline(
+static cl::opt<std::string> UserDefinedPassPipeline(
     "sbvec-passes", cl::init(DefaultPipelineMagicStr), cl::Hidden,
     cl::desc("Comma-separated list of vectorizer passes. If not set "
              "we run the predefined pipeline."));
@@ -31,14 +31,14 @@ static void registerAllRegionPasses(sandboxir::PassRegistry &PR) {
   PR.registerPass(std::make_unique<sandboxir::NullPass>());
 }
 
-static sandboxir::RegionPassManager &
-parseAndCreatePassPipeline(sandboxir::PassRegistry &PR, StringRef Pipeline) {
+/// Adds to `RPM` a sequence of passes described by `Pipeline` from the `PR`
+/// pass registry.
+static void parseAndCreatePassPipeline(RegionPassManager &RPM, PassRegistry &PR,
+                                       StringRef Pipeline) {
   static constexpr const char EndToken = '\0';
   // Add EndToken to the end to ease parsing.
   std::string PipelineStr = std::string(Pipeline) + EndToken;
   int FlagBeginIdx = 0;
-  auto &RPM = static_cast<sandboxir::RegionPassManager &>(
-      PR.registerPass(std::make_unique<sandboxir::RegionPassManager>("rpm")));
 
   for (auto [Idx, C] : enumerate(PipelineStr)) {
     // Keep moving Idx until we find the end of the pass name.
@@ -59,21 +59,17 @@ parseAndCreatePassPipeline(sandboxir::PassRegistry &PR, StringRef Pipeline) {
     // registerAllRegionPasses only registers regions passes.
     RPM.addPass(static_cast<sandboxir::RegionPass *>(Pass));
   }
-  return RPM;
 }
 
-BottomUpVec::BottomUpVec() : FunctionPass("bottom-up-vec") {
+BottomUpVec::BottomUpVec() : FunctionPass("bottom-up-vec"), RPM("rpm") {
   registerAllRegionPasses(PR);
 
   // Create a pipeline to be run on each Region created by BottomUpVec.
   if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
-    // Create the default pass pipeline.
-    RPM = &static_cast<sandboxir::RegionPassManager &>(PR.registerPass(
-        std::make_unique<sandboxir::FunctionPassManager>("rpm")));
-    // TODO: Add passes to the default pipeline.
+    // TODO: Add default passes to RPM.
   } else {
     // Create the user-defined pipeline.
-    RPM = &parseAndCreatePassPipeline(PR, UserDefinedPassPipeline);
+    parseAndCreatePassPipeline(RPM, PR, UserDefinedPassPipeline);
   }
 }
 
@@ -115,7 +111,7 @@ void BottomUpVec::tryVectorize(ArrayRef<Value *> Bndl) { vectorizeRec(Bndl); }
 
 bool BottomUpVec::runOnFunction(Function &F) {
   if (PrintPassPipeline) {
-    RPM->printPipeline(outs());
+    RPM.printPipeline(outs());
     return false;
   }
 

From 1772712a6bdad26aa80f246e98446111ed5f763b Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Mon, 7 Oct 2024 18:17:45 -0700
Subject: [PATCH 5/8] Remove PassRegistry, use a .def file for vectorizer
 RegionPasses.

PassManagers now own the passes they contain, and pipeline parsing is
done by a standalone function in BottomUpVec.cpp.
---
 llvm/include/llvm/SandboxIR/PassManager.h     | 42 ++---------
 .../SandboxVectorizer/Passes/BottomUpVec.h    |  4 --
 llvm/lib/SandboxIR/PassManager.cpp            | 44 ++----------
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 28 ++++----
 .../SandboxVectorizer/Passes/PassRegistry.def | 20 ++++++
 llvm/unittests/SandboxIR/PassTest.cpp         | 72 ++-----------------
 6 files changed, 48 insertions(+), 162 deletions(-)
 create mode 100644 llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index c14fc1c6d3dba..d8a1be2101825 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/llvm/include/llvm/SandboxIR/PassManager.h
@@ -18,6 +18,8 @@
 #ifndef LLVM_SANDBOXIR_PASSMANAGER_H
 #define LLVM_SANDBOXIR_PASSMANAGER_H
 
+#include <memory>
+
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/SandboxIR/Pass.h"
@@ -32,7 +34,7 @@ template <typename ParentPass, typename ContainedPass>
 class PassManager : public ParentPass {
 protected:
   /// The list of passes that this pass manager will run.
-  SmallVector<ContainedPass *> Passes;
+  SmallVector<std::unique_ptr<ContainedPass>> Passes;
 
   PassManager(StringRef Name) : ParentPass(Name) {}
   PassManager(const PassManager &) = delete;
@@ -42,16 +44,16 @@ class PassManager : public ParentPass {
 
 public:
   /// Adds \p Pass to the pass pipeline.
-  void addPass(ContainedPass *Pass) {
+  void addPass(std::unique_ptr<ContainedPass> Pass) {
     // TODO: Check that Pass's class type works with this PassManager type.
-    Passes.push_back(Pass);
+    Passes.push_back(std::move(Pass));
   }
 #ifndef NDEBUG
   void print(raw_ostream &OS) const override {
     OS << this->getName();
     OS << "(";
     // TODO: This should call Pass->print(OS) because Pass may be a PM.
-    interleave(Passes, OS, [&OS](auto *Pass) { OS << Pass->getName(); }, ",");
+    interleave(Passes, OS, [&OS](auto &Pass) { OS << Pass->getName(); }, ",");
     OS << ")";
   }
   LLVM_DUMP_METHOD void dump() const override {
@@ -80,38 +82,6 @@ class RegionPassManager final : public PassManager<RegionPass, RegionPass> {
   bool runOnRegion(Region &R) final;
 };
 
-/// Owns the passes and provides an API to get a pass by its name.
-class PassRegistry {
-  SmallVector<std::unique_ptr<Pass>, 8> Passes;
-  DenseMap<StringRef, Pass *> NameToPassMap;
-
-public:
-  static constexpr const char PassDelimToken = ',';
-  PassRegistry() = default;
-  /// Registers \p PassPtr and takes ownership.
-  Pass &registerPass(std::unique_ptr<Pass> &&PassPtr) {
-    auto &PassRef = *PassPtr.get();
-    NameToPassMap[PassRef.getName()] = &PassRef;
-    Passes.push_back(std::move(PassPtr));
-    return PassRef;
-  }
-  /// \Returns the pass with name \p Name, or null if not registered.
-  Pass *getPassByName(StringRef Name) const {
-    auto It = NameToPassMap.find(Name);
-    return It != NameToPassMap.end() ? It->second : nullptr;
-  }
-  /// Creates a pass pipeline and returns the first pass manager.
-  FunctionPassManager &parseAndCreatePassPipeline(StringRef Pipeline);
-
-#ifndef NDEBUG
-  void print(raw_ostream &OS) const {
-    for (const auto &PassPtr : Passes)
-      OS << PassPtr->getName() << "\n";
-  }
-  LLVM_DUMP_METHOD void dump() const;
-#endif
-};
-
 } // namespace llvm::sandboxir
 
 #endif // LLVM_SANDBOXIR_PASSMANAGER_H
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
index ae06a85b2b819..02abdf0a1ef0d 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h
@@ -28,10 +28,6 @@ class BottomUpVec final : public FunctionPass {
   void vectorizeRec(ArrayRef<Value *> Bndl);
   void tryVectorize(ArrayRef<Value *> Seeds);
 
-  // Used to build a RegionPass pipeline to be run on Regions created by the
-  // bottom-up vectorization pass. Owns the sub-passes.
-  PassRegistry PR;
-
   // The PM containing the pipeline of region passes.
   RegionPassManager RPM;
 
diff --git a/llvm/lib/SandboxIR/PassManager.cpp b/llvm/lib/SandboxIR/PassManager.cpp
index 95bc5e56bb3ec..3a1cfa1d367a2 100644
--- a/llvm/lib/SandboxIR/PassManager.cpp
+++ b/llvm/lib/SandboxIR/PassManager.cpp
@@ -8,11 +8,11 @@
 
 #include "llvm/SandboxIR/PassManager.h"
 
-using namespace llvm::sandboxir;
+namespace llvm::sandboxir {
 
 bool FunctionPassManager::runOnFunction(Function &F) {
   bool Change = false;
-  for (FunctionPass *Pass : Passes) {
+  for (auto &Pass : Passes) {
     Change |= Pass->runOnFunction(F);
     // TODO: run the verifier.
   }
@@ -22,7 +22,7 @@ bool FunctionPassManager::runOnFunction(Function &F) {
 
 bool RegionPassManager::runOnRegion(Region &R) {
   bool Change = false;
-  for (RegionPass *Pass : Passes) {
+  for (auto &Pass : Passes) {
     Change |= Pass->runOnRegion(R);
     // TODO: run the verifier.
   }
@@ -30,40 +30,4 @@ bool RegionPassManager::runOnRegion(Region &R) {
   return Change;
 }
 
-FunctionPassManager &
-PassRegistry::parseAndCreatePassPipeline(StringRef Pipeline) {
-  static constexpr const char EndToken = '\0';
-  // Add EndToken to the end to ease parsing.
-  std::string PipelineStr = std::string(Pipeline) + EndToken;
-  int FlagBeginIdx = 0;
-  // Start with a FunctionPassManager.
-  auto &InitialPM = static_cast<FunctionPassManager &>(
-      registerPass(std::make_unique<FunctionPassManager>("init-fpm")));
-
-  for (auto [Idx, C] : enumerate(PipelineStr)) {
-    // Keep moving Idx until we find the end of the pass name.
-    bool FoundDelim = C == EndToken || C == PassDelimToken;
-    if (!FoundDelim)
-      continue;
-    unsigned Sz = Idx - FlagBeginIdx;
-    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
-    FlagBeginIdx = Idx + 1;
-
-    // Get the pass that corresponds to PassName and add it to the pass manager.
-    auto *Pass = getPassByName(PassName);
-    if (Pass == nullptr) {
-      errs() << "Pass '" << PassName << "' not registered!\n";
-      exit(1);
-    }
-    // TODO: This is safe for now, but would require proper upcasting once we
-    // add more Pass sub-classes.
-    InitialPM.addPass(static_cast<FunctionPass *>(Pass));
-  }
-  return InitialPM;
-}
-#ifndef NDEBUG
-void PassRegistry::dump() const {
-  print(dbgs());
-  dbgs() << "\n";
-}
-#endif // NDEBUG
+} // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 5c9b7639008b7..a40e8e0b23168 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -27,22 +27,26 @@ static cl::opt<std::string> UserDefinedPassPipeline(
     cl::desc("Comma-separated list of vectorizer passes. If not set "
              "we run the predefined pipeline."));
 
-static void registerAllRegionPasses(sandboxir::PassRegistry &PR) {
-  PR.registerPass(std::make_unique<sandboxir::NullPass>());
+static std::unique_ptr<RegionPass> createRegionPass(StringRef Name) {
+#define REGION_PASS(NAME, CREATE_PASS)                                         \
+  if (Name == NAME)                                                            \
+    return std::make_unique<decltype(CREATE_PASS)>(CREATE_PASS);
+#include "PassRegistry.def"
+  return nullptr;
 }
 
-/// Adds to `RPM` a sequence of passes described by `Pipeline` from the `PR`
-/// pass registry.
-static void parseAndCreatePassPipeline(RegionPassManager &RPM, PassRegistry &PR,
+/// Adds to `RPM` a sequence of region passes described by `Pipeline`.
+static void parseAndCreatePassPipeline(RegionPassManager &RPM,
                                        StringRef Pipeline) {
   static constexpr const char EndToken = '\0';
+  static constexpr const char PassDelimToken = ',';
   // Add EndToken to the end to ease parsing.
   std::string PipelineStr = std::string(Pipeline) + EndToken;
   int FlagBeginIdx = 0;
 
   for (auto [Idx, C] : enumerate(PipelineStr)) {
     // Keep moving Idx until we find the end of the pass name.
-    bool FoundDelim = C == EndToken || C == PR.PassDelimToken;
+    bool FoundDelim = C == EndToken || C == PassDelimToken;
     if (!FoundDelim)
       continue;
     unsigned Sz = Idx - FlagBeginIdx;
@@ -50,26 +54,22 @@ static void parseAndCreatePassPipeline(RegionPassManager &RPM, PassRegistry &PR,
     FlagBeginIdx = Idx + 1;
 
     // Get the pass that corresponds to PassName and add it to the pass manager.
-    auto *Pass = PR.getPassByName(PassName);
-    if (Pass == nullptr) {
+    auto RegionPass = createRegionPass(PassName);
+    if (RegionPass == nullptr) {
       errs() << "Pass '" << PassName << "' not registered!\n";
       exit(1);
     }
-    // TODO: Add a type check here. The downcast is correct as long as
-    // registerAllRegionPasses only registers regions passes.
-    RPM.addPass(static_cast<sandboxir::RegionPass *>(Pass));
+    RPM.addPass(std::move(RegionPass));
   }
 }
 
 BottomUpVec::BottomUpVec() : FunctionPass("bottom-up-vec"), RPM("rpm") {
-  registerAllRegionPasses(PR);
-
   // Create a pipeline to be run on each Region created by BottomUpVec.
   if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
     // TODO: Add default passes to RPM.
   } else {
     // Create the user-defined pipeline.
-    parseAndCreatePassPipeline(RPM, PR, UserDefinedPassPipeline);
+    parseAndCreatePassPipeline(RPM, UserDefinedPassPipeline);
   }
 }
 
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
new file mode 100644
index 0000000000000..ce5ab868187ea
--- /dev/null
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
@@ -0,0 +1,20 @@
+//===- PassRegistry.def - Registry of passes --------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is used as the registry of sub-passes that are part of the
+// SandboxVectorizer pass.
+//
+//===----------------------------------------------------------------------===//
+
+// NOTE: NO INCLUDE GUARD DESIRED!
+
+#ifndef REGION_PASS
+#define REGION_PASS(NAME, CREATE_PASS)
+#endif
+
+REGION_PASS("null", NullPass())
diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp
index b380ae9fd475a..a9f2ceecbb9bc 100644
--- a/llvm/unittests/SandboxIR/PassTest.cpp
+++ b/llvm/unittests/SandboxIR/PassTest.cpp
@@ -180,12 +180,10 @@ define void @foo() {
   };
   unsigned BBCnt1 = 0;
   unsigned BBCnt2 = 0;
-  TestPass1 TPass1(BBCnt1);
-  TestPass2 TPass2(BBCnt2);
 
   FunctionPassManager FPM("test-fpm");
-  FPM.addPass(&TPass1);
-  FPM.addPass(&TPass2);
+  FPM.addPass(std::make_unique<TestPass1>(BBCnt1));
+  FPM.addPass(std::make_unique<TestPass2>(BBCnt2));
   // Check runOnFunction().
   FPM.runOnFunction(*F);
   EXPECT_EQ(BBCnt1, 1u);
@@ -238,12 +236,10 @@ define i8 @foo(i8 %v0, i8 %v1) {
   };
   unsigned InstCount1 = 0;
   unsigned InstCount2 = 0;
-  TestPass1 TPass1(InstCount1);
-  TestPass2 TPass2(InstCount2);
 
   RegionPassManager RPM("test-rpm");
-  RPM.addPass(&TPass1);
-  RPM.addPass(&TPass2);
+  RPM.addPass(std::make_unique<TestPass1>(InstCount1));
+  RPM.addPass(std::make_unique<TestPass2>(InstCount2));
   // Check runOnRegion().
   llvm::SmallVector<std::unique_ptr<Region>> Regions =
       Region::createRegionsFromMD(*F);
@@ -259,63 +255,3 @@ define i8 @foo(i8 %v0, i8 %v1) {
   EXPECT_EQ(Buff, "test-rpm(test-pass1,test-pass2)");
 #endif // NDEBUG
 }
-
-TEST_F(PassTest, PassRegistry) {
-  class TestPass1 final : public FunctionPass {
-  public:
-    TestPass1() : FunctionPass("test-pass1") {}
-    bool runOnFunction(Function &F) final { return false; }
-  };
-  class TestPass2 final : public FunctionPass {
-  public:
-    TestPass2() : FunctionPass("test-pass2") {}
-    bool runOnFunction(Function &F) final { return false; }
-  };
-
-  PassRegistry Registry;
-  auto &TP1 = Registry.registerPass(std::make_unique<TestPass1>());
-  auto &TP2 = Registry.registerPass(std::make_unique<TestPass2>());
-
-  // Check getPassByName().
-  EXPECT_EQ(Registry.getPassByName("test-pass1"), &TP1);
-  EXPECT_EQ(Registry.getPassByName("test-pass2"), &TP2);
-
-#ifndef NDEBUG
-  // Check print().
-  std::string Buff;
-  llvm::raw_string_ostream SS(Buff);
-  Registry.print(SS);
-  EXPECT_EQ(Buff, "test-pass1\ntest-pass2\n");
-#endif // NDEBUG
-}
-
-TEST_F(PassTest, ParsePassPipeline) {
-  class TestPass1 final : public FunctionPass {
-  public:
-    TestPass1() : FunctionPass("test-pass1") {}
-    bool runOnFunction(Function &F) final { return false; }
-  };
-  class TestPass2 final : public FunctionPass {
-  public:
-    TestPass2() : FunctionPass("test-pass2") {}
-    bool runOnFunction(Function &F) final { return false; }
-  };
-
-  PassRegistry Registry;
-  Registry.registerPass(std::make_unique<TestPass1>());
-  Registry.registerPass(std::make_unique<TestPass2>());
-
-  [[maybe_unused]] auto &FPM =
-      Registry.parseAndCreatePassPipeline("test-pass1,test-pass2,test-pass1");
-#ifndef NDEBUG
-  std::string Buff;
-  llvm::raw_string_ostream SS(Buff);
-  FPM.print(SS);
-  EXPECT_EQ(Buff, "init-fpm(test-pass1,test-pass2,test-pass1)");
-#endif // NDEBUG
-
-  EXPECT_DEATH(Registry.parseAndCreatePassPipeline("bad-pass-name"),
-               ".*not registered.*");
-  EXPECT_DEATH(Registry.parseAndCreatePassPipeline(""), ".*not registered.*");
-  EXPECT_DEATH(Registry.parseAndCreatePassPipeline(","), ".*not registered.*");
-}

From fe5f7d8d6802da6cac815beb24587244920f2d09 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Tue, 8 Oct 2024 14:29:43 -0700
Subject: [PATCH 6/8] Added  cleanup in PassRegistry.def

---
 .../Vectorize/SandboxVectorizer/Passes/PassRegistry.def         | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
index ce5ab868187ea..bbb0dcba1ea51 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
@@ -18,3 +18,5 @@
 #endif
 
 REGION_PASS("null", NullPass())
+
+#undef REGION_PASS

From 196b167b22e61bf37b115cd26d455c382a569d79 Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Tue, 8 Oct 2024 16:32:35 -0700
Subject: [PATCH 7/8] Moved pipeline creation to a PassManager method.

The new method receives a CreatePass function as an argument to avoid
the PassManager class having to "know" about existing passes in, for
example, SandboxVectorizer.

Also added back a unit test similar to the now deleted PassRegistry test
that checks PassManager::setPassPipeline.
---
 llvm/include/llvm/SandboxIR/PassManager.h     | 38 +++++++++++++
 .../SandboxVectorizer/Passes/BottomUpVec.cpp  | 30 +---------
 llvm/unittests/SandboxIR/PassTest.cpp         | 55 +++++++++++++++++++
 3 files changed, 94 insertions(+), 29 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index d8a1be2101825..e5a9540af071e 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/llvm/include/llvm/SandboxIR/PassManager.h
@@ -48,6 +48,44 @@ class PassManager : public ParentPass {
     // TODO: Check that Pass's class type works with this PassManager type.
     Passes.push_back(std::move(Pass));
   }
+
+  using CreatePassFunc =
+      std::function<std::unique_ptr<ContainedPass>(StringRef)>;
+
+  /// Parses \p Pipeline as a comma-separated sequence of pass names and sets
+  /// the pass pipeline, using \p CreatePass to instantiate passes by name.
+  ///
+  /// After calling this function, the PassManager contains only the specified
+  /// pipeline, any previously added passes are cleared.
+  void setPassPipeline(StringRef Pipeline, CreatePassFunc CreatePass) {
+    static constexpr const char EndToken = '\0';
+    static constexpr const char PassDelimToken = ',';
+
+    Passes.clear();
+    // Add EndToken to the end to ease parsing.
+    std::string PipelineStr = std::string(Pipeline) + EndToken;
+    int FlagBeginIdx = 0;
+
+    for (auto [Idx, C] : enumerate(PipelineStr)) {
+      // Keep moving Idx until we find the end of the pass name.
+      bool FoundDelim = C == EndToken || C == PassDelimToken;
+      if (!FoundDelim)
+        continue;
+      unsigned Sz = Idx - FlagBeginIdx;
+      std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
+      FlagBeginIdx = Idx + 1;
+
+      // Get the pass that corresponds to PassName and add it to the pass
+      // manager.
+      auto Pass = CreatePass(PassName);
+      if (Pass == nullptr) {
+        errs() << "Pass '" << PassName << "' not registered!\n";
+        exit(1);
+      }
+      addPass(std::move(Pass));
+    }
+  }
+
 #ifndef NDEBUG
   void print(raw_ostream &OS) const override {
     OS << this->getName();
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index a40e8e0b23168..77198f932a3ec 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -35,41 +35,13 @@ static std::unique_ptr<RegionPass> createRegionPass(StringRef Name) {
   return nullptr;
 }
 
-/// Adds to `RPM` a sequence of region passes described by `Pipeline`.
-static void parseAndCreatePassPipeline(RegionPassManager &RPM,
-                                       StringRef Pipeline) {
-  static constexpr const char EndToken = '\0';
-  static constexpr const char PassDelimToken = ',';
-  // Add EndToken to the end to ease parsing.
-  std::string PipelineStr = std::string(Pipeline) + EndToken;
-  int FlagBeginIdx = 0;
-
-  for (auto [Idx, C] : enumerate(PipelineStr)) {
-    // Keep moving Idx until we find the end of the pass name.
-    bool FoundDelim = C == EndToken || C == PassDelimToken;
-    if (!FoundDelim)
-      continue;
-    unsigned Sz = Idx - FlagBeginIdx;
-    std::string PassName(&PipelineStr[FlagBeginIdx], Sz);
-    FlagBeginIdx = Idx + 1;
-
-    // Get the pass that corresponds to PassName and add it to the pass manager.
-    auto RegionPass = createRegionPass(PassName);
-    if (RegionPass == nullptr) {
-      errs() << "Pass '" << PassName << "' not registered!\n";
-      exit(1);
-    }
-    RPM.addPass(std::move(RegionPass));
-  }
-}
-
 BottomUpVec::BottomUpVec() : FunctionPass("bottom-up-vec"), RPM("rpm") {
   // Create a pipeline to be run on each Region created by BottomUpVec.
   if (UserDefinedPassPipeline == DefaultPipelineMagicStr) {
     // TODO: Add default passes to RPM.
   } else {
     // Create the user-defined pipeline.
-    parseAndCreatePassPipeline(RPM, UserDefinedPassPipeline);
+    RPM.setPassPipeline(UserDefinedPassPipeline, createRegionPass);
   }
 }
 
diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp
index a9f2ceecbb9bc..979a8e1ad3053 100644
--- a/llvm/unittests/SandboxIR/PassTest.cpp
+++ b/llvm/unittests/SandboxIR/PassTest.cpp
@@ -255,3 +255,58 @@ define i8 @foo(i8 %v0, i8 %v1) {
   EXPECT_EQ(Buff, "test-rpm(test-pass1,test-pass2)");
 #endif // NDEBUG
 }
+
+TEST_F(PassTest, SetPassPipeline) {
+  auto *F = parseFunction(R"IR(
+define void @f() {
+  ret void
+}
+)IR",
+                          "f");
+  class FooPass final : public FunctionPass {
+    std::string &Str;
+
+  public:
+    FooPass(std::string &Str) : FunctionPass("foo-pass"), Str(Str) {}
+    bool runOnFunction(Function &F) final {
+      Str += "foo";
+      return false;
+    }
+  };
+  class BarPass final : public FunctionPass {
+    std::string &Str;
+
+  public:
+    BarPass(std::string &Str) : FunctionPass("bar-pass"), Str(Str) {}
+    bool runOnFunction(Function &F) final {
+      Str += "bar";
+      return false;
+    }
+  };
+
+  std::string Str;
+  auto CreatePass =
+      [&Str](llvm::StringRef Name) -> std::unique_ptr<FunctionPass> {
+    if (Name == "foo")
+      return std::make_unique<FooPass>(Str);
+    if (Name == "bar")
+      return std::make_unique<BarPass>(Str);
+    return nullptr;
+  };
+
+  FunctionPassManager FPM("test-fpm");
+  FPM.setPassPipeline("foo,bar,foo", CreatePass);
+  FPM.runOnFunction(*F);
+  EXPECT_EQ(Str, "foobarfoo");
+
+  // Check that a second call to setPassPipeline clears the previous pipeline.
+  Str.clear();
+  FPM.setPassPipeline("bar,bar,foo", CreatePass);
+  FPM.runOnFunction(*F);
+  EXPECT_EQ(Str, "barbarfoo");
+
+  EXPECT_DEATH(FPM.setPassPipeline("bad-pass-name", CreatePass),
+               ".*not registered.*");
+  EXPECT_DEATH(FPM.setPassPipeline("", CreatePass), ".*not registered.*");
+  EXPECT_DEATH(FPM.setPassPipeline(",", CreatePass), ".*not registered.*");
+}

From 314bf8f549c890da981e72f0190787dc8a0b176e Mon Sep 17 00:00:00 2001
From: Jorge Gorbe Moya <jgorbe@google.com>
Date: Wed, 9 Oct 2024 10:13:12 -0700
Subject: [PATCH 8/8] Assert on multiple calls to setPassPipeline. Update unit
 test.

---
 llvm/include/llvm/SandboxIR/PassManager.h |  3 ++-
 llvm/unittests/SandboxIR/PassTest.cpp     | 19 +++++++++++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/include/llvm/SandboxIR/PassManager.h b/llvm/include/llvm/SandboxIR/PassManager.h
index e5a9540af071e..247c43615f576 100644
--- a/llvm/include/llvm/SandboxIR/PassManager.h
+++ b/llvm/include/llvm/SandboxIR/PassManager.h
@@ -61,7 +61,8 @@ class PassManager : public ParentPass {
     static constexpr const char EndToken = '\0';
     static constexpr const char PassDelimToken = ',';
 
-    Passes.clear();
+    assert(Passes.empty() &&
+           "setPassPipeline called on a non-empty sandboxir::PassManager");
     // Add EndToken to the end to ease parsing.
     std::string PipelineStr = std::string(Pipeline) + EndToken;
     int FlagBeginIdx = 0;
diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp
index 979a8e1ad3053..ae7284ecf2deb 100644
--- a/llvm/unittests/SandboxIR/PassTest.cpp
+++ b/llvm/unittests/SandboxIR/PassTest.cpp
@@ -299,14 +299,17 @@ define void @f() {
   FPM.runOnFunction(*F);
   EXPECT_EQ(Str, "foobarfoo");
 
-  // Check that a second call to setPassPipeline clears the previous pipeline.
-  Str.clear();
-  FPM.setPassPipeline("bar,bar,foo", CreatePass);
-  FPM.runOnFunction(*F);
-  EXPECT_EQ(Str, "barbarfoo");
+  // A second call to setPassPipeline will trigger an assertion in debug mode.
+#ifndef NDEBUG
+  EXPECT_DEATH(FPM.setPassPipeline("bar,bar,foo", CreatePass),
+               "setPassPipeline called on a non-empty sandboxir::PassManager");
+#endif
 
-  EXPECT_DEATH(FPM.setPassPipeline("bad-pass-name", CreatePass),
+  // Fresh PM for the death tests so they die from bad pipeline strings, rather
+  // than from multiple setPassPipeline calls.
+  FunctionPassManager FPM2("test-fpm");
+  EXPECT_DEATH(FPM2.setPassPipeline("bad-pass-name", CreatePass),
                ".*not registered.*");
-  EXPECT_DEATH(FPM.setPassPipeline("", CreatePass), ".*not registered.*");
-  EXPECT_DEATH(FPM.setPassPipeline(",", CreatePass), ".*not registered.*");
+  EXPECT_DEATH(FPM2.setPassPipeline("", CreatePass), ".*not registered.*");
+  EXPECT_DEATH(FPM2.setPassPipeline(",", CreatePass), ".*not registered.*");
 }