[SandboxVec][SeedCollector] Reject non-simple memory ops for memory seeds

[Sterling-Augustine]

Load/Store isSimple is a necessary condition for VectorSeeds, but not sufficient, so reverse the condition and return value, and continue the check. Add relevant tests.

[llvmbot]

@llvm/pr-subscribers-vectorizers

@llvm/pr-subscribers-llvm-transforms

Author: None (Sterling-Augustine)

Changes

Load/Store isSimple is a necessary condition for VectorSeeds, but not sufficient, so reverse the condition and return value, and continue the check. Add relevant tests.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp (+2-2)
• (modified) llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp (+53-4)

diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
index 17544afcc185fd..6ea34c5e0598df 100644
--- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
+++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
@@ -140,8 +140,8 @@ LLVM_DUMP_METHOD void SeedContainer::dump() const { print(dbgs()); }
 #endif // NDEBUG
 
 template <typename LoadOrStoreT> static bool isValidMemSeed(LoadOrStoreT *LSI) {
-  if (LSI->isSimple())
-    return true;
+  if (!LSI->isSimple())
+    return false;
   auto *Ty = Utils::getExpectedType(LSI);
   // Omit types that are architecturally unvectorizable
   if (Ty->isX86_FP80Ty() || Ty->isPPC_FP128Ty())
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
index 95a8f66ac1e662..149be97b2b216d 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
@@ -394,12 +394,16 @@ define void @foo(ptr noalias %ptr, float %val) {
 
 TEST_F(SeedBundleTest, VectorStores) {
   parseIR(C, R"IR(
-define void @foo(ptr noalias %ptr, <2 x float> %val) {
+define void @foo(ptr noalias %ptr, <2 x float> %val0, i64 %val1) {
 bb:
   %ptr0 = getelementptr float, ptr %ptr, i32 0
   %ptr1 = getelementptr float, ptr %ptr, i32 1
-  store <2 x float> %val, ptr %ptr1
-  store <2 x float> %val, ptr %ptr0
+  %ptr2 = getelementptr i64, ptr %ptr, i32 2
+  store <2 x float> %val0, ptr %ptr1
+  store <2 x float> %val0, ptr %ptr0
+  store atomic i64 %val1, ptr %ptr2 unordered, align 8
+  store volatile i64 %val1, ptr %ptr2
+
   ret void
 }
 )IR");
@@ -418,7 +422,7 @@ define void @foo(ptr noalias %ptr, <2 x float> %val) {
   sandboxir::SeedCollector SC(&*BB, SE);
 
   // Find the stores
-  auto It = std::next(BB->begin(), 2);
+  auto It = std::next(BB->begin(), 3);
   // StX with X as the order by offset in memory
   auto *St1 = &*It++;
   auto *St0 = &*It++;
@@ -426,6 +430,8 @@ define void @foo(ptr noalias %ptr, <2 x float> %val) {
   auto StoreSeedsRange = SC.getStoreSeeds();
   EXPECT_EQ(range_size(StoreSeedsRange), 1u);
   auto &SB = *StoreSeedsRange.begin();
+  // isValidMemSeed check: The atomic and volatile stores should not
+  // be included in the bundle, but the vector stores should be.
   ExpectThatElementsAre(SB, {St0, St1});
 }
 
@@ -466,5 +472,48 @@ define void @foo(ptr noalias %ptr, float %v, <2 x float> %val) {
   auto StoreSeedsRange = SC.getStoreSeeds();
   EXPECT_EQ(range_size(StoreSeedsRange), 1u);
   auto &SB = *StoreSeedsRange.begin();
+  // isValidMemSeedCheck here: all of the three stores should be included.
   ExpectThatElementsAre(SB, {St0, St1, St3});
 }
+
+TEST_F(SeedBundleTest, VectorLoads) {
+  parseIR(C, R"IR(
+define void @foo(ptr noalias %ptr, <2 x float> %val0) {
+bb:
+  %ptr0 = getelementptr float, ptr %ptr, i32 0
+  %ptr1 = getelementptr float, ptr %ptr, i32 1
+  %r0 = load <2 x float>, ptr %ptr0
+  %r1 = load <2 x float>, ptr %ptr1
+  %r2 = load atomic i64, ptr %ptr0 unordered, align 8
+  %r3 = load volatile i64, ptr %ptr1
+
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+  DominatorTree DT(LLVMF);
+  TargetLibraryInfoImpl TLII;
+  TargetLibraryInfo TLI(TLII);
+  DataLayout DL(M->getDataLayout());
+  LoopInfo LI(DT);
+  AssumptionCache AC(LLVMF);
+  ScalarEvolution SE(LLVMF, TLI, AC, DT, LI);
+
+  sandboxir::Context Ctx(C);
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto BB = F.begin();
+  sandboxir::SeedCollector SC(&*BB, SE);
+
+  // Find the stores
+  auto It = std::next(BB->begin(), 2);
+  // StX with X as the order by offset in memory
+  auto *Ld0 = &*It++;
+  auto *Ld1 = &*It++;
+
+  auto LoadSeedsRange = SC.getLoadSeeds();
+  EXPECT_EQ(range_size(LoadSeedsRange), 1u);
+  auto &SB = *LoadSeedsRange.begin();
+  // isValidMemSeed check: The atomic and volatile stores should not
+  // be included in the bundle, but the vector stores should be.
+  ExpectThatElementsAre(SB, {Ld0, Ld1});
+}

[vporpo]

Please add a [SandboxVec][SeedCollector] or similar prefix to the title.

[slackito]

Is there any test for cases where the LoadOrStoreT is simple, but one of the later checks fails?

[Sterling-Augustine]

added

[vporpo]

A couple of comments, looks good otherwise.

[vporpo]

I would change this to auto *Ld0 = cast<sandboxir::LoadInst>(&*It++); just to make sure that the iterator is actually pointing to a load.

[Sterling-Augustine]

done

[vporpo]

same here.

[Sterling-Augustine]

done