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

116

When does this happen?

119

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?

132

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.

624

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

116

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

119

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.

132

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

624

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

123

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

5

Add baz code here too?

27

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

123

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

27

makes sense. updated the tests.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

123

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/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>
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));
 }
 
-bool unlikelyExecuted(BasicBlock &BB) {
+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 TrueWeight, FalseWeight;
+  if (!extractBranchWeights(*CondBr, TrueWeight, FalseWeight))
+    return;
+
+  if (TrueWeight + FalseWeight == 0)
+    return;
+
+  if (BranchProbability(TrueWeight, TrueWeight + FalseWeight) <= BranchProbability(1, 1000))
+    AnnotatedColdBlocks.insert(CondBr->getSuccessor(0));
+
+  if (BranchProbability(FalseWeight, TrueWeight + FalseWeight) <= BranchProbability(1, 1000))
+    AnnotatedColdBlocks.insert(CondBr->getSuccessor(1));
+}
+
+bool unlikelyExecuted(BasicBlock &BB, SmallPtrSetImpl<BasicBlock *> &AnnotatedColdBlocks) {
   // Exception handling blocks are unlikely executed.
   if (BB.isEHPad() || isa<ResumeInst>(BB.getTerminator()))
     return true;
 
+  analyzeProfMetadata(BB, AnnotatedColdBlocks);
+  if (AnnotatedColdBlocks.count(&BB))
+    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;
 

 };
 } // namespace
 
 bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
   bool Changed = false;
 
   // The set of cold blocks.
   SmallPtrSet<BasicBlock *, 4> ColdBlocks;
+  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);
bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
 
   // 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));
+                (EnableStaticAnalysis && unlikelyExecuted(*BB, AnnotatedColdBlocks));
     if (!Cold)
       continue;
 
     LLVM_DEBUG({
       dbgs() << "Found a cold block:\n";
       BB->dump();
     });
 

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

+; RUN: opt -S -passes=hotcoldsplit < %s 2>&1  | FileCheck %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]] {
+;    cold("same");
+;    cold("same");
+;  } else {
+;    hot("different");
+;  }
+;}
+
+; CHECK: internal void @_Z3fooii.cold.1() #[[ATTR0:[0-9]+]]
+; CHECK-NEXT: newFuncRoot
+; CHECK: tail call noundef i32 @_Z4coldPKc
+; CHECK: tail call noundef i32 @_Z4coldPKc
+; 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 @_Z3fooii(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 !6
+
+if.then:
+  %call = tail call noundef i32 @_Z4coldPKc(ptr noundef nonnull @.str)
+  %call1 = tail call noundef i32 @_Z4coldPKc(ptr noundef nonnull @.str)
+  br label %if.end
+
+if.else:
+  %call2 = tail call noundef i32 @_Z3hotPKc(ptr noundef nonnull @.str.1)
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare noundef i32 @_Z4coldPKc(ptr noundef) local_unnamed_addr #1
+
+declare noundef i32 @_Z3hotPKc(ptr noundef) local_unnamed_addr #1
+
+!llvm.linker.options = !{}
+!llvm.module.flags = !{!0, !1, !2, !3, !4}
+!llvm.ident = !{!5}
+
+!0 = !{i32 7, !"Dwarf Version", i32 4}
+!1 = !{i32 1, !"wchar_size", i32 4}
+!2 = !{i32 8, !"PIC Level", i32 2}
+!3 = !{i32 7, !"PIE Level", i32 2}
+!4 = !{i32 7, !"uwtable", i32 2}
+!5 = !{!"clang version 18.0.0 (https://github.com/llvm/llvm-project.git 733edde00582c2a4a6266a4cd2bed7015f4246ce)"}
+!6 = !{!"branch_weights", i32 1, i32 2000}