[llvm] [SandboxVec][SeedCollection] Support pass argument for enabling different types (PR #155079)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

<details>
<summary>Changes</summary>

With this patch you can call the SeedCollection pass with the argument "enable-diff-types" to get it to collect seeds of different types. This will be used in a follow-up patch.

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


5 Files Affected:

- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h (+5) 
- (modified) llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h (+29) 
- (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp (+14-2) 
- (added) llvm/test/Transforms/SandboxVectorizer/seed_collection.ll (+19) 
- (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp (+27) 


``````````diff
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h
index 286d971ff4851..f55dbf753dbcc 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.h
@@ -25,6 +25,11 @@ class SeedCollection final : public FunctionPass {
 
   /// The PM containing the pipeline of region passes.
   RegionPassManager RPM;
+  /// The auxiliary argument passed to the pass that tells us that we should
+  /// collect seeds of different types.
+  static constexpr StringRef DiffTypesArgStr = "enable-diff-types";
+  /// Collect seeds of different types.
+  bool AllowDiffTypes = false;
 
 public:
   SeedCollection(StringRef Pipeline);
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h
index d32bfbaf7a4c8..76752872f4d5c 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h
@@ -262,6 +262,35 @@ class VecUtils {
     return Pack;
   }
 
+  static constexpr const char AuxArgBeginToken = '(';
+  static constexpr const char AuxArgEndToken = ')';
+  /// \Returns the auxiliary pass argument by parsing \p Args. The auxiliary
+  /// pass argument is enclosed in parentheses and should be before any other
+  /// argument.
+  static StringRef getAuxPassArg(StringRef Args) {
+    if (Args.empty())
+      return Args;
+    if (Args[0] != AuxArgBeginToken)
+      return StringRef();
+    // We found the Begin token, so look for the End token.
+    size_t EndIdx = Args.find(AuxArgEndToken);
+    if (EndIdx == StringRef::npos) {
+      errs() << "Missing '" << AuxArgEndToken << "' in '" << Args << "' !\n";
+      exit(1);
+    }
+    assert(EndIdx >= 1 && "Expected at index 1 or later!");
+    return Args.substr(1, EndIdx - 1);
+  }
+
+  /// \Returns \p Args with the auxiliary argument removed.
+  /// For example: "(foo)bar1,bar2" returns "bar1,bar2".
+  static StringRef stripAuxPassArg(StringRef Args) {
+    StringRef Aux = getAuxPassArg(Args);
+    if (Aux.empty())
+      return Args;
+    return Args.slice(Aux.size() + 2, StringRef::npos);
+  }
+
 #ifndef NDEBUG
   /// Helper dump function for debugging.
   LLVM_DUMP_METHOD static void dump(ArrayRef<Value *> Bndl);
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
index ddb4a1e154a18..e87e6426d3be6 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/SeedCollection.cpp
@@ -32,9 +32,20 @@ cl::opt<std::string> CollectSeeds(
              "list of '" StoreSeedsDef "' and '" LoadSeedsDef "'."));
 
 namespace sandboxir {
+
 SeedCollection::SeedCollection(StringRef Pipeline)
     : FunctionPass("seed-collection"),
-      RPM("rpm", Pipeline, SandboxVectorizerPassBuilder::createRegionPass) {}
+      RPM("rpm", VecUtils::stripAuxPassArg(Pipeline),
+          SandboxVectorizerPassBuilder::createRegionPass) {
+  StringRef AuxArg = VecUtils::getAuxPassArg(Pipeline);
+  if (!AuxArg.empty()) {
+    if (DiffTypesArgStr != DiffTypesArgStr) {
+      errs() << "SeedCollection only supports '" << DiffTypesArgStr << "' !\n";
+      exit(1);
+    }
+    AllowDiffTypes = true;
+  }
+}
 
 bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
   bool Change = false;
@@ -50,7 +61,8 @@ bool SeedCollection::runOnFunction(Function &F, const Analyses &A) {
 
   // TODO: Start from innermost BBs first
   for (auto &BB : F) {
-    SeedCollector SC(&BB, A.getScalarEvolution(), CollectStores, CollectLoads);
+    SeedCollector SC(&BB, A.getScalarEvolution(), CollectStores, CollectLoads,
+                     AllowDiffTypes);
     for (SeedBundle &Seeds : SC.getStoreSeeds()) {
       unsigned ElmBits =
           Utils::getNumBits(VecUtils::getElementType(Utils::getExpectedType(
diff --git a/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll b/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll
new file mode 100644
index 0000000000000..0fc1fa04a081f
--- /dev/null
+++ b/llvm/test/Transforms/SandboxVectorizer/seed_collection.ll
@@ -0,0 +1,19 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=sandbox-vectorizer -sbvec-vec-reg-bits=1024 -disable-output -sbvec-passes="seed-collection<(enable-diff-types)print-region>" %s  | FileCheck %s
+
+define void @collect_seeds_of_diff_types(ptr %ptr) {
+; CHECK: Aux:
+; CHECK-NEXT:  store i8 %ld0, ptr %ptr0, {{.*}}
+; CHECK-NEXT:  store i16 %ld1, ptr %ptr1, {{.*}}
+; CHECK-NEXT:  store i8 %ld2, ptr %ptr2, {{.*}}
+  %ptr0 = getelementptr i8, ptr %ptr, i32 0
+  %ptr1 = getelementptr i16, ptr %ptr, i32 1
+  %ptr2 = getelementptr i8, ptr %ptr, i32 3
+  %ld0 = load i8, ptr %ptr0
+  %ld1 = load i16, ptr %ptr1
+  %ld2 = load i8, ptr %ptr2
+  store i8 %ld0, ptr %ptr0
+  store i16 %ld1, ptr %ptr1
+  store i8 %ld2, ptr %ptr2
+  ret void
+}
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp
index 2bfea6908305c..566c31446ec20 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp
@@ -617,3 +617,30 @@ define void @foo(i8 %v0, i8 %v1) {
       EXPECT_FALSE(sandboxir::VecUtils::matchPack(NotPack));
   }
 }
+
+TEST_F(VecUtilsTest, GetAuxPassArg) {
+  // Check no aux argument.
+  EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("no aux arg"), "");
+  // Check an illegal aux argument.
+  EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("illegal (arg) other"), "");
+  // Check a valid argument.
+  EXPECT_EQ(sandboxir::VecUtils::getAuxPassArg("(some arg)other stuff"),
+            "some arg");
+  // Missing token.
+  EXPECT_DEBUG_DEATH(sandboxir::VecUtils::getAuxPassArg("(arg other"),
+                     "Missing.*");
+}
+
+TEST_F(VecUtilsTest, StripAuxPassArg) {
+  // Check no aux argument.
+  EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("no aux arg"), "no aux arg");
+  // Check a legal aux argument.
+  EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("(legal aux arg)foo bar"),
+            "foo bar");
+  // Check an illegal aux argument.
+  EXPECT_EQ(sandboxir::VecUtils::stripAuxPassArg("illegal (arg) other"),
+            "illegal (arg) other");
+  // Missing token.
+  EXPECT_DEBUG_DEATH(sandboxir::VecUtils::stripAuxPassArg("(arg other"),
+                     "Missing.*");
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/155079