[llvm] [AArch64][SelectionDAG] Improve codegen for insert_into_scalable(zext_fixed_length_vector) (PR #192405)

Sushant Gokhale via llvm-commits llvm-commits at lists.llvm.org
Thu Apr 16 11:17:43 PDT 2026


https://github.com/sushgokh updated https://github.com/llvm/llvm-project/pull/192405

>From 1a0d1bddfa50e3bcff48117039ef284051bc86b7 Mon Sep 17 00:00:00 2001
From: sgokhale <sgokhale at nvidia.com>
Date: Wed, 15 Apr 2026 23:46:10 -0700
Subject: [PATCH 1/4] [AArch64][NFC] Add tests for  zext(icmp)

When these cases can generate SVE instructions, we get unoptimized codegen in some cases.

The inefficiency would be addressed in immediate patch.
---
 .../CodeGen/AArch64/sve-icmp-sext-zext.ll     | 142 ++++++++++++++++++
 1 file changed, 142 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll

diff --git a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
new file mode 100644
index 0000000000000..330deff930711
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
@@ -0,0 +1,142 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
+; RUN: llc -mtriple=aarch64 -mattr=+sve2 --aarch64-sve-vector-bits-min=256 -o - < %s | FileCheck %s
+
+define <vscale x 16 x i8> @zext_slt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+; CHECK-LABEL: zext_slt_test1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp slt <vscale x 16 x i8> %a, %b
+  %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
+  ret <vscale x 16 x i8> %2
+}
+
+define <vscale x 16 x i8> @zext_slt_test2(<32 x i8> %a, <32 x i8> %b)
+; CHECK-LABEL: zext_slt_test2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
+; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
+; CHECK-NEXT:    ptrue p0.b, vl32
+; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp slt <32 x i8> %a, %b
+  %2 = zext <32 x i1> %1 to <32 x i8>
+  %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
+  ret <vscale x 16 x i8> %3
+}
+
+define <vscale x 16 x i8> @zext_ult_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+; CHECK-LABEL: zext_ult_test1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp ult <vscale x 16 x i8> %a, %b
+  %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
+  ret <vscale x 16 x i8> %2
+}
+
+define <vscale x 16 x i8> @zext_ult_test2(<32 x i8> %a, <32 x i8> %b)
+; CHECK-LABEL: zext_ult_test2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
+; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
+; CHECK-NEXT:    ptrue p0.b, vl32
+; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp ult <32 x i8> %a, %b
+  %2 = zext <32 x i1> %1 to <32 x i8>
+  %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
+  ret <vscale x 16 x i8> %3
+}
+
+define <vscale x 16 x i8> @zext_sgt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+; CHECK-LABEL: zext_sgt_test1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z1.b
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp sgt <vscale x 16 x i8> %a, %b
+  %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
+  ret <vscale x 16 x i8> %2
+}
+
+define <vscale x 16 x i8> @zext_sgt_test2(<32 x i8> %a, <32 x i8> %b)
+; CHECK-LABEL: zext_sgt_test2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    splice z2.b, p0, { z2.b, z3.b }
+; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
+; CHECK-NEXT:    ptrue p0.b, vl32
+; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z2.b
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp sgt <32 x i8> %a, %b
+  %2 = zext <32 x i1> %1 to <32 x i8>
+  %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
+  ret <vscale x 16 x i8> %3
+}
+
+define <vscale x 16 x i8> @zext_ugt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
+; CHECK-LABEL: zext_ugt_test1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.b
+; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z1.b
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp ugt <vscale x 16 x i8> %a, %b
+  %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
+  ret <vscale x 16 x i8> %2
+}
+
+define <vscale x 16 x i8> @zext_ugt_test2(<32 x i8> %a, <32 x i8> %b)
+; CHECK-LABEL: zext_ugt_test2:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.b, vl16
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    splice z2.b, p0, { z2.b, z3.b }
+; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
+; CHECK-NEXT:    ptrue p0.b, vl32
+; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z2.b
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    ret
+{
+  %1 = icmp ugt <32 x i8> %a, %b
+  %2 = zext <32 x i1> %1 to <32 x i8>
+  %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
+  ret <vscale x 16 x i8> %3
+}

>From 1fbef322f87f57ff8d3e075bff1fc64ce49c7b72 Mon Sep 17 00:00:00 2001
From: sgokhale <sgokhale at nvidia.com>
Date: Wed, 15 Apr 2026 23:51:46 -0700
Subject: [PATCH 2/4] [AArch64][SelectionDAG] Improve codegen for
 insert_into_scalable(zext_fixed_length_vector)

For the above code, what we get before isel is:
`and(splat(1), sext(setcc_merge_zero)) `

We can perform the following folding for better codegen:
`and(splat(1), sext(setcc_merge_zero)) -> zext(setcc_merge_zero)`

Acts  as enabler for PR #192052 in some of the cases
---
 .../lib/Target/AArch64/AArch64ISelLowering.cpp | 18 ++++++++++++++++++
 .../test/CodeGen/AArch64/sve-icmp-sext-zext.ll | 12 ++++--------
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9b34d9b385b4e..3cd23cdf68fc7 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -20942,12 +20942,30 @@ static SDValue performReinterpretCastCombine(SDNode *N) {
   return SDValue();
 }
 
+// and(splat(1), sext(setcc_merge_zero)) -> zext(setcc_merge_zero)
+SDValue performSVEAndSplatSetCCMergeZeroCombine(SDNode *N, SelectionDAG &DAG) {
+  SDLoc DL(N);
+  SDValue SplatOp = N->getOperand(0);
+  SDValue NonSplatOp = N->getOperand(1);
+  if (NonSplatOp.getOpcode() == ISD::SPLAT_VECTOR)
+    std::swap(SplatOp, NonSplatOp);
+  SDValue Compare = NonSplatOp.getOperand(0);
+  if (SplatOp.getOpcode() != ISD::SPLAT_VECTOR ||
+      NonSplatOp.getOpcode() != ISD::SIGN_EXTEND ||
+      Compare.getOpcode() != AArch64ISD::SETCC_MERGE_ZERO)
+    return SDValue();
+  return DAG.getNode(ISD::ZERO_EXTEND, DL, N->getValueType(0), Compare);
+}
+
 static SDValue performSVEAndCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue Src = N->getOperand(0);
   unsigned Opc = Src->getOpcode();
 
+  if (SDValue R = performSVEAndSplatSetCCMergeZeroCombine(N, DAG))
+    return R;
+
   // Zero/any extend of an unsigned unpack
   if (Opc == AArch64ISD::UUNPKHI || Opc == AArch64ISD::UUNPKLO) {
     SDValue UnpkOp = Src->getOperand(0);
diff --git a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
index 330deff930711..c140deedbb6cf 100644
--- a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
+++ b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
@@ -26,8 +26,7 @@ define <vscale x 16 x i8> @zext_slt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
 {
   %1 = icmp slt <32 x i8> %a, %b
@@ -61,8 +60,7 @@ define <vscale x 16 x i8> @zext_ult_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
 {
   %1 = icmp ult <32 x i8> %a, %b
@@ -96,8 +94,7 @@ define <vscale x 16 x i8> @zext_sgt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
 {
   %1 = icmp sgt <32 x i8> %a, %b
@@ -131,8 +128,7 @@ define <vscale x 16 x i8> @zext_ugt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
 {
   %1 = icmp ugt <32 x i8> %a, %b

>From d771549e85d6e4b2bd9708a0347e88f198a33b07 Mon Sep 17 00:00:00 2001
From: sgokhale <sgokhale at nvidia.com>
Date: Thu, 16 Apr 2026 04:13:45 -0700
Subject: [PATCH 3/4] Add more test cases

---
 .../Target/AArch64/AArch64ISelLowering.cpp    |  18 ---
 .../CodeGen/AArch64/sve-icmp-sext-zext.ll     | 148 ++++++++++++++----
 2 files changed, 120 insertions(+), 46 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 3cd23cdf68fc7..9b34d9b385b4e 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -20942,30 +20942,12 @@ static SDValue performReinterpretCastCombine(SDNode *N) {
   return SDValue();
 }
 
-// and(splat(1), sext(setcc_merge_zero)) -> zext(setcc_merge_zero)
-SDValue performSVEAndSplatSetCCMergeZeroCombine(SDNode *N, SelectionDAG &DAG) {
-  SDLoc DL(N);
-  SDValue SplatOp = N->getOperand(0);
-  SDValue NonSplatOp = N->getOperand(1);
-  if (NonSplatOp.getOpcode() == ISD::SPLAT_VECTOR)
-    std::swap(SplatOp, NonSplatOp);
-  SDValue Compare = NonSplatOp.getOperand(0);
-  if (SplatOp.getOpcode() != ISD::SPLAT_VECTOR ||
-      NonSplatOp.getOpcode() != ISD::SIGN_EXTEND ||
-      Compare.getOpcode() != AArch64ISD::SETCC_MERGE_ZERO)
-    return SDValue();
-  return DAG.getNode(ISD::ZERO_EXTEND, DL, N->getValueType(0), Compare);
-}
-
 static SDValue performSVEAndCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue Src = N->getOperand(0);
   unsigned Opc = Src->getOpcode();
 
-  if (SDValue R = performSVEAndSplatSetCCMergeZeroCombine(N, DAG))
-    return R;
-
   // Zero/any extend of an unsigned unpack
   if (Opc == AArch64ISD::UUNPKHI || Opc == AArch64ISD::UUNPKLO) {
     SDValue UnpkOp = Src->getOperand(0);
diff --git a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
index c140deedbb6cf..0fcb4ab388818 100644
--- a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
+++ b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
@@ -1,21 +1,20 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 6
 ; RUN: llc -mtriple=aarch64 -mattr=+sve2 --aarch64-sve-vector-bits-min=256 -o - < %s | FileCheck %s
 
-define <vscale x 16 x i8> @zext_slt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
-; CHECK-LABEL: zext_slt_test1:
+define <vscale x 16 x i8> @zext_slt_nxv16i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: zext_slt_nxv16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.b
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
 ; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp slt <vscale x 16 x i8> %a, %b
   %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
   ret <vscale x 16 x i8> %2
 }
 
-define <vscale x 16 x i8> @zext_slt_test2(<32 x i8> %a, <32 x i8> %b)
-; CHECK-LABEL: zext_slt_test2:
+define <vscale x 16 x i8> @zext_slt_v32i8(<32 x i8> %a, <32 x i8> %b) {
+; CHECK-LABEL: zext_slt_v32i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
 ; CHECK-NEXT:    ptrue p0.b, vl16
@@ -26,30 +25,29 @@ define <vscale x 16 x i8> @zext_slt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp slt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
   %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
   ret <vscale x 16 x i8> %3
 }
 
-define <vscale x 16 x i8> @zext_ult_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
-; CHECK-LABEL: zext_ult_test1:
+define <vscale x 16 x i8> @zext_ult_nxv16i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: zext_ult_nxv16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.b
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
 ; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp ult <vscale x 16 x i8> %a, %b
   %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
   ret <vscale x 16 x i8> %2
 }
 
-define <vscale x 16 x i8> @zext_ult_test2(<32 x i8> %a, <32 x i8> %b)
-; CHECK-LABEL: zext_ult_test2:
+define <vscale x 16 x i8> @zext_ult_v32i8(<32 x i8> %a, <32 x i8> %b) {
+; CHECK-LABEL: zext_ult_v32i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
 ; CHECK-NEXT:    ptrue p0.b, vl16
@@ -60,30 +58,29 @@ define <vscale x 16 x i8> @zext_ult_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp ult <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
   %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
   ret <vscale x 16 x i8> %3
 }
 
-define <vscale x 16 x i8> @zext_sgt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
-; CHECK-LABEL: zext_sgt_test1:
+define <vscale x 16 x i8> @zext_sgt_nxv16i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: zext_sgt_nxv16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.b
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z1.b
 ; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp sgt <vscale x 16 x i8> %a, %b
   %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
   ret <vscale x 16 x i8> %2
 }
 
-define <vscale x 16 x i8> @zext_sgt_test2(<32 x i8> %a, <32 x i8> %b)
-; CHECK-LABEL: zext_sgt_test2:
+define <vscale x 16 x i8> @zext_sgt_v32i8(<32 x i8> %a, <32 x i8> %b) {
+; CHECK-LABEL: zext_sgt_v32i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
 ; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
@@ -94,30 +91,29 @@ define <vscale x 16 x i8> @zext_sgt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp sgt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
   %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
   ret <vscale x 16 x i8> %3
 }
 
-define <vscale x 16 x i8> @zext_ugt_test1(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
-; CHECK-LABEL: zext_ugt_test1:
+define <vscale x 16 x i8> @zext_ugt_nxv16i8(<vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
+; CHECK-LABEL: zext_ugt_nxv16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.b
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z1.b
 ; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp ugt <vscale x 16 x i8> %a, %b
   %2 = zext <vscale x 16 x i1> %1 to <vscale x 16 x i8>
   ret <vscale x 16 x i8> %2
 }
 
-define <vscale x 16 x i8> @zext_ugt_test2(<32 x i8> %a, <32 x i8> %b)
-; CHECK-LABEL: zext_ugt_test2:
+define <vscale x 16 x i8> @zext_ugt_v32i8(<32 x i8> %a, <32 x i8> %b) {
+; CHECK-LABEL: zext_ugt_v32i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
 ; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
@@ -128,11 +124,107 @@ define <vscale x 16 x i8> @zext_ugt_test2(<32 x i8> %a, <32 x i8> %b)
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.b, z0.b, #0x1
 ; CHECK-NEXT:    ret
-{
   %1 = icmp ugt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
   %3 = tail call <vscale x 16 x i8> @llvm.vector.insert(<vscale x 16 x i8> poison, <32 x i8> %2, i64 0)
   ret <vscale x 16 x i8> %3
 }
+
+define <vscale x 2 x i64> @zext_slt_nxv2i64(<vscale x 2 x i64> %a, <vscale x 2 x i64> %b) {
+; CHECK-LABEL: zext_slt_nxv2i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.d
+; CHECK-NEXT:    cmpgt p1.d, p0/z, z1.d, z0.d
+; CHECK-NEXT:    mov z0.d, p1/z, #1 // =0x1
+; CHECK-NEXT:    ret
+  %1 = icmp slt <vscale x 2 x i64> %a, %b
+  %2 = zext <vscale x 2 x i1> %1 to <vscale x 2 x i64>
+  ret <vscale x 2 x i64> %2
+}
+
+define <vscale x 2 x i64> @zext_slt_v2i64(<2 x i64> %a, <2 x i64> %b) {
+; CHECK-LABEL: zext_slt_v2i64:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmgt v0.2d, v1.2d, v0.2d
+; CHECK-NEXT:    and z0.d, z0.d, #0x1
+; CHECK-NEXT:    ret
+  %1 = icmp slt <2 x i64> %a, %b
+  %2 = zext <2 x i1> %1 to <2 x i64>
+  %3 = tail call <vscale x 2 x i64> @llvm.vector.insert(<vscale x 2 x i64> poison, <2 x i64> %2, i64 0)
+  ret <vscale x 2 x i64> %3
+}
+
+define <8 x i32> @zext_slt_v8i32(<8 x i32> %a, <8 x i32> %b) {
+; CHECK-LABEL: zext_slt_v8i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.s, vl4
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    splice z0.s, p0, { z0.s, z1.s }
+; CHECK-NEXT:    splice z1.s, p0, { z2.s, z3.s }
+; CHECK-NEXT:    ptrue p0.s, vl8
+; CHECK-NEXT:    cmpgt p1.s, p0/z, z1.s, z0.s
+; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    movprfx z1, z0
+; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z1
+; CHECK-NEXT:    ret
+  %1 = icmp slt <8 x i32> %a, %b
+  %2 = zext <8 x i1> %1 to <8 x i32>
+  ret <8 x i32> %2
+}
+
+define <8 x i32> @zext_sle_v8i32(<8 x i32> %a, <8 x i32> %b) {
+; CHECK-LABEL: zext_sle_v8i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.s, vl4
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    splice z0.s, p0, { z0.s, z1.s }
+; CHECK-NEXT:    splice z1.s, p0, { z2.s, z3.s }
+; CHECK-NEXT:    ptrue p0.s, vl8
+; CHECK-NEXT:    cmpge p1.s, p0/z, z1.s, z0.s
+; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    movprfx z1, z0
+; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z1
+; CHECK-NEXT:    ret
+  %1 = icmp sle <8 x i32> %a, %b
+  %2 = zext <8 x i1> %1 to <8 x i32>
+  ret <8 x i32> %2
+}
+
+define <8 x i32> @zext_sge_v8i32(<8 x i32> %a, <8 x i32> %b) {
+; CHECK-LABEL: zext_sge_v8i32:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    // kill: def $q3 killed $q3 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    ptrue p0.s, vl4
+; CHECK-NEXT:    // kill: def $q2 killed $q2 killed $z2_z3 def $z2_z3
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0_z1 def $z0_z1
+; CHECK-NEXT:    splice z2.s, p0, { z2.s, z3.s }
+; CHECK-NEXT:    splice z0.s, p0, { z0.s, z1.s }
+; CHECK-NEXT:    ptrue p0.s, vl8
+; CHECK-NEXT:    cmpge p1.s, p0/z, z0.s, z2.s
+; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
+; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    movprfx z1, z0
+; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
+; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
+; CHECK-NEXT:    // kill: def $q1 killed $q1 killed $z1
+; CHECK-NEXT:    ret
+  %1 = icmp sge <8 x i32> %a, %b
+  %2 = zext <8 x i1> %1 to <8 x i32>
+  ret <8 x i32> %2
+}

>From 312353ad26c85ae315d8e9b5d384b4cc6f540e87 Mon Sep 17 00:00:00 2001
From: sgokhale <sgokhale at nvidia.com>
Date: Thu, 16 Apr 2026 11:15:18 -0700
Subject: [PATCH 4/4] Amend patch as per review comments

---
 .../Target/AArch64/AArch64ISelLowering.cpp    | 24 +++++++++++++++++++
 .../CodeGen/AArch64/sve-icmp-sext-zext.ll     | 21 ++++++----------
 2 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 9b34d9b385b4e..cc34f4d1cfe83 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -20942,12 +20942,36 @@ static SDValue performReinterpretCastCombine(SDNode *N) {
   return SDValue();
 }
 
+// and(sext(Op), splat(1)) -> zext(Op)
+static SDValue
+performSVEAndSplatSExtCombine(SDNode *N, TargetLowering::DAGCombinerInfo &DCI) {
+  if (DCI.isBeforeLegalizeOps())
+    return SDValue();
+
+  SDLoc DL(N);
+  SelectionDAG &DAG = DCI.DAG;
+  SDValue ExtendOp = N->getOperand(0);
+  SDValue SplatOp = N->getOperand(1);
+
+  if (SplatOp.getOpcode() != ISD::SPLAT_VECTOR ||
+      !isOneConstant(SplatOp.getOperand(0)) ||
+      ExtendOp.getOpcode() != ISD::SIGN_EXTEND)
+    return SDValue();
+
+  SDValue ZExtOp = ExtendOp.getOperand(0);
+
+  return DAG.getNode(ISD::ZERO_EXTEND, DL, N->getValueType(0), ZExtOp);
+}
+
 static SDValue performSVEAndCombine(SDNode *N,
                                     TargetLowering::DAGCombinerInfo &DCI) {
   SelectionDAG &DAG = DCI.DAG;
   SDValue Src = N->getOperand(0);
   unsigned Opc = Src->getOpcode();
 
+  if (SDValue R = performSVEAndSplatSExtCombine(N, DCI))
+    return R;
+
   // Zero/any extend of an unsigned unpack
   if (Opc == AArch64ISD::UUNPKHI || Opc == AArch64ISD::UUNPKLO) {
     SDValue UnpkOp = Src->getOperand(0);
diff --git a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
index 0fcb4ab388818..f17b69630c59d 100644
--- a/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
+++ b/llvm/test/CodeGen/AArch64/sve-icmp-sext-zext.ll
@@ -25,8 +25,7 @@ define <vscale x 16 x i8> @zext_slt_v32i8(<32 x i8> %a, <32 x i8> %b) {
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
   %1 = icmp slt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
@@ -58,8 +57,7 @@ define <vscale x 16 x i8> @zext_ult_v32i8(<32 x i8> %a, <32 x i8> %b) {
 ; CHECK-NEXT:    splice z1.b, p0, { z2.b, z3.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z1.b, z0.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
   %1 = icmp ult <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
@@ -91,8 +89,7 @@ define <vscale x 16 x i8> @zext_sgt_v32i8(<32 x i8> %a, <32 x i8> %b) {
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmpgt p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
   %1 = icmp sgt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
@@ -124,8 +121,7 @@ define <vscale x 16 x i8> @zext_ugt_v32i8(<32 x i8> %a, <32 x i8> %b) {
 ; CHECK-NEXT:    splice z0.b, p0, { z0.b, z1.b }
 ; CHECK-NEXT:    ptrue p0.b, vl32
 ; CHECK-NEXT:    cmphi p1.b, p0/z, z0.b, z2.b
-; CHECK-NEXT:    mov z0.b, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.b, z0.b, #0x1
+; CHECK-NEXT:    mov z0.b, p1/z, #1 // =0x1
 ; CHECK-NEXT:    ret
   %1 = icmp ugt <32 x i8> %a, %b
   %2 = zext <32 x i1> %1 to <32 x i8>
@@ -169,8 +165,7 @@ define <8 x i32> @zext_slt_v8i32(<8 x i32> %a, <8 x i32> %b) {
 ; CHECK-NEXT:    splice z1.s, p0, { z2.s, z3.s }
 ; CHECK-NEXT:    ptrue p0.s, vl8
 ; CHECK-NEXT:    cmpgt p1.s, p0/z, z1.s, z0.s
-; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    mov z0.s, p1/z, #1 // =0x1
 ; CHECK-NEXT:    movprfx z1, z0
 ; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
 ; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
@@ -193,8 +188,7 @@ define <8 x i32> @zext_sle_v8i32(<8 x i32> %a, <8 x i32> %b) {
 ; CHECK-NEXT:    splice z1.s, p0, { z2.s, z3.s }
 ; CHECK-NEXT:    ptrue p0.s, vl8
 ; CHECK-NEXT:    cmpge p1.s, p0/z, z1.s, z0.s
-; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    mov z0.s, p1/z, #1 // =0x1
 ; CHECK-NEXT:    movprfx z1, z0
 ; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
 ; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0
@@ -217,8 +211,7 @@ define <8 x i32> @zext_sge_v8i32(<8 x i32> %a, <8 x i32> %b) {
 ; CHECK-NEXT:    splice z0.s, p0, { z0.s, z1.s }
 ; CHECK-NEXT:    ptrue p0.s, vl8
 ; CHECK-NEXT:    cmpge p1.s, p0/z, z0.s, z2.s
-; CHECK-NEXT:    mov z0.s, p1/z, #-1 // =0xffffffffffffffff
-; CHECK-NEXT:    and z0.s, z0.s, #0x1
+; CHECK-NEXT:    mov z0.s, p1/z, #1 // =0x1
 ; CHECK-NEXT:    movprfx z1, z0
 ; CHECK-NEXT:    ext z1.b, z1.b, z0.b, #16
 ; CHECK-NEXT:    // kill: def $q0 killed $q0 killed $z0



More information about the llvm-commits mailing list