[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

1322

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

2721

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

2721

good catch.

[mingmingl]

plan changes as suggested.

llvm/lib/Transforms/Utils/Local.cpp

2721

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

2721

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

2721

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

#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);
+  /// Returns nullptr if merge is not implementated for the given types
+  /// instructions and !prof metadata. Currently only implemented with direct
+  /// callsites with branch weights (used only in SamplePGO as documented in
+  /// https://llvm.org/docs/BranchWeightMetadata.html#callinst). Pass in both
+  /// instructions and nodes. Instruction information (e.g., instruction type)
+  /// helps interpret profiles and make implementation clearer.
+  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;
 }
 
+MDNode *MDNode::getMergedProfMetadata(MDNode *A, MDNode *B,
+                                      const Instruction *AInstr,
+                                      const Instruction *BInstr) {
+  if (!A)
+    return B;
+  if (!B)
+    return A;
+
+  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);
+  // Proceed if Both ACall and BCall are direct callsites.
+  // The rest of the cases are not implemented but could be added
+  // when there are use cases.
+  if (!(ACall && BCall && ACall->getCalledFunction() &&
+        BCall->getCalledFunction()))
+    return nullptr;
+
+  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));
+    if (AInstrWeight && BInstrWeight) {
+      return MDNode::get(Ctx,
+                         {MDHelper.createString("branch_weights"),
+                          MDHelper.createConstant(ConstantInt::get(
+                              Type::getInt64Ty(Ctx),
+                              SaturatingAdd(AInstrWeight->getZExtValue(),
+                                            BInstrWeight->getZExtValue())))});
+    }
+  }
+  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:
+        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/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}
+;.