[PGO]Implement metadata combine for 'branch_weights' of direct callsites when none of the instructions folds the rest away

[mingmingl]

Summary

• Merge cases are added for simplify-cfg {sink,hoist}, based on https://gcc.godbolt.org/z/avGvc38W7 and https://gcc.godbolt.org/z/dbWbjGhaE
• When one instruction folds the others in, do not update branch_weights with sum (see test/Transforms/GVN/calls-readonly.ll)

[nikic]

llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-in-hoist.ll

2

You need --check-globals to check the resulting metadata. Also, it should be possible to reduce these tests to basically an if/else with a call each.

[mingmingl]

resolve comments.

[davidxl]

llvm/include/llvm/IR/Metadata.h

1327

Have an one line summary of what the interface does here. Move the detailed documentation to the function definition.

llvm/lib/IR/Metadata.cpp

1075

Why not naming it 'mergeDirectCallProfMetadata'? Even if there is a need to merge indirect call profiles, the implementation will be quite different and may require a different method for it.

1078

Combine these two into if (! (A&&B)) return;

1101

how about == 2?

1116

why is the check needed?

[mingmingl]

resolve comments.

[davidxl]

lgtm

[mingmingl]

llvm/lib/IR/Metadata.cpp

1075

Even if there is a need to merge indirect call profiles, the implementation will be quite different and may require a different method for it.

This is true. Here either the caller 'combineMetadata' or callee ('getMergedProfMetadata' or 'mergeDirectCallProfMetadata') needs to look at 'AInstr' or 'BInstr' to decide if it's the implemented case, and it looks cleaner (in the caller) by putting details in the callee 'getMergedProfMetadata'

Added a private method 'mergeDirectCallProfMetadata' to make future expansions easier.

1101

'== 2' works for 'callinstr' with 'branch_weights' (e.g., after seeing AProfName.equals("branch_weights") && BProfName.equals("branch_weights")) -> yet !prof for direct callsites might be expanded in the future (like invoke prof could have both branch weights and value profiles) so here make assertions closer what LLVM verifier already verifies.

1116

Changed it to assert since verifier already has Check(mdconst::dyn_extract<ConstantInt>(MDO), "!prof brunch_weights operand is not a const int");

llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-in-hoist.ll

2

thanks! Reduced the test case and added '-check-globals'

[nikic]

llvm/lib/Transforms/Utils/Local.cpp

2713

This should probably be behind a DoesKMove check? If we're CSEing with a dominating call, it doesn't make sense to add the branch weights, as the second call is being removed.

I'd suggest adding a test with a readonly/nounwind/willreturn call being GVNed to cover that case.

[davidxl]

llvm/lib/Transforms/Utils/Local.cpp

2713

good catch.

[mingmingl]

plan changes as suggested.

llvm/lib/Transforms/Utils/Local.cpp

2713

thanks for the catch!

Just to confirm I'm getting the point (not particular to a specific pass yet), in https://gcc.godbolt.org/z/nGa6YqWdP, two calls to 'func1' in line 11 and line14 should keep one branch weight (rather than the sum of two). So that is the case when unifying (not merge) is correct.

[nikic]

llvm/lib/Transforms/Utils/Local.cpp

2713

Yeah, that's the case I had in mind. After the transform only the first call will be executed, so it would make sense to me to only keep its call count.

[mingmingl]

gate the sum of branch weights by DoesKMove and add test case when one callsite folds away (GVN'ed the other)

[mingmingl]

llvm/lib/Transforms/Utils/Local.cpp

2713

https://gcc.godbolt.org/z/9cnd3qEce is a test case when EarlyCSE CSE'ed call to 'func1', while turns out the current implementation doesn't invoke merge-metadata around where calls are CSE'ed (probably goes to a follow-up patch)

Modified existing test case in GVN.

llvm/include/llvm/IR/Metadata.h

private:
   static void dispatchRecalculateHash(NodeTy *, std::false_type) {}
   template <class NodeTy>
   static void dispatchResetHash(NodeTy *N, std::true_type) {
     N->setHash(0);
   }
   template <class NodeTy>
   static void dispatchResetHash(NodeTy *, std::false_type) {}
 
+  /// Merge branch weights from two direct callsites.
+  static MDNode *mergeDirectCallProfMetadata(MDNode *A, MDNode *B,
+                                             const Instruction *AInstr,
+                                             const Instruction *BInstr);
+
 public:
   using op_iterator = const MDOperand *;
   using op_range = iterator_range<op_iterator>;
 
   op_iterator op_begin() const {
     return const_cast<MDNode *>(this)->mutable_begin();
   }
 
#include "llvm/IR/Metadata.def"
   /// Methods for metadata merging.
   static MDNode *concatenate(MDNode *A, MDNode *B);
   static MDNode *intersect(MDNode *A, MDNode *B);
   static MDNode *getMostGenericTBAA(MDNode *A, MDNode *B);
   static MDNode *getMostGenericFPMath(MDNode *A, MDNode *B);
   static MDNode *getMostGenericRange(MDNode *A, MDNode *B);
   static MDNode *getMostGenericAliasScope(MDNode *A, MDNode *B);
   static MDNode *getMostGenericAlignmentOrDereferenceable(MDNode *A, MDNode *B);
+  /// Merge !prof metadata from two instructions.
+  /// Currently only implemented with direct callsites with branch weights.
+  static MDNode *getMergedProfMetadata(MDNode *A, MDNode *B,
+                                       const Instruction *AInstr,
+                                       const Instruction *BInstr);
 };
 
 /// Tuple of metadata.
 ///
 /// This is the simple \a MDNode arbitrary tuple.  Nodes are uniqued by
 /// default based on their operands.
 class MDTuple : public MDNode {
   friend class LLVMContextImpl;

llvm/lib/IR/Metadata.cpp

MDNode *MDNode::getMostGenericFPMath(MDNode *A, MDNode *B) {
 
   APFloat AVal = mdconst::extract<ConstantFP>(A->getOperand(0))->getValueAPF();
   APFloat BVal = mdconst::extract<ConstantFP>(B->getOperand(0))->getValueAPF();
   if (AVal < BVal)
     return A;
   return B;
 }
 
+// Call instructions with branch weights are only used in SamplePGO as
+// documented in
+/// https://llvm.org/docs/BranchWeightMetadata.html#callinst).
+MDNode *MDNode::mergeDirectCallProfMetadata(MDNode *A, MDNode *B,
+                                            const Instruction *AInstr,
+                                            const Instruction *BInstr) {
+  assert(A && B && AInstr && BInstr && "Caller should guarantee");
+  auto &Ctx = AInstr->getContext();
+  MDBuilder MDHelper(Ctx);
+
+  // LLVM IR verifier verifies !prof metadata has at least 2 operands.
+  assert(A->getNumOperands() >= 2 && B->getNumOperands() >= 2 &&
+         "!prof annotations should have no less than 2 operands");
+  MDString *AMDS = dyn_cast<MDString>(A->getOperand(0));
+  MDString *BMDS = dyn_cast<MDString>(B->getOperand(0));
+  // LLVM IR verfier verifies first operand is MDString.
+  assert(AMDS != nullptr && BMDS != nullptr &&
+         "first operand should be a non-null MDString");
+  StringRef AProfName = AMDS->getString();
+  StringRef BProfName = BMDS->getString();
+  if (AProfName.equals("branch_weights") &&
+      BProfName.equals("branch_weights")) {
+    ConstantInt *AInstrWeight =
+        mdconst::dyn_extract<ConstantInt>(A->getOperand(1));
+    ConstantInt *BInstrWeight =
+        mdconst::dyn_extract<ConstantInt>(B->getOperand(1));
+    assert(AInstrWeight && BInstrWeight && "verified by LLVM verifier");
+    return MDNode::get(Ctx,
+                       {MDHelper.createString("branch_weights"),
+                        MDHelper.createConstant(ConstantInt::get(
+                            Type::getInt64Ty(Ctx),
+                            SaturatingAdd(AInstrWeight->getZExtValue(),
+                                          BInstrWeight->getZExtValue())))});
+  }
+  return nullptr;
+}
+
+// Pass in both instructions and nodes. Instruction information (e.g.,
+// instruction type) helps interpret profiles and make implementation clearer.
+MDNode *MDNode::getMergedProfMetadata(MDNode *A, MDNode *B,
+                                      const Instruction *AInstr,
+                                      const Instruction *BInstr) {
+  if (!(A && B)) {
+    return A ? A : B;
+  }
+
+  assert(AInstr->getMetadata(LLVMContext::MD_prof) == A &&
+         "Caller should guarantee");
+  assert(BInstr->getMetadata(LLVMContext::MD_prof) == B &&
+         "Caller should guarantee");
+
+  const CallInst *ACall = dyn_cast<CallInst>(AInstr);
+  const CallInst *BCall = dyn_cast<CallInst>(BInstr);
+
+  // Both ACall and BCall are direct callsites.
+  if (ACall && BCall && ACall->getCalledFunction() &&
+      BCall->getCalledFunction())
+    return mergeDirectCallProfMetadata(A, B, AInstr, BInstr);
+
+  // The rest of the cases are not implemented but could be added
+  // when there are use cases.
+  return nullptr;
+}
+
 static bool isContiguous(const ConstantRange &A, const ConstantRange &B) {
   return A.getUpper() == B.getLower() || A.getLower() == B.getUpper();
 }
 
 static bool canBeMerged(const ConstantRange &A, const ConstantRange &B) {
   return !A.intersectWith(B).isEmptySet() || isContiguous(A, B);
 }
 

llvm/lib/Transforms/Utils/Local.cpp

switch (Kind) {
         // If K does move, keep noundef if it is present in both instructions.
         if (DoesKMove)
           K->setMetadata(Kind, JMD);
         break;
       case LLVMContext::MD_nontemporal:
         // Preserve !nontemporal if it is present on both instructions.
         K->setMetadata(Kind, JMD);
         break;
+      case LLVMContext::MD_prof:
+        if (DoesKMove)
+          K->setMetadata(Kind, MDNode::getMergedProfMetadata(KMD, JMD, K, J));
+        break;
     }
   }
   // Set !invariant.group from J if J has it. If both instructions have it
   // then we will just pick it from J - even when they are different.
   // Also make sure that K is load or store - f.e. combining bitcast with load
   // could produce bitcast with invariant.group metadata, which is invalid.
   // FIXME: we should try to preserve both invariant.group md if they are
   // different, but right now instruction can only have one invariant.group.
unsigned KnownIDs[] = {LLVMContext::MD_tbaa,
                          LLVMContext::MD_invariant_load,
                          LLVMContext::MD_nonnull,
                          LLVMContext::MD_invariant_group,
                          LLVMContext::MD_align,
                          LLVMContext::MD_dereferenceable,
                          LLVMContext::MD_dereferenceable_or_null,
                          LLVMContext::MD_access_group,
                          LLVMContext::MD_preserve_access_index,
+                         LLVMContext::MD_prof,
                          LLVMContext::MD_nontemporal,
                          LLVMContext::MD_noundef};
   combineMetadata(K, J, KnownIDs, KDominatesJ);
 }
 
 void llvm::copyMetadataForLoad(LoadInst &Dest, const LoadInst &Source) {
   SmallVector<std::pair<unsigned, MDNode *>, 8> MD;
   Source.getAllMetadata(MD);

llvm/test/Transforms/GVN/calls-readonly.ll

 ; RUN: opt < %s -passes=gvn -S | FileCheck %s
 ; Should delete the second call to strlen even though the intervening strchr call exists.
 
 target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:128:128"
 target triple = "i386-apple-darwin7"
 
 define ptr @test(ptr %P, ptr %Q, i32 %x, i32 %y) nounwind readonly {
 entry:
-  %0 = tail call i32 @strlen(ptr %P)              ; <i32> [#uses=2]
+  %0 = tail call i32 @strlen(ptr %P), !prof !0    ; <i32> [#uses=2]
   %1 = icmp eq i32 %0, 0                          ; <i1> [#uses=1]
   br i1 %1, label %bb, label %bb1
 
 bb:                                               ; preds = %entry
   %2 = sdiv i32 %x, %y                            ; <i32> [#uses=1]
   br label %bb1
 
 bb1:                                              ; preds = %bb, %entry
   %x_addr.0 = phi i32 [ %2, %bb ], [ %x, %entry ] ; <i32> [#uses=1]
   %3 = tail call ptr @strchr(ptr %Q, i32 97)      ; <ptr> [#uses=1]
-  %4 = tail call i32 @strlen(ptr %P)              ; <i32> [#uses=1]
+  %4 = tail call i32 @strlen(ptr %P) , !prof !1   ; <i32> [#uses=1]
   %5 = add i32 %x_addr.0, %0                      ; <i32> [#uses=1]
   %.sum = sub i32 %5, %4                          ; <i32> [#uses=1]
   %6 = getelementptr i8, ptr %3, i32 %.sum            ; <ptr> [#uses=1]
   ret ptr %6
 }
 
 ; CHECK: define ptr @test(ptr %P, ptr %Q, i32 %x, i32 %y) #0 {
 ; CHECK: entry:
-; CHECK-NEXT:   %0 = tail call i32 @strlen(ptr %P)
+; CHECK-NEXT:   %0 = tail call i32 @strlen(ptr %P), !prof !0
 ; CHECK-NEXT:   %1 = icmp eq i32 %0, 0
 ; CHECK-NEXT:   br i1 %1, label %bb, label %bb1
 ; CHECK: bb:
 ; CHECK-NEXT:   %2 = sdiv i32 %x, %y
 ; CHECK-NEXT:   br label %bb1
 ; CHECK: bb1:
 ; CHECK-NEXT:   %x_addr.0 = phi i32 [ %2, %bb ], [ %x, %entry ]
 ; CHECK-NEXT:   %3 = tail call ptr @strchr(ptr %Q, i32 97)
 ; CHECK-NEXT:   %4 = add i32 %x_addr.0, %0
 ; CHECK-NEXT:   %5 = getelementptr i8, ptr %3, i32 %x_addr.0
 ; CHECK-NEXT:   ret ptr %5
 ; CHECK: }
 
 declare i32 @strlen(ptr) nounwind readonly
 
 declare ptr @strchr(ptr, i32) nounwind readonly
+
+!0 = !{!"branch_weights", i32 95}
+!1 = !{!"branch_weights", i32 95}

llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-in-hoist.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
+; RUN: opt < %s -passes='simplifycfg<no-sink-common-insts;hoist-common-insts>' -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=HOIST
+
+; Test case based on C++ code with manualy annotated !prof metadata.
+; This is to test that when calls to 'func1' from 'if.then' block
+; and 'if.else' block are hoisted, the branch_weights are merged and
+; attached to merged call rather than dropped.
+;
+; int func1(int a, int b) ;
+; int func2(int a, int b) ;
+
+; int func(int a, int b, bool c) {
+;    int sum= 0;
+;    if(c) {
+;        sum += func1(a, b);
+;    } else {
+;        sum += func1(a, b);
+;        sum -= func2(a, b);
+;    }
+;    return sum;
+; }
+define i32 @_Z4funciib(i32 %a, i32 %b, i1 %c) {
+; HOIST-LABEL: define i32 @_Z4funciib
+; HOIST-SAME: (i32 [[A:%.*]], i32 [[B:%.*]], i1 [[C:%.*]]) {
+; HOIST-NEXT:  entry:
+; HOIST-NEXT:    [[CALL:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[B]]), !prof [[PROF0:![0-9]+]]
+; HOIST-NEXT:    br i1 [[C]], label [[IF_END:%.*]], label [[IF_ELSE:%.*]]
+; HOIST:       if.else:
+; HOIST-NEXT:    [[CALL3:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B]])
+; HOIST-NEXT:    [[SUB:%.*]] = sub i32 [[CALL]], [[CALL3]]
+; HOIST-NEXT:    br label [[IF_END]]
+; HOIST:       if.end:
+; HOIST-NEXT:    [[SUM_0:%.*]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[CALL]], [[ENTRY:%.*]] ]
+; HOIST-NEXT:    ret i32 [[SUM_0]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !0
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  %call1 = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !1
+  %call3 = tail call i32 @_Z5func2ii(i32 %a, i32 %b)
+  %sub = sub i32 %call1, %call3
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  %sum.0 = phi i32 [ %call, %if.then ], [ %sub, %if.else ]
+  ret i32 %sum.0
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+!1 = !{!"branch_weights", i32 90}
+;.
+; HOIST: [[PROF0]] = !{!"branch_weights", i64 100}
+;.

llvm/test/Transforms/SimplifyCFG/merge-direct-call-branch-weights-in-sink.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals --version 2
+; RUN: opt < %s -passes='simplifycfg<sink-common-insts;no-hoist-common-insts>' -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s --check-prefix=SINK
+
+
+; Test case based on the following C++ code with manualy annotated !prof metadata.
+; This is to test that when calls to 'func1' from 'if.then' and 'if.else' are
+; sinked, the branch weights are merged and attached to sinked call.
+;
+; int func1(int a, int b) ;
+; int func2(int a, int b) ;
+
+; int func(int a, int b, bool c) {
+;    int sum = 0;
+;    if (c) {
+;        sum += func1(a,b);
+;    } else {
+;        b -= func2(a,b);
+;        sum += func1(a,b);
+;    }
+;    return sum;
+; }
+
+define i32 @_Z4funciib(i32 %a, i32 %b, i1 %c) {
+; SINK-LABEL: define i32 @_Z4funciib
+; SINK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]], i1 [[C:%.*]]) {
+; SINK-NEXT:  entry:
+; SINK-NEXT:    br i1 [[C]], label [[IF_END:%.*]], label [[IF_ELSE:%.*]]
+; SINK:       if.else:
+; SINK-NEXT:    [[CALL1:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B]])
+; SINK-NEXT:    [[SUB:%.*]] = sub i32 [[B]], [[CALL1]]
+; SINK-NEXT:    br label [[IF_END]]
+; SINK:       if.end:
+; SINK-NEXT:    [[SUB_SINK:%.*]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[B]], [[ENTRY:%.*]] ]
+; SINK-NEXT:    [[CALL2:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[SUB_SINK]]), !prof [[PROF0:![0-9]+]]
+; SINK-NEXT:    ret i32 [[CALL2]]
+;
+entry:
+  br i1 %c, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !0
+  br label %if.end
+
+if.else:                                          ; preds = %entry
+  %call1 = tail call i32 @_Z5func2ii(i32 %a, i32 %b)
+  %sub = sub i32 %b, %call1
+  %call2 = tail call i32 @_Z5func1ii(i32 %a, i32 %sub), !prof !1
+  br label %if.end
+
+if.end:                                           ; preds = %if.else, %if.then
+  %sum.0 = phi i32 [ %call, %if.then ], [ %call2, %if.else ]
+  ret i32 %sum.0
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+!1 = !{!"branch_weights", i32 90}
+;.
+; SINK: [[PROF0]] = !{!"branch_weights", i64 100}
+;.