From 5bc9ea6eeebf89c57c5e5fe41466d8bc90597593 Mon Sep 17 00:00:00 2001 From: Sterling Augustine Date: Tue, 8 Oct 2024 13:43:10 -0700 Subject: [PATCH 1/2] [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 460e3f675fa79..5b6fff05d2137 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 class MemSeedBundle : public SeedBundle { +public: + explicit MemSeedBundle(SmallVector &&SV, ScalarEvolution &SE) + : SeedBundle(std::move(SV)) { + assert(all_of(Seeds, [](auto *S) { return isa(S); }) && + "Expected Load or Store instructions!"); + auto Cmp = [&SE](Instruction *I0, Instruction *I1) { + return Utils::atLowerAddress(cast(I0), + cast(I1), SE); + }; + std::sort(Seeds.begin(), Seeds.end(), Cmp); + } + explicit MemSeedBundle(LoadOrStoreT *MemI) : SeedBundle(MemI) { + assert(isa(MemI) && "Expected Load or Store!"); + } + void insert(sandboxir::Instruction *I, ScalarEvolution &SE) { + assert(isa(I) && "Expected a Store or a Load!"); + auto Cmp = [&SE](Instruction *I0, Instruction *I1) { + return Utils::atLowerAddress(cast(I0), + cast(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; +using LoadSeedBundle = MemSeedBundle; + } // 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 36400afeaf4c5..dd41b0a660509 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(&*It++); + auto *S1 = cast(&*It++); + auto *S2 = cast(&*It++); + auto *S3 = cast(&*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(&*It++); + auto *L1 = cast(&*It++); + auto *L2 = cast(&*It++); + auto *L3 = cast(&*It++); + SmallVector 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)); +} From b06d064b2b6a147d71b4d2d644d8216be6ca698a Mon Sep 17 00:00:00 2001 From: Sterling Augustine Date: Tue, 8 Oct 2024 15:03:55 -0700 Subject: [PATCH 2/2] Add static asserts and comment --- .../Vectorize/SandboxVectorizer/SeedCollector.h | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h index 5b6fff05d2137..619c2147f2e5c 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/SeedCollector.h @@ -124,10 +124,16 @@ class SeedBundle { #endif // NDEBUG }; +/// Specialization of SeedBundle for memory access instructions. Keeps +/// seeds sorted in ascending memory order, which is convenient for slicing +/// these bundles into vectorizable groups. template class MemSeedBundle : public SeedBundle { public: explicit MemSeedBundle(SmallVector &&SV, ScalarEvolution &SE) : SeedBundle(std::move(SV)) { + static_assert(std::is_same::value || + std::is_same::value, + "Expected LoadInst or StoreInst!"); assert(all_of(Seeds, [](auto *S) { return isa(S); }) && "Expected Load or Store instructions!"); auto Cmp = [&SE](Instruction *I0, Instruction *I1) { @@ -137,6 +143,9 @@ template class MemSeedBundle : public SeedBundle { std::sort(Seeds.begin(), Seeds.end(), Cmp); } explicit MemSeedBundle(LoadOrStoreT *MemI) : SeedBundle(MemI) { + static_assert(std::is_same::value || + std::is_same::value, + "Expected LoadInst or StoreInst!"); assert(isa(MemI) && "Expected Load or Store!"); } void insert(sandboxir::Instruction *I, ScalarEvolution &SE) {