[llvm] [AArch64][SelectionDAG] Mask for SUBS with multiple users cannot be elided (PR #90911)

[Weihang Fan via llvm-commits]

https://github.com/weihangf-apple updated https://github.com/llvm/llvm-project/pull/90911

>From b70dd1f3c3710b11fd6f3569e428bc970326b1be Mon Sep 17 00:00:00 2001
From: Weihang Fan <weihang_fan at apple.com>
Date: Wed, 24 Apr 2024 15:52:27 -0700
Subject: [PATCH] [AArch64][SelectionDAG] Mask for SUBS with multiple users
 cannot be elided

---
 .../Target/AArch64/AArch64ISelLowering.cpp    |  3 ++-
 llvm/test/CodeGen/AArch64/and-mask-removal.ll | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index b27d204f3dded0..855573fdab9358 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -22833,7 +22833,8 @@ SDValue performCONDCombine(SDNode *N,
   SDNode *SubsNode = N->getOperand(CmpIndex).getNode();
   unsigned CondOpcode = SubsNode->getOpcode();
 
-  if (CondOpcode != AArch64ISD::SUBS || SubsNode->hasAnyUseOfValue(0))
+  if (CondOpcode != AArch64ISD::SUBS || SubsNode->hasAnyUseOfValue(0) ||
+      !SubsNode->hasOneUse())
     return SDValue();
 
   // There is a SUBS feeding this condition. Is it fed by a mask we can
diff --git a/llvm/test/CodeGen/AArch64/and-mask-removal.ll b/llvm/test/CodeGen/AArch64/and-mask-removal.ll
index a8a59f1591268f..493d503de2cc13 100644
--- a/llvm/test/CodeGen/AArch64/and-mask-removal.ll
+++ b/llvm/test/CodeGen/AArch64/and-mask-removal.ll
@@ -526,4 +526,26 @@ define i64 @pr58109b(i8 signext %0, i64 %a, i64 %b) {
   ret i64 %4
 }
 
+define i64 @test_2_selects(i8 zeroext %a) {
+; CHECK-LABEL: test_2_selects:
+; CHECK:       ; %bb.0:
+; CHECK-NEXT:    add w9, w0, #24
+; CHECK-NEXT:    mov w8, #131
+; CHECK-NEXT:    and w9, w9, #0xff
+; CHECK-NEXT:    cmp w9, #81
+; CHECK-NEXT:    mov w9, #57
+; CHECK-NEXT:    csel x8, x8, xzr, lo
+; CHECK-NEXT:    csel x9, xzr, x9, eq
+; CHECK-NEXT:    add x0, x8, x9
+; CHECK-NEXT:    ret
+  %1 = add i8 %a, 24
+  %2 = zext i8 %1 to i64
+  %3 = icmp ult i8 %1, 81
+  %4 = select i1 %3, i64 131, i64 0
+  %5 = icmp eq i8 %1, 81
+  %6 = select i1 %5, i64 0, i64 57
+  %7 = add i64 %4, %6
+  ret i64 %7
+}
+
 declare i8 @llvm.usub.sat.i8(i8, i8) #0