diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h index 8133e01734ea7..c1195141cb54c 100644 --- a/llvm/include/llvm/SandboxIR/Region.h +++ b/llvm/include/llvm/SandboxIR/Region.h @@ -6,18 +6,55 @@ // //===----------------------------------------------------------------------===// -#ifndef LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_REGION_H -#define LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_REGION_H +#ifndef LLVM_SANDBOXIR_REGION_H +#define LLVM_SANDBOXIR_REGION_H #include #include "llvm/ADT/SetVector.h" #include "llvm/ADT/iterator_range.h" +#include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/SandboxIR/Instruction.h" #include "llvm/Support/raw_ostream.h" namespace llvm::sandboxir { +class Region; + +class ScoreBoard { + const Region &Rgn; + TargetTransformInfo &TTI; + constexpr static TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; + /// The cost of all instructions added to the region. + InstructionCost AfterCost = 0; + /// The cost of all instructions that got removed and replaced by new ones. + InstructionCost BeforeCost = 0; + /// Helper for both add() and remove(). \Returns the TTI cost of \p I. + InstructionCost getCost(Instruction *I) const; + /// No need to allow copies. + ScoreBoard(const ScoreBoard &) = delete; + const ScoreBoard &operator=(const ScoreBoard &) = delete; + +public: + ScoreBoard(Region &Rgn, TargetTransformInfo &TTI) : Rgn(Rgn), TTI(TTI) {} + /// Mark \p I as a newly added instruction to the region. + void add(Instruction *I) { AfterCost += getCost(I); } + /// Mark \p I as a deleted instruction from the region. + void remove(Instruction *I); + /// \Returns the cost of the newly added instructions. + InstructionCost getAfterCost() const { return AfterCost; } + /// \Returns the cost of the Removed instructions. + InstructionCost getBeforeCost() const { return BeforeCost; } + +#ifndef NDEBUG + void dump(raw_ostream &OS) const { + OS << "BeforeCost: " << BeforeCost << "\n"; + OS << "AfterCost: " << AfterCost << "\n"; + } + LLVM_DUMP_METHOD void dump() const; +#endif // NDEBUG +}; + /// The main job of the Region is to point to new instructions generated by /// vectorization passes. It is the unit that RegionPasses operate on with their /// runOnRegion() function. @@ -62,6 +99,8 @@ class Region { static constexpr const char *RegionStr = "sandboxregion"; Context &Ctx; + /// Keeps track of cost of instructions added and removed. + ScoreBoard Scoreboard; /// ID (for later deregistration) of the "create instruction" callback. Context::CallbackID CreateInstCB; @@ -72,7 +111,7 @@ class Region { // TODO: Add a way to encode/decode region info to/from metadata. public: - Region(Context &Ctx); + Region(Context &Ctx, TargetTransformInfo &TTI); ~Region(); Context &getContext() const { return Ctx; } @@ -91,7 +130,10 @@ class Region { iterator end() { return Insts.end(); } iterator_range insts() { return make_range(begin(), end()); } - static SmallVector> createRegionsFromMD(Function &F); + static SmallVector> + createRegionsFromMD(Function &F, TargetTransformInfo &TTI); + /// \Returns the ScoreBoard data structure that keeps track of instr costs. + const ScoreBoard &getScoreboard() const { return Scoreboard; } #ifndef NDEBUG /// This is an expensive check, meant for testing. @@ -109,4 +151,4 @@ class Region { } // namespace llvm::sandboxir -#endif // LLVM_TRANSFORMS_VECTORIZE_SANDBOXVECTORIZER_REGION_H +#endif // LLVM_SANDBOXIR_REGION_H diff --git a/llvm/include/llvm/SandboxIR/Value.h b/llvm/include/llvm/SandboxIR/Value.h index 243195f4c1c4b..28e33ca0f2312 100644 --- a/llvm/include/llvm/SandboxIR/Value.h +++ b/llvm/include/llvm/SandboxIR/Value.h @@ -167,6 +167,7 @@ class Value { // Region needs to manipulate metadata in the underlying LLVM Value, we don't // expose metadata in sandboxir. friend class Region; + friend class ScoreBoard; // Needs access to `Val` for the instruction cost. /// All values point to the context. Context &Ctx; diff --git a/llvm/lib/SandboxIR/Region.cpp b/llvm/lib/SandboxIR/Region.cpp index 1455012440f90..8c84d0c46fa10 100644 --- a/llvm/lib/SandboxIR/Region.cpp +++ b/llvm/lib/SandboxIR/Region.cpp @@ -11,7 +11,29 @@ namespace llvm::sandboxir { -Region::Region(Context &Ctx) : Ctx(Ctx) { +InstructionCost ScoreBoard::getCost(Instruction *I) const { + auto *LLVMI = cast(I->Val); + SmallVector Operands(LLVMI->operands()); + return TTI.getInstructionCost(LLVMI, Operands, CostKind); +} + +void ScoreBoard::remove(Instruction *I) { + auto Cost = getCost(I); + if (Rgn.contains(I)) + // If `I` is one the newly added ones, then we should adjust `AfterCost` + AfterCost -= Cost; + else + // If `I` is one of the original instructions (outside the region) then it + // is part of the original code, so adjust `BeforeCost`. + BeforeCost += Cost; +} + +#ifndef NDEBUG +void ScoreBoard::dump() const { dump(dbgs()); } +#endif + +Region::Region(Context &Ctx, TargetTransformInfo &TTI) + : Ctx(Ctx), Scoreboard(*this, TTI) { LLVMContext &LLVMCtx = Ctx.LLVMCtx; auto *RegionStrMD = MDString::get(LLVMCtx, RegionStr); RegionMDN = MDNode::getDistinct(LLVMCtx, {RegionStrMD}); @@ -31,9 +53,15 @@ void Region::add(Instruction *I) { Insts.insert(I); // TODO: Consider tagging instructions lazily. cast(I->Val)->setMetadata(MDKind, RegionMDN); + // Keep track of the instruction cost. + Scoreboard.add(I); } void Region::remove(Instruction *I) { + // Keep track of the instruction cost. This need to be done *before* we remove + // `I` from the region. + Scoreboard.remove(I); + Insts.remove(I); cast(I->Val)->setMetadata(MDKind, nullptr); } @@ -58,7 +86,8 @@ void Region::dump() const { } #endif // NDEBUG -SmallVector> Region::createRegionsFromMD(Function &F) { +SmallVector> +Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) { SmallVector> Regions; DenseMap MDNToRegion; auto &Ctx = F.getContext(); @@ -68,7 +97,7 @@ SmallVector> Region::createRegionsFromMD(Function &F) { Region *R = nullptr; auto It = MDNToRegion.find(MDN); if (It == MDNToRegion.end()) { - Regions.push_back(std::make_unique(Ctx)); + Regions.push_back(std::make_unique(Ctx, TTI)); R = Regions.back().get(); MDNToRegion[MDN] = R; } else { diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.cpp index 8e3f5b77429c5..121a195f45ee4 100644 --- a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.cpp +++ b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.cpp @@ -19,7 +19,7 @@ RegionsFromMetadata::RegionsFromMetadata(StringRef Pipeline) bool RegionsFromMetadata::runOnFunction(Function &F, const Analyses &A) { SmallVector> Regions = - sandboxir::Region::createRegionsFromMD(F); + sandboxir::Region::createRegionsFromMD(F, A.getTTI()); for (auto &R : Regions) { RPM.runOnRegion(*R, A); } diff --git a/llvm/unittests/SandboxIR/PassTest.cpp b/llvm/unittests/SandboxIR/PassTest.cpp index 751aedefd8fe2..19fce94563e48 100644 --- a/llvm/unittests/SandboxIR/PassTest.cpp +++ b/llvm/unittests/SandboxIR/PassTest.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "llvm/SandboxIR/Pass.h" +#include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/AsmParser/Parser.h" #include "llvm/IR/Module.h" #include "llvm/SandboxIR/Constant.h" @@ -23,10 +24,13 @@ struct PassTest : public testing::Test { llvm::LLVMContext LLVMCtx; std::unique_ptr LLVMM; std::unique_ptr Ctx; + std::unique_ptr TTI; Function *parseFunction(const char *IR, const char *FuncName) { llvm::SMDiagnostic Err; LLVMM = parseAssemblyString(IR, Err, LLVMCtx); + TTI = std::make_unique(LLVMM->getDataLayout()); + if (!LLVMM) Err.print("PassTest", llvm::errs()); Ctx = std::make_unique(LLVMCtx); @@ -119,7 +123,7 @@ define i8 @foo(i8 %v0, i8 %v1) { EXPECT_EQ(TPass.getName(), "test-pass"); // Check runOnRegion(); llvm::SmallVector> Regions = - Region::createRegionsFromMD(*F); + Region::createRegionsFromMD(*F, *TTI); ASSERT_EQ(Regions.size(), 1u); TPass.runOnRegion(*Regions[0], Analyses::emptyForTesting()); EXPECT_EQ(InstCount, 2u); @@ -242,7 +246,7 @@ define i8 @foo(i8 %v0, i8 %v1) { RPM.addPass(std::make_unique(InstCount2)); // Check runOnRegion(). llvm::SmallVector> Regions = - Region::createRegionsFromMD(*F); + Region::createRegionsFromMD(*F, *TTI); ASSERT_EQ(Regions.size(), 1u); RPM.runOnRegion(*Regions[0], Analyses::emptyForTesting()); EXPECT_EQ(InstCount1, 2u); diff --git a/llvm/unittests/SandboxIR/RegionTest.cpp b/llvm/unittests/SandboxIR/RegionTest.cpp index 47368f93a32c0..1ee72d127daa4 100644 --- a/llvm/unittests/SandboxIR/RegionTest.cpp +++ b/llvm/unittests/SandboxIR/RegionTest.cpp @@ -7,6 +7,7 @@ //===----------------------------------------------------------------------===// #include "llvm/SandboxIR/Region.h" +#include "llvm/Analysis/TargetTransformInfo.h" #include "llvm/AsmParser/Parser.h" #include "llvm/SandboxIR/Context.h" #include "llvm/SandboxIR/Function.h" @@ -20,10 +21,12 @@ using namespace llvm; struct RegionTest : public testing::Test { LLVMContext C; std::unique_ptr M; + std::unique_ptr TTI; void parseIR(LLVMContext &C, const char *IR) { SMDiagnostic Err; M = parseAssemblyString(IR, Err, C); + TTI = std::make_unique(M->getDataLayout()); if (!M) Err.print("RegionTest", errs()); } @@ -45,7 +48,7 @@ define i8 @foo(i8 %v0, i8 %v1) { auto *T0 = cast(&*It++); auto *T1 = cast(&*It++); auto *Ret = cast(&*It++); - sandboxir::Region Rgn(Ctx); + sandboxir::Region Rgn(Ctx, *TTI); // Check getContext. EXPECT_EQ(&Ctx, &Rgn.getContext()); @@ -73,7 +76,7 @@ define i8 @foo(i8 %v0, i8 %v1) { #ifndef NDEBUG // Check equality comparison. Insert in reverse order into `Other` to check // that comparison is order-independent. - sandboxir::Region Other(Ctx); + sandboxir::Region Other(Ctx, *TTI); Other.add(Ret); EXPECT_NE(Rgn, Other); Other.add(T1); @@ -98,7 +101,7 @@ define i8 @foo(i8 %v0, i8 %v1, ptr %ptr) { auto *T0 = cast(&*It++); auto *T1 = cast(&*It++); auto *Ret = cast(&*It++); - sandboxir::Region Rgn(Ctx); + sandboxir::Region Rgn(Ctx, *TTI); Rgn.add(T0); Rgn.add(T1); @@ -134,7 +137,7 @@ define i8 @foo(i8 %v0, i8 %v1) { auto *T2 = cast(&*It++); SmallVector> Regions = - sandboxir::Region::createRegionsFromMD(*F); + sandboxir::Region::createRegionsFromMD(*F, *TTI); EXPECT_THAT(Regions[0]->insts(), testing::UnorderedElementsAre(T0)); EXPECT_THAT(Regions[1]->insts(), testing::UnorderedElementsAre(T1, T2)); } @@ -160,7 +163,7 @@ define i8 @foo(i8 %v0, i8 %v1) { auto *T2 = cast(&*It++); SmallVector> Regions = - sandboxir::Region::createRegionsFromMD(*F); + sandboxir::Region::createRegionsFromMD(*F, *TTI); EXPECT_THAT(Regions[0]->insts(), testing::UnorderedElementsAre(T0, T2)); } @@ -182,9 +185,9 @@ define i8 @foo(i8 %v0, i8 %v1) { [[maybe_unused]] auto *T1 = cast(&*It++); auto *T2 = cast(&*It++); [[maybe_unused]] auto *Ret = cast(&*It++); - sandboxir::Region Rgn(Ctx); + sandboxir::Region Rgn(Ctx, *TTI); Rgn.add(T0); - sandboxir::Region Rgn2(Ctx); + sandboxir::Region Rgn2(Ctx, *TTI); Rgn2.add(T2); std::string output; @@ -226,14 +229,66 @@ define i8 @foo(i8 %v0, i8 %v1) { auto *T0 = cast(&*It++); auto *T1 = cast(&*It++); - sandboxir::Region Rgn(Ctx); + sandboxir::Region Rgn(Ctx, *TTI); Rgn.add(T0); Rgn.add(T1); SmallVector> Regions = - sandboxir::Region::createRegionsFromMD(*F); + sandboxir::Region::createRegionsFromMD(*F, *TTI); ASSERT_EQ(1U, Regions.size()); #ifndef NDEBUG EXPECT_EQ(Rgn, *Regions[0].get()); #endif } + +TEST_F(RegionTest, RegionCost) { + parseIR(C, R"IR( +define void @foo(i8 %v0, i8 %v1, i8 %v2) { + %add0 = add i8 %v0, 1 + %add1 = add i8 %v1, 2 + %add2 = add i8 %v2, 3 + ret void +} +)IR"); + llvm::Function *LLVMF = &*M->getFunction("foo"); + auto *LLVMBB = &*LLVMF->begin(); + auto LLVMIt = LLVMBB->begin(); + auto *LLVMAdd0 = &*LLVMIt++; + auto *LLVMAdd1 = &*LLVMIt++; + auto *LLVMAdd2 = &*LLVMIt++; + + 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++); + auto *Add2 = cast(&*It++); + + sandboxir::Region Rgn(Ctx, *TTI); + const auto &SB = Rgn.getScoreboard(); + EXPECT_EQ(SB.getAfterCost(), 0); + EXPECT_EQ(SB.getBeforeCost(), 0); + + auto GetCost = [this](llvm::Instruction *LLVMI) { + constexpr static TTI::TargetCostKind CostKind = TTI::TCK_RecipThroughput; + SmallVector Operands(LLVMI->operands()); + return TTI->getInstructionCost(LLVMI, Operands, CostKind); + }; + // Add `Add0` to the region, should be counted in "After". + Rgn.add(Add0); + EXPECT_EQ(SB.getBeforeCost(), 0); + EXPECT_EQ(SB.getAfterCost(), GetCost(LLVMAdd0)); + // Same for `Add1`. + Rgn.add(Add1); + EXPECT_EQ(SB.getBeforeCost(), 0); + EXPECT_EQ(SB.getAfterCost(), GetCost(LLVMAdd0) + GetCost(LLVMAdd1)); + // Remove `Add0`, should be subtracted from "After". + Rgn.remove(Add0); + EXPECT_EQ(SB.getBeforeCost(), 0); + EXPECT_EQ(SB.getAfterCost(), GetCost(LLVMAdd1)); + // Remove `Add2` which was never in the region, should counted in "Before". + Rgn.remove(Add2); + EXPECT_EQ(SB.getBeforeCost(), GetCost(LLVMAdd2)); + EXPECT_EQ(SB.getAfterCost(), GetCost(LLVMAdd1)); +}