diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h index 6b56f348f328c..c8ca65a155704 100644 --- a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h +++ b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h @@ -130,6 +130,10 @@ class SchedBundle { N->clearSchedBundle(); } bool empty() const { return Nodes.empty(); } + /// Singleton bundles are created when scheduling instructions temporarily to + /// fill in the schedule until we schedule the vector bundle. These are + /// non-vector bundles containing just a single instruction. + bool isSingleton() const { return Nodes.size() == 1u; } DGNode *back() const { return Nodes.back(); } using iterator = ContainerTy::iterator; using const_iterator = ContainerTy::const_iterator; @@ -187,10 +191,12 @@ class Scheduler { /// The scheduling state of the instructions in the bundle. enum class BndlSchedState { NoneScheduled, ///> No instruction in the bundle was previously scheduled. - PartiallyOrDifferentlyScheduled, ///> Only some of the instrs in the bundle - /// were previously scheduled, or all of - /// them were but not in the same - /// SchedBundle. + AlreadyScheduled, ///> At least one instruction in the bundle belongs to a + /// different non-singleton scheduling bundle. + TemporarilyScheduled, ///> Instructions were temporarily scheduled as + /// singleton bundles or some of them were not + /// scheduled at all. None of them were in a vector + ///(non-singleton) bundle. FullyScheduled, ///> All instrs in the bundle were previously scheduled and /// were in the same SchedBundle. }; @@ -243,6 +249,11 @@ class Scheduler { class SchedulerInternalsAttorney { public: static DependencyGraph &getDAG(Scheduler &Sched) { return Sched.DAG; } + using BndlSchedState = Scheduler::BndlSchedState; + static BndlSchedState getBndlSchedState(const Scheduler &Sched, + ArrayRef Instrs) { + return Sched.getBndlSchedState(Instrs); + } }; } // namespace llvm::sandboxir diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp index 3e37e07aabc5c..ad46683d95063 100644 --- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp +++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Scheduler.cpp @@ -161,30 +161,36 @@ bool Scheduler::tryScheduleUntil(ArrayRef Instrs) { Scheduler::BndlSchedState Scheduler::getBndlSchedState(ArrayRef Instrs) const { assert(!Instrs.empty() && "Expected non-empty bundle"); - bool PartiallyScheduled = false; - bool FullyScheduled = true; - for (auto *I : Instrs) { + auto *N0 = DAG.getNode(Instrs[0]); + auto *SB0 = N0 != nullptr ? N0->getSchedBundle() : nullptr; + bool AllUnscheduled = SB0 == nullptr; + bool FullyScheduled = SB0 != nullptr && !SB0->isSingleton(); + for (auto *I : drop_begin(Instrs)) { auto *N = DAG.getNode(I); - if (N != nullptr && N->scheduled()) - PartiallyScheduled = true; - else - FullyScheduled = false; - } - if (FullyScheduled) { - // If not all instrs in the bundle are in the same SchedBundle then this - // should be considered as partially-scheduled, because we will need to - // re-schedule. - SchedBundle *SB = DAG.getNode(Instrs[0])->getSchedBundle(); - assert(SB != nullptr && "FullyScheduled assumes that there is an SB!"); - if (any_of(drop_begin(Instrs), [this, SB](sandboxir::Value *SBV) { - return DAG.getNode(cast(SBV)) - ->getSchedBundle() != SB; - })) + auto *SB = N != nullptr ? N->getSchedBundle() : nullptr; + if (SB != nullptr) { + // We found a scheduled instr, so there is now way all are unscheduled. + AllUnscheduled = false; + if (SB->isSingleton()) { + // We found an instruction in a temporarily scheduled singleton. There + // is no way that all instructions are scheduled in the same bundle. + FullyScheduled = false; + } + } + + if (SB != SB0) { + // Either one of SB, SB0 is null, or they are in different bundles, so + // Instrs are definitely not in the same vector bundle. FullyScheduled = false; + // One of SB, SB0 are in a vector bundle and they differ. + if ((SB != nullptr && !SB->isSingleton()) || + (SB0 != nullptr && !SB0->isSingleton())) + return BndlSchedState::AlreadyScheduled; + } } - return FullyScheduled ? BndlSchedState::FullyScheduled - : PartiallyScheduled ? BndlSchedState::PartiallyOrDifferentlyScheduled - : BndlSchedState::NoneScheduled; + return AllUnscheduled ? BndlSchedState::NoneScheduled + : FullyScheduled ? BndlSchedState::FullyScheduled + : BndlSchedState::TemporarilyScheduled; } void Scheduler::trimSchedule(ArrayRef Instrs) { @@ -203,13 +209,14 @@ void Scheduler::trimSchedule(ArrayRef Instrs) { // Instruction *TopI = &*ScheduleTopItOpt.value(); Instruction *LowestI = VecUtils::getLowest(Instrs); - // Destroy the schedule bundles from LowestI all the way to the top. + // Destroy the singleton schedule bundles from LowestI all the way to the top. for (auto *I = LowestI, *E = TopI->getPrevNode(); I != E; I = I->getPrevNode()) { auto *N = DAG.getNode(I); if (N == nullptr) continue; - if (auto *SB = N->getSchedBundle()) + auto *SB = N->getSchedBundle(); + if (SB->isSingleton()) eraseBundle(SB); } // The DAG Nodes contain state like the number of UnscheduledSuccs and the @@ -259,7 +266,12 @@ bool Scheduler::trySchedule(ArrayRef Instrs) { case BndlSchedState::FullyScheduled: // Nothing to do. return true; - case BndlSchedState::PartiallyOrDifferentlyScheduled: + case BndlSchedState::AlreadyScheduled: + // Instructions are part of a different vector schedule, so we can't + // schedule \p Instrs in the same bundle (without destroying the existing + // schedule). + return false; + case BndlSchedState::TemporarilyScheduled: // If one or more instrs are already scheduled we need to destroy the // top-most part of the schedule that includes the instrs in the bundle and // re-schedule. diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll index fc5795708c7d8..531ed8cb618fc 100644 --- a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll +++ b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll @@ -386,3 +386,33 @@ define void @vecInstrsPlacement(ptr %ptr0) { store double %add1, ptr %ptr1 ret void } + +; During the bottom-up traversal we form bundle {ldA0,ldA1} but later when we +; visit the RHS operands of the additions we try to form {ldA1,ldA2} +; which is not allowed. +define void @instrsInMultipleBundles(ptr noalias %ptr) { +; CHECK-LABEL: define void @instrsInMultipleBundles( +; CHECK-SAME: ptr noalias [[PTR:%.*]]) { +; CHECK-NEXT: [[GEP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 0 +; CHECK-NEXT: [[GEP2:%.*]] = getelementptr i8, ptr [[PTR]], i64 2 +; CHECK-NEXT: [[LDA2:%.*]] = load i8, ptr [[GEP2]], align 1 +; CHECK-NEXT: [[VECL:%.*]] = load <2 x i8>, ptr [[GEP0]], align 1 +; CHECK-NEXT: [[VEXT:%.*]] = extractelement <2 x i8> [[VECL]], i32 1 +; CHECK-NEXT: [[VINS:%.*]] = insertelement <2 x i8> poison, i8 [[VEXT]], i32 0 +; CHECK-NEXT: [[VINS1:%.*]] = insertelement <2 x i8> [[VINS]], i8 [[LDA2]], i32 1 +; CHECK-NEXT: [[VEC:%.*]] = add <2 x i8> [[VECL]], [[VINS1]] +; CHECK-NEXT: store <2 x i8> [[VEC]], ptr [[GEP0]], align 1 +; CHECK-NEXT: ret void +; + %gep0 = getelementptr i8, ptr %ptr, i64 0 + %gep1 = getelementptr i8, ptr %ptr, i64 1 + %gep2 = getelementptr i8, ptr %ptr, i64 2 + %ldA0 = load i8, ptr %gep0 + %ldA1 = load i8, ptr %gep1 + %ldA2 = load i8, ptr %gep2 + %add0 = add i8 %ldA0, %ldA1 + %add1 = add i8 %ldA1, %ldA2 + store i8 %add0, ptr %gep0 + store i8 %add1, ptr %gep1 + ret void +} diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp index 97724100ba341..854d2bcde9537 100644 --- a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp +++ b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/SchedulerTest.cpp @@ -210,9 +210,10 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) { EXPECT_TRUE(Sched.trySchedule({L0, L1})); } -TEST_F(SchedulerTest, RescheduleAlreadyScheduled) { +TEST_F(SchedulerTest, TrimSchedule) { parseIR(C, R"IR( -define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) { +define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %arg) { + %zext = zext i8 0 to i32 %ld0 = load i8, ptr %ptr0 %ld1 = load i8, ptr %ptr1 %add0 = add i8 %ld0, %ld0 @@ -227,6 +228,7 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) { auto *F = Ctx.createFunction(LLVMF); auto *BB = &*F->begin(); auto It = BB->begin(); + auto *Z = cast(&*It++); auto *L0 = cast(&*It++); auto *L1 = cast(&*It++); auto *Add0 = cast(&*It++); @@ -240,10 +242,224 @@ define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) { EXPECT_TRUE(Sched.trySchedule({S0, S1})); EXPECT_TRUE(Sched.trySchedule({L0, L1})); // At this point Add0 and Add1 should have been individually scheduled - // as single bundles. + // as singleton bundles, but {S0,S1} and {L0,L1} as vector bundles. // Check if rescheduling works. EXPECT_TRUE(Sched.trySchedule({Add0, Add1})); + // These should fail because {L0,L1} is a vector bundle. + EXPECT_FALSE(Sched.trySchedule({L0, Z})); + EXPECT_FALSE(Sched.trySchedule({L1, Z})); + // This should succeed because it matches the original vec bundle. + EXPECT_TRUE(Sched.trySchedule({L0, L1})); +} + +// Test that an instruction can't belong in two bundles! +TEST_F(SchedulerTest, CheckBundles) { + parseIR(C, R"IR( +define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) { + %L0 = load i8, ptr %ptr0 + %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles! + %L2 = load i8, ptr %ptr2 + %add0 = add i8 %L0, %L1 + %add1 = add i8 %L1, %L2 + store i8 %add0, ptr %ptr0 + store i8 %add1, ptr %ptr1 + ret void +} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + sandboxir::Context Ctx(C); + auto *F = Ctx.createFunction(LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + auto *L0 = cast(&*It++); + auto *L1 = cast(&*It++); + auto *L2 = cast(&*It++); + auto *Add0 = cast(&*It++); + auto *Add1 = cast(&*It++); + auto *S0 = cast(&*It++); + auto *S1 = cast(&*It++); + + sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx); + EXPECT_TRUE(Sched.trySchedule({S0, S1})); + EXPECT_TRUE(Sched.trySchedule({Add0, Add1})); + EXPECT_TRUE(Sched.trySchedule({L0, L1})); + // This should fail because L1 is already part of {L0,L1} + EXPECT_FALSE(Sched.trySchedule({L1, L2})); + EXPECT_FALSE(Sched.trySchedule({L2, L1})); +} + +// Try schedule a bundle {L1,L2} where L1 is already scheduled in {L0,L1} +// but L2 is not in the DAG at all +TEST_F(SchedulerTest, CheckBundles2) { + parseIR(C, R"IR( +define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) { + %L2 = load i8, ptr %ptr2 ; This is not in the DAG + %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles! + %L0 = load i8, ptr %ptr0 + %add1 = add i8 %L1, %L2 + %add0 = add i8 %L0, %L1 + store i8 %add1, ptr %ptr1 + store i8 %add0, ptr %ptr0 + ret void +} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + sandboxir::Context Ctx(C); + auto *F = Ctx.createFunction(LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + auto *L2 = cast(&*It++); + auto *L1 = cast(&*It++); + auto *L0 = cast(&*It++); + auto *Add1 = cast(&*It++); + auto *Add0 = cast(&*It++); + auto *S1 = cast(&*It++); + auto *S0 = cast(&*It++); + + sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx); + EXPECT_TRUE(Sched.trySchedule({S0, S1})); + EXPECT_TRUE(Sched.trySchedule({Add0, Add1})); + EXPECT_TRUE(Sched.trySchedule({L0, L1})); + // This should fail because L1 is already part of {L0,L1}. + EXPECT_FALSE(Sched.trySchedule({L1, L2})); + EXPECT_FALSE(Sched.trySchedule({L2, L1})); +} + +// Try schedule a bundle {L1,L2} where L1 is already scheduled in {L0,L1} +// but L2 is in the DAG but isn't scheduled. +TEST_F(SchedulerTest, CheckBundles3) { + parseIR(C, R"IR( +define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) { + %L2 = load i8, ptr %ptr2 ; This is not in the DAG + %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles! + %L0 = load i8, ptr %ptr0 + %add1 = add i8 %L1, %L2 + %add0 = add i8 %L0, %L1 + store i8 %add1, ptr %ptr1 + store i8 %add0, ptr %ptr0 + ret void +} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + sandboxir::Context Ctx(C); + auto *F = Ctx.createFunction(LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + auto *L2 = cast(&*It++); + auto *L1 = cast(&*It++); + auto *L0 = cast(&*It++); + auto *Add1 = cast(&*It++); + auto *Add0 = cast(&*It++); + auto *S1 = cast(&*It++); + auto *S0 = cast(&*It++); + + sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx); + EXPECT_TRUE(Sched.trySchedule({S0, S1})); + EXPECT_TRUE(Sched.trySchedule({Add0, Add1})); + EXPECT_TRUE(Sched.trySchedule({L0, L1})); + // Add L2 to the DAG, but don't schedule it. + auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched); + DAG.extend(L2); + // This should fail because L1 is already part of {L0,L1}. + EXPECT_FALSE(Sched.trySchedule({L1, L2})); + EXPECT_FALSE(Sched.trySchedule({L2, L1})); +} + +// Check that Scheduler::getBndlSchedState() works correctly. +TEST_F(SchedulerTest, GetBndlSchedState) { + parseIR(C, R"IR( +define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) { + %L2 = load i8, ptr %ptr2 ; This is not in the DAG + %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles! + %L0 = load i8, ptr %ptr0 + %add1 = add i8 %L1, %L2 + %add0 = add i8 %L0, %L1 + store i8 %add1, ptr %ptr1 + store i8 %add0, ptr %ptr0 + ret void +} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + sandboxir::Context Ctx(C); + auto *F = Ctx.createFunction(LLVMF); + auto *BB = &*F->begin(); + auto It = BB->begin(); + auto *L2 = cast(&*It++); + auto *L1 = cast(&*It++); + auto *L0 = cast(&*It++); + auto *Add1 = cast(&*It++); + auto *Add0 = cast(&*It++); + auto *S1 = cast(&*It++); + auto *S0 = cast(&*It++); + + sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx); + auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched); + auto GetBndlSchedState = [&Sched](ArrayRef Instrs) { + return sandboxir::SchedulerInternalsAttorney::getBndlSchedState(Sched, + Instrs); + }; + using BndlSchedState = sandboxir::SchedulerInternalsAttorney::BndlSchedState; + // Check when instructions are not in the DAG. + EXPECT_EQ(GetBndlSchedState({S0}), BndlSchedState::NoneScheduled); + EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled); + EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled); + // Check when instructions are in the DAG. + DAG.extend({S0, S1}); + EXPECT_EQ(GetBndlSchedState({S0}), BndlSchedState::NoneScheduled); + EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled); + EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled); + // One instruction in the DAG and the other not in the DAG. + EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::NoneScheduled); + + // Check with scheduled instructions. + Sched.clear(); // Manually extending the DAG messes with the scheduler. + EXPECT_TRUE(Sched.trySchedule({S0, S1})); + // Check fully scheduled. + EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::FullyScheduled); + // Check scheduled + not in DAG. + EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, S0}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, S1}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, Add1}), BndlSchedState::NoneScheduled); + // Extend DAG such that Add0 and Add1 are in the DAG but are not scheduled. + DAG.extend({Add0, Add1}); + // Check both in DAG but not scheduled. + EXPECT_EQ(GetBndlSchedState({Add0, Add1}), BndlSchedState::NoneScheduled); + // Check scheduled + in DAG but not scheduled. + EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, S0}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, S1}), BndlSchedState::AlreadyScheduled); + + Sched.clear(); // Manually extending the DAG messes with the scheduler. + // Schedule instructions towards the top so that intermediate instructions + // (namely Add0, Add1) get temporarily scheduled in singleton bundles. + EXPECT_TRUE(Sched.trySchedule({S0, S1})); EXPECT_TRUE(Sched.trySchedule({L0, L1})); + // Check fully scheduled. + EXPECT_EQ(GetBndlSchedState({L0, L1}), BndlSchedState::FullyScheduled); + // Check both singletons. + EXPECT_EQ(GetBndlSchedState({Add0, Add1}), + BndlSchedState::TemporarilyScheduled); + // Check single singleton. + EXPECT_EQ(GetBndlSchedState({Add0}), BndlSchedState::TemporarilyScheduled); + EXPECT_EQ(GetBndlSchedState({Add1}), BndlSchedState::TemporarilyScheduled); + // Check singleton + scheduled. + EXPECT_EQ(GetBndlSchedState({L0, S1}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({S1, L0}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({L0, Add1}), BndlSchedState::AlreadyScheduled); + EXPECT_EQ(GetBndlSchedState({Add1, L0}), BndlSchedState::AlreadyScheduled); + // Check singleton + not in DAG. + EXPECT_EQ(GetBndlSchedState({Add1, L2}), + BndlSchedState::TemporarilyScheduled); + EXPECT_EQ(GetBndlSchedState({L2, Add0}), + BndlSchedState::TemporarilyScheduled); + + // Check duplicates. + // TODO: Should duplicates be allowed? + EXPECT_EQ(GetBndlSchedState({L2, L2}), BndlSchedState::NoneScheduled); + EXPECT_EQ(GetBndlSchedState({S0, S0}), BndlSchedState::FullyScheduled); + EXPECT_EQ(GetBndlSchedState({Add0, Add1}), + BndlSchedState::TemporarilyScheduled); } // Check scheduling in the following order: {A0,A1},{B0,B1},{C0,C1},{D0,D1}