[SandboxIR][Region] Add cost modeling to the region

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

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;

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 {

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);
   }

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);

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));
+}