[FuncSpec] Estimate dead blocks more accurately.

[labrinea]

Summary

Currently we only consider basic blocks with a unique predecessor when
estimating the size of dead code. However, we could expand to this to
consider blocks with a back-edge, or blocks preceded by dead blocks.

[ChuanqiXu]

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

109–117

My instinct reaction to the function is "doesn't DeadBlocks contain DeadBB"?

[labrinea]

So generally, when we encounter root dead blocks, then their predecessor will not be in DeadBlocks. Successors of known dead blocks do have at least one predecessor in DeadBlocks (the one we followed while traversing the graph).

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

109–117

Not yet. If isBlockExecutable and canEliminateSuccessor, then we push it to the work list. Blocks are marked dead after being popped from the work list.

158

Here, BB, which is the predecessor of SuccBB, is indeed already in DeadBlocks.

247

Here, BB corresponds to the not taken branch. If I.getParent is the predecessor of BB (despite not being in DeadBlocks), we can eliminate BB.

265

Here, Succ corresponds to the not taken branch. If I.getParent is the predecessor of Succ (despite not being in DeadBlocks), we can eliminate Succ.

[labrinea]

ping?

[ChuanqiXu]

LGTM with nit.

Sorry that I imaged I commented on this : (

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

109–117

Got it.

nit: Then I feel slightly better to make it a static private member function of InstCostVisitor. Then the reader may be less confused.

llvm/include/llvm/Transforms/IPO/FunctionSpecialization.h

public:
   Bonus getUserBonus(Instruction *User, Value *Use = nullptr,
                      Constant *C = nullptr);
 
   Bonus getBonusFromPendingPHIs();
 
 private:
   friend class InstVisitor<InstCostVisitor, Constant *>;
 
+  Cost estimateBasicBlocks(SmallVectorImpl<BasicBlock *> &WorkList);
   Cost estimateSwitchInst(SwitchInst &I);
   Cost estimateBranchInst(BranchInst &I);
 
   Constant *visitInstruction(Instruction &I) { return nullptr; }
   Constant *visitPHINode(PHINode &I);
   Constant *visitFreezeInst(FreezeInst &I);
   Constant *visitCallBase(CallBase &I);
   Constant *visitLoadInst(LoadInst &I);

llvm/lib/Transforms/IPO/FunctionSpecialization.cpp

static cl::opt<unsigned> MaxClones(
     "The maximum number of clones allowed for a single function "
     "specialization"));
 
 static cl::opt<unsigned> MaxIncomingPhiValues(
     "funcspec-max-incoming-phi-values", cl::init(4), cl::Hidden, cl::desc(
     "The maximum number of incoming values a PHI node can have to be "
     "considered during the specialization bonus estimation"));
 
+static cl::opt<unsigned> MaxBlockPredecessors(
+    "funcspec-max-block-predecessors", cl::init(2), cl::Hidden, cl::desc(
+    "The maximum number of predecessors a basic block can have to be "
+    "considered during the estimation of dead code"));
+
 static cl::opt<unsigned> MinFunctionSize(
     "funcspec-min-function-size", cl::init(100), cl::Hidden, cl::desc(
     "Don't specialize functions that have less than this number of "
     "instructions"));
 
 static cl::opt<bool> SpecializeOnAddress(
     "funcspec-on-address", cl::init(false), cl::Hidden, cl::desc(
     "Enable function specialization on the address of global values"));
 
 // Disabled by default as it can significantly increase compilation times.
 //
 // https://llvm-compile-time-tracker.com
 // https://github.com/nikic/llvm-compile-time-tracker
 static cl::opt<bool> SpecializeLiteralConstant(
     "funcspec-for-literal-constant", cl::init(false), cl::Hidden, cl::desc(
     "Enable specialization of functions that take a literal constant as an "
     "argument"));
 
+static bool canEliminateSuccessor(BasicBlock *DeadBB, BasicBlock *SuccBB,
+                                  DenseSet<BasicBlock *> &DeadBlocks) {
+  unsigned I = 0;
+  return all_of(predecessors(SuccBB),
+    [&I, DeadBB, SuccBB, &DeadBlocks] (BasicBlock *PredBB) {
+    return I++ < MaxBlockPredecessors &&
+      (PredBB == DeadBB || PredBB == SuccBB || DeadBlocks.contains(PredBB));
+  });
+}
+
 // Estimates the codesize savings due to dead code after constant propagation.
 // \p WorkList represents the basic blocks of a specialization which will
 // eventually become dead once we replace instructions that are known to be
 // constants. The successors of such blocks are added to the list as long as
 // the \p Solver found they were executable prior to specialization, and only
 // if they have a unique predecessor.
-static Cost estimateBasicBlocks(SmallVectorImpl<BasicBlock *> &WorkList,
-                                DenseSet<BasicBlock *> &DeadBlocks,
-                                ConstMap &KnownConstants, SCCPSolver &Solver,
-                                TargetTransformInfo &TTI) {
+Cost InstCostVisitor::estimateBasicBlocks(
+                          SmallVectorImpl<BasicBlock *> &WorkList) {
   Cost CodeSize = 0;
   // Accumulate the instruction cost of each basic block weighted by frequency.
   while (!WorkList.empty()) {
     BasicBlock *BB = WorkList.pop_back_val();
 
     // These blocks are considered dead as far as the InstCostVisitor
     // is concerned. They haven't been proven dead yet by the Solver,
     // but may become if we propagate the specialization arguments.
for (Instruction &I : *BB) {
       LLVM_DEBUG(dbgs() << "FnSpecialization:     CodeSize " << C
                         << " for user " << I << "\n");
       CodeSize += C;
     }
 
     // Keep adding dead successors to the list as long as they are
     // executable and they have a unique predecessor.
     for (BasicBlock *SuccBB : successors(BB))
-      if (Solver.isBlockExecutable(SuccBB) &&
-          SuccBB->getUniquePredecessor() == BB)
+      if (isBlockExecutable(SuccBB) &&
+          canEliminateSuccessor(BB, SuccBB, DeadBlocks))
         WorkList.push_back(SuccBB);
   }
   return CodeSize;
 }
 
 static Constant *findConstantFor(Value *V, ConstMap &KnownConstants) {
   if (auto *C = dyn_cast<Constant>(V))
     return C;
Bonus InstCostVisitor::getUserBonus(Instruction *User, Value *Use, Constant *C) {
   if (auto *I = dyn_cast<SwitchInst>(User)) {
     CodeSize = estimateSwitchInst(*I);
   } else if (auto *I = dyn_cast<BranchInst>(User)) {
     CodeSize = estimateBranchInst(*I);
   } else {
     C = visit(*User);
     if (!C)
       return {0, 0};
-    KnownConstants.insert({User, C});
   }
 
+  // Even though it doesn't make sense to bind switch and branch instructions
+  // with a constant, unlike any other instruction type, it prevents estimating
+  // their bonus multiple times.
+  KnownConstants.insert({User, C});
+
   CodeSize += TTI.getInstructionCost(User, TargetTransformInfo::TCK_CodeSize);
 
   uint64_t Weight = BFI.getBlockFreq(User->getParent()).getFrequency() /
                     BFI.getEntryFreq();
 
   Cost Latency = Weight *
       TTI.getInstructionCost(User, TargetTransformInfo::TCK_Latency);
 
Cost InstCostVisitor::estimateSwitchInst(SwitchInst &I) {
 
   BasicBlock *Succ = I.findCaseValue(C)->getCaseSuccessor();
   // Initialize the worklist with the dead basic blocks. These are the
   // destination labels which are different from the one corresponding
   // to \p C. They should be executable and have a unique predecessor.
   SmallVector<BasicBlock *> WorkList;
   for (const auto &Case : I.cases()) {
     BasicBlock *BB = Case.getCaseSuccessor();
-    if (BB == Succ || !Solver.isBlockExecutable(BB) ||
-        BB->getUniquePredecessor() != I.getParent())
-      continue;
+    if (BB != Succ && isBlockExecutable(BB) &&
+        canEliminateSuccessor(I.getParent(), BB, DeadBlocks))
       WorkList.push_back(BB);
   }
 
-  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, TTI);
+  return estimateBasicBlocks(WorkList);
 }
 
 Cost InstCostVisitor::estimateBranchInst(BranchInst &I) {
   assert(LastVisited != KnownConstants.end() && "Invalid iterator!");
 
   if (I.getCondition() != LastVisited->first)
     return 0;
 
   BasicBlock *Succ = I.getSuccessor(LastVisited->second->isOneValue());
   // Initialize the worklist with the dead successor as long as
   // it is executable and has a unique predecessor.
   SmallVector<BasicBlock *> WorkList;
-  if (Solver.isBlockExecutable(Succ) &&
-      Succ->getUniquePredecessor() == I.getParent())
+  if (isBlockExecutable(Succ) &&
+      canEliminateSuccessor(I.getParent(), Succ, DeadBlocks))
     WorkList.push_back(Succ);
 
-  return estimateBasicBlocks(WorkList, DeadBlocks, KnownConstants, Solver, TTI);
+  return estimateBasicBlocks(WorkList);
 }
 
 Constant *InstCostVisitor::visitPHINode(PHINode &I) {
   if (I.getNumIncomingValues() > MaxIncomingPhiValues)
     return nullptr;
 
   bool Inserted = VisitedPHIs.insert(&I).second;
   Constant *Const = nullptr;

llvm/unittests/Transforms/IPO/FunctionSpecializationTest.cpp

 }
 
 TEST_F(FunctionSpecializationTest, BranchInst) {
   const char *ModuleString = R"(
     define void @foo(i32 %a, i32 %b, i1 %cond) {
     entry:
       br label %loop
     loop:
-      br i1 %cond, label %bb0, label %bb2
+      br i1 %cond, label %bb0, label %bb3
     bb0:
       %0 = mul i32 %a, 2
       %1 = sub i32 6, 5
-      br label %bb1
+      br i1 %cond, label %bb1, label %bb2
     bb1:
       %2 = add i32 %0, %b
       %3 = sdiv i32 8, 2
-      br label %loop
+      br label %bb2
     bb2:
+      br label %loop
+    bb3:
       ret void
     }
   )";
 
   Module &M = parseModule(ModuleString);
   Function *F = M.getFunction("foo");
   FunctionSpecializer Specializer = getSpecializerFor(F);
   InstCostVisitor Visitor = Specializer.getInstCostVisitorFor(F);
 
   Constant *One = ConstantInt::get(IntegerType::getInt32Ty(M.getContext()), 1);
   Constant *False = ConstantInt::getFalse(M.getContext());
 
   auto FuncIter = F->begin();
   BasicBlock &Loop = *++FuncIter;
   BasicBlock &BB0 = *++FuncIter;
   BasicBlock &BB1 = *++FuncIter;
+  BasicBlock &BB2 = *++FuncIter;
 
   Instruction &Branch = Loop.front();
   Instruction &Mul = BB0.front();
   Instruction &Sub = *++BB0.begin();
-  Instruction &BrBB1 = BB0.back();
+  Instruction &BrBB1BB2 = BB0.back();
   Instruction &Add = BB1.front();
   Instruction &Sdiv = *++BB1.begin();
-  Instruction &BrLoop = BB1.back();
+  Instruction &BrBB2 = BB1.back();
+  Instruction &BrLoop = BB2.front();
 
   // mul
   Bonus Ref = getInstCost(Mul);
   Bonus Test = Specializer.getSpecializationBonus(F->getArg(0), One, Visitor);
   EXPECT_EQ(Test, Ref);
   EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // add
   Ref = getInstCost(Add);
   Test = Specializer.getSpecializationBonus(F->getArg(1), One, Visitor);
   EXPECT_EQ(Test, Ref);
   EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 
   // branch + sub + br + sdiv + br
   Ref = getInstCost(Branch) +
         getInstCost(Sub, /*SizeOnly =*/ true) +
-        getInstCost(BrBB1, /*SizeOnly =*/ true) +
+        getInstCost(BrBB1BB2) +
         getInstCost(Sdiv, /*SizeOnly =*/ true) +
+        getInstCost(BrBB2, /*SizeOnly =*/ true) +
         getInstCost(BrLoop, /*SizeOnly =*/ true);
   Test = Specializer.getSpecializationBonus(F->getArg(2), False, Visitor);
   EXPECT_EQ(Test, Ref);
   EXPECT_TRUE(Test.CodeSize > 0 && Test.Latency > 0);
 }
 
 TEST_F(FunctionSpecializationTest, Misc) {
   const char *ModuleString = R"(