[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.

Author: vporpo

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