[SandboxVec] Add BottomUpVec test flag to build regions from metadata.

[slackito]

This allows us to write lit tests for region passes where regions are specified by metadata in textual IR. See added test file for an example.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-vectorizers

Author: Jorge Gorbe Moya (slackito)

Changes

This allows us to write lit tests for region passes where regions are specified by metadata in textual IR. See added test file for an example.

---

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

4 Files Affected:

• (added) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCountPass.h (+24)
• (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp (+15)
• (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def (+1)
• (added) llvm/test/Transforms/SandboxVectorizer/regions-from-metadata.ll (+16)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCountPass.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCountPass.h
new file mode 100644
index 00000000000000..d8cfdcb672ccce
--- /dev/null
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCountPass.h
@@ -0,0 +1,24 @@
+#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_PRINTINSTRUCTIONCOUNTPASS_H
+#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_PASSES_PRINTINSTRUCTIONCOUNTPASS_H
+
+#include "llvm/SandboxIR/Pass.h"
+
+namespace llvm::sandboxir {
+
+class Region;
+
+/// 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 PrintInstructionCountPass final : public RegionPass {
+public:
+  PrintInstructionCountPass() : 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
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
index 77198f932a3ec0..7743f5c50ff913 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -10,8 +10,10 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/SandboxIR/Function.h"
 #include "llvm/SandboxIR/Instruction.h"
+#include "llvm/SandboxIR/Region.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/NullPass.h"
+#include "llvm/Transforms/Vectorize/SandboxVectorizer/Passes/PrintInstructionCountPass.h"
 
 namespace llvm::sandboxir {
 
@@ -19,6 +21,11 @@ static cl::opt<bool>
     PrintPassPipeline("sbvec-print-pass-pipeline", cl::init(false), cl::Hidden,
                       cl::desc("Prints the pass pipeline and returns."));
 
+static cl::opt<bool> UseRegionsFromMetadata(
+    "sbvec-use-regions-from-metadata", cl::init(false), cl::Hidden,
+    cl::desc("Skips bottom-up vectorization, builds regions from metadata "
+             "already present in the IR and runs the region pass pipeline."));
+
 /// A magic string for the default pass pipeline.
 static const char *DefaultPipelineMagicStr = "*";
 
@@ -86,6 +93,14 @@ bool BottomUpVec::runOnFunction(Function &F) {
     RPM.printPipeline(outs());
     return false;
   }
+  if (UseRegionsFromMetadata) {
+    SmallVector<std::unique_ptr<Region>> Regions =
+        Region::createRegionsFromMD(F);
+    for (auto &R : Regions) {
+      RPM.runOnRegion(*R);
+    }
+    return false;
+  }
 
   Change = false;
   // TODO: Start from innermost BBs first
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
index bbb0dcba1ea516..c0ad1710d06b91 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/PassRegistry.def
@@ -18,5 +18,6 @@
 #endif
 
 REGION_PASS("null", NullPass())
+REGION_PASS("print-instruction-count", PrintInstructionCountPass())
 
 #undef REGION_PASS
diff --git a/llvm/test/Transforms/SandboxVectorizer/regions-from-metadata.ll b/llvm/test/Transforms/SandboxVectorizer/regions-from-metadata.ll
new file mode 100644
index 00000000000000..3874876996f8c0
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/regions-from-metadata.ll
@@ -0,0 +1,16 @@
+; RUN: opt -disable-output --passes=sandbox-vectorizer \
+; RUN:    -sbvec-passes=print-instruction-count \
+; RUN:    -sbvec-use-regions-from-metadata %s | FileCheck %s
+
+define i8 @foo(i8 %v0, i8 %v1) {
+  %t0 = add i8 %v0, 1, !sandboxvec !0
+  %t1 = add i8 %t0, %v1, !sandboxvec !1
+  %t2 = add i8 %t1, %v1, !sandboxvec !1
+  ret i8 %t2
+}
+
+!0 = distinct !{!"sandboxregion"}
+!1 = distinct !{!"sandboxregion"}
+
+; CHECK: InstructionCount: 1
+; CHECK: InstructionCount: 2

[aeubanks]

I thought we were going to have a separate function pass for the purposes of testing, rather than hacking this into BoUpVec

printing instruction count is good

[slackito]

I thought we were going to have a separate function pass for the purposes of testing, rather than hacking this into BoUpVec

I wasn't feeling comfortable with the idea of adding a new pass to the overall list just to test the internals of the vectorizer.

[slackito]

Chatted offline about moving it to the main SandboxVectorizer pass rather than BottomUpVec. What do you think about this?

[aeubanks]

can this be a reference?

[tmsri]

I am not following the various discussions. Why can't this be another function pass called "TestsFuncPass" which then runs the various regions?

[slackito]

It can. I'm just having some trouble figuring out what we want.