From c8f73c415649f3d8bc7676533af74baf09c4d881 Mon Sep 17 00:00:00 2001 From: Vasileios Porpodas Date: Fri, 24 Jan 2025 12:46:34 -0800 Subject: [PATCH] [SandboxIR][Region] Add cost modeling to the region This patch implements cost modeling for Region. All instructions that are added or removed get their cost counted in the Scoreboard. This is used for checking if the region before or after a transformation is more profitable. --- llvm/include/llvm/SandboxIR/Region.h | 52 +++++++++++-- llvm/include/llvm/SandboxIR/Value.h | 1 + llvm/lib/SandboxIR/Region.cpp | 35 ++++++++- .../Passes/RegionsFromMetadata.cpp | 2 +- llvm/unittests/SandboxIR/PassTest.cpp | 8 +- llvm/unittests/SandboxIR/RegionTest.cpp | 73 ++++++++++++++++--- 6 files changed, 151 insertions(+), 20 deletions(-) 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)); +}