[llvm] [AArch64] Add SVE shuffle optimization pass (PR #193951)

[Sander de Smalen via llvm-commits]

================
@@ -0,0 +1,602 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -O3 -mtriple=aarch64-linux-gnu -mattr=+sve < %s | FileCheck %s
+
+define void @zext_nxv8i16_to_nxv8i64_deinterleave_in_loop(ptr %src, ptr %dst, <vscale x 8 x i1> %mask) #0 {
+; CHECK-LABEL: zext_nxv8i16_to_nxv8i64_deinterleave_in_loop:
+; CHECK:       // %bb.0: // %entry
+; CHECK-NEXT:    index z7.d, #0, #4
+; CHECK-NEXT:    mov z4.d, #0xffffffffffff0000
+; CHECK-NEXT:    mov z5.d, #0xffffffffffff0001
+; CHECK-NEXT:    mov x8, #-65534 // =0xffffffffffff0002
+; CHECK-NEXT:    mov z24.d, #0xffffffffffff0003
+; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    mov z6.d, x8
+; CHECK-NEXT:    movi v1.2d, #0000000000000000
+; CHECK-NEXT:    mov w8, #2048 // =0x800
+; CHECK-NEXT:    movi v2.2d, #0000000000000000
+; CHECK-NEXT:    movi v3.2d, #0000000000000000
+; CHECK-NEXT:    cntd x9
+; CHECK-NEXT:    add z4.d, z7.d, z4.d
+; CHECK-NEXT:    add z5.d, z7.d, z5.d
+; CHECK-NEXT:    rdvl x10, #1
+; CHECK-NEXT:    add z6.d, z7.d, z6.d
+; CHECK-NEXT:    add z7.d, z7.d, z24.d
+; CHECK-NEXT:  .LBB0_1: // %loop
+; CHECK-NEXT:    // =>This Inner Loop Header: Depth=1
+; CHECK-NEXT:    ld1h { z24.h }, p0/z, [x0]
+; CHECK-NEXT:    subs x8, x8, x9
+; CHECK-NEXT:    add x0, x0, x10
+; CHECK-NEXT:    tbl z25.h, { z24.h }, z4.h
+; CHECK-NEXT:    tbl z26.h, { z24.h }, z5.h
+; CHECK-NEXT:    tbl z27.h, { z24.h }, z6.h
+; CHECK-NEXT:    tbl z24.h, { z24.h }, z7.h
+; CHECK-NEXT:    add z0.d, z0.d, z25.d
+; CHECK-NEXT:    add z1.d, z1.d, z26.d
+; CHECK-NEXT:    add z2.d, z2.d, z27.d
+; CHECK-NEXT:    add z3.d, z3.d, z24.d
+; CHECK-NEXT:    b.ne .LBB0_1
+; CHECK-NEXT:  // %bb.2: // %exit
+; CHECK-NEXT:    str z0, [x1]
+; CHECK-NEXT:    str z1, [x1, #1, mul vl]
+; CHECK-NEXT:    str z2, [x1, #2, mul vl]
+; CHECK-NEXT:    str z3, [x1, #3, mul vl]
+; CHECK-NEXT:    ret
+entry:
+  %vscale = tail call i64 @llvm.vscale.i64()
+  %stride = shl nuw nsw i64 %vscale, 1
+  br label %loop
+
+loop:
+  %iv = phi i64 [ 0, %entry ], [ %iv.next, %loop ]
+  %acc.b.i64 = phi <vscale x 2 x i64> [ splat(i64 0), %entry ], [ %add.b.i64, %loop ]
+  %acc.g.i64 = phi <vscale x 2 x i64> [ splat(i64 0), %entry ], [ %add.g.i64, %loop ]
+  %acc.r.i64 = phi <vscale x 2 x i64> [ splat(i64 0), %entry ], [ %add.r.i64, %loop ]
+  %acc.a.i64 = phi <vscale x 2 x i64> [ splat(i64 0), %entry ], [ %add.a.i64, %loop ]
+  %src.gep = getelementptr inbounds nuw [4 x i16], ptr %src, i64 %iv
+  %bgra = call <vscale x 8 x i16> @llvm.masked.load(ptr %src.gep, <vscale x 8 x i1> %mask, <vscale x 8 x i16> zeroinitializer)
+  %deinterleave = tail call { <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16> } @llvm.vector.deinterleave4(<vscale x 8 x i16> %bgra)
+  %b.i16 = extractvalue { <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16> } %deinterleave, 0
+  %g.i16 = extractvalue { <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16> } %deinterleave, 1
+  %r.i16 = extractvalue { <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16> } %deinterleave, 2
+  %a.i16 = extractvalue { <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16>, <vscale x 2 x i16> } %deinterleave, 3
+  %b.i64 = zext <vscale x 2 x i16> %b.i16 to <vscale x 2 x i64>
+  %g.i64 = zext <vscale x 2 x i16> %g.i16 to <vscale x 2 x i64>
+  %r.i64 = zext <vscale x 2 x i16> %r.i16 to <vscale x 2 x i64>
+  %a.i64 = zext <vscale x 2 x i16> %a.i16 to <vscale x 2 x i64>
+  %add.b.i64 = add <vscale x 2 x i64> %acc.b.i64, %b.i64
+  %add.g.i64 = add <vscale x 2 x i64> %acc.g.i64, %g.i64
+  %add.r.i64 = add <vscale x 2 x i64> %acc.r.i64, %r.i64
+  %add.a.i64 = add <vscale x 2 x i64> %acc.a.i64, %a.i64
+  %iv.next = add nuw i64 %iv, %stride
+  %ec = icmp eq i64 %iv.next, 2048
+  br i1 %ec, label %exit, label %loop
+
+exit:
+  store <vscale x 2 x i64> %add.b.i64, ptr %dst
+  %g.i64.gep = getelementptr <vscale x 2 x i64>, ptr %dst, i64 1
+  store <vscale x 2 x i64> %add.g.i64, ptr %g.i64.gep
+  %r.i64.gep = getelementptr <vscale x 2 x i64>, ptr %dst, i64 2
+  store <vscale x 2 x i64> %add.r.i64, ptr %r.i64.gep
+  %a.i64.gep = getelementptr <vscale x 2 x i64>, ptr %dst, i64 3
+  store <vscale x 2 x i64> %add.a.i64, ptr %a.i64.gep
+  ret void
+}
+
+;; TODO: Do we want to perform the sext equivalent? Requires a splat of the
----------------
sdesmalen-arm wrote:

That seems like a sensible follow-up to me, and yes it could use the two-source TBL for this case.

https://github.com/llvm/llvm-project/pull/193951