[SandboxVec] Add region-from-bbs helper pass (#130153)

RegionFromBBs is a helper Sandbox IR function pass that builds a region
for each BB. This helps test region passes in isolation without relying
on other parts of the vectorizer, which is especially useful for
stress-testing.

7 files changed, 108 insertions, 0 deletions

diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index 1a05194..14f35c9c 100644
--- a/llvm/include/llvm/SandboxIR/Region.h
+++ b/llvm/include/llvm/SandboxIR/Region.h
@@ -120,6 +120,7 @@ class Region {
   void remove(Instruction *I);
   friend class Context; // The callbacks need to call add() and remove().
   friend class RegionInternalsAttorney; // For unit tests.
+  friend class RegionsFromBBs;          // For add().
 
   /// Set \p I as the \p Idx'th element in the auxiliary vector.
   /// NOTE: This is for internal use, it does not set the metadata.
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.h
new file mode 100644
index 0000000..66c9cd4
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.h
@@ -0,0 +1,38 @@
+//===- RegionsFromBBs.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 one region per BB 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 other vectorizer components.
+//
+
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMBBS_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMBBS_H
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/SandboxIR/Pass.h"
+#include "llvm/SandboxIR/PassManager.h"
+
+namespace llvm::sandboxir {
+
+class RegionsFromBBs final : public FunctionPass {
+  // The PM containing the pipeline of region passes.
+  RegionPassManager RPM;
+
+public:
+  RegionsFromBBs(StringRef Pipeline);
+  bool runOnFunction(Function &F, const Analyses &A) final;
+  void printPipeline(raw_ostream &OS) const final {
+    OS << getName() << "\n";
+    RPM.printPipeline(OS);
+  }
+};
+
+} // namespace llvm::sandboxir
+
+#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_REGIONSFROMBBS_H
diff --git a/llvm/lib/Transforms/Vectorize/CMakeLists.txt b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
index 38670ba..f8d08b9 100644
--- a/llvm/lib/Transforms/Vectorize/CMakeLists.txt
+++ b/llvm/lib/Transforms/Vectorize/CMakeLists.txt
@@ -8,6 +8,7 @@ add_llvm_component_library(LLVMVectorize
   SandboxVectorizer/Interval.cpp
   SandboxVectorizer/Legality.cpp
   SandboxVectorizer/Passes/BottomUpVec.cpp
+  SandboxVectorizer/Passes/RegionsFromBBs.cpp
   SandboxVectorizer/Passes/RegionsFromMetadata.cpp
   SandboxVectorizer/Passes/SeedCollection.cpp
   SandboxVectorizer/Passes/TransactionAcceptOrRevert.cpp
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
index f745073..c525608 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
@@ -31,6 +31,7 @@ REGION_PASS("bottom-up-vec", ::llvm::sandboxir::BottomUpVec)
 #endif
 
 FUNCTION_PASS_WITH_PARAMS("seed-collection", ::llvm::sandboxir::SeedCollection)
+FUNCTION_PASS_WITH_PARAMS("regions-from-bbs", ::llvm::sandboxir::RegionsFromBBs)
 FUNCTION_PASS_WITH_PARAMS("regions-from-metadata", ::llvm::sandboxir::RegionsFromMetadata)
 
 #undef FUNCTION_PASS_WITH_PARAMS
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.cpp
new file mode 100644
index 0000000..6801524
--- /dev/null
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.cpp
@@ -0,0 +1,35 @@
+//===- RegionsFromBBs.cpp - A helper to test RegionPasses -----------------===//
+//
+// 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/Passes/RegionsFromBBs.h"
+#include "llvm/SandboxIR/Function.h"
+#include "llvm/SandboxIR/Region.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.h"
+
+namespace llvm::sandboxir {
+
+RegionsFromBBs::RegionsFromBBs(StringRef Pipeline)
+    : FunctionPass("regions-from-bbs"),
+      RPM("rpm", Pipeline, SandboxVectorizerPassBuilder::createRegionPass) {}
+
+bool RegionsFromBBs::runOnFunction(Function &F, const Analyses &A) {
+  SmallVector<std::unique_ptr<Region>, 16> Regions;
+  // Create a region for each BB.
+  for (BasicBlock &BB : F) {
+    Regions.push_back(std::make_unique<Region>(F.getContext(), A.getTTI()));
+    auto &RgnPtr = Regions.back();
+    for (Instruction &I : BB)
+      RgnPtr->add(&I);
+  }
+  // For each region run the region pass pipeline.
+  for (auto &RgnPtr : Regions)
+    RPM.runOnRegion(*RgnPtr, A);
+  return false;
+}
+
+} // namespace llvm::sandboxir
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
index 389f9cc..013ccf6 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SandboxVectorizerPassBuilder.cpp
@@ -3,6 +3,7 @@
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCount.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromBBs.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/TransactionAcceptOrRevert.h"
diff --git a/llvm/test/Transforms/SandboxVectorizer/regions-from-bbs.ll b/llvm/test/Transforms/SandboxVectorizer/regions-from-bbs.ll
new file mode 100644
index 0000000..065baad
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/regions-from-bbs.ll
@@ -0,0 +1,31 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt --passes=sandbox-vectorizer -sbvec-passes='regions-from-bbs<null>' %s -S | FileCheck %s
+
+define void @foo(i8 %v) {
+; CHECK-LABEL: define void @foo(
+; CHECK-SAME: i8 [[V:%.*]]) {
+; CHECK-NEXT:  [[BBA:.*:]]
+; CHECK-NEXT:    [[ADD0:%.*]] = add i8 [[V]], 0, !sandboxvec [[META0:![0-9]+]]
+; CHECK-NEXT:    br label %[[BBB:.*]], !sandboxvec [[META0]]
+; CHECK:       [[BBB]]:
+; CHECK-NEXT:    [[ADD1:%.*]] = add i8 [[V]], 1, !sandboxvec [[META1:![0-9]+]]
+; CHECK-NEXT:    br label %[[BBC:.*]], !sandboxvec [[META1]]
+; CHECK:       [[BBC]]:
+; CHECK-NEXT:    ret void, !sandboxvec [[META2:![0-9]+]]
+;
+bbA:
+  %add0 = add i8 %v, 0
+  br label %bbB
+
+bbB:
+  %add1 = add i8 %v, 1
+  br label %bbC
+
+bbC:
+  ret void
+}
+;.
+; CHECK: [[META0]] = distinct !{!"sandboxregion"}
+; CHECK: [[META1]] = distinct !{!"sandboxregion"}
+; CHECK: [[META2]] = distinct !{!"sandboxregion"}
+;.