[llvm] Fix unassigned add handling in aarch64 (PR #86636)

[Yusuke Abe via llvm-commits]

https://github.com/chansuke created https://github.com/llvm/llvm-project/pull/86636

Fixes #84746 

>From 427f8b097c88c3a8698463aa06f3a5ba83461545 Mon Sep 17 00:00:00 2001
From: chansuke <moonset20 at gmail.com>
Date: Mon, 25 Mar 2024 18:27:53 +0900
Subject: [PATCH] Fix unassigned add handling in aarch64

---
 llvm/lib/Target/AArch64/AArch64ISelLowering.cpp     | 10 +++++++---
 .../CodeGen/AArch64/aarch64-non-negative-args.ll    | 13 +++++++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/aarch64-non-negative-args.ll

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 7fab274ab957c8..e604d31db35ea2 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -5711,13 +5711,17 @@ SDValue AArch64TargetLowering::LowerINTRINSIC_WO_CHAIN(SDValue Op,
   case Intrinsic::aarch64_neon_urhadd:
   case Intrinsic::aarch64_neon_shadd:
   case Intrinsic::aarch64_neon_uhadd: {
+    bool IsUnsignedAdd = IntNo == Intrinsic::aarch64_neon_uhadd;
     bool IsSignedAdd = (IntNo == Intrinsic::aarch64_neon_srhadd ||
                         IntNo == Intrinsic::aarch64_neon_shadd);
     bool IsRoundingAdd = (IntNo == Intrinsic::aarch64_neon_srhadd ||
                           IntNo == Intrinsic::aarch64_neon_urhadd);
-    unsigned Opcode = IsSignedAdd
-                          ? (IsRoundingAdd ? ISD::AVGCEILS : ISD::AVGFLOORS)
-                          : (IsRoundingAdd ? ISD::AVGCEILU : ISD::AVGFLOORU);
+    unsigned Opcode;
+    if (IsUnsignedAdd) {
+      Opcode = IsRoundingAdd ? ISD::AVGCEILU : ISD::AVGFLOORU;
+    } else {
+      Opcode = IsRoundingAdd ? ISD::AVGCEILS : ISD::AVGFLOORS;
+    }
     return DAG.getNode(Opcode, dl, Op.getValueType(), Op.getOperand(1),
                        Op.getOperand(2));
   }
diff --git a/llvm/test/CodeGen/AArch64/aarch64-non-negative-args.ll b/llvm/test/CodeGen/AArch64/aarch64-non-negative-args.ll
new file mode 100644
index 00000000000000..07f445803469ef
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/aarch64-non-negative-args.ll
@@ -0,0 +1,13 @@
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu < %s | FileCheck %s
+
+define <8 x i16> @uhadd_in_disguise(<8 x i16> %a0, <8 x i16> %a1) {
+; CHECK-LABEL: uhadd_in_disguise:
+; CHECK:       and
+; CHECK-NEXT:  uhadd
+; CHECK-NEXT:  ret
+  %m0 = and <8 x i16> %a0, <i16 510, i16 510, i16 510, i16 510, i16 510, i16 510, i16 510, i16 510>
+  %m1 = and <8 x i16> %a1, <i16 510, i16 510, i16 510, i16 510, i16 510, i16 510, i16 510, i16 510>
+  %r = call <8 x i16> @llvm.aarch64.neon.shadd.v8i16(<8 x i16> %m0, <8 x i16> %m1)
+  ret <8 x i16> %r
+}
+declare <8 x i16> @llvm.aarch64.neon.shadd.v8i16(<8 x i16>, <8 x i16>)