[SandboxVec] Add pass to create Regions from metadata. Generalize SandboxVec pass pipelines.

llvm/include/llvm/SandboxIR/Pass.h

@@ -43,7 +43,7 @@ class Pass {
   LLVM_DUMP_METHOD virtual void dump() const;
 #endif
   /// Similar to print() but adds a newline. Used for testing.
-  void printPipeline(raw_ostream &OS) const { OS << Name << "\n"; }
+  virtual void printPipeline(raw_ostream &OS) const { OS << Name << "\n"; }
 };
 
 /// A pass that runs on a sandbox::Function.

llvm/include/llvm/SandboxIR/PassManager.h

@@ -50,40 +50,130 @@ class PassManager : public ParentPass {
   }
 
   using CreatePassFunc =
-      std::function<std::unique_ptr<ContainedPass>(StringRef)>;
+      std::function<std::unique_ptr<ContainedPass>(StringRef, 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.
+  /// Passes can have arguments, for example:
+  ///   "pass1<arg1,arg2>,pass2,pass3<arg3,arg4>"
+  ///
+  /// The arguments between angle brackets are treated as a mostly opaque string
+  /// and each pass is responsible for parsing its arguments. The exception to
+  /// this are nested angle brackets, which must match pair-wise to allow
+  /// arguments to contain nested pipelines, like:
+  ///
+  ///   "pass1<subpass1,subpass2<arg1,arg2>,subpass3>"
+  ///
+  /// An empty args string is treated the same as no args, so "pass" and
+  /// "pass<>" are equivalent.
   void setPassPipeline(StringRef Pipeline, CreatePassFunc CreatePass) {
     static constexpr const char EndToken = '\0';
+    static constexpr const char BeginArgsToken = '<';
+    static constexpr const char EndArgsToken = '>';
     static constexpr const char PassDelimToken = ',';
 
     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;
 
-    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;
+    // Accept an empty pipeline as a special case. This can be useful, for
+    // example, to test conversion to SandboxIR without running any passes on
+    // it.
+    if (Pipeline.empty())
+      return;
 
+    // Add EndToken to the end to ease parsing.
+    std::string PipelineStr = std::string(Pipeline) + EndToken;
+    Pipeline = StringRef(PipelineStr);
+
+    enum {
+      ScanName,  // reading a pass name
+      ScanArgs,  // reading a list of args
+      ArgsEnded, // read the last '>' in an args list, must read delimiter next
+    } State;
+    State = ScanName;
+    int PassBeginIdx = 0;
+    int ArgsBeginIdx;
+    StringRef PassName;
+    StringRef PassArgs;
+    int NestedArgs = 0;
+
+    auto AddPass = [this, CreatePass](StringRef PassName, StringRef PassArgs) {
+      if (PassName.empty()) {
+        errs() << "Found empty pass name.\n";
+        exit(1);
+      }
       // Get the pass that corresponds to PassName and add it to the pass
       // manager.
-      auto Pass = CreatePass(PassName);
+      auto Pass = CreatePass(PassName, PassArgs);
       if (Pass == nullptr) {
         errs() << "Pass '" << PassName << "' not registered!\n";
         exit(1);
       }
       addPass(std::move(Pass));
+    };
+    for (auto [Idx, C] : enumerate(Pipeline)) {
+      switch (State) {
+      case ScanName:
+        if (C == BeginArgsToken) {
+          // Save pass name for later and begin scanning args.
+          PassName = Pipeline.slice(PassBeginIdx, Idx);
+          ArgsBeginIdx = Idx + 1;
+          ++NestedArgs;
+          State = ScanArgs;
+          break;
+        }
+        if (C == EndArgsToken) {
+          errs() << "Unexpected '>' in pass pipeline.\n";
+          exit(1);
+        }
+        if (C == EndToken || C == PassDelimToken) {
+          // Delimiter found, add the pass (with empty args), stay in the
+          // ScanName state.
+          AddPass(Pipeline.slice(PassBeginIdx, Idx), StringRef());
+          PassBeginIdx = Idx + 1;
+        }
+        break;
+      case ScanArgs:
+        // While scanning args, we only care about making sure nesting of angle
+        // brackets is correct.
+        if (C == BeginArgsToken) {
+          ++NestedArgs;
+          break;
+        }
+        if (C == EndArgsToken) {
+          --NestedArgs;
+          if (NestedArgs == 0) {
+            // Done scanning args.
+            PassArgs = Pipeline.slice(ArgsBeginIdx, Idx);
+            AddPass(PassName, PassArgs);
+            State = ArgsEnded;
+          } else if (NestedArgs < 0) {
+            errs() << "Unbalanced '>' in pass pipeline.\n";
+            exit(1);
+          }
+          break;
+        }
+        if (C == EndToken) {
+          errs() << "Missing '>' in pass pipeline. End-of-string reached while "
+                    "reading arguments for pass '"
+                 << PassName << "'.\n";
+          exit(1);
+        }
+        break;
+      case ArgsEnded:
+        // Once we're done scanning args, only a delimiter is valid. This avoids
+        // accepting strings like "foo<args><more-args>" or "foo<args>bar".
+        if (C == EndToken || C == PassDelimToken) {
+          PassBeginIdx = Idx + 1;
+          State = ScanName;
+        } else {
+          errs() << "Expected delimiter or end-of-string after pass "
+                    "arguments.\n";
+          exit(1);
+        }
+        break;
+      }
     }
   }
 
@@ -101,7 +191,7 @@ class PassManager : public ParentPass {
   }
 #endif
   /// Similar to print() but prints one pass per line. Used for testing.
-  void printPipeline(raw_ostream &OS) const {
+  void printPipeline(raw_ostream &OS) const override {
     OS << this->getName() << "\n";
     for (const auto &PassPtr : Passes)
       PassPtr->printPipeline(OS);

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h

@@ -13,15 +13,15 @@
 #define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_BOTTOMUPVEC_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/SandboxIR/Constant.h"
 #include "llvm/SandboxIR/Pass.h"
 #include "llvm/SandboxIR/PassManager.h"
+#include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Legality.h"
 
 namespace llvm::sandboxir {
 
-class RegionPassManager;
-
 class BottomUpVec final : public FunctionPass {
   bool Change = false;
   LegalityAnalysis Legality;
@@ -32,8 +32,12 @@ class BottomUpVec final : public FunctionPass {
   RegionPassManager RPM;
 
 public:
-  BottomUpVec();
+  BottomUpVec(StringRef Pipeline);
   bool runOnFunction(Function &F) final;
+  void printPipeline(raw_ostream &OS) const final {
+    OS << getName() << "\n";
+    RPM.printPipeline(OS);
+  }
 };
 
 } // namespace llvm::sandboxir

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCount.h

@@ -0,0 +1,23 @@
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_PRINTINSTRUCTIONCOUNT_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_PRINTINSTRUCTIONCOUNT_H
+
+#include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/Region.h"
+
+namespace llvm::sandboxir {
+
+/// A Region pass that prints the instruction count for the region to stdout.
+/// Used to test -sbvec-passes while we don't have any actual optimization
+/// passes.
+class PrintInstructionCount final : public RegionPass {
+public:
+  PrintInstructionCount() : RegionPass("null") {}
+  bool runOnRegion(Region &R) final {
+    outs() << "InstructionCount: " << std::distance(R.begin(), R.end()) << "\n";
+    return false;
+  }
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_PRINTINSTRUCTIONCOUNTPASS_H

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.h

@@ -0,0 +1,34 @@
+//===- RegionsFromMetadata.h ------------------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A SandboxIR function pass that builds regions from IR metadata and then runs
+// a pipeline of region passes on them. This is useful to test region passes in
+// isolation without relying on the output of the bottom-up vectorizer.
+//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMMETADATA_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMMETADATA_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/PassManager.h"
+
+namespace llvm::sandboxir {
+
+class RegionsFromMetadata final : public FunctionPass {
+  // The PM containing the pipeline of region passes.
+  RegionPassManager RPM;
+
+public:
+  RegionsFromMetadata(StringRef Pipeline);
+  bool runOnFunction(Function &F) final;
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMMETADATA_H

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizer.h

@@ -11,7 +11,7 @@
 #include <memory>
 
 #include "llvm/IR/PassManager.h"
-#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
+#include "llvm/SandboxIR/PassManager.h"
 
 namespace llvm {
 
@@ -20,8 +20,8 @@ class TargetTransformInfo;
 class SandboxVectorizerPass : public PassInfoMixin<SandboxVectorizerPass> {
   TargetTransformInfo *TTI = nullptr;
 
-  // The main vectorizer pass.
-  sandboxir::BottomUpVec BottomUpVecPass;
+  // A pipeline of SandboxIR function passes run by the vectorizer.
+  sandboxir::FunctionPassManager FPM;
 
   bool runImpl(Function &F);
 

llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.h

@@ -0,0 +1,32 @@
+//===- SandboxVectorizerPassBuilder.h ---------------------------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Utility functions so passes with sub-pipelines can create SandboxVectorizer
+// passes without replicating the same logic in each pass.
+//
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SANDBOXVECTORIZERPASSBUILDER_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SANDBOXVECTORIZERPASSBUILDER_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/SandboxIR/Pass.h"
+
+#include <memory>
+
+namespace llvm::sandboxir {
+
+class SandboxVectorizerPassBuilder {
+public:
+  static std::unique_ptr<FunctionPass> createFunctionPass(StringRef Name,
+                                                          StringRef Args);
+  static std::unique_ptr<RegionPass> createRegionPass(StringRef Name,
+                                                      StringRef Args);
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SANDBOXVECTORIZERPASSBUILDER_H

llvm/lib/Transforms/Vectorize/CMakeLists.txt

@@ -6,7 +6,9 @@ add_llvm_component_library(LLVMVectorize
   SandboxVectorizer/DependencyGraph.cpp
   SandboxVectorizer/Interval.cpp
   SandboxVectorizer/Passes/BottomUpVec.cpp
+  SandboxVectorizer/Passes/RegionsFromMetadata.cpp
   SandboxVectorizer/SandboxVectorizer.cpp
+  SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
   SandboxVectorizer/SeedCollector.cpp
   SLPVectorizer.cpp
   Vectorize.cpp

llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp

@@ -7,42 +7,17 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
+
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/SandboxIR/Function.h"
 #include "llvm/SandboxIR/Instruction.h"
-#include "llvm/Support/CommandLine.h"
-#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.h"
 
 namespace llvm::sandboxir {
 
-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.
-static const char *DefaultPipelineMagicStr = "*";
-
-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."));
-
-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;
-}
-
-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.
-    RPM.setPassPipeline(UserDefinedPassPipeline, createRegionPass);
-  }
+BottomUpVec::BottomUpVec(StringRef Pipeline)
+    : FunctionPass("bottom-up-vec"), RPM("rpm") {
+  RPM.setPassPipeline(Pipeline, SandboxVectorizerPassBuilder::createRegionPass);
 }
 
 // TODO: This is a temporary function that returns some seeds.
@@ -82,11 +57,6 @@ 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) {

llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def

@@ -14,9 +14,19 @@
 // NOTE: NO INCLUDE GUARD DESIRED!
 
 #ifndef REGION_PASS
-#define REGION_PASS(NAME, CREATE_PASS)
+#define REGION_PASS(NAME, CLASS_NAME)
 #endif
 
-REGION_PASS("null", NullPass())
+REGION_PASS("null", ::llvm::sandboxir::NullPass)
+REGION_PASS("print-instruction-count", ::llvm::sandboxir::PrintInstructionCount)
 
 #undef REGION_PASS
+
+#ifndef FUNCTION_PASS_WITH_PARAMS
+#define FUNCTION_PASS_WITH_PARAMS(NAME, CLASS_NAME)
+#endif
+
+FUNCTION_PASS_WITH_PARAMS("bottom-up-vec", ::llvm::sandboxir::BottomUpVec)
+FUNCTION_PASS_WITH_PARAMS("regions-from-metadata", ::llvm::sandboxir::RegionsFromMetadata)
+
+#undef FUNCTION_PASS_WITH_PARAMS