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

[via llvm-commits]

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

In DAGCombiner, the `performCONDCombine` function attempts to remove AND instructions in front of SUBS (cmp) instructions for which the AND is transparent. The rules for that are correct, but it fails to take into account the case where the SUBS instruction has multiple users with different condition codes for comparison and simply removes the AND for all of them. This causes a miscompilation in the attached test case.

>From feb0276e6df4da9f0513a4992da5d54ccf183ba8 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