diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h index 14f35c9c4d8a9..f86199ab6c228 100644 --- a/llvm/include/llvm/SandboxIR/Region.h +++ b/llvm/include/llvm/SandboxIR/Region.h @@ -114,8 +114,18 @@ class Region { /// ID (for later deregistration) of the "erase instruction" callback. Context::CallbackID EraseInstCB; - /// Adds I to the set. - void add(Instruction *I); + /// Adds \p I to the set but also don't track the instruction's score if \p + /// IgnoreCost is true. Only to be used when adding an instruction to the + /// auxiliary vector. + /// NOTE: When an instruction is added to the region we track it cost in the + /// scoreboard, which currently resides in the region class. However, when we + /// add an instruction to the auxiliary vector it does get tagged as being a + /// member of the region (for ownership reasons), but its cost does not get + /// counted because the instruction hasn't been added in the "normal" way. + void addImpl(Instruction *I, bool IgnoreCost); + /// Adds I to the set. This is the main API for adding an instruction to the + /// region. + void add(Instruction *I) { addImpl(I, /*IgnoreCost=*/false); } /// Removes I from the set. void remove(Instruction *I); friend class Context; // The callbacks need to call add() and remove(). diff --git a/llvm/lib/SandboxIR/Region.cpp b/llvm/lib/SandboxIR/Region.cpp index 2eb84bd72ed00..bd79a97cdd0a9 100644 --- a/llvm/lib/SandboxIR/Region.cpp +++ b/llvm/lib/SandboxIR/Region.cpp @@ -51,12 +51,13 @@ Region::~Region() { Ctx.unregisterEraseInstrCallback(EraseInstCB); } -void Region::add(Instruction *I) { +void Region::addImpl(Instruction *I, bool IgnoreCost) { Insts.insert(I); // TODO: Consider tagging instructions lazily. cast(I->Val)->setMetadata(MDKind, RegionMDN); - // Keep track of the instruction cost. - Scoreboard.add(I); + if (!IgnoreCost) + // Keep track of the instruction cost. + Scoreboard.add(I); } void Region::setAux(ArrayRef Aux) { @@ -69,6 +70,8 @@ void Region::setAux(ArrayRef Aux) { "Instruction already in Aux!"); cast(I->Val)->setMetadata( AuxMDKind, MDNode::get(LLVMCtx, ConstantAsMetadata::get(IdxC))); + // Aux instrs should always be in a region. + addImpl(I, /*DontTrackCost=*/true); } } @@ -84,6 +87,8 @@ void Region::setAux(unsigned Idx, Instruction *I) { Aux[Idx] = nullptr; } Aux[Idx] = I; + // Aux instrs should always be in a region. + addImpl(I, /*DontTrackCost=*/true); } void Region::dropAuxMetadata(Instruction *I) { @@ -151,8 +156,8 @@ Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) { for (BasicBlock &BB : F) { for (Instruction &Inst : BB) { auto *LLVMI = cast(Inst.Val); + Region *R = nullptr; if (auto *MDN = LLVMI->getMetadata(MDKind)) { - Region *R = nullptr; auto [It, Inserted] = MDNToRegion.try_emplace(MDN); if (Inserted) { Regions.push_back(std::make_unique(Ctx, TTI)); @@ -161,14 +166,17 @@ Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) { } else { R = It->second; } - R->add(&Inst); - - if (auto *AuxMDN = LLVMI->getMetadata(AuxMDKind)) { - llvm::Constant *IdxC = - dyn_cast(AuxMDN->getOperand(0))->getValue(); - auto Idx = cast(IdxC)->getSExtValue(); - R->setAux(Idx, &Inst); + R->addImpl(&Inst, /*IgnoreCost=*/true); + } + if (auto *AuxMDN = LLVMI->getMetadata(AuxMDKind)) { + llvm::Constant *IdxC = + dyn_cast(AuxMDN->getOperand(0))->getValue(); + auto Idx = cast(IdxC)->getSExtValue(); + if (R == nullptr) { + errs() << "No region specified for Aux: '" << *LLVMI << "'\n"; + exit(1); } + R->setAux(Idx, &Inst); } } } diff --git a/llvm/unittests/SandboxIR/RegionTest.cpp b/llvm/unittests/SandboxIR/RegionTest.cpp index c3390dfcdbc8e..7d998b147e37c 100644 --- a/llvm/unittests/SandboxIR/RegionTest.cpp +++ b/llvm/unittests/SandboxIR/RegionTest.cpp @@ -429,6 +429,22 @@ define void @foo(i8 %v) { } } +TEST_F(RegionTest, AuxWithoutRegion) { + parseIR(C, R"IR( +define void @foo(i8 %v) { + %Add0 = add i8 %v, 0, !sandboxaux !0 + ret void +} +!0 = !{i32 0} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + sandboxir::Context Ctx(C); + auto *F = Ctx.createFunction(LLVMF); +#ifndef NDEBUG + EXPECT_DEATH(sandboxir::Region::createRegionsFromMD(*F, *TTI), "No region.*"); +#endif +} + TEST_F(RegionTest, AuxRoundTrip) { parseIR(C, R"IR( define i8 @foo(i8 %v0, i8 %v1) { @@ -461,3 +477,36 @@ define i8 @foo(i8 %v0, i8 %v1) { #endif EXPECT_THAT(Rgn.getAux(), testing::ElementsAre(T1, T0)); } + +// Same as before but only add instructions to aux. They should get added too +// the region too automatically. +TEST_F(RegionTest, AuxOnlyRoundTrip) { + parseIR(C, R"IR( +define void @foo(i8 %v) { + %add0 = add i8 %v, 0 + %add1 = add i8 %v, 1 + 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 *Add0 = cast(&*It++); + auto *Add1 = cast(&*It++); + + sandboxir::Region Rgn(Ctx, *TTI); +#ifndef NDEBUG + EXPECT_DEATH(Rgn.setAux({Add0, Add0}), ".*already.*"); +#endif + Rgn.setAux({Add1, Add0}); + + SmallVector> Regions = + sandboxir::Region::createRegionsFromMD(*F, *TTI); + ASSERT_EQ(1U, Regions.size()); +#ifndef NDEBUG + EXPECT_EQ(Rgn, *Regions[0].get()); +#endif + EXPECT_THAT(Rgn.getAux(), testing::ElementsAre(Add1, Add0)); +}