add profile unknown only if the enclosing function has profile counts

Author: alanzhao1

llvm/include/llvm/IR/ProfDataUtils.h

@@ -185,6 +185,14 @@ inline uint32_t scaleBranchCount(uint64_t Count, uint64_t Scale) {
 LLVM_ABI void setExplicitlyUnknownBranchWeights(Instruction &I,
                                                 StringRef PassName);
 
+/// Like setExplicitlyUnknownBranchWeights(...), but only sets unknown branch
+/// weights in the new instruction if the parent function of the original
+/// instruction has function counts. This is to not confuse users by injecting
+/// profile data into non-profiled functions.
+LLVM_ABI void setExplicitlyUnknownBranchWeightsIfProfiled(Instruction &New,
+                                                          Instruction &Original,
+                                                          StringRef PassName);
+
 /// Analogous to setExplicitlyUnknownBranchWeights, but for functions and their
 /// entry counts.
 LLVM_ABI void setExplicitlyUnknownFunctionEntryCount(Function &F,

llvm/lib/IR/ProfDataUtils.cpp

@@ -252,6 +252,16 @@ void setExplicitlyUnknownBranchWeights(Instruction &I, StringRef PassName) {
                    MDB.createString(PassName)}));
 }
 
+void setExplicitlyUnknownBranchWeightsIfProfiled(Instruction &New,
+                                                 Instruction &Original,
+                                                 StringRef PassName) {
+  Function *F = Original.getFunction();
+  assert(F && "instruction does not belong to a function!");
+  std::optional<Function::ProfileCount> EC = F->getEntryCount();
+  if (EC && EC->getCount() > 0)
+    setExplicitlyUnknownBranchWeights(New, PassName);
+}
+
 void setExplicitlyUnknownFunctionEntryCount(Function &F, StringRef PassName) {
   MDBuilder MDB(F.getContext());
   F.setMetadata(

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

@@ -888,7 +888,7 @@ Instruction *InstCombinerImpl::foldAddWithConstant(BinaryOperator &Add) {
       X->getType()->getScalarSizeInBits() == 1)
     SI = SelectInst::Create(X, InstCombiner::SubOne(Op1C), Op1);
   if (SI) {
-    setExplicitlyUnknownBranchWeights(*SI, DEBUG_TYPE);
+    setExplicitlyUnknownBranchWeightsIfProfiled(*SI, Add, DEBUG_TYPE);
     return SI;
   }
 

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -1255,7 +1255,7 @@ Instruction *InstCombinerImpl::visitShl(BinaryOperator &I) {
     if (match(Op0, m_ZExt(m_Value(X))) && X->getType()->isIntOrIntVectorTy(1)) {
       auto *NewC = Builder.CreateShl(ConstantInt::get(Ty, 1), C1);
       auto *SI = SelectInst::Create(X, NewC, ConstantInt::getNullValue(Ty));
-      setExplicitlyUnknownBranchWeights(*SI, DEBUG_TYPE);
+      setExplicitlyUnknownBranchWeightsIfProfiled(*SI, I, DEBUG_TYPE);
       return SI;
     }
   }

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -1737,7 +1737,7 @@ Instruction *InstCombinerImpl::foldBinopOfSextBoolToSelect(BinaryOperator &BO) {
   Value *TVal = Builder.CreateBinOp(BO.getOpcode(), Ones, C);
   Value *FVal = Builder.CreateBinOp(BO.getOpcode(), Zero, C);
   SelectInst *SI = SelectInst::Create(X, TVal, FVal);
-  setExplicitlyUnknownBranchWeights(*SI, DEBUG_TYPE);
+  setExplicitlyUnknownBranchWeightsIfProfiled(*SI, BO, DEBUG_TYPE);
   return SI;
 }
 

llvm/test/Transforms/InstCombine/preserve-profile.ll

@@ -46,9 +46,9 @@ define i32 @NegBin(i1 %C) !prof !0 {
   ret i32 %V
 }
 
-define i32 @select_C_minus_1_or_C_from_bool(i1 %x) {
+define i32 @select_C_minus_1_or_C_from_bool(i1 %x) !prof !0 {
 ; CHECK-LABEL: define i32 @select_C_minus_1_or_C_from_bool(
-; CHECK-SAME: i1 [[X:%.*]]) {
+; CHECK-SAME: i1 [[X:%.*]]) !prof [[PROF0]] {
 ; CHECK-NEXT:    [[ADD:%.*]] = select i1 [[X]], i32 41, i32 42, !prof [[PROF2:![0-9]+]]
 ; CHECK-NEXT:    ret i32 [[ADD]]
 ;
@@ -57,9 +57,9 @@ define i32 @select_C_minus_1_or_C_from_bool(i1 %x) {
   ret i32 %add
 }
 
-define i5 @and_add(i1 %x, i1 %y) {
+define i5 @and_add(i1 %x, i1 %y) !prof !0 {
 ; CHECK-LABEL: define i5 @and_add(
-; CHECK-SAME: i1 [[X:%.*]], i1 [[Y:%.*]]) {
+; CHECK-SAME: i1 [[X:%.*]], i1 [[Y:%.*]]) !prof [[PROF0]] {
 ; CHECK-NEXT:    [[TMP1:%.*]] = xor i1 [[X]], true
 ; CHECK-NEXT:    [[TMP2:%.*]] = and i1 [[Y]], [[TMP1]]
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[TMP2]], i5 -2, i5 0, !prof [[PROF2]]
@@ -72,9 +72,9 @@ define i5 @and_add(i1 %x, i1 %y) {
   ret i5 %r
 }
 
-define i32 @add_zext_zext_i1(i1 %a) {
+define i32 @add_zext_zext_i1(i1 %a) !prof !0 {
 ; CHECK-LABEL: define i32 @add_zext_zext_i1(
-; CHECK-SAME: i1 [[A:%.*]]) {
+; CHECK-SAME: i1 [[A:%.*]]) !prof [[PROF0]] {
 ; CHECK-NEXT:    [[ADD:%.*]] = select i1 [[A]], i32 2, i32 0, !prof [[PROF2]]
 ; CHECK-NEXT:    ret i32 [[ADD]]
 ;
@@ -83,6 +83,17 @@ define i32 @add_zext_zext_i1(i1 %a) {
   ret i32 %add
 }
 
+define i32 @no_count_no_branch_weights(i1 %a) {
+; CHECK-LABEL: define i32 @no_count_no_branch_weights(
+; CHECK-SAME: i1 [[A:%.*]]) {
+; CHECK-NEXT:    [[ADD:%.*]] = select i1 [[A]], i32 2, i32 0
+; CHECK-NEXT:    ret i32 [[ADD]]
+;
+  %zext = zext i1 %a to i32
+  %add = add i32 %zext, %zext
+  ret i32 %add
+}
+
 
 !0 = !{!"function_entry_count", i64 1000}
 !1 = !{!"branch_weights", i32 2, i32 3}