Statically analyze likely and unlikely blocks based on metadata

[hiraditya]

Summary

The builtin_expect(), and C++20's likely, unlikely attributes assign branch_weights to annotated branches.

This patch adds the the ability to query branch !prof metadata and improve static analysis based on that.

https://github.com/llvm/llvm-project/issues/64998

[tejohnson]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

143

I'm confused about how this works. It appears that analyzeProfMetadata analyzes the successors of BB, but then we query the information for the current BB. Ah, it looks like this works because this is called in a reversed PO traversal of the BBs?

It would be good to add some comments here. Otherwise it looks a little confusing as the back to back call and set lookup make it appear as though the call is doing the work for the set lookup in the same invocation, but really it is doing the work for a later invocation, with unstated traversal ordering assumptions.

147

When does this happen?

150

Should this use BranchProbability::getBranchProbability instead of the constructor directly?

Also, is the (1,1000) threshold related to or the same as the branch probability assigned for builtin_expect and likely/unlikely?

654–655

Should the static analysis i.e. the branch probabilities only be looked at if there is no BFI aka profile information? I think if there is profile info that the creation of the BFI/PSI has already taken expect and unlikely/likely attributes into account?

llvm/test/Transforms/HotColdSplit/split-static-profile.ll

35

Missing "CHECK:"?

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

143

yes it is due to RPOT. i'have added comments.

147

i see this in https://llvm.org/doxygen/JumpThreading_8cpp_source.html so added as precaution. although i'm not 100% sure about it.

150

builtins adds 1/2000 but i think that is too low. i've modified this to reflect builtins but added a flag for users to provide as they see fit.

654–655

that is possible for unlikely/likely attributes true. but unlikelyExecuted does more analysis taking into account ehpad, noreturn functions etc. i'll refactor this code a little.

[tejohnson]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

128

Is there an internal option that determines the value used there that can be used in the code here? Looks like it is controlled by LikelyBranchWeight - can that be made a global option and referenced here so that they don't get out of sync?

llvm/test/Transforms/HotColdSplit/split-static-profile.ll

4

Add baz code here too?

26

I think these cases should also be the same in the CHECK-PROB case - probably want a common check label to check the common cases between the two invocations, and a different check label just for the different parts.

Also, probably want to ensure that the hot calls/funcs are all above the cold calls/funcs - the checks below don't guarantee that.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

128

found one. it has a higher probability but that is what other parts of llvm is using so i'll use the same here.

llvm/test/Transforms/HotColdSplit/split-static-profile.ll

26

makes sense. updated the tests.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

128

TTI.getPredictableBranchThreshold is what i used. it seems this is the recommended API to use for branch threshold (https://llvm.org/doxygen/LowerExpectIntrinsic_8cpp_source.html)

[tejohnson]

lgtm wth a couple of small test tweaks.

llvm/test/Transforms/HotColdSplit/split-static-profile.ll

5

Suggest moving this comment to the top of file

99

Suggest moving these CHECK-PROB lines up earlier, to where you have the CHECK-NOOUTLINE-BAZ-NOT, and the checks for the other 2 functions, so that all the checks are consolidated.

103

Probably check the attributes once with CHECK-OUTLINE, which covers both cases (I think the CHECK-OUTLINE,CHECK-NOOUTLINE-BAZ case is not checking the attributes currently).

llvm/include/llvm/Transforms/IPO/HotColdSplitting.h

HotColdSplitting(ProfileSummaryInfo *ProfSI,
                    function_ref<TargetTransformInfo &(Function &)> GTTI,
                    std::function<OptimizationRemarkEmitter &(Function &)> *GORE,
                    function_ref<AssumptionCache *(Function &)> LAC)
       : PSI(ProfSI), GetBFI(GBFI), GetTTI(GTTI), GetORE(GORE), LookupAC(LAC) {}
   bool run(Module &M);
 
 private:
   bool isFunctionCold(const Function &F) const;
+  bool isBasicBlockCold(BasicBlock* BB,
+                        SmallPtrSetImpl<BasicBlock *> &ColdBlocks,
+                        SmallPtrSetImpl<BasicBlock *> &AnnotatedColdBlocks,
+                        BlockFrequencyInfo *BFI) const;
   bool shouldOutlineFrom(const Function &F) const;
   bool outlineColdRegions(Function &F, bool HasProfileSummary);
   Function *extractColdRegion(const BlockSequence &Region,
                               const CodeExtractorAnalysisCache &CEAC,
                               DominatorTree &DT, BlockFrequencyInfo *BFI,
                               TargetTransformInfo &TTI,
                               OptimizationRemarkEmitter &ORE,
                               AssumptionCache *AC, unsigned Count);

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PassManager.h"
+#include "llvm/IR/ProfDataUtils.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include <algorithm>
ColdSectionName("hotcoldsplit-cold-section-name", cl::init("__llvm_cold"),
                     cl::Hidden,
                     cl::desc("Name for the section containing cold functions "
                              "extracted by hot-cold splitting."));
 
 static cl::opt<int> MaxParametersForSplit(
     "hotcoldsplit-max-params", cl::init(4), cl::Hidden,
     cl::desc("Maximum number of parameters for a split function"));
 
+static cl::opt<int> ColdBranchProbDenom(
+    "hotcoldsplit-cold-probability-denom", cl::init(2000), cl::Hidden,
+    cl::desc("Divisor of cold branch probability."
+             "BranchProbability = 1/ColdBranchProbDenom"));
+
 namespace {
 // Same as blockEndsInUnreachable in CodeGen/BranchFolding.cpp. Do not modify
 // this function unless you modify the MBB version as well.
 //
 /// A no successor, non-return block probably ends in unreachable and is cold.
 /// Also consider a block that ends in an indirect branch to be a return block,
 /// since many targets use plain indirect branches to return.
 bool blockEndsInUnreachable(const BasicBlock &BB) {
   if (!succ_empty(&BB))
     return false;
   if (BB.empty())
     return true;
   const Instruction *I = BB.getTerminator();
   return !(isa<ReturnInst>(I) || isa<IndirectBrInst>(I));
 }
 
+void analyzeProfMetadata(BasicBlock *BB, SmallPtrSetImpl<BasicBlock *> &AnnotatedColdBlocks) {
+  // TODO: Handle branches with > 2 successors.
+  BranchInst *CondBr = dyn_cast<BranchInst>(BB->getTerminator());
+  if (!CondBr)
+    return;
+
+  uint64_t TrueWt, FalseWt;
+  if (!extractBranchWeights(*CondBr, TrueWt, FalseWt))
+    return;
+
+  auto SumWt = TrueWt + FalseWt;
+  if (SumWt == 0)
+    return;
+
+  auto TrueProb = BranchProbability::getBranchProbability(TrueWt, SumWt);
+  auto FalseProb = BranchProbability::getBranchProbability(FalseWt, SumWt);
+
+  // [[unlikely]] and builtin_expect assigns 1/2000 as cold branch probability.
+  auto ColdProb = BranchProbability(1, std::min(2000, ColdBranchProbDenom.getValue()));
+  if (TrueProb <= ColdProb)
+    AnnotatedColdBlocks.insert(CondBr->getSuccessor(0));
+
+  if (FalseProb <= ColdProb)
+    AnnotatedColdBlocks.insert(CondBr->getSuccessor(1));
+}
+
 bool unlikelyExecuted(BasicBlock &BB) {
   // Exception handling blocks are unlikely executed.
   if (BB.isEHPad() || isa<ResumeInst>(BB.getTerminator()))
     return true;
 
   // The block is cold if it calls/invokes a cold function. However, do not
   // mark sanitizer traps as cold.
   for (Instruction &I : BB)
     if (auto *CB = dyn_cast<CallBase>(&I))
       if (CB->hasFnAttr(Attribute::Cold) &&
           !CB->getMetadata(LLVMContext::MD_nosanitize))
         return true;
 
   // The block is cold if it has an unreachable terminator, unless it's
   // preceded by a call to a (possibly warm) noreturn call (e.g. longjmp).
   if (blockEndsInUnreachable(BB)) {
     if (auto *CI =
             dyn_cast_or_null<CallInst>(BB.getTerminator()->getPrevNode()))
       if (CI->hasFnAttr(Attribute::NoReturn))
         return false;
     return true;
   }
if (F.getCallingConv() == CallingConv::Cold)
     return true;
 
   if (PSI->isFunctionEntryCold(&F))
     return true;
 
   return false;
 }
 
+bool HotColdSplitting::isBasicBlockCold(BasicBlock *BB,
+                          SmallPtrSetImpl<BasicBlock *> &ColdBlocks,
+                          SmallPtrSetImpl<BasicBlock *> &AnnotatedColdBlocks,
+                          BlockFrequencyInfo *BFI) const {
+  // This block is already part of some outlining region.
+  if (ColdBlocks.count(BB))
+    return true;
+
+  if (BFI) {
+    if (PSI->isColdBlock(BB, BFI))
+      return true;
+  } else {
+    // Find cold blocks of successors of BB during a reverse postorder traversal.
+    analyzeProfMetadata(BB, AnnotatedColdBlocks);
+
+    // A statically cold BB would be known before it is visited
+    // because the prof-data of incoming edges are 'analyzed' as part of RPOT.
+    if (AnnotatedColdBlocks.count(BB))
+      return true;
+  }
+
+  if (EnableStaticAnalysis && unlikelyExecuted(*BB))
+    return true;
+  
+  return false;
+}
+
 // Returns false if the function should not be considered for hot-cold split
 // optimization.
 bool HotColdSplitting::shouldOutlineFrom(const Function &F) const {
   if (F.hasFnAttribute(Attribute::AlwaysInline))
     return false;
 
   if (F.hasFnAttribute(Attribute::NoInline))
     return false;

 } // namespace
 
 bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
   bool Changed = false;
 
   // The set of cold blocks.
   SmallPtrSet<BasicBlock *, 4> ColdBlocks;
 
+  // Set of cold blocks obtained with RPOT.
+  SmallPtrSet<BasicBlock *, 4> AnnotatedColdBlocks;
+
   // The worklist of non-intersecting regions left to outline.
   SmallVector<OutliningRegion, 2> OutliningWorklist;
 
   // Set up an RPO traversal. Experimentally, this performs better (outlines
   // more) than a PO traversal, because we prevent region overlap by keeping
   // the first region to contain a block.
   ReversePostOrderTraversal<Function *> RPOT(&F);
 
   // Calculate domtrees lazily. This reduces compile-time significantly.
   std::unique_ptr<DominatorTree> DT;
   std::unique_ptr<PostDominatorTree> PDT;
 
   // Calculate BFI lazily (it's only used to query ProfileSummaryInfo). This
   // reduces compile-time significantly. TODO: When we *do* use BFI, we should
   // be able to salvage its domtrees instead of recomputing them.
   BlockFrequencyInfo *BFI = nullptr;
   if (HasProfileSummary)
     BFI = GetBFI(F);
 
   TargetTransformInfo &TTI = GetTTI(F);
   OptimizationRemarkEmitter &ORE = (*GetORE)(F);
   AssumptionCache *AC = LookupAC(F);
 
   // Find all cold regions.
   for (BasicBlock *BB : RPOT) {
-    // This block is already part of some outlining region.
-    if (ColdBlocks.count(BB))
-      continue;
-
-    bool Cold = (BFI && PSI->isColdBlock(BB, BFI)) ||
-                (EnableStaticAnalysis && unlikelyExecuted(*BB));
-    if (!Cold)
+    if (!isBasicBlockCold(BB, ColdBlocks, AnnotatedColdBlocks, BFI))
       continue;
 
     LLVM_DEBUG({
       dbgs() << "Found a cold block:\n";
       BB->dump();
     });
 
     if (!DT)

llvm/test/Transforms/HotColdSplit/split-static-profile.ll

+; RUN: opt -S -passes=hotcoldsplit < %s | FileCheck %s
+; RUN: opt -S -passes=hotcoldsplit -hotcoldsplit-cold-probability-denom=1000 < %s | FileCheck --check-prefix=CHECK-PROB %s
+
+; Check that the unlikely branch is outlined.
+; int cold(const char*);
+; int hot(const char*);
+; void foo(int a, int b) {
+;   if (a == b) [[unlikely]] { // Should be outlined.
+;     cold("same");
+;     cold("same");
+;   } else {
+;     hot("different");
+;   }
+; }
+
+; void bar(int a, int b) {
+;   if (a == b) [[likely]] {
+;     hot("same");
+;   } else { // Should be outlined.
+;     cold("different");
+;     cold("different");
+;   }
+; }
+
+
+; CHECK: internal void @foo.cold.1() #[[ATTR0:[0-9]+]]
+; CHECK-NEXT: newFuncRoot
+; CHECK: tail call noundef i32 @cold
+; CHECK: tail call noundef i32 @cold
+
+; CHECK: internal void @bar.cold.1() #[[ATTR0:[0-9]+]]
+; CHECK-NEXT: newFuncRoot
+; CHECK: tail call noundef i32 @cold
+; CHECK: tail call noundef i32 @cold
+
+; CHECK-NOT: internal void @baz.cold.1()
+; CHECK: attributes #[[ATTR0]] = { cold minsize }
+
+source_filename = "/app/example.cpp"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+@.str = private unnamed_addr constant [5 x i8] c"same\00", align 1
+@.str.1 = private unnamed_addr constant [10 x i8] c"different\00", align 1
+
+define dso_local void @foo(i32 noundef %a, i32 noundef %b) local_unnamed_addr {
+entry:
+  %cmp = icmp eq i32 %a, %b
+  br i1 %cmp, label %if.then, label %if.else, !prof !1
+
+if.then:
+  %call = tail call noundef i32 @cold(ptr noundef nonnull @.str)
+  %call1 = tail call noundef i32 @cold(ptr noundef nonnull @.str)
+  br label %if.end
+
+if.else:
+  %call2 = tail call noundef i32 @hot(ptr noundef nonnull @.str.1)
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare noundef i32 @cold(ptr noundef) local_unnamed_addr #1
+
+declare noundef i32 @hot(ptr noundef) local_unnamed_addr #1
+
+define dso_local void @bar(i32 noundef %a, i32 noundef %b) local_unnamed_addr {
+entry:
+  %cmp = icmp eq i32 %a, %b
+  br i1 %cmp, label %if.then, label %if.else, !prof !2
+
+if.then:
+  %call = tail call noundef i32 @hot(ptr noundef nonnull @.str)
+  br label %if.end
+
+if.else:
+  %call1 = tail call noundef i32 @cold(ptr noundef nonnull @.str.1)
+  %call2 = tail call noundef i32 @cold(ptr noundef nonnull @.str.1)
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+; CHECK-PROB: internal void @baz.cold.1() #[[ATTR0:[0-9]+]]
+; CHECK-PROB-NEXT: newFuncRoot
+; CHECK-PROB: tail call noundef i32 @cold
+; CHECK-PROB: tail call noundef i32 @cold
+; CHECK-PROB: attributes #[[ATTR0]] = { cold minsize }
+define dso_local void @baz(i32 noundef %a, i32 noundef %b) local_unnamed_addr {
+entry:
+  %cmp = icmp eq i32 %a, %b
+  br i1 %cmp, label %if.then, label %if.else, !prof !3
+
+if.then:
+  %call = tail call noundef i32 @hot(ptr noundef nonnull @.str)
+  br label %if.end
+
+if.else:
+  %call1 = tail call noundef i32 @cold(ptr noundef nonnull @.str.1)
+  %call2 = tail call noundef i32 @cold(ptr noundef nonnull @.str.1)
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+!1 = !{!"branch_weights", i32 1, i32 2000}
+!2 = !{!"branch_weights", i32 2000, i32 1}
+!3 = !{!"branch_weights", i32 1000, i32 1}
 No newline at end of file