[llvm] [SandboxVectorizer] Add MemSeed bundle types (PR #111584)

[via llvm-commits]

https://github.com/Sterling-Augustine created https://github.com/llvm/llvm-project/pull/111584

None

>From 5bc9ea6eeebf89c57c5e5fe41466d8bc90597593 Mon Sep 17 00:00:00 2001
From: Sterling Augustine <saugustine at google.com>
Date: Tue, 8 Oct 2024 13:43:10 -0700
Subject: [PATCH] [SandboxVectorizer] Add MemSeed bundle types

---
 .../SandboxVectorizer/SeedCollector.h         | 30 ++++++++
 .../SandboxVectorizer/SeedCollectorTest.cpp   | 73 +++++++++++++++++++
 2 files changed, 103 insertions(+)

diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
index 460e3f675fa797..5b6fff05d2137e 100644
--- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
+++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h
@@ -123,5 +123,35 @@ class SeedBundle {
   }
 #endif // NDEBUG
 };
+
+template <typename LoadOrStoreT> class MemSeedBundle : public SeedBundle {
+public:
+  explicit MemSeedBundle(SmallVector<Instruction *> &&SV, ScalarEvolution &SE)
+      : SeedBundle(std::move(SV)) {
+    assert(all_of(Seeds, [](auto *S) { return isa<LoadOrStoreT>(S); }) &&
+           "Expected Load or Store instructions!");
+    auto Cmp = [&SE](Instruction *I0, Instruction *I1) {
+      return Utils::atLowerAddress(cast<LoadOrStoreT>(I0),
+                                   cast<LoadOrStoreT>(I1), SE);
+    };
+    std::sort(Seeds.begin(), Seeds.end(), Cmp);
+  }
+  explicit MemSeedBundle(LoadOrStoreT *MemI) : SeedBundle(MemI) {
+    assert(isa<LoadOrStoreT>(MemI) && "Expected Load or Store!");
+  }
+  void insert(sandboxir::Instruction *I, ScalarEvolution &SE) {
+    assert(isa<LoadOrStoreT>(I) && "Expected a Store or a Load!");
+    auto Cmp = [&SE](Instruction *I0, Instruction *I1) {
+      return Utils::atLowerAddress(cast<LoadOrStoreT>(I0),
+                                   cast<LoadOrStoreT>(I1), SE);
+    };
+    // Find the first element after I in mem. Then insert I before it.
+    insertAt(std::upper_bound(begin(), end(), I, Cmp), I);
+  }
+};
+
+using StoreSeedBundle = MemSeedBundle<sandboxir::StoreInst>;
+using LoadSeedBundle = MemSeedBundle<sandboxir::LoadInst>;
+
 } // namespace llvm::sandboxir
 #endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_SEEDCOLLECTOR_H
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
index 36400afeaf4c59..dd41b0a6605095 100644
--- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
+++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SeedCollectorTest.cpp
@@ -7,7 +7,13 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h"
+#include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BasicAliasAnalysis.h"
+#include "llvm/Analysis/LoopInfo.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/SandboxIR/Function.h"
 #include "llvm/SandboxIR/Instruction.h"
 #include "llvm/Support/SourceMgr.h"
@@ -26,6 +32,12 @@ struct SeedBundleTest : public testing::Test {
     if (!M)
       Err.print("LegalityTest", errs());
   }
+  BasicBlock *getBasicBlockByName(Function &F, StringRef Name) {
+    for (BasicBlock &BB : F)
+      if (BB.getName() == Name)
+        return &BB;
+    llvm_unreachable("Expected to find basic block!");
+  }
 };
 
 TEST_F(SeedBundleTest, SeedBundle) {
@@ -123,3 +135,64 @@ define void @foo(float %v0, i32 %i0, i16 %i1, i8 %i2) {
                              /* ForcePowerOf2 */ true);
   EXPECT_EQ(Slice4.size(), 0u);
 }
+
+TEST_F(SeedBundleTest, MemSeedBundle) {
+  parseIR(C, R"IR(
+define void @foo(ptr %ptrA, float %val, ptr %ptr) {
+bb:
+  %gep0 = getelementptr float, ptr %ptr, i32 0
+  %gep1 = getelementptr float, ptr %ptr, i32 1
+  %gep2 = getelementptr float, ptr %ptr, i32 3
+  %gep3 = getelementptr float, ptr %ptr, i32 4
+  store float %val, ptr %gep0
+  store float %val, ptr %gep1
+  store float %val, ptr %gep2
+  store float %val, ptr %gep3
+
+  load float, ptr %gep0
+  load float, ptr %gep1
+  load float, ptr %gep2
+  load float, ptr %gep3
+
+  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();
+  auto It = std::next(BB->begin(), 4);
+  auto *S0 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S1 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S2 = cast<sandboxir::StoreInst>(&*It++);
+  auto *S3 = cast<sandboxir::StoreInst>(&*It++);
+
+  // Single instruction constructor; test insert out of memory order
+  sandboxir::StoreSeedBundle SB(S3);
+  SB.insert(S1, SE);
+  SB.insert(S2, SE);
+  SB.insert(S0, SE);
+  EXPECT_THAT(SB, testing::ElementsAre(S0, S1, S2, S3));
+
+  // Instruction list constructor; test list out of order
+  auto *L0 = cast<sandboxir::LoadInst>(&*It++);
+  auto *L1 = cast<sandboxir::LoadInst>(&*It++);
+  auto *L2 = cast<sandboxir::LoadInst>(&*It++);
+  auto *L3 = cast<sandboxir::LoadInst>(&*It++);
+  SmallVector<sandboxir::Instruction *> Loads;
+  Loads.push_back(L1);
+  Loads.push_back(L3);
+  Loads.push_back(L2);
+  Loads.push_back(L0);
+  sandboxir::LoadSeedBundle LB(std::move(Loads), SE);
+  EXPECT_THAT(LB, testing::ElementsAre(L0, L1, L2, L3));
+}