[AArch64] Combine shift and truncate into deinterleave.

This is an alternative approach to #213252, from which I've
taken some of the tests.

Co-authored by Jacob Crawley

Author: sdesmalen-arm

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -26632,6 +26632,27 @@ static SDValue performUzpCombine(SDNode *N, SelectionDAG &DAG,
     }
   }
 
+  // uzp1(nvcast(x >> k), nvcast(y >> k)) -> uzp2(nvcast(x), nvcast(y))
+  // where 'k' is 'sizeof(eltty(x))/2'.
+  if (SDValue PreCastOp0 = isNVCastToHalfWidthElements(Op0)) {
+    if (SDValue PreCastOp1 = isNVCastToHalfWidthElements(Op1)) {
+      if (PreCastOp0.getOpcode() == ISD::SRL &&
+          PreCastOp1.getOpcode() == ISD::SRL &&
+          PreCastOp0.getOperand(1) == PreCastOp1.getOperand(1)) {
+        ConstantSDNode *ShiftAmount =
+            isConstOrConstSplat(PreCastOp0.getOperand(1), /*AllowUndef=*/false,
+                                /*AllowTruncate=*/true);
+        if (ShiftAmount &&
+            ShiftAmount->getAsZExtVal() == ResVT.getScalarSizeInBits())
+          return DAG.getNode(AArch64ISD::UZP2, DL, ResVT,
+                             DAG.getNode(AArch64ISD::NVCAST, DL, ResVT,
+                                         PreCastOp0.getOperand(0)),
+                             DAG.getNode(AArch64ISD::NVCAST, DL, ResVT,
+                                         PreCastOp1.getOperand(0)));
+      }
+    }
+  }
+
   // These optimizations only work on little endian.
   if (!DAG.getDataLayout().isLittleEndian())
     return SDValue();

llvm/test/CodeGen/AArch64/sve-shift-trunc-combine.ll

@@ -9,23 +9,19 @@
 define void @lshr_trunc_i16_load_scalable(ptr %src, ptr %dst) #0 {
 ; CHECK-LABEL: lshr_trunc_i16_load_scalable:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr z0, [x0]
-; CHECK-NEXT:    ldr z1, [x0, #1, mul vl]
-; CHECK-NEXT:    lsr z1.h, z1.h, #8
-; CHECK-NEXT:    lsr z0.h, z0.h, #8
-; CHECK-NEXT:    uzp1 z0.b, z0.b, z1.b
+; CHECK-NEXT:    ldr z0, [x0, #1, mul vl]
+; CHECK-NEXT:    ldr z1, [x0]
+; CHECK-NEXT:    uzp2 z0.b, z1.b, z0.b
 ; CHECK-NEXT:    str z0, [x1]
 ; CHECK-NEXT:    ret
 ;
 ; BE-LABEL: lshr_trunc_i16_load_scalable:
 ; BE:       // %bb.0:
 ; BE-NEXT:    ptrue p0.h
-; BE-NEXT:    ld1h { z0.h }, p0/z, [x0]
-; BE-NEXT:    ld1h { z1.h }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1h { z0.h }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1h { z1.h }, p0/z, [x0]
 ; BE-NEXT:    ptrue p0.b
-; BE-NEXT:    lsr z1.h, z1.h, #8
-; BE-NEXT:    lsr z0.h, z0.h, #8
-; BE-NEXT:    uzp1 z0.b, z0.b, z1.b
+; BE-NEXT:    uzp2 z0.b, z1.b, z0.b
 ; BE-NEXT:    st1b { z0.b }, p0, [x1]
 ; BE-NEXT:    ret
   %x = load <vscale x 16 x i16>, ptr %src, align 2
@@ -61,23 +57,19 @@ define void @lshr_trunc_i16_load_fixed_length(ptr %src, ptr %dst) #0 {
 define void @ashr_trunc_i16_load(ptr %src, ptr %dst)#0 {
 ; CHECK-LABEL: ashr_trunc_i16_load:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr z0, [x0]
-; CHECK-NEXT:    ldr z1, [x0, #1, mul vl]
-; CHECK-NEXT:    lsr z1.h, z1.h, #8
-; CHECK-NEXT:    lsr z0.h, z0.h, #8
-; CHECK-NEXT:    uzp1 z0.b, z0.b, z1.b
+; CHECK-NEXT:    ldr z0, [x0, #1, mul vl]
+; CHECK-NEXT:    ldr z1, [x0]
+; CHECK-NEXT:    uzp2 z0.b, z1.b, z0.b
 ; CHECK-NEXT:    str z0, [x1]
 ; CHECK-NEXT:    ret
 ;
 ; BE-LABEL: ashr_trunc_i16_load:
 ; BE:       // %bb.0:
 ; BE-NEXT:    ptrue p0.h
-; BE-NEXT:    ld1h { z0.h }, p0/z, [x0]
-; BE-NEXT:    ld1h { z1.h }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1h { z0.h }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1h { z1.h }, p0/z, [x0]
 ; BE-NEXT:    ptrue p0.b
-; BE-NEXT:    lsr z1.h, z1.h, #8
-; BE-NEXT:    lsr z0.h, z0.h, #8
-; BE-NEXT:    uzp1 z0.b, z0.b, z1.b
+; BE-NEXT:    uzp2 z0.b, z1.b, z0.b
 ; BE-NEXT:    st1b { z0.b }, p0, [x1]
 ; BE-NEXT:    ret
   %x = load <vscale x 16 x i16>, ptr %src, align 2
@@ -90,23 +82,19 @@ define void @ashr_trunc_i16_load(ptr %src, ptr %dst)#0 {
 define void @lshr_trunc_i32_load(ptr %src, ptr %dst) #0 {
 ; CHECK-LABEL: lshr_trunc_i32_load:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr z0, [x0]
-; CHECK-NEXT:    ldr z1, [x0, #1, mul vl]
-; CHECK-NEXT:    lsr z1.s, z1.s, #16
-; CHECK-NEXT:    lsr z0.s, z0.s, #16
-; CHECK-NEXT:    uzp1 z0.h, z0.h, z1.h
+; CHECK-NEXT:    ldr z0, [x0, #1, mul vl]
+; CHECK-NEXT:    ldr z1, [x0]
+; CHECK-NEXT:    uzp2 z0.h, z1.h, z0.h
 ; CHECK-NEXT:    str z0, [x1]
 ; CHECK-NEXT:    ret
 ;
 ; BE-LABEL: lshr_trunc_i32_load:
 ; BE:       // %bb.0:
 ; BE-NEXT:    ptrue p0.s
-; BE-NEXT:    ld1w { z0.s }, p0/z, [x0]
-; BE-NEXT:    ld1w { z1.s }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1w { z0.s }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1w { z1.s }, p0/z, [x0]
 ; BE-NEXT:    ptrue p0.h
-; BE-NEXT:    lsr z1.s, z1.s, #16
-; BE-NEXT:    lsr z0.s, z0.s, #16
-; BE-NEXT:    uzp1 z0.h, z0.h, z1.h
+; BE-NEXT:    uzp2 z0.h, z1.h, z0.h
 ; BE-NEXT:    st1h { z0.h }, p0, [x1]
 ; BE-NEXT:    ret
   %x = load <vscale x 8 x i32>, ptr %src, align 2
@@ -119,23 +107,19 @@ define void @lshr_trunc_i32_load(ptr %src, ptr %dst) #0 {
 define void @lshr_trunc_i64_load(ptr %src, ptr %dst) #0 {
 ; CHECK-LABEL: lshr_trunc_i64_load:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ldr z0, [x0]
-; CHECK-NEXT:    ldr z1, [x0, #1, mul vl]
-; CHECK-NEXT:    lsr z1.d, z1.d, #32
-; CHECK-NEXT:    lsr z0.d, z0.d, #32
-; CHECK-NEXT:    uzp1 z0.s, z0.s, z1.s
+; CHECK-NEXT:    ldr z0, [x0, #1, mul vl]
+; CHECK-NEXT:    ldr z1, [x0]
+; CHECK-NEXT:    uzp2 z0.s, z1.s, z0.s
 ; CHECK-NEXT:    str z0, [x1]
 ; CHECK-NEXT:    ret
 ;
 ; BE-LABEL: lshr_trunc_i64_load:
 ; BE:       // %bb.0:
 ; BE-NEXT:    ptrue p0.d
-; BE-NEXT:    ld1d { z0.d }, p0/z, [x0]
-; BE-NEXT:    ld1d { z1.d }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1d { z0.d }, p0/z, [x0, #1, mul vl]
+; BE-NEXT:    ld1d { z1.d }, p0/z, [x0]
 ; BE-NEXT:    ptrue p0.s
-; BE-NEXT:    lsr z1.d, z1.d, #32
-; BE-NEXT:    lsr z0.d, z0.d, #32
-; BE-NEXT:    uzp1 z0.s, z0.s, z1.s
+; BE-NEXT:    uzp2 z0.s, z1.s, z0.s
 ; BE-NEXT:    st1w { z0.s }, p0, [x1]
 ; BE-NEXT:    ret
   %x = load <vscale x 4 x i64>, ptr %src, align 2