[llvm] [InstCombine] Propagate flags in `foldSelectICmpAndBinOp` (PR #127437)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/127437

It is always safe to add poison-generating flags for `BinOp Y, Identity`.
Proof: https://alive2.llvm.org/ce/z/8BLEpq

Then we can propagate flags from one of the arms:
```
select Cond, Y, (BinOp flags Y, Z) ->
select Cond, (BinOp flags Y, Identity), (BinOp flags Y, Z) ->
BinOp flags Y, (select Cond, Identity, Z)
```
This patch is proposed to avoid information loss caused by https://github.com/llvm/llvm-project/pull/127390.



>From 97950e8874ce3faa68e5e955bf30d3f6964575e4 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 17 Feb 2025 13:08:59 +0800
Subject: [PATCH 1/2] [InstCombine] Add pre-commit tests. NFC.

---
 .../InstCombine/select-with-bitwise-ops.ll    | 42 +++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll b/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
index 7c100f579399d..c895946421f0f 100644
--- a/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
+++ b/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
@@ -20,6 +20,34 @@ define i32 @select_icmp_eq_and_1_0_or_2(i32 %x, i32 %y) {
   ret i32 %select
 }
 
+define i32 @select_icmp_eq_and_1_0_or_2_disjoint(i32 %x, i32 %y) {
+; CHECK-LABEL: @select_icmp_eq_and_1_0_or_2_disjoint(
+; CHECK-NEXT:    [[AND:%.*]] = shl i32 [[X:%.*]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 2
+; CHECK-NEXT:    [[SELECT:%.*]] = or i32 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i32 [[SELECT]]
+;
+  %and = and i32 %x, 1
+  %cmp = icmp eq i32 %and, 0
+  %or = or disjoint i32 %y, 2
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
+define i32 @select_icmp_eq_and_1_0_add_2_nsw_nuw(i32 %x, i32 %y) {
+; CHECK-LABEL: @select_icmp_eq_and_1_0_add_2_nsw_nuw(
+; CHECK-NEXT:    [[AND:%.*]] = shl i32 [[X:%.*]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 2
+; CHECK-NEXT:    [[SELECT:%.*]] = add i32 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i32 [[SELECT]]
+;
+  %and = and i32 %x, 1
+  %cmp = icmp eq i32 %and, 0
+  %or = add nsw nuw i32 %y, 2
+  %select = select i1 %cmp, i32 %y, i32 %or
+  ret i32 %select
+}
+
 define <2 x i32> @select_icmp_eq_and_1_0_or_2_vec(<2 x i32> %x, <2 x i32> %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_1_0_or_2_vec(
 ; CHECK-NEXT:    [[AND:%.*]] = shl <2 x i32> [[X:%.*]], splat (i32 1)
@@ -1696,6 +1724,20 @@ define i8 @select_icmp_eq_and_1_0_lshr_fv(i8 %x, i8 %y) {
   ret i8 %select
 }
 
+define i8 @select_icmp_eq_and_1_0_lshr_exact_fv(i8 %x, i8 %y) {
+; CHECK-LABEL: @select_icmp_eq_and_1_0_lshr_exact_fv(
+; CHECK-NEXT:    [[AND:%.*]] = shl i8 [[X:%.*]], 1
+; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[AND]], 2
+; CHECK-NEXT:    [[SELECT:%.*]] = lshr i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    ret i8 [[SELECT]]
+;
+  %and = and i8 %x, 1
+  %cmp = icmp eq i8 %and, 0
+  %blshr = lshr exact i8 %y, 2
+  %select = select i1 %cmp, i8 %y, i8 %blshr
+  ret i8 %select
+}
+
 define i8 @select_icmp_eq_and_1_0_lshr_tv(i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_1_0_lshr_tv(
 ; CHECK-NEXT:    [[AND:%.*]] = shl i8 [[X:%.*]], 1

>From c8f73c1335ca71247f65b91a1b22dca4316ddee5 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Mon, 17 Feb 2025 13:21:15 +0800
Subject: [PATCH 2/2] [InstCombine] Propagate flags in `foldSelectICmpAndBinOp`

---
 llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp       | 5 ++++-
 llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll | 6 +++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 2e14145aef884..cf38fc5f058f2 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -829,7 +829,10 @@ static Value *foldSelectICmpAndBinOp(const ICmpInst *IC, Value *TrueVal,
   if (NeedXor)
     V = Builder.CreateXor(V, *C2);
 
-  return Builder.CreateBinOp(BinOp->getOpcode(), Y, V);
+  auto *Res = Builder.CreateBinOp(BinOp->getOpcode(), Y, V);
+  if (auto *BO = dyn_cast<BinaryOperator>(Res))
+    BO->copyIRFlags(BinOp);
+  return Res;
 }
 
 /// Canonicalize a set or clear of a masked set of constant bits to
diff --git a/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll b/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
index c895946421f0f..67dec9178eeca 100644
--- a/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
+++ b/llvm/test/Transforms/InstCombine/select-with-bitwise-ops.ll
@@ -24,7 +24,7 @@ define i32 @select_icmp_eq_and_1_0_or_2_disjoint(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_1_0_or_2_disjoint(
 ; CHECK-NEXT:    [[AND:%.*]] = shl i32 [[X:%.*]], 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 2
-; CHECK-NEXT:    [[SELECT:%.*]] = or i32 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[SELECT:%.*]] = or disjoint i32 [[Y:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %and = and i32 %x, 1
@@ -38,7 +38,7 @@ define i32 @select_icmp_eq_and_1_0_add_2_nsw_nuw(i32 %x, i32 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_1_0_add_2_nsw_nuw(
 ; CHECK-NEXT:    [[AND:%.*]] = shl i32 [[X:%.*]], 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[AND]], 2
-; CHECK-NEXT:    [[SELECT:%.*]] = add i32 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[SELECT:%.*]] = add nuw nsw i32 [[Y:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i32 [[SELECT]]
 ;
   %and = and i32 %x, 1
@@ -1728,7 +1728,7 @@ define i8 @select_icmp_eq_and_1_0_lshr_exact_fv(i8 %x, i8 %y) {
 ; CHECK-LABEL: @select_icmp_eq_and_1_0_lshr_exact_fv(
 ; CHECK-NEXT:    [[AND:%.*]] = shl i8 [[X:%.*]], 1
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[AND]], 2
-; CHECK-NEXT:    [[SELECT:%.*]] = lshr i8 [[Y:%.*]], [[TMP1]]
+; CHECK-NEXT:    [[SELECT:%.*]] = lshr exact i8 [[Y:%.*]], [[TMP1]]
 ; CHECK-NEXT:    ret i8 [[SELECT]]
 ;
   %and = and i8 %x, 1