Skip to content

Conversation

vporpo
Copy link
Contributor

@vporpo vporpo commented Aug 23, 2025

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.

@llvmbot
Copy link
Member

llvmbot commented Aug 23, 2025

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: vporpo (vporpo)

Changes

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 --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.*");
+}

…rent types

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants