[llvm] 4b209c5 - [SandboxIR][Region] Add cost modeling to the region (#124354)
via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 24 14:28:59 PST 2025
Author: vporpo
Date: 2025-01-24T14:28:55-08:00
New Revision: 4b209c5d87c8b8eb4bbf2750ea9daa5927a13699
URL: https://github.com/llvm/llvm-project/commit/4b209c5d87c8b8eb4bbf2750ea9daa5927a13699
DIFF: https://github.com/llvm/llvm-project/commit/4b209c5d87c8b8eb4bbf2750ea9daa5927a13699.diff
LOG: [SandboxIR][Region] Add cost modeling to the region (#124354)
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.
Added:
Modified:
llvm/include/llvm/SandboxIR/Region.h
llvm/include/llvm/SandboxIR/Value.h
llvm/lib/SandboxIR/Region.cpp
llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/RegionsFromMetadata.cpp
llvm/unittests/SandboxIR/PassTest.cpp
llvm/unittests/SandboxIR/RegionTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/SandboxIR/Region.h b/llvm/include/llvm/SandboxIR/Region.h
index 8133e01734ea78..c1195141cb54ce 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 <memory>
#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<iterator> insts() { return make_range(begin(), end()); }
- static SmallVector<std::unique_ptr<Region>> createRegionsFromMD(Function &F);
+ static SmallVector<std::unique_ptr<Region>>
+ 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 243195f4c1c4bd..28e33ca0f23125 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 1455012440f90c..8c84d0c46fa10a 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<llvm::Instruction>(I->Val);
+ SmallVector<const llvm::Value *> 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<llvm::Instruction>(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<llvm::Instruction>(I->Val)->setMetadata(MDKind, nullptr);
}
@@ -58,7 +86,8 @@ void Region::dump() const {
}
#endif // NDEBUG
-SmallVector<std::unique_ptr<Region>> Region::createRegionsFromMD(Function &F) {
+SmallVector<std::unique_ptr<Region>>
+Region::createRegionsFromMD(Function &F, TargetTransformInfo &TTI) {
SmallVector<std::unique_ptr<Region>> Regions;
DenseMap<MDNode *, Region *> MDNToRegion;
auto &Ctx = F.getContext();
@@ -68,7 +97,7 @@ SmallVector<std::unique_ptr<Region>> Region::createRegionsFromMD(Function &F) {
Region *R = nullptr;
auto It = MDNToRegion.find(MDN);
if (It == MDNToRegion.end()) {
- Regions.push_back(std::make_unique<Region>(Ctx));
+ Regions.push_back(std::make_unique<Region>(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 8e3f5b77429c5a..121a195f45ee46 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<std::unique_ptr<sandboxir::Region>> 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 751aedefd8fe2d..19fce94563e485 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<llvm::Module> LLVMM;
std::unique_ptr<Context> Ctx;
+ std::unique_ptr<llvm::TargetTransformInfo> TTI;
Function *parseFunction(const char *IR, const char *FuncName) {
llvm::SMDiagnostic Err;
LLVMM = parseAssemblyString(IR, Err, LLVMCtx);
+ TTI = std::make_unique<llvm::TargetTransformInfo>(LLVMM->getDataLayout());
+
if (!LLVMM)
Err.print("PassTest", llvm::errs());
Ctx = std::make_unique<Context>(LLVMCtx);
@@ -119,7 +123,7 @@ define i8 @foo(i8 %v0, i8 %v1) {
EXPECT_EQ(TPass.getName(), "test-pass");
// Check runOnRegion();
llvm::SmallVector<std::unique_ptr<Region>> 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<TestPass2>(InstCount2));
// Check runOnRegion().
llvm::SmallVector<std::unique_ptr<Region>> 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 47368f93a32c0c..1ee72d127daa4d 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<Module> M;
+ std::unique_ptr<TargetTransformInfo> TTI;
void parseIR(LLVMContext &C, const char *IR) {
SMDiagnostic Err;
M = parseAssemblyString(IR, Err, C);
+ TTI = std::make_unique<TargetTransformInfo>(M->getDataLayout());
if (!M)
Err.print("RegionTest", errs());
}
@@ -45,7 +48,7 @@ define i8 @foo(i8 %v0, i8 %v1) {
auto *T0 = cast<sandboxir::Instruction>(&*It++);
auto *T1 = cast<sandboxir::Instruction>(&*It++);
auto *Ret = cast<sandboxir::Instruction>(&*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<sandboxir::Instruction>(&*It++);
auto *T1 = cast<sandboxir::Instruction>(&*It++);
auto *Ret = cast<sandboxir::Instruction>(&*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<sandboxir::Instruction>(&*It++);
SmallVector<std::unique_ptr<sandboxir::Region>> 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<sandboxir::Instruction>(&*It++);
SmallVector<std::unique_ptr<sandboxir::Region>> 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<sandboxir::Instruction>(&*It++);
auto *T2 = cast<sandboxir::Instruction>(&*It++);
[[maybe_unused]] auto *Ret = cast<sandboxir::Instruction>(&*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<sandboxir::Instruction>(&*It++);
auto *T1 = cast<sandboxir::Instruction>(&*It++);
- sandboxir::Region Rgn(Ctx);
+ sandboxir::Region Rgn(Ctx, *TTI);
Rgn.add(T0);
Rgn.add(T1);
SmallVector<std::unique_ptr<sandboxir::Region>> 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<sandboxir::Instruction>(&*It++);
+ auto *Add1 = cast<sandboxir::Instruction>(&*It++);
+ auto *Add2 = cast<sandboxir::Instruction>(&*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<const llvm::Value *> 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));
+}
More information about the llvm-commits
mailing list