[llvm] c0ee8e2 - [SandboxVec][SeedCollector] Reject non-simple memory ops for memory seeds (#116891)

[via llvm-commits]

Author: Sterling-Augustine
Date: 2024-11-20T11:53:41-08:00
New Revision: c0ee8e22f4093ea1fda42cc037d50cb4619e1445

URL: https://github.com/llvm/llvm-project/commit/c0ee8e22f4093ea1fda42cc037d50cb4619e1445
DIFF: https://github.com/llvm/llvm-project/commit/c0ee8e22f4093ea1fda42cc037d50cb4619e1445.diff

LOG: [SandboxVec][SeedCollector] Reject non-simple memory ops for memory seeds (#116891)

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.

Added: 
    

Modified: 
    llvm/lib/Transforms/Vectorize/SandboxVectorizer/SeedCollector.cpp
    llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp

Removed: 
    


################################################################################
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..99a13801c7c335 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,50 @@ 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
+  %r4 = load void()*, 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 loads
+  auto It = std::next(BB->begin(), 2);
+  // StX with X as the order by offset in memory
+  auto *Ld0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *Ld1 = cast<sandboxir::LoadInst>(&*It++);
+
+  auto LoadSeedsRange = SC.getLoadSeeds();
+  EXPECT_EQ(range_size(LoadSeedsRange), 2u);
+  auto &SB = *LoadSeedsRange.begin();
+  // isValidMemSeed check: The atomic and volatile loads should not
+  // be included in the bundle, the vector stores should be, but the
+  // void-typed load should not.
+  ExpectThatElementsAre(SB, {Ld0, Ld1});
+}