[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

1

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

1

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 implementation for the type of instrutions
+  /// and !prof metadata. 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;
 }
 
+MDNode *MDNode::getMergedProfMetadata(MDNode *A, MDNode *B,
+                                      const Instruction *AInstr,
+                                      const Instruction *BInstr) {
+  if (!A)
+    return B;
+  if (!B)
+    return A;
+
+  const CallInst *ACall = dyn_cast<CallInst>(AInstr);
+  const CallInst *BCall = dyn_cast<CallInst>(BInstr);
+  if (!(ACall && BCall && ACall->getCalledFunction() &&
+        BCall->getCalledFunction()))
+    return nullptr;
+
+  auto &Ctx = AInstr->getContext();
+  MDBuilder MDHelper(Ctx);
+
+  // Both ACall and BCall are direct callsites.
+  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));
+  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), 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: --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
+
+define dso_local i32 @_Z4funciiii(i32 %a, i32 %b){
+; HOIST-LABEL: define dso_local i32 @_Z4funciiii
+; HOIST-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) {
+; HOIST-NEXT:  entry:
+; HOIST-NEXT:    br label [[FOR_BODY:%.*]]
+; HOIST:       for.cond.cleanup:
+; HOIST-NEXT:    [[SUM_1_LCSSA:%.*]] = phi i32 [ [[SUM_1:%.*]], [[FOR_INC:%.*]] ]
+; HOIST-NEXT:    ret i32 [[SUM_1_LCSSA]]
+; HOIST:       for.body:
+; HOIST-NEXT:    [[I_016:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_INC]] ]
+; HOIST-NEXT:    [[SUM_015:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_1]], [[FOR_INC]] ]
+; HOIST-NEXT:    [[REM_LHS_TRUNC:%.*]] = trunc i32 [[I_016]] to i16
+; HOIST-NEXT:    [[REM14:%.*]] = urem i16 [[REM_LHS_TRUNC]], 100
+; HOIST-NEXT:    [[CMP1:%.*]] = icmp eq i16 [[REM14]], 0
+; HOIST-NEXT:    [[CALL:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[B]]), !prof [[PROF0:![0-9]+]]
+; HOIST-NEXT:    [[ADD:%.*]] = add nsw i32 [[CALL]], [[SUM_015]]
+; HOIST-NEXT:    br i1 [[CMP1]], label [[FOR_INC]], label [[IF_ELSE:%.*]]
+; HOIST:       if.else:
+; HOIST-NEXT:    [[CALL4:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B]])
+; HOIST-NEXT:    [[SUB:%.*]] = sub i32 [[ADD]], [[CALL4]]
+; HOIST-NEXT:    br label [[FOR_INC]]
+; HOIST:       for.inc:
+; HOIST-NEXT:    [[SUM_1]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[ADD]], [[FOR_BODY]] ]
+; HOIST-NEXT:    [[INC]] = add nuw nsw i32 [[I_016]], 1
+; HOIST-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; HOIST-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc
+  %sum.1.lcssa = phi i32 [ %sum.1, %for.inc ]
+  ret i32 %sum.1.lcssa
+
+for.body:                                         ; preds = %entry, %for.inc
+  %i.016 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  %sum.015 = phi i32 [ 0, %entry ], [ %sum.1, %for.inc ]
+  %rem.lhs.trunc = trunc i32 %i.016 to i16
+  %rem14 = urem i16 %rem.lhs.trunc, 100
+  %cmp1 = icmp eq i16 %rem14, 0
+  br i1 %cmp1, label %if.then, label %if.else
+
+if.then:                                          ; preds = %for.body
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !0
+  %add = add nsw i32 %call, %sum.015
+  br label %for.inc
+
+if.else:                                          ; preds = %for.body
+  %call2 = tail call i32 @_Z5func1ii(i32 %a, i32 %b), !prof !1
+  %add3 = add nsw i32 %call2, %sum.015
+  %call4 = tail call i32 @_Z5func2ii(i32 %a, i32 %b)
+  %sub = sub i32 %add3, %call4
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.then, %if.else
+  %sum.1 = phi i32 [ %add, %if.then ], [ %sub, %if.else ]
+  %inc = add nuw nsw i32 %i.016, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+!1 = !{!"branch_weights", i32 990}

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: --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=HOIST
+
+define  i32 @_Z4funcii(i32 %a, i32 %b) {
+; HOIST-LABEL: define i32 @_Z4funcii
+; HOIST-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) {
+; HOIST-NEXT:  entry:
+; HOIST-NEXT:    br label [[FOR_BODY:%.*]]
+; HOIST:       for.cond.cleanup:
+; HOIST-NEXT:    [[SUM_1_LCSSA:%.*]] = phi i32 [ [[SUM_1:%.*]], [[FOR_INC:%.*]] ]
+; HOIST-NEXT:    ret i32 [[SUM_1_LCSSA]]
+; HOIST:       for.body:
+; HOIST-NEXT:    [[I_017:%.*]] = phi i32 [ 0, [[ENTRY:%.*]] ], [ [[INC:%.*]], [[FOR_INC]] ]
+; HOIST-NEXT:    [[SUM_016:%.*]] = phi i32 [ 0, [[ENTRY]] ], [ [[SUM_1]], [[FOR_INC]] ]
+; HOIST-NEXT:    [[B_ADDR_015:%.*]] = phi i32 [ [[B]], [[ENTRY]] ], [ [[SUB_SINK:%.*]], [[FOR_INC]] ]
+; HOIST-NEXT:    [[REM_LHS_TRUNC:%.*]] = trunc i32 [[I_017]] to i16
+; HOIST-NEXT:    [[REM14:%.*]] = urem i16 [[REM_LHS_TRUNC]], 100
+; HOIST-NEXT:    [[CMP1:%.*]] = icmp eq i16 [[REM14]], 0
+; HOIST-NEXT:    br i1 [[CMP1]], label [[FOR_INC]], label [[IF_ELSE:%.*]]
+; HOIST:       if.else:
+; HOIST-NEXT:    [[CALL2:%.*]] = tail call i32 @_Z5func2ii(i32 [[A]], i32 [[B_ADDR_015]])
+; HOIST-NEXT:    [[SUB:%.*]] = sub nsw i32 [[B_ADDR_015]], [[CALL2]]
+; HOIST-NEXT:    br label [[FOR_INC]]
+; HOIST:       for.inc:
+; HOIST-NEXT:    [[SUB_SINK]] = phi i32 [ [[SUB]], [[IF_ELSE]] ], [ [[B_ADDR_015]], [[FOR_BODY]] ]
+; HOIST-NEXT:    [[CALL3:%.*]] = tail call i32 @_Z5func1ii(i32 [[A]], i32 [[SUB_SINK]]), !prof [[PROF0:![0-9]+]]
+; HOIST-NEXT:    [[SUM_1]] = add nsw i32 [[CALL3]], [[SUM_016]]
+; HOIST-NEXT:    [[INC]] = add nuw nsw i32 [[I_017]], 1
+; HOIST-NEXT:    [[EXITCOND_NOT:%.*]] = icmp eq i32 [[INC]], 1000
+; HOIST-NEXT:    br i1 [[EXITCOND_NOT]], label [[FOR_COND_CLEANUP:%.*]], label [[FOR_BODY]]
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.inc
+  %sum.1.lcssa = phi i32 [ %sum.1, %for.inc ]
+  ret i32 %sum.1.lcssa
+
+for.body:                                         ; preds = %entry, %for.inc
+  %i.017 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
+  %sum.016 = phi i32 [ 0, %entry ], [ %sum.1, %for.inc ]
+  %b.addr.015 = phi i32 [ %b, %entry ], [ %b.addr.1, %for.inc ]
+  %rem.lhs.trunc = trunc i32 %i.017 to i16
+  %rem14 = urem i16 %rem.lhs.trunc, 100
+  %cmp1 = icmp eq i16 %rem14, 0
+  br i1 %cmp1, label %if.then, label %if.else
+
+if.then:                                          ; preds = %for.body
+  %call = tail call i32 @_Z5func1ii(i32 %a, i32 %b.addr.015), !prof !0
+  br label %for.inc
+
+if.else:                                          ; preds = %for.body
+  %call2 = tail call i32 @_Z5func2ii(i32 %a, i32 %b.addr.015)
+  %sub = sub nsw i32 %b.addr.015, %call2
+  %call3 = tail call i32 @_Z5func1ii(i32 %a, i32 %sub), !prof !1
+  br label %for.inc
+
+for.inc:                                          ; preds = %if.then, %if.else
+  %b.addr.1 = phi i32 [ %b.addr.015, %if.then ], [ %sub, %if.else ]
+  %call.pn = phi i32 [ %call, %if.then ], [ %call3, %if.else ]
+  %sum.1 = add nsw i32 %call.pn, %sum.016
+  %inc = add nuw nsw i32 %i.017, 1
+  %exitcond.not = icmp eq i32 %inc, 1000
+  br i1 %exitcond.not, label %for.cond.cleanup, label %for.body
+}
+
+declare i32 @_Z5func1ii(i32, i32)
+
+declare i32 @_Z5func2ii(i32, i32)
+
+!0 = !{!"branch_weights", i32 10}
+!1 = !{!"branch_weights", i32 990}