[llvm] [InstCombine] Keep branch weights when folding a shift through a select (PR #227987)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Alok Kumar Sharma (alokkrsharma)

<details>
<summary>Changes</summary>

Pulling a binop out of a select and through a constant shift rebuilds the select, and the new one was losing the original branch weights. Copy !prof so the weights stay attached to the same condition.

---
Full diff: https://github.com/llvm/llvm-project/pull/227987.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp (+8-2) 
- (modified) llvm/test/Transforms/InstCombine/pull-conditional-binop-through-shift.ll (+33) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index 97964cebb50b8..07c66045ae8b9 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -994,7 +994,10 @@ Instruction *InstCombinerImpl::FoldShiftByConstant(Value *Op0, Constant *C1,
 
       Value *NewShift = Builder.CreateBinOp(I.getOpcode(), FalseVal, C1);
       Value *NewOp = Builder.CreateBinOp(TBO->getOpcode(), NewShift, NewRHS);
-      return SelectInst::Create(Cond, NewOp, NewShift);
+      auto *NewSel = SelectInst::Create(Cond, NewOp, NewShift);
+      if (auto *SI = dyn_cast<SelectInst>(Op0))
+        NewSel->copyMetadata(*SI, {LLVMContext::MD_prof});
+      return NewSel;
     }
   }
 
@@ -1011,7 +1014,10 @@ Instruction *InstCombinerImpl::FoldShiftByConstant(Value *Op0, Constant *C1,
 
       Value *NewShift = Builder.CreateBinOp(I.getOpcode(), TrueVal, C1);
       Value *NewOp = Builder.CreateBinOp(FBO->getOpcode(), NewShift, NewRHS);
-      return SelectInst::Create(Cond, NewShift, NewOp);
+      auto *NewSel = SelectInst::Create(Cond, NewShift, NewOp);
+      if (auto *SI = dyn_cast<SelectInst>(Op0))
+        NewSel->copyMetadata(*SI, {LLVMContext::MD_prof});
+      return NewSel;
     }
   }
 
diff --git a/llvm/test/Transforms/InstCombine/pull-conditional-binop-through-shift.ll b/llvm/test/Transforms/InstCombine/pull-conditional-binop-through-shift.ll
index 5d8becae50023..cc3180fba8a35 100644
--- a/llvm/test/Transforms/InstCombine/pull-conditional-binop-through-shift.ll
+++ b/llvm/test/Transforms/InstCombine/pull-conditional-binop-through-shift.ll
@@ -306,3 +306,36 @@ define i32 @add_nosignbit_select_ashr(i32 %x, i1 %cond) {
   %r = ashr i32 %t1, 8
   ret i32 %r
 }
+
+; Branch weights stay on the new select, and stay attached to the same
+; condition arms. The binop may be on either arm.
+
+define i32 @and_select_shl_prof(i32 %x, i1 %cond) {
+; CHECK-LABEL: @and_select_shl_prof(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 8
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -16777216
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[COND:%.*]], i32 [[TMP2]], i32 [[TMP1]], !prof [[PROF0:![0-9]+]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %t0 = and i32 %x, 4294901760 ; 0xFFFF0000
+  %t1 = select i1 %cond, i32 %t0, i32 %x, !prof !0
+  %r = shl i32 %t1, 8
+  ret i32 %r
+}
+
+define i32 @and_select_shl_prof_binop_on_false(i32 %x, i1 %cond) {
+; CHECK-LABEL: @and_select_shl_prof_binop_on_false(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl i32 [[X:%.*]], 8
+; CHECK-NEXT:    [[TMP2:%.*]] = and i32 [[TMP1]], -16777216
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[COND:%.*]], i32 [[TMP1]], i32 [[TMP2]], !prof [[PROF0]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %t0 = and i32 %x, 4294901760 ; 0xFFFF0000
+  %t1 = select i1 %cond, i32 %x, i32 %t0, !prof !0
+  %r = shl i32 %t1, 8
+  ret i32 %r
+}
+
+!0 = !{!"branch_weights", i32 1, i32 99}
+
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 1, i32 99}

``````````

</details>


https://github.com/llvm/llvm-project/pull/227987