[llvm] [InstCombine] avoid extra instructions in foldSelectICmpAnd (PR #127398)

[Andreas Jonson via llvm-commits]

https://github.com/andjo403 updated https://github.com/llvm/llvm-project/pull/127398

>From fe78e0f0eaf1849a98c902c22bc86c5b0167a724 Mon Sep 17 00:00:00 2001
From: Andreas Jonson <andjo403 at hotmail.com>
Date: Sun, 16 Feb 2025 14:07:55 +0100
Subject: [PATCH] [InstCombine] avoid extra instructions in foldSelectICmpAnd

---
 .../InstCombine/InstCombineSelect.cpp         | 14 +++++++----
 .../Transforms/InstCombine/select-icmp-and.ll | 23 ++++++++-----------
 2 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 2e14145aef884..3989577ff2667 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -205,11 +205,15 @@ static Value *foldSelectICmpAnd(SelectInst &Sel, ICmpInst *Cmp,
   unsigned ValZeros = ValC.logBase2();
   unsigned AndZeros = AndMask.logBase2();
   bool ShouldNotVal = !TC.isZero();
-
-  // If we would need to create an 'and' + 'shift' + 'xor' to replace a 'select'
-  // + 'icmp', then this transformation would result in more instructions and
-  // potentially interfere with other folding.
-  if (CreateAnd && ShouldNotVal && ValZeros != AndZeros)
+  bool NeedShift = ValZeros != AndZeros;
+  bool NeedZExtTrunc =
+      SelType->getScalarSizeInBits() != V->getType()->getScalarSizeInBits();
+
+  // If we would need to create an 'and' + 'shift' + 'xor' + cast to replace
+  // a 'select' + 'icmp', then this transformation would result in more
+  // instructions and potentially interfere with other folding.
+  if (CreateAnd + ShouldNotVal + NeedShift + NeedZExtTrunc >
+      1 + Cmp->hasOneUse())
     return nullptr;
 
   // Insert the 'and' instruction on the input to the truncate.
diff --git a/llvm/test/Transforms/InstCombine/select-icmp-and.ll b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
index c854a770aa74c..b9128fc2b3727 100644
--- a/llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ b/llvm/test/Transforms/InstCombine/select-icmp-and.ll
@@ -391,9 +391,8 @@ define i32 @test15e_extra_use(i32 %X) {
 ;; (a & 128) ? 256 : 0
 define i32 @test15e_zext(i8 %X) {
 ; CHECK-LABEL: @test15e_zext(
-; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[X:%.*]], -128
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i32
-; CHECK-NEXT:    [[T3:%.*]] = shl nuw nsw i32 [[TMP2]], 1
+; CHECK-NEXT:    [[T2_NOT:%.*]] = icmp sgt i8 [[X:%.*]], -1
+; CHECK-NEXT:    [[T3:%.*]] = select i1 [[T2_NOT]], i32 0, i32 256
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
   %t1 = and i8 %X, 128
@@ -406,9 +405,7 @@ define i32 @test15e_zext(i8 %X) {
 define i32 @test15e_zext_extra_use(i8 %X) {
 ; CHECK-LABEL: @test15e_zext_extra_use(
 ; CHECK-NEXT:    [[T2:%.*]] = icmp slt i8 [[X:%.*]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = and i8 [[X]], -128
-; CHECK-NEXT:    [[TMP2:%.*]] = zext i8 [[TMP1]] to i32
-; CHECK-NEXT:    [[T3:%.*]] = shl nuw nsw i32 [[TMP2]], 1
+; CHECK-NEXT:    [[T3:%.*]] = select i1 [[T2]], i32 256, i32 0
 ; CHECK-NEXT:    call void @use1(i1 [[T2]])
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
@@ -438,8 +435,7 @@ define i32 @test15f_extra_use(i32 %X) {
 ; CHECK-LABEL: @test15f_extra_use(
 ; CHECK-NEXT:    [[T1:%.*]] = and i32 [[X:%.*]], 128
 ; CHECK-NEXT:    [[T2:%.*]] = icmp ne i32 [[T1]], 0
-; CHECK-NEXT:    [[TMP1:%.*]] = shl nuw nsw i32 [[T1]], 1
-; CHECK-NEXT:    [[T3:%.*]] = xor i32 [[TMP1]], 256
+; CHECK-NEXT:    [[T3:%.*]] = select i1 [[T2]], i32 0, i32 256
 ; CHECK-NEXT:    call void @use1(i1 [[T2]])
 ; CHECK-NEXT:    ret i32 [[T3]]
 ;
@@ -453,10 +449,9 @@ define i32 @test15f_extra_use(i32 %X) {
 ;; (a & 128) ? 0 : 256
 define i16 @test15f_trunc(i32 %X) {
 ; CHECK-LABEL: @test15f_trunc(
-; CHECK-NEXT:    [[TMP1:%.*]] = trunc i32 [[X:%.*]] to i16
-; CHECK-NEXT:    [[TMP2:%.*]] = shl i16 [[TMP1]], 1
-; CHECK-NEXT:    [[TMP3:%.*]] = and i16 [[TMP2]], 256
-; CHECK-NEXT:    [[T3:%.*]] = xor i16 [[TMP3]], 256
+; CHECK-NEXT:    [[T1:%.*]] = and i32 [[X:%.*]], 128
+; CHECK-NEXT:    [[T2_NOT:%.*]] = icmp eq i32 [[T1]], 0
+; CHECK-NEXT:    [[T3:%.*]] = select i1 [[T2_NOT]], i16 256, i16 0
 ; CHECK-NEXT:    ret i16 [[T3]]
 ;
   %t1 = and i32 %X, 128
@@ -799,7 +794,9 @@ define i8 @select_bittest_to_xor(i8 %x) {
 ; CHECK-LABEL: @select_bittest_to_xor(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp sgt i8 [[X:%.*]], -1
 ; CHECK-NEXT:    call void @use1(i1 [[CMP]])
-; CHECK-NEXT:    [[MASKSEL:%.*]] = xor i8 [[X]], -128
+; CHECK-NEXT:    [[AND:%.*]] = and i8 [[X]], 127
+; CHECK-NEXT:    [[MASKSEL1:%.*]] = select i1 [[CMP]], i8 -128, i8 0
+; CHECK-NEXT:    [[MASKSEL:%.*]] = or disjoint i8 [[AND]], [[MASKSEL1]]
 ; CHECK-NEXT:    ret i8 [[MASKSEL]]
 ;
   %cmp = icmp sgt i8 %x, -1