[InstCombine] Preserve profile data with select instructions and binary operators (#158375)

Tracking issue: #147390

Author: alanzhao1

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -144,6 +144,8 @@ static cl::opt<unsigned>
 MaxArraySize("instcombine-maxarray-size", cl::init(1024),
              cl::desc("Maximum array size considered when doing a combine"));
 
+extern cl::opt<bool> ProfcheckDisableMetadataFixes;
+
 // FIXME: Remove this flag when it is no longer necessary to convert
 // llvm.dbg.declare to avoid inaccurate debug info. Setting this to false
 // increases variable availability at the cost of accuracy. Variables that
@@ -1361,6 +1363,10 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
   if (!LHSIsSelect && !RHSIsSelect)
     return nullptr;
 
+  SelectInst *SI = ProfcheckDisableMetadataFixes
+                       ? nullptr
+                       : cast<SelectInst>(LHSIsSelect ? LHS : RHS);
+
   FastMathFlags FMF;
   BuilderTy::FastMathFlagGuard Guard(Builder);
   if (isa<FPMathOperator>(&I)) {
@@ -1381,15 +1387,14 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
     // We need an 'add' and exactly 1 arm of the select to have been simplified.
     if (Opcode != Instruction::Add || (!True && !False) || (True && False))
       return nullptr;
-
     Value *N;
     if (True && match(FVal, m_Neg(m_Value(N)))) {
       Value *Sub = Builder.CreateSub(Z, N);
-      return Builder.CreateSelect(Cond, True, Sub, I.getName());
+      return Builder.CreateSelect(Cond, True, Sub, I.getName(), SI);
     }
     if (False && match(TVal, m_Neg(m_Value(N)))) {
       Value *Sub = Builder.CreateSub(Z, N);
-      return Builder.CreateSelect(Cond, Sub, False, I.getName());
+      return Builder.CreateSelect(Cond, Sub, False, I.getName(), SI);
     }
     return nullptr;
   };
@@ -1425,9 +1430,9 @@ Value *InstCombinerImpl::SimplifySelectsFeedingBinaryOp(BinaryOperator &I,
   if (!True || !False)
     return nullptr;
 
-  Value *SI = Builder.CreateSelect(Cond, True, False);
-  SI->takeName(&I);
-  return SI;
+  Value *NewSI = Builder.CreateSelect(Cond, True, False, I.getName(), SI);
+  NewSI->takeName(&I);
+  return NewSI;
 }
 
 /// Freely adapt every user of V as-if V was changed to !V.

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

@@ -0,0 +1,54 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 6
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i32 @LHSBin(i1 %C) !prof !0 {
+; CHECK-LABEL: define i32 @LHSBin(
+; CHECK-SAME: i1 [[C:%.*]]) !prof [[PROF0:![0-9]+]] {
+; CHECK-NEXT:    [[V:%.*]] = select i1 [[C]], i32 1010, i32 20, !prof [[PROF1:![0-9]+]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+  %A = select i1 %C, i32 1000, i32 10, !prof !1
+  %V = add i32 %A, 10
+  ret i32 %V
+}
+
+define i32 @RHSBin(i1 %C) !prof !0 {
+; CHECK-LABEL: define i32 @RHSBin(
+; CHECK-SAME: i1 [[C:%.*]]) !prof [[PROF0]] {
+; CHECK-NEXT:    [[V:%.*]] = select i1 [[C]], i32 1010, i32 20, !prof [[PROF1]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+  %A = select i1 %C, i32 1000, i32 10, !prof !1
+  %V = add i32 10, %A
+  ret i32 %V;
+}
+
+define i32 @BothBin(i1 %C) !prof !0 {
+; CHECK-LABEL: define i32 @BothBin(
+; CHECK-SAME: i1 [[C:%.*]]) !prof [[PROF0]] {
+; CHECK-NEXT:    [[V:%.*]] = select i1 [[C]], i32 2000, i32 20, !prof [[PROF1]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+  %A = select i1 %C, i32 1000, i32 10, !prof !1
+  %B = select i1 %C, i32 1000, i32 10, !prof !1
+  %V = add i32 %A, %B
+  ret i32 %V;
+}
+
+define i32 @NegBin(i1 %C) !prof !0 {
+; CHECK-LABEL: define i32 @NegBin(
+; CHECK-SAME: i1 [[C:%.*]]) !prof [[PROF0]] {
+; CHECK-NEXT:    [[V:%.*]] = select i1 [[C]], i32 1010, i32 0, !prof [[PROF1]]
+; CHECK-NEXT:    ret i32 [[V]]
+;
+  %A = select i1 %C, i32 1000, i32 -10, !prof !1
+  %V = add i32 %A, 10
+  ret i32 %V
+}
+
+!0 = !{!"function_entry_count", i64 1000}
+!1 = !{!"branch_weights", i32 2, i32 3}
+;.
+; CHECK: [[PROF0]] = !{!"function_entry_count", i64 1000}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 2, i32 3}
+;.