882 lines · cpp
1//===- SchedulerTest.cpp --------------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "llvm/Transforms/Vectorize/SandboxVectorizer/Scheduler.h"10#include "llvm/ADT/SmallVector.h"11#include "llvm/Analysis/AliasAnalysis.h"12#include "llvm/Analysis/AssumptionCache.h"13#include "llvm/Analysis/BasicAliasAnalysis.h"14#include "llvm/Analysis/TargetLibraryInfo.h"15#include "llvm/AsmParser/Parser.h"16#include "llvm/IR/Dominators.h"17#include "llvm/SandboxIR/Context.h"18#include "llvm/SandboxIR/Function.h"19#include "llvm/SandboxIR/Instruction.h"20#include "llvm/Support/SourceMgr.h"21#include "gmock/gmock-matchers.h"22#include "gtest/gtest.h"23 24using namespace llvm;25 26struct SchedulerTest : public testing::Test {27 LLVMContext C;28 std::unique_ptr<Module> M;29 std::unique_ptr<AssumptionCache> AC;30 std::unique_ptr<DominatorTree> DT;31 std::unique_ptr<BasicAAResult> BAA;32 std::unique_ptr<AAResults> AA;33 std::unique_ptr<TargetLibraryInfoImpl> TLII;34 std::unique_ptr<TargetLibraryInfo> TLI;35 36 void parseIR(LLVMContext &C, const char *IR) {37 SMDiagnostic Err;38 M = parseAssemblyString(IR, Err, C);39 if (!M) {40 Err.print("SchedulerTest", errs());41 return;42 }43 44 TLII = std::make_unique<TargetLibraryInfoImpl>(M->getTargetTriple());45 TLI = std::make_unique<TargetLibraryInfo>(*TLII);46 }47 48 AAResults &getAA(llvm::Function &LLVMF) {49 AA = std::make_unique<AAResults>(*TLI);50 AC = std::make_unique<AssumptionCache>(LLVMF);51 DT = std::make_unique<DominatorTree>(LLVMF);52 BAA = std::make_unique<BasicAAResult>(M->getDataLayout(), LLVMF, *TLI, *AC,53 DT.get());54 AA->addAAResult(*BAA);55 return *AA;56 }57};58 59static sandboxir::BasicBlock *getBasicBlockByName(sandboxir::Function *F,60 StringRef Name) {61 for (sandboxir::BasicBlock &BB : *F)62 if (BB.getName() == Name)63 return &BB;64 llvm_unreachable("Expected to find basic block!");65}66 67TEST_F(SchedulerTest, SchedBundle) {68 parseIR(C, R"IR(69define void @foo(ptr %ptr, i8 %v0, i8 %v1) {70 store i8 %v0, ptr %ptr71 %other = add i8 %v0, %v172 store i8 %v1, ptr %ptr73 ret void74}75)IR");76 llvm::Function *LLVMF = &*M->getFunction("foo");77 sandboxir::Context Ctx(C);78 auto *F = Ctx.createFunction(LLVMF);79 auto *BB = &*F->begin();80 auto It = BB->begin();81 auto *S0 = cast<sandboxir::StoreInst>(&*It++);82 auto *Other = &*It++;83 auto *S1 = cast<sandboxir::StoreInst>(&*It++);84 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);85 86 sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);87 DAG.extend({&*BB->begin(), BB->getTerminator()});88 auto *SN0 = DAG.getNode(S0);89 auto *SN1 = DAG.getNode(S1);90 sandboxir::SchedBundle Bndl({SN0, SN1});91 92 // Check getTop().93 EXPECT_EQ(Bndl.getTop(), SN0);94 // Check getBot().95 EXPECT_EQ(Bndl.getBot(), SN1);96 // Check cluster().97 Bndl.cluster(S1->getIterator());98 {99 auto It = BB->begin();100 EXPECT_EQ(&*It++, Other);101 EXPECT_EQ(&*It++, S0);102 EXPECT_EQ(&*It++, S1);103 EXPECT_EQ(&*It++, Ret);104 S0->moveBefore(Other);105 }106 107 Bndl.cluster(S0->getIterator());108 {109 auto It = BB->begin();110 EXPECT_EQ(&*It++, S0);111 EXPECT_EQ(&*It++, S1);112 EXPECT_EQ(&*It++, Other);113 EXPECT_EQ(&*It++, Ret);114 S1->moveAfter(Other);115 }116 117 Bndl.cluster(Other->getIterator());118 {119 auto It = BB->begin();120 EXPECT_EQ(&*It++, S0);121 EXPECT_EQ(&*It++, S1);122 EXPECT_EQ(&*It++, Other);123 EXPECT_EQ(&*It++, Ret);124 S1->moveAfter(Other);125 }126 127 Bndl.cluster(Ret->getIterator());128 {129 auto It = BB->begin();130 EXPECT_EQ(&*It++, Other);131 EXPECT_EQ(&*It++, S0);132 EXPECT_EQ(&*It++, S1);133 EXPECT_EQ(&*It++, Ret);134 Other->moveBefore(S1);135 }136 137 Bndl.cluster(BB->end());138 {139 auto It = BB->begin();140 EXPECT_EQ(&*It++, Other);141 EXPECT_EQ(&*It++, Ret);142 EXPECT_EQ(&*It++, S0);143 EXPECT_EQ(&*It++, S1);144 Ret->moveAfter(S1);145 Other->moveAfter(S0);146 }147 // Check iterators.148 EXPECT_THAT(Bndl, testing::ElementsAre(SN0, SN1));149 EXPECT_THAT((const sandboxir::SchedBundle &)Bndl,150 testing::ElementsAre(SN0, SN1));151}152 153// Check that when we erase a DAG node its SchedBundle gets updated.154TEST_F(SchedulerTest, SchedBundleEraseDGNode) {155 parseIR(C, R"IR(156define void @foo(ptr %ptr, i8 %v0, i8 %v1, i8 %v2, i8 %v3) {157 store i8 %v0, ptr %ptr158 store i8 %v1, ptr %ptr159 store i8 %v2, ptr %ptr160 store i8 %v3, ptr %ptr161 ret void162}163)IR");164 llvm::Function *LLVMF = &*M->getFunction("foo");165 sandboxir::Context Ctx(C);166 auto *F = Ctx.createFunction(LLVMF);167 auto *BB = &*F->begin();168 auto It = BB->begin();169 auto *S0 = cast<sandboxir::StoreInst>(&*It++);170 auto *S1 = cast<sandboxir::StoreInst>(&*It++);171 auto *S2 = cast<sandboxir::StoreInst>(&*It++);172 auto *S3 = cast<sandboxir::StoreInst>(&*It++);173 174 sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);175 DAG.extend({&*BB->begin(), BB->getTerminator()});176 auto *SN0 = DAG.getNode(S0);177 auto *SN1 = DAG.getNode(S1);178 auto *SN2 = DAG.getNode(S2);179 auto *SN3 = DAG.getNode(S3);180 {181 // Check the common case, when the bundle contains unique nodes.182 sandboxir::SchedBundle Bndl({SN0, SN1});183 S0->eraseFromParent();184 EXPECT_THAT(Bndl, testing::ElementsAre(SN1));185 }186 {187 // Check corner case when the node appears more than once.188 sandboxir::SchedBundle Bndl({SN2, SN3, SN2});189 S2->eraseFromParent();190 EXPECT_THAT(Bndl, testing::ElementsAre(SN3));191 }192}193 194// Check that assigning a bundle to a DAG Node that is already assigned to a195// bundle, removes the node from the old bundle.196TEST_F(SchedulerTest, SchedBundleReassign) {197 parseIR(C, R"IR(198define void @foo(ptr %ptr, i8 %v0, i8 %v1, i8 %v2) {199 store i8 %v0, ptr %ptr200 store i8 %v1, ptr %ptr201 store i8 %v2, ptr %ptr202 ret void203}204)IR");205 llvm::Function *LLVMF = &*M->getFunction("foo");206 sandboxir::Context Ctx(C);207 auto *F = Ctx.createFunction(LLVMF);208 auto *BB = &*F->begin();209 auto It = BB->begin();210 auto *S0 = cast<sandboxir::StoreInst>(&*It++);211 auto *S1 = cast<sandboxir::StoreInst>(&*It++);212 auto *S2 = cast<sandboxir::StoreInst>(&*It++);213 214 sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);215 DAG.extend({&*BB->begin(), BB->getTerminator()});216 auto *SN0 = DAG.getNode(S0);217 auto *SN1 = DAG.getNode(S1);218 auto *SN2 = DAG.getNode(S2);219 sandboxir::SchedBundle BndlOld({SN0, SN1});220 sandboxir::SchedBundle BndlNew({SN0, SN2});221 EXPECT_THAT(BndlOld, testing::ElementsAre(SN1));222 EXPECT_THAT(BndlNew, testing::ElementsAre(SN0, SN2));223}224 225TEST_F(SchedulerTest, Basic) {226 parseIR(C, R"IR(227define void @foo(ptr %ptr, i8 %v0, i8 %v1) {228 store i8 %v0, ptr %ptr229 store i8 %v1, ptr %ptr230 ret void231}232)IR");233 llvm::Function *LLVMF = &*M->getFunction("foo");234 sandboxir::Context Ctx(C);235 auto *F = Ctx.createFunction(LLVMF);236 auto *BB = &*F->begin();237 auto It = BB->begin();238 auto *S0 = cast<sandboxir::StoreInst>(&*It++);239 auto *S1 = cast<sandboxir::StoreInst>(&*It++);240 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);241 242 {243 // Schedule all instructions in sequence.244 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);245 EXPECT_TRUE(Sched.trySchedule({Ret}));246 EXPECT_TRUE(Sched.trySchedule({S1}));247 EXPECT_TRUE(Sched.trySchedule({S0}));248 }249 {250 // Skip instructions.251 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);252 EXPECT_TRUE(Sched.trySchedule({Ret}));253 EXPECT_TRUE(Sched.trySchedule({S0}));254 }255 {256 // Try invalid scheduling. Dependency S0->S1.257 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);258 EXPECT_TRUE(Sched.trySchedule({Ret}));259 EXPECT_FALSE(Sched.trySchedule({S0, S1}));260 }261}262 263TEST_F(SchedulerTest, Bundles) {264 parseIR(C, R"IR(265define void @foo(ptr noalias %ptr0, ptr noalias %ptr1) {266 %ld0 = load i8, ptr %ptr0267 %ld1 = load i8, ptr %ptr1268 store i8 %ld0, ptr %ptr0269 store i8 %ld1, ptr %ptr1270 ret void271}272)IR");273 llvm::Function *LLVMF = &*M->getFunction("foo");274 sandboxir::Context Ctx(C);275 auto *F = Ctx.createFunction(LLVMF);276 auto *BB = &*F->begin();277 auto It = BB->begin();278 auto *L0 = cast<sandboxir::LoadInst>(&*It++);279 auto *L1 = cast<sandboxir::LoadInst>(&*It++);280 auto *S0 = cast<sandboxir::StoreInst>(&*It++);281 auto *S1 = cast<sandboxir::StoreInst>(&*It++);282 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);283 284 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);285 EXPECT_TRUE(Sched.trySchedule({Ret}));286 EXPECT_TRUE(Sched.trySchedule({S0, S1}));287 EXPECT_TRUE(Sched.trySchedule({L0, L1}));288}289 290TEST_F(SchedulerTest, TrimSchedule) {291 parseIR(C, R"IR(292define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %arg) {293 %zext = zext i8 0 to i32294 %ld0 = load i8, ptr %ptr0295 %ld1 = load i8, ptr %ptr1296 %add0 = add i8 %ld0, %ld0297 %add1 = add i8 %ld1, %ld1298 store i8 %add0, ptr %ptr0299 store i8 %add1, ptr %ptr1300 ret void301}302)IR");303 llvm::Function *LLVMF = &*M->getFunction("foo");304 sandboxir::Context Ctx(C);305 auto *F = Ctx.createFunction(LLVMF);306 auto *BB = &*F->begin();307 auto It = BB->begin();308 auto *Z = cast<sandboxir::CastInst>(&*It++);309 auto *L0 = cast<sandboxir::LoadInst>(&*It++);310 auto *L1 = cast<sandboxir::LoadInst>(&*It++);311 auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);312 auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);313 auto *S0 = cast<sandboxir::StoreInst>(&*It++);314 auto *S1 = cast<sandboxir::StoreInst>(&*It++);315 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);316 317 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);318 EXPECT_TRUE(Sched.trySchedule({Ret}));319 EXPECT_TRUE(Sched.trySchedule({S0, S1}));320 EXPECT_TRUE(Sched.trySchedule({L0, L1}));321 // At this point Add0 and Add1 should have been individually scheduled322 // as singleton bundles, but {S0,S1} and {L0,L1} as vector bundles.323 // Check if rescheduling works.324 EXPECT_TRUE(Sched.trySchedule({Add0, Add1}));325 // These should fail because {L0,L1} is a vector bundle.326 EXPECT_FALSE(Sched.trySchedule({L0, Z}));327 EXPECT_FALSE(Sched.trySchedule({L1, Z}));328 // This should succeed because it matches the original vec bundle.329 EXPECT_TRUE(Sched.trySchedule({L0, L1}));330}331 332// Make sure that instructions in SchedBundles are always scheduled333// back-to-back334TEST_F(SchedulerTest, SchedBundleBackToBack) {335 parseIR(C, R"IR(336define void @foo(ptr %ptr, i16 %arg) {337 %gep0 = getelementptr i32, ptr %ptr, i64 0338 %gep1 = getelementptr i32, ptr %ptr, i64 1339 %zextX = zext i16 0 to i32340 %zext1 = zext i16 0 to i32341 %zext0 = zext i16 %arg to i32342 %shl1 = shl i32 %zextX, 0343 %shl0 = shl i32 %zext1, 0344 %sub1 = sub i32 %zext1, %shl1345 %sub0 = sub i32 %zext0, %shl0346 store i32 %sub1, ptr %gep1347 store i32 %sub0, ptr %gep0348 ret void349})IR");350 llvm::Function *LLVMF = &*M->getFunction("foo");351 sandboxir::Context Ctx(C);352 auto *F = Ctx.createFunction(LLVMF);353 auto *BB = &*F->begin();354 auto It = BB->begin();355 [[maybe_unused]] auto *Gep0 = cast<sandboxir::GetElementPtrInst>(&*It++);356 [[maybe_unused]] auto *Gep1 = cast<sandboxir::GetElementPtrInst>(&*It++);357 [[maybe_unused]] auto *ZextX = cast<sandboxir::CastInst>(&*It++);358 auto *Zext1 = cast<sandboxir::CastInst>(&*It++);359 auto *Zext0 = cast<sandboxir::CastInst>(&*It++);360 auto *Shl1 = cast<sandboxir::BinaryOperator>(&*It++);361 auto *Shl0 = cast<sandboxir::BinaryOperator>(&*It++);362 [[maybe_unused]] auto *Sub1 = cast<sandboxir::BinaryOperator>(&*It++);363 [[maybe_unused]] auto *Sub0 = cast<sandboxir::BinaryOperator>(&*It++);364 auto *S0 = cast<sandboxir::StoreInst>(&*It++);365 auto *S1 = cast<sandboxir::StoreInst>(&*It++);366 367 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);368 EXPECT_TRUE(Sched.trySchedule({S0, S1}));369 EXPECT_TRUE(Sched.trySchedule({Zext0, Zext1}));370 EXPECT_TRUE(Sched.trySchedule({Shl0, Shl1}));371 auto BackToBack = [](sandboxir::Instruction *I1, sandboxir::Instruction *I2) {372 return I1->getNextNode() == I2 || I2->getNextNode() == I1;373 };374 EXPECT_TRUE(BackToBack(S0, S1));375 EXPECT_TRUE(BackToBack(Zext0, Zext1));376 EXPECT_TRUE(BackToBack(Shl0, Shl1));377}378 379// Test that an instruction can't belong in two bundles!380TEST_F(SchedulerTest, CheckBundles) {381 parseIR(C, R"IR(382define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) {383 %L0 = load i8, ptr %ptr0384 %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles!385 %L2 = load i8, ptr %ptr2386 %add0 = add i8 %L0, %L1387 %add1 = add i8 %L1, %L2388 store i8 %add0, ptr %ptr0389 store i8 %add1, ptr %ptr1390 ret void391}392)IR");393 llvm::Function *LLVMF = &*M->getFunction("foo");394 sandboxir::Context Ctx(C);395 auto *F = Ctx.createFunction(LLVMF);396 auto *BB = &*F->begin();397 auto It = BB->begin();398 auto *L0 = cast<sandboxir::LoadInst>(&*It++);399 auto *L1 = cast<sandboxir::LoadInst>(&*It++);400 auto *L2 = cast<sandboxir::LoadInst>(&*It++);401 auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);402 auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);403 auto *S0 = cast<sandboxir::StoreInst>(&*It++);404 auto *S1 = cast<sandboxir::StoreInst>(&*It++);405 406 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);407 EXPECT_TRUE(Sched.trySchedule({S0, S1}));408 EXPECT_TRUE(Sched.trySchedule({Add0, Add1}));409 EXPECT_TRUE(Sched.trySchedule({L0, L1}));410 // This should fail because L1 is already part of {L0,L1}411 EXPECT_FALSE(Sched.trySchedule({L1, L2}));412 EXPECT_FALSE(Sched.trySchedule({L2, L1}));413}414 415// Try schedule a bundle {L1,L2} where L1 is already scheduled in {L0,L1}416// but L2 is not in the DAG at all417TEST_F(SchedulerTest, CheckBundles2) {418 parseIR(C, R"IR(419define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) {420 %L2 = load i8, ptr %ptr2 ; This is not in the DAG421 %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles!422 %L0 = load i8, ptr %ptr0423 %add1 = add i8 %L1, %L2424 %add0 = add i8 %L0, %L1425 store i8 %add1, ptr %ptr1426 store i8 %add0, ptr %ptr0427 ret void428}429)IR");430 llvm::Function *LLVMF = &*M->getFunction("foo");431 sandboxir::Context Ctx(C);432 auto *F = Ctx.createFunction(LLVMF);433 auto *BB = &*F->begin();434 auto It = BB->begin();435 auto *L2 = cast<sandboxir::LoadInst>(&*It++);436 auto *L1 = cast<sandboxir::LoadInst>(&*It++);437 auto *L0 = cast<sandboxir::LoadInst>(&*It++);438 auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);439 auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);440 auto *S1 = cast<sandboxir::StoreInst>(&*It++);441 auto *S0 = cast<sandboxir::StoreInst>(&*It++);442 443 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);444 EXPECT_TRUE(Sched.trySchedule({S0, S1}));445 EXPECT_TRUE(Sched.trySchedule({Add0, Add1}));446 EXPECT_TRUE(Sched.trySchedule({L0, L1}));447 // This should fail because L1 is already part of {L0,L1}.448 EXPECT_FALSE(Sched.trySchedule({L1, L2}));449 EXPECT_FALSE(Sched.trySchedule({L2, L1}));450}451 452// Try schedule a bundle {L1,L2} where L1 is already scheduled in {L0,L1}453// but L2 is in the DAG but isn't scheduled.454TEST_F(SchedulerTest, CheckBundles3) {455 parseIR(C, R"IR(456define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) {457 %L2 = load i8, ptr %ptr2 ; This is not in the DAG458 %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles!459 %L0 = load i8, ptr %ptr0460 %add1 = add i8 %L1, %L2461 %add0 = add i8 %L0, %L1462 store i8 %add1, ptr %ptr1463 store i8 %add0, ptr %ptr0464 ret void465}466)IR");467 llvm::Function *LLVMF = &*M->getFunction("foo");468 sandboxir::Context Ctx(C);469 auto *F = Ctx.createFunction(LLVMF);470 auto *BB = &*F->begin();471 auto It = BB->begin();472 auto *L2 = cast<sandboxir::LoadInst>(&*It++);473 auto *L1 = cast<sandboxir::LoadInst>(&*It++);474 auto *L0 = cast<sandboxir::LoadInst>(&*It++);475 auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);476 auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);477 auto *S1 = cast<sandboxir::StoreInst>(&*It++);478 auto *S0 = cast<sandboxir::StoreInst>(&*It++);479 480 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);481 EXPECT_TRUE(Sched.trySchedule({S0, S1}));482 EXPECT_TRUE(Sched.trySchedule({Add0, Add1}));483 EXPECT_TRUE(Sched.trySchedule({L0, L1}));484 // Add L2 to the DAG, but don't schedule it.485 auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);486 DAG.extend(L2);487 // This should fail because L1 is already part of {L0,L1}.488 EXPECT_FALSE(Sched.trySchedule({L1, L2}));489 EXPECT_FALSE(Sched.trySchedule({L2, L1}));490}491 492// Check that Scheduler::getBndlSchedState() works correctly.493TEST_F(SchedulerTest, GetBndlSchedState) {494 parseIR(C, R"IR(495define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, ptr noalias %ptr2) {496 %L2 = load i8, ptr %ptr2 ; This is not in the DAG497 %L1 = load i8, ptr %ptr1 ; This belongs in 2 bundles!498 %L0 = load i8, ptr %ptr0499 %add1 = add i8 %L1, %L2500 %add0 = add i8 %L0, %L1501 store i8 %add1, ptr %ptr1502 store i8 %add0, ptr %ptr0503 ret void504}505)IR");506 llvm::Function *LLVMF = &*M->getFunction("foo");507 sandboxir::Context Ctx(C);508 auto *F = Ctx.createFunction(LLVMF);509 auto *BB = &*F->begin();510 auto It = BB->begin();511 auto *L2 = cast<sandboxir::LoadInst>(&*It++);512 auto *L1 = cast<sandboxir::LoadInst>(&*It++);513 auto *L0 = cast<sandboxir::LoadInst>(&*It++);514 auto *Add1 = cast<sandboxir::BinaryOperator>(&*It++);515 auto *Add0 = cast<sandboxir::BinaryOperator>(&*It++);516 auto *S1 = cast<sandboxir::StoreInst>(&*It++);517 auto *S0 = cast<sandboxir::StoreInst>(&*It++);518 519 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);520 auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);521 auto GetBndlSchedState = [&Sched](ArrayRef<sandboxir::Instruction *> Instrs) {522 return sandboxir::SchedulerInternalsAttorney::getBndlSchedState(Sched,523 Instrs);524 };525 using BndlSchedState = sandboxir::SchedulerInternalsAttorney::BndlSchedState;526 // Check when instructions are not in the DAG.527 EXPECT_EQ(GetBndlSchedState({S0}), BndlSchedState::NoneScheduled);528 EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled);529 EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled);530 // Check when instructions are in the DAG.531 DAG.extend({S0, S1});532 EXPECT_EQ(GetBndlSchedState({S0}), BndlSchedState::NoneScheduled);533 EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled);534 EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::NoneScheduled);535 // One instruction in the DAG and the other not in the DAG.536 EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::NoneScheduled);537 538 // Check with scheduled instructions.539 Sched.clear(); // Manually extending the DAG messes with the scheduler.540 EXPECT_TRUE(Sched.trySchedule({S0, S1}));541 // Check fully scheduled.542 EXPECT_EQ(GetBndlSchedState({S0, S1}), BndlSchedState::FullyScheduled);543 // Check scheduled + not in DAG.544 EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::AlreadyScheduled);545 EXPECT_EQ(GetBndlSchedState({Add0, S0}), BndlSchedState::AlreadyScheduled);546 EXPECT_EQ(GetBndlSchedState({Add0, S1}), BndlSchedState::AlreadyScheduled);547 EXPECT_EQ(GetBndlSchedState({Add0, Add1}), BndlSchedState::NoneScheduled);548 // Extend DAG such that Add0 and Add1 are in the DAG but are not scheduled.549 DAG.extend({Add0, Add1});550 // Check both in DAG but not scheduled.551 EXPECT_EQ(GetBndlSchedState({Add0, Add1}), BndlSchedState::NoneScheduled);552 // Check scheduled + in DAG but not scheduled.553 EXPECT_EQ(GetBndlSchedState({S0, Add0}), BndlSchedState::AlreadyScheduled);554 EXPECT_EQ(GetBndlSchedState({Add0, S0}), BndlSchedState::AlreadyScheduled);555 EXPECT_EQ(GetBndlSchedState({Add0, S1}), BndlSchedState::AlreadyScheduled);556 557 Sched.clear(); // Manually extending the DAG messes with the scheduler.558 // Schedule instructions towards the top so that intermediate instructions559 // (namely Add0, Add1) get temporarily scheduled in singleton bundles.560 EXPECT_TRUE(Sched.trySchedule({S0, S1}));561 EXPECT_TRUE(Sched.trySchedule({L0, L1}));562 // Check fully scheduled.563 EXPECT_EQ(GetBndlSchedState({L0, L1}), BndlSchedState::FullyScheduled);564 // Check both singletons.565 EXPECT_EQ(GetBndlSchedState({Add0, Add1}),566 BndlSchedState::TemporarilyScheduled);567 // Check single singleton.568 EXPECT_EQ(GetBndlSchedState({Add0}), BndlSchedState::TemporarilyScheduled);569 EXPECT_EQ(GetBndlSchedState({Add1}), BndlSchedState::TemporarilyScheduled);570 // Check singleton + scheduled.571 EXPECT_EQ(GetBndlSchedState({L0, S1}), BndlSchedState::AlreadyScheduled);572 EXPECT_EQ(GetBndlSchedState({S1, L0}), BndlSchedState::AlreadyScheduled);573 EXPECT_EQ(GetBndlSchedState({L0, Add1}), BndlSchedState::AlreadyScheduled);574 EXPECT_EQ(GetBndlSchedState({Add1, L0}), BndlSchedState::AlreadyScheduled);575 // Check singleton + not in DAG.576 EXPECT_EQ(GetBndlSchedState({Add1, L2}),577 BndlSchedState::TemporarilyScheduled);578 EXPECT_EQ(GetBndlSchedState({L2, Add0}),579 BndlSchedState::TemporarilyScheduled);580 581 // Check duplicates.582 // TODO: Should duplicates be allowed?583 EXPECT_EQ(GetBndlSchedState({L2, L2}), BndlSchedState::NoneScheduled);584 EXPECT_EQ(GetBndlSchedState({S0, S0}), BndlSchedState::FullyScheduled);585 EXPECT_EQ(GetBndlSchedState({Add0, Add1}),586 BndlSchedState::TemporarilyScheduled);587}588 589// Check scheduling in the following order: {A0,A1},{B0,B1},{C0,C1},{D0,D1}590// assuming program order: B0,B1,C0,C1,D0,D1,E0,D1.591// This will effectively schedule nodes below already scheduled nodes, which592// can expose issues in the code that adds nodes to the ready list.593// For example, we schedule {D0,D1} while {C0,C1} are scheduled and there is594// a dependency D0->C0 and D1->C1.595//596// {A0,A1} {B0,B1} {C0,C1} {D0,D1}597// B0,B1 | S598// |\ |599// | C0,C1 | | S | S600// | | \ | |601// | | D0,D1 | | S602// | / |603// A0,A1 | S | S604// +------------------------+605// | Legend |: DAG |606// | S: Scheduled |607TEST_F(SchedulerTest, ScheduledPredecessors) {608 parseIR(C, R"IR(609define void @foo(ptr noalias %ptrA0, ptr noalias %ptrA1,610 ptr noalias %ptrB0, ptr noalias %ptrB1,611 ptr noalias %ptrD0, ptr noalias %ptrD1) {612 %B0 = load i8, ptr %ptrB0613 %B1 = load i8, ptr %ptrB1614 %C0 = add i8 %B0, 0615 %C1 = add i8 %B1, 1616 store i8 %C0, ptr %ptrD0617 store i8 %C1, ptr %ptrD1618 store i8 %B0, ptr %ptrA0619 store i8 %B1, ptr %ptrA1620 ret void621}622)IR");623 llvm::Function *LLVMF = &*M->getFunction("foo");624 sandboxir::Context Ctx(C);625 auto *F = Ctx.createFunction(LLVMF);626 auto *BB = &*F->begin();627 auto It = BB->begin();628 auto *B1 = cast<sandboxir::LoadInst>(&*It++);629 auto *B0 = cast<sandboxir::LoadInst>(&*It++);630 auto *C1 = cast<sandboxir::BinaryOperator>(&*It++);631 auto *C0 = cast<sandboxir::BinaryOperator>(&*It++);632 auto *D1 = cast<sandboxir::StoreInst>(&*It++);633 auto *D0 = cast<sandboxir::StoreInst>(&*It++);634 auto *A1 = cast<sandboxir::StoreInst>(&*It++);635 auto *A0 = cast<sandboxir::StoreInst>(&*It++);636 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);637 (void)Ret;638 639 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);640 EXPECT_TRUE(Sched.trySchedule({A0, A1}));641 // NOTE: We schedule the intermediate nodes between {A0,A1} and {B0,B1} by642 // hand one by one to make sure they are scheduled in that order because643 // the scheduler may reorder them a bit if we let it do it.644 EXPECT_TRUE(Sched.trySchedule(D0));645 EXPECT_TRUE(Sched.trySchedule(D1));646 EXPECT_TRUE(Sched.trySchedule(C0));647 EXPECT_TRUE(Sched.trySchedule(C1));648 EXPECT_TRUE(Sched.trySchedule({B0, B1}));649 // At this point all nodes must have been scheduled from B0,B1 to A0,A1.650 // The ones in between are scheduled as single-instruction nodes.651 // So when we attempt to schedule {C0,C1} we will need to reschedule.652 // At this point we will trim the schedule from {C0,C1} upwards.653 EXPECT_TRUE(Sched.trySchedule({C0, C1}));654 // Now the schedule should only contain {C0,C1} which should be marked as655 // "scheduled".656 // {D0,D1} are below {C0,C1}, so we grow the DAG downwards, while657 // {C0,C1} are marked as "scheduled" above them.658 EXPECT_TRUE(Sched.trySchedule({D0, D1}));659}660 661TEST_F(SchedulerTest, DontCrossBBs) {662 parseIR(C, R"IR(663define void @foo(ptr noalias %ptr0, ptr noalias %ptr1, i8 %v0, i8 %v1) {664bb0:665 %add0 = add i8 %v0, 0666 %add1 = add i8 %v1, 1667 br label %bb1668 669bb1:670 store i8 %add0, ptr %ptr0671 store i8 %add1, ptr %ptr1672 ret void673}674)IR");675 llvm::Function *LLVMF = &*M->getFunction("foo");676 sandboxir::Context Ctx(C);677 auto *F = Ctx.createFunction(LLVMF);678 auto *BB0 = getBasicBlockByName(F, "bb0");679 auto *BB1 = getBasicBlockByName(F, "bb1");680 auto It = BB0->begin();681 auto *Add0 = &*It++;682 auto *Add1 = &*It++;683 684 It = BB1->begin();685 auto *S0 = cast<sandboxir::StoreInst>(&*It++);686 auto *S1 = cast<sandboxir::StoreInst>(&*It++);687 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);688 689 {690 // Schedule bottom-up691 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);692 EXPECT_TRUE(Sched.trySchedule({Ret}));693 EXPECT_TRUE(Sched.trySchedule({S0, S1}));694 // Scheduling across blocks should fail.695 EXPECT_FALSE(Sched.trySchedule({Add0, Add1}));696 }697 {698 // Schedule top-down699 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);700 EXPECT_TRUE(Sched.trySchedule({Add0, Add1}));701 // Scheduling across blocks should fail.702 EXPECT_FALSE(Sched.trySchedule({S0, S1}));703 }704}705 706TEST_F(SchedulerTest, NotifyCreateInst) {707 parseIR(C, R"IR(708define void @foo(ptr noalias %ptr, ptr noalias %ptr1, ptr noalias %ptr2) {709 %ld0 = load i8, ptr %ptr710 store i8 %ld0, ptr %ptr711 ret void712}713)IR");714 llvm::Function *LLVMF = &*M->getFunction("foo");715 sandboxir::Context Ctx(C);716 auto *F = Ctx.createFunction(LLVMF);717 auto *BB = &*F->begin();718 auto It = BB->begin();719 auto *L0 = cast<sandboxir::LoadInst>(&*It++);720 auto *S0 = cast<sandboxir::StoreInst>(&*It++);721 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);722 auto *Ptr1 = F->getArg(1);723 auto *Ptr2 = F->getArg(2);724 725 sandboxir::Scheduler Sched(getAA(*LLVMF), Ctx);726 // Schedule Ret and S0. The top of schedule should be at S0.727 EXPECT_TRUE(Sched.trySchedule({Ret}));728 EXPECT_TRUE(Sched.trySchedule({S0}));729 auto &DAG = sandboxir::SchedulerInternalsAttorney::getDAG(Sched);730 DAG.extend({L0});731 auto *L0N = DAG.getNode(L0);732 EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 0u);733 // We should have DAG nodes for all instructions at this point734 735 // Now create a new instruction below S0.736 sandboxir::StoreInst *NewS1 =737 sandboxir::StoreInst::create(L0, Ptr1, Align(8), Ret->getIterator(),738 /*IsVolatile=*/false, Ctx);739 // Check that it is marked as "scheduled".740 auto *NewS1N = DAG.getNode(NewS1);741 EXPECT_TRUE(NewS1N->scheduled());742 // Check that L0's UnscheduledSuccs are still == 0 since NewS1 is "scheduled".743 EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 0u);744 745 // Now create a new instruction above S0.746 sandboxir::StoreInst *NewS2 =747 sandboxir::StoreInst::create(L0, Ptr2, Align(8), S0->getIterator(),748 /*IsVolatile=*/false, Ctx);749 // Check that it is not marked as "scheduled".750 auto *NewS2N = DAG.getNode(NewS2);751 EXPECT_FALSE(NewS2N->scheduled());752 // Check that L0's UnscheduledSuccs got updated because of NewS2.753 EXPECT_EQ(L0N->getNumUnscheduledSuccs(), 1u);754 755 sandboxir::ReadyListContainer ReadyList;756 // Check empty().757 EXPECT_TRUE(ReadyList.empty());758}759 760TEST_F(SchedulerTest, ReadyList) {761 parseIR(C, R"IR(762define void @foo(ptr %ptr) {763 %ld0 = load i8, ptr %ptr764 store i8 %ld0, ptr %ptr765 ret void766}767)IR");768 llvm::Function *LLVMF = &*M->getFunction("foo");769 sandboxir::Context Ctx(C);770 auto *F = Ctx.createFunction(LLVMF);771 auto *BB = &*F->begin();772 auto It = BB->begin();773 auto *L0 = cast<sandboxir::LoadInst>(&*It++);774 auto *S0 = cast<sandboxir::StoreInst>(&*It++);775 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);776 777 sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);778 DAG.extend({&*BB->begin(), BB->getTerminator()});779 auto *L0N = DAG.getNode(L0);780 auto *S0N = DAG.getNode(S0);781 auto *RetN = DAG.getNode(Ret);782 783 sandboxir::ReadyListContainer ReadyList;784 // Check empty().785 EXPECT_TRUE(ReadyList.empty());786 // Check insert(), pop().787 ReadyList.insert(L0N);788 EXPECT_FALSE(ReadyList.empty());789 EXPECT_EQ(ReadyList.pop(), L0N);790 // Check clear().791 ReadyList.insert(L0N);792 EXPECT_FALSE(ReadyList.empty());793 ReadyList.clear();794 EXPECT_TRUE(ReadyList.empty());795 // Check remove().796 EXPECT_TRUE(ReadyList.empty());797 ReadyList.remove(L0N); // Removing a non-existing node should be valid.798 ReadyList.insert(L0N);799 ReadyList.insert(S0N);800 ReadyList.insert(RetN);801 ReadyList.remove(S0N);802 DenseSet<sandboxir::DGNode *> Nodes;803 Nodes.insert(ReadyList.pop());804 Nodes.insert(ReadyList.pop());805 EXPECT_TRUE(ReadyList.empty());806 EXPECT_THAT(Nodes, testing::UnorderedElementsAre(L0N, RetN));807}808 809TEST_F(SchedulerTest, ReadyListPriorities) {810 parseIR(C, R"IR(811define void @foo(ptr %ptr) {812bb0:813 br label %bb1814 815bb1:816 %phi0 = phi i8 [0, %bb0], [1, %bb1]817 %phi1 = phi i8 [0, %bb0], [1, %bb1]818 %ld0 = load i8, ptr %ptr819 store i8 %ld0, ptr %ptr820 ret void821}822)IR");823 llvm::Function *LLVMF = &*M->getFunction("foo");824 sandboxir::Context Ctx(C);825 auto *F = Ctx.createFunction(LLVMF);826 auto *BB1 = getBasicBlockByName(F, "bb1");827 auto It = BB1->begin();828 auto *Phi0 = cast<sandboxir::PHINode>(&*It++);829 auto *Phi1 = cast<sandboxir::PHINode>(&*It++);830 auto *L0 = cast<sandboxir::LoadInst>(&*It++);831 auto *S0 = cast<sandboxir::StoreInst>(&*It++);832 auto *Ret = cast<sandboxir::ReturnInst>(&*It++);833 834 sandboxir::DependencyGraph DAG(getAA(*LLVMF), Ctx);835 DAG.extend({&*BB1->begin(), BB1->getTerminator()});836 auto *Phi0N = DAG.getNode(Phi0);837 auto *Phi1N = DAG.getNode(Phi1);838 auto *L0N = DAG.getNode(L0);839 auto *S0N = DAG.getNode(S0);840 auto *RetN = DAG.getNode(Ret);841 842 sandboxir::ReadyListContainer ReadyList;843 // Check PHI vs non-PHI.844 ReadyList.insert(S0N);845 ReadyList.insert(Phi0N);846 EXPECT_EQ(ReadyList.pop(), Phi0N);847 EXPECT_EQ(ReadyList.pop(), S0N);848 ReadyList.insert(Phi0N);849 ReadyList.insert(S0N);850 EXPECT_EQ(ReadyList.pop(), Phi0N);851 EXPECT_EQ(ReadyList.pop(), S0N);852 // Check PHI vs terminator.853 ReadyList.insert(RetN);854 ReadyList.insert(Phi1N);855 EXPECT_EQ(ReadyList.pop(), Phi1N);856 EXPECT_EQ(ReadyList.pop(), RetN);857 ReadyList.insert(Phi1N);858 ReadyList.insert(RetN);859 EXPECT_EQ(ReadyList.pop(), Phi1N);860 EXPECT_EQ(ReadyList.pop(), RetN);861 // Check terminator vs non-terminator.862 ReadyList.insert(RetN);863 ReadyList.insert(L0N);864 EXPECT_EQ(ReadyList.pop(), L0N);865 EXPECT_EQ(ReadyList.pop(), RetN);866 ReadyList.insert(L0N);867 ReadyList.insert(RetN);868 EXPECT_EQ(ReadyList.pop(), L0N);869 EXPECT_EQ(ReadyList.pop(), RetN);870 // Check all, program order.871 ReadyList.insert(RetN);872 ReadyList.insert(L0N);873 ReadyList.insert(Phi1N);874 ReadyList.insert(S0N);875 ReadyList.insert(Phi0N);876 EXPECT_EQ(ReadyList.pop(), Phi0N);877 EXPECT_EQ(ReadyList.pop(), Phi1N);878 EXPECT_EQ(ReadyList.pop(), L0N);879 EXPECT_EQ(ReadyList.pop(), S0N);880 EXPECT_EQ(ReadyList.pop(), RetN);881}882