xref: /aosp_15_r20/external/XNNPACK/src/qu8-gemm/gen/4x16c4-minmax-rndnu-aarch64-neondot-ld128.S (revision 4bdc94577ba0e567308109d787f7fec7b531ce36)

// Auto-generated file. Do not edit!
//   Template: src/qu8-gemm/4x16c4-aarch64-neondot-ld128.S.in
//   Generator: tools/xngen
//
// Copyright 2020 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.


#include <xnnpack/assembly.h>

# void xnn_qu8_gemm_minmax_rndnu_ukernel_4x16c4__aarch64_neondot_ld128(
#     size_t mr,                 x0
#     size_t nc,                 x1
#     size_t kc,                 x2 / x0
#     const int8_t* restrict a,  x3
#     size_t a_stride,           x4
#     const void* restrict w,    x5
#     int8_t* restrict c,        x6
#     size_t cm_stride,          x7
#     size_t cn_stride,          [sp] -> x12
#     const union xnn_qu8_conv_minmax_params params)  [sp + 8] -> x11

# params structure is 20 bytes
#  struct {
#    uint8_t kernel_zero_point[4];
#    int32_t right_pre_shift;
#    int32_t multiplier;
#    int32_t right_post_shift;
#    int16_t output_zero_point;
#    int8_t output_min;
#    int8_t output_max;
#  } rndnu_neon;

# d8-d15, x19-x30 need to be preserved if used. x18 is reserved by the OS.

# Register usage
# A0  x3 v0
# A1 x15 v1
# A2 x13 v2
# A3  x4 v3
# B   x5 v4  v5  v6  v7
# C0  x6 v16 v20 v24 v28
# C1  x8 v17 v21 v25 v29
# C2  x9 v18 v22 v26 v30
# C3  x7 v19 v23 v27 v31
# zero_point v8 v12 v13 v14 v15
# unused v9 v10 v11

BEGIN_FUNCTION xnn_qu8_gemm_minmax_rndnu_ukernel_4x16c4__aarch64_neondot_ld128

        # Clamp A and C pointers
        CMP     x0, 2                   // if mr < 2
        LDP     x12, x11, [sp]          // cn_stride, params
        ADD     x2, x2, 3               // kc = (kc + 3) & ~3
        ADD     x15, x3, x4             // a1 = a0 + a_stride
        ADD     x8, x6, x7              // c1 = c0 + cm_stride

        # Save d8,d12-d15 on stack
        STR     d8,  [sp, -48]!
        CSEL    x15, x3, x15, LO        //   a1 = a0
        CSEL    x8, x6,  x8, LO         //   c1 = c0
        BIC     x2, x2, 3

        STP     d12, d13, [sp, 16]
        ADD     x13, x15, x4            // a2 = a1 + a_stride
        ADD     x9,  x8, x7             // c2 = c1 + cm_stride
                                        // if mr <= 2
        CSEL    x13, x15, x13, LS       //   a2 = a1
        CSEL    x9,  x8,  x9, LS        //   c2 = c1

        STP     d14, d15, [sp, 32]
        CMP     x0, 4                   // if mr < 4
        ADD     x4, x13, x4             // a3 = a2 + a_stride
        ADD     x7,  x9, x7             // c3 = c2 + cm_stride

        LD1R    {v8.4s}, [x11], 4       // kernel_zero_point

        CSEL    x4, x13, x4, LO         //   a3 = a2
        CSEL    x7,  x9, x7, LO         //   c3 = c2

        .p2align 3
0:
        # Load initial bias from w into accumulators
        LDP     q16, q20, [x5], 32

        MOVI    v12.4s, 0
        MOVI    v13.4s, 0
        MOVI    v14.4s, 0
        MOVI    v15.4s, 0

        MOV     v17.16b, v16.16b
        MOV     v18.16b, v16.16b
        LDP     q24, q28, [x5], 32
        MOV     v19.16b, v16.16b
        MOV     v21.16b, v20.16b
        SUBS    x0, x2, 16              // k = kc - 16
        MOV     v22.16b, v20.16b
        MOV     v23.16b, v20.16b
        MOV     v25.16b, v24.16b
        MOV     v26.16b, v24.16b
        MOV     v27.16b, v24.16b
        MOV     v29.16b, v28.16b
        MOV     v30.16b, v28.16b
        MOV     v31.16b, v28.16b

        # Is there at least 16 bytes?
        B.LO    3f

        # Main loop - 16 bytes of A
        .p2align 3
1:
        LDR     q0,  [x3], 16
        LDR     q4,  [x5], 16
        LDR     q1, [x15], 16
        LDR     q2, [x13], 16
        LDR     q3,  [x4], 16
        LDR     q5,  [x5], 16

        UDOT    v12.4s, v8.16b,  v0.16b  // update zero point
        UDOT    v13.4s, v8.16b,  v1.16b
        UDOT    v14.4s, v8.16b,  v2.16b
        UDOT    v15.4s, v8.16b,  v3.16b

        UDOT    v16.4s, v4.16b,  v0.4b[0]
        UDOT    v17.4s, v4.16b,  v1.4b[0]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[0]
        UDOT    v19.4s, v4.16b,  v3.4b[0]
        UDOT    v20.4s, v5.16b,  v0.4b[0]
        UDOT    v21.4s, v5.16b,  v1.4b[0]
        UDOT    v22.4s, v5.16b,  v2.4b[0]
        UDOT    v23.4s, v5.16b,  v3.4b[0]
        UDOT    v24.4s, v6.16b, v0.4b[0]
        UDOT    v25.4s, v6.16b, v1.4b[0]
        LDP     q4, q5, [x5], 32
        UDOT    v26.4s, v6.16b, v2.4b[0]
        UDOT    v27.4s, v6.16b, v3.4b[0]
        UDOT    v28.4s, v7.16b, v0.4b[0]
        UDOT    v29.4s, v7.16b, v1.4b[0]
        UDOT    v30.4s, v7.16b, v2.4b[0]
        UDOT    v31.4s, v7.16b, v3.4b[0]

        UDOT    v16.4s, v4.16b,  v0.4b[1]
        UDOT    v17.4s, v4.16b,  v1.4b[1]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[1]
        UDOT    v19.4s, v4.16b,  v3.4b[1]
        UDOT    v20.4s, v5.16b,  v0.4b[1]
        UDOT    v21.4s, v5.16b,  v1.4b[1]
        UDOT    v22.4s, v5.16b,  v2.4b[1]
        UDOT    v23.4s, v5.16b,  v3.4b[1]
        UDOT    v24.4s, v6.16b,  v0.4b[1]
        UDOT    v25.4s, v6.16b,  v1.4b[1]
        LDP     q4, q5, [x5], 32
        UDOT    v26.4s, v6.16b,  v2.4b[1]
        UDOT    v27.4s, v6.16b,  v3.4b[1]
        UDOT    v28.4s, v7.16b,  v0.4b[1]
        UDOT    v29.4s, v7.16b,  v1.4b[1]
        UDOT    v30.4s, v7.16b,  v2.4b[1]
        UDOT    v31.4s, v7.16b,  v3.4b[1]

        UDOT    v16.4s, v4.16b,  v0.4b[2]
        UDOT    v17.4s, v4.16b,  v1.4b[2]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[2]
        UDOT    v19.4s, v4.16b,  v3.4b[2]
        UDOT    v20.4s, v5.16b,  v0.4b[2]
        UDOT    v21.4s, v5.16b,  v1.4b[2]
        UDOT    v22.4s, v5.16b,  v2.4b[2]
        UDOT    v23.4s, v5.16b,  v3.4b[2]
        UDOT    v24.4s, v6.16b,  v0.4b[2]
        UDOT    v25.4s, v6.16b,  v1.4b[2]
        LDP     q4, q5, [x5], 32
        UDOT    v26.4s, v6.16b,  v2.4b[2]
        UDOT    v27.4s, v6.16b,  v3.4b[2]
        UDOT    v28.4s, v7.16b,  v0.4b[2]
        UDOT    v29.4s, v7.16b,  v1.4b[2]
        UDOT    v30.4s, v7.16b,  v2.4b[2]
        UDOT    v31.4s, v7.16b,  v3.4b[2]

        UDOT    v16.4s, v4.16b,  v0.4b[3]
        UDOT    v17.4s, v4.16b,  v1.4b[3]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[3]
        UDOT    v19.4s, v4.16b,  v3.4b[3]
        UDOT    v20.4s, v5.16b,  v0.4b[3]
        UDOT    v21.4s, v5.16b,  v1.4b[3]
        UDOT    v22.4s, v5.16b,  v2.4b[3]
        UDOT    v23.4s, v5.16b,  v3.4b[3]
        UDOT    v24.4s, v6.16b,  v0.4b[3]
        UDOT    v25.4s, v6.16b,  v1.4b[3]
        UDOT    v26.4s, v6.16b,  v2.4b[3]
        UDOT    v27.4s, v6.16b,  v3.4b[3]
        SUBS    x0, x0, 16
        UDOT    v28.4s, v7.16b,  v0.4b[3]
        UDOT    v29.4s, v7.16b,  v1.4b[3]
        UDOT    v30.4s, v7.16b,  v2.4b[3]
        UDOT    v31.4s, v7.16b,  v3.4b[3]
        B.HS    1b

        # Is there a remainder?- 4 to 12 bytes of A
        TST     x0, 15
        B.NE    3f

2:
        ADDP    v0.4s, v12.4s, v12.4s
        ADDP    v1.4s, v13.4s, v13.4s
        ADDP    v2.4s, v14.4s, v14.4s
        ADDP    v3.4s, v15.4s, v15.4s
        ADDP    v12.4s, v0.4s, v0.4s
        ADDP    v13.4s, v1.4s, v1.4s
        ADDP    v14.4s, v2.4s, v2.4s
        ADDP    v15.4s, v3.4s, v3.4s

        # Subtract zero point from accumulators
        SUB     v16.4s, v16.4s, v12.4s
        SUB     v17.4s, v17.4s, v13.4s
        SUB     v18.4s, v18.4s, v14.4s
        SUB     v19.4s, v19.4s, v15.4s
        SUB     v20.4s, v20.4s, v12.4s
        SUB     v21.4s, v21.4s, v13.4s
        SUB     v22.4s, v22.4s, v14.4s
        SUB     v23.4s, v23.4s, v15.4s
        SUB     v24.4s, v24.4s, v12.4s
        SUB     v25.4s, v25.4s, v13.4s
        SUB     v26.4s, v26.4s, v14.4s
        SUB     v27.4s, v27.4s, v15.4s
        SUB     v28.4s, v28.4s, v12.4s
        SUB     v29.4s, v29.4s, v13.4s
        SUB     v30.4s, v30.4s, v14.4s
        SUB     v31.4s, v31.4s, v15.4s

        # Apply params - preshift, scale, postshift, bias and clamp
        LD1R    {v4.4s}, [x11], 4
        SSHL    v16.4s, v16.4s, v4.4s   // shift to upper bits
        SSHL    v17.4s, v17.4s, v4.4s
        SSHL    v18.4s, v18.4s, v4.4s
        SSHL    v19.4s, v19.4s, v4.4s
        SSHL    v20.4s, v20.4s, v4.4s
        SSHL    v21.4s, v21.4s, v4.4s
        SSHL    v22.4s, v22.4s, v4.4s
        SSHL    v23.4s, v23.4s, v4.4s
        LD1R    {v5.4s}, [x11], 4
        SSHL    v24.4s, v24.4s, v4.4s
        SSHL    v25.4s, v25.4s, v4.4s
        SSHL    v26.4s, v26.4s, v4.4s
        SSHL    v27.4s, v27.4s, v4.4s
        SSHL    v28.4s, v28.4s, v4.4s
        SSHL    v29.4s, v29.4s, v4.4s
        SSHL    v30.4s, v30.4s, v4.4s
        SSHL    v31.4s, v31.4s, v4.4s
        LD1R    {v6.4s}, [x11], 4
        SQDMULH v16.4s, v16.4s, v5.4s   // scale without rounding
        SQDMULH v17.4s, v17.4s, v5.4s
        SQDMULH v18.4s, v18.4s, v5.4s
        SQDMULH v19.4s, v19.4s, v5.4s
        SQDMULH v20.4s, v20.4s, v5.4s
        SQDMULH v21.4s, v21.4s, v5.4s
        SQDMULH v22.4s, v22.4s, v5.4s
        SQDMULH v23.4s, v23.4s, v5.4s
        SQDMULH v24.4s, v24.4s, v5.4s
        SQDMULH v25.4s, v25.4s, v5.4s
        SQDMULH v26.4s, v26.4s, v5.4s
        SQDMULH v27.4s, v27.4s, v5.4s
        SQDMULH v28.4s, v28.4s, v5.4s
        SQDMULH v29.4s, v29.4s, v5.4s
        SQDMULH v30.4s, v30.4s, v5.4s
        SQDMULH v31.4s, v31.4s, v5.4s
        SRSHL   v16.4s, v16.4s, v6.4s   // signed rounding shift left
        SRSHL   v17.4s, v17.4s, v6.4s
        SRSHL   v18.4s, v18.4s, v6.4s
        SRSHL   v19.4s, v19.4s, v6.4s
        SRSHL   v20.4s, v20.4s, v6.4s
        SRSHL   v21.4s, v21.4s, v6.4s
        SRSHL   v22.4s, v22.4s, v6.4s
        SRSHL   v23.4s, v23.4s, v6.4s
        SRSHL   v24.4s, v24.4s, v6.4s
        SRSHL   v25.4s, v25.4s, v6.4s
        SRSHL   v26.4s, v26.4s, v6.4s
        SRSHL   v27.4s, v27.4s, v6.4s
        SRSHL   v28.4s, v28.4s, v6.4s
        SRSHL   v29.4s, v29.4s, v6.4s
        SRSHL   v30.4s, v30.4s, v6.4s
        SRSHL   v31.4s, v31.4s, v6.4s

        SQXTN   v16.4h, v16.4s
        SQXTN   v17.4h, v17.4s
        SQXTN   v18.4h, v18.4s
        SQXTN   v19.4h, v19.4s
        SQXTN   v24.4h, v24.4s
        SQXTN   v25.4h, v25.4s
        SQXTN   v26.4h, v26.4s
        SQXTN   v27.4h, v27.4s
        LD1R    {v6.8h}, [x11], 2       // add bias

        SQXTN2  v16.8h, v20.4s
        SQXTN2  v17.8h, v21.4s
        SQXTN2  v18.8h, v22.4s
        SQXTN2  v19.8h, v23.4s
        SQXTN2  v24.8h, v28.4s
        SQXTN2  v25.8h, v29.4s
        SQXTN2  v26.8h, v30.4s
        SQXTN2  v27.8h, v31.4s

        SQADD   v16.8h, v16.8h, v6.8h
        SQADD   v17.8h, v17.8h, v6.8h
        SQADD   v18.8h, v18.8h, v6.8h
        SQADD   v19.8h, v19.8h, v6.8h
        SQADD   v24.8h, v24.8h, v6.8h
        SQADD   v25.8h, v25.8h, v6.8h
        SQADD   v26.8h, v26.8h, v6.8h
        SQADD   v27.8h, v27.8h, v6.8h
        LD1R    {v4.16b}, [x11], 1      // clamp min value

        SQXTUN  v0.8b, v16.8h
        SQXTUN  v1.8b, v17.8h
        SQXTUN  v2.8b, v18.8h
        SQXTUN  v3.8b, v19.8h
        LD1R    {v5.16b}, [x11]         // clamp max value
        SQXTUN2 v0.16b, v24.8h
        SQXTUN2 v1.16b, v25.8h
        SQXTUN2 v2.16b, v26.8h
        SQXTUN2 v3.16b, v27.8h

        SUB     x11, x11, 15             // rewind params pointer

        UMAX    v0.16b, v0.16b, v4.16b
        UMAX    v1.16b, v1.16b, v4.16b
        UMAX    v2.16b, v2.16b, v4.16b
        UMAX    v3.16b, v3.16b, v4.16b
        SUBS    x1, x1, 16
        UMIN    v0.16b, v0.16b, v5.16b
        UMIN    v1.16b, v1.16b, v5.16b
        UMIN    v2.16b, v2.16b, v5.16b
        UMIN    v3.16b, v3.16b, v5.16b
        B.LO    5f

        # Store full 4 x 16
        ST1     {v0.16b}, [x6], x12
        SUB     x3,  x3, x2             // a0 -= kc
        ST1     {v1.16b}, [x8], x12
        SUB     x15, x15, x2            // a1 -= kc
        ST1     {v2.16b}, [x9], x12
        SUB     x13, x13, x2            // a2 -= kc
        ST1     {v3.16b}, [x7], x12
        SUB     x4,  x4, x2             // a3 -= kc
        B.NE    0b

        # Restore d8,d12-d15 from stack
        LDP     d14, d15, [sp, 32]
        LDP     d12, d13, [sp, 16]
        LDR     d8,  [sp], 48
        RET

        # Remainder- 8 bytes of A
        .p2align 3
3:
        # Is there a remainder?- 8 bytes of A
        TBZ     x0, 3, 4f

        LDR     d0,  [x3], 8
        LDR     q4,  [x5], 16
        LDR     d1, [x15], 8
        LDR     d2, [x13], 8
        LDR     d3,  [x4], 8
        LDR     q5,  [x5], 16

        UDOT    v12.4s, v8.16b,  v0.16b  // update zero point
        UDOT    v13.4s, v8.16b,  v1.16b
        UDOT    v14.4s, v8.16b,  v2.16b
        UDOT    v15.4s, v8.16b,  v3.16b

        UDOT    v16.4s, v4.16b,  v0.4b[0]
        UDOT    v17.4s, v4.16b,  v1.4b[0]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[0]
        UDOT    v19.4s, v4.16b,  v3.4b[0]
        UDOT    v20.4s, v5.16b,  v0.4b[0]
        UDOT    v21.4s, v5.16b,  v1.4b[0]
        UDOT    v22.4s, v5.16b,  v2.4b[0]
        UDOT    v23.4s, v5.16b,  v3.4b[0]
        UDOT    v24.4s, v6.16b, v0.4b[0]
        UDOT    v25.4s, v6.16b, v1.4b[0]
        LDP     q4, q5, [x5], 32
        UDOT    v26.4s, v6.16b, v2.4b[0]
        UDOT    v27.4s, v6.16b, v3.4b[0]
        UDOT    v28.4s, v7.16b, v0.4b[0]
        UDOT    v29.4s, v7.16b, v1.4b[0]
        UDOT    v30.4s, v7.16b, v2.4b[0]
        UDOT    v31.4s, v7.16b, v3.4b[0]
        UDOT    v16.4s, v4.16b,  v0.4b[1]
        UDOT    v17.4s, v4.16b,  v1.4b[1]
        LDP     q6, q7, [x5], 32
        UDOT    v18.4s, v4.16b,  v2.4b[1]
        UDOT    v19.4s, v4.16b,  v3.4b[1]
        UDOT    v20.4s, v5.16b,  v0.4b[1]
        UDOT    v21.4s, v5.16b,  v1.4b[1]
        UDOT    v22.4s, v5.16b,  v2.4b[1]
        UDOT    v23.4s, v5.16b,  v3.4b[1]
        UDOT    v24.4s, v6.16b,  v0.4b[1]
        UDOT    v25.4s, v6.16b,  v1.4b[1]
        UDOT    v26.4s, v6.16b,  v2.4b[1]
        UDOT    v27.4s, v6.16b,  v3.4b[1]
        UDOT    v28.4s, v7.16b,  v0.4b[1]
        UDOT    v29.4s, v7.16b,  v1.4b[1]
        UDOT    v30.4s, v7.16b,  v2.4b[1]
        UDOT    v31.4s, v7.16b,  v3.4b[1]
        # Is there a remainder?- 4 bytes of A
        TBZ     x0, 2, 2b

        # Remainder- 4 bytes of A
4:
        LDR     s0,  [x3], 4
        LDR     q4,  [x5], 16
        LDR     s1, [x15], 4
        LDR     s2, [x13], 4
        LDR     s3,  [x4], 4
        LDR     q5, [x5], 16

        UDOT    v12.4s, v8.16b,  v0.16b  // update zero point
        UDOT    v13.4s, v8.16b,  v1.16b
        UDOT    v14.4s, v8.16b,  v2.16b
        UDOT    v15.4s, v8.16b,  v3.16b

        UDOT    v16.4s, v4.16b,  v0.4b[0]
        UDOT    v17.4s, v4.16b,  v1.4b[0]
        UDOT    v18.4s, v4.16b,  v2.4b[0]
        UDOT    v19.4s, v4.16b,  v3.4b[0]
        LDP     q6, q7, [x5], 32
        UDOT    v20.4s, v5.16b,  v0.4b[0]
        UDOT    v21.4s, v5.16b,  v1.4b[0]
        UDOT    v22.4s, v5.16b,  v2.4b[0]
        UDOT    v23.4s, v5.16b,  v3.4b[0]
        UDOT    v24.4s, v6.16b, v0.4b[0]
        UDOT    v25.4s, v6.16b, v1.4b[0]
        UDOT    v26.4s, v6.16b, v2.4b[0]
        UDOT    v27.4s, v6.16b, v3.4b[0]
        UDOT    v28.4s, v7.16b, v0.4b[0]
        UDOT    v29.4s, v7.16b, v1.4b[0]
        UDOT    v30.4s, v7.16b, v2.4b[0]
        UDOT    v31.4s, v7.16b, v3.4b[0]
        B       2b

        # Store odd width
        .p2align 3
5:
        TBZ     x1, 3, 6f
        STR     d0, [x6], 8
        STR     d1, [x8], 8
        DUP     d0, v0.d[1]
        DUP     d1, v1.d[1]
        STR     d2, [x9], 8
        STR     d3, [x7], 8
        DUP     d2, v2.d[1]
        DUP     d3, v3.d[1]
6:
        TBZ     x1, 2, 7f
        STR     s0, [x6], 4
        STR     s1, [x8], 4
        DUP     s0, v0.s[1]
        DUP     s1, v1.s[1]
        STR     s2, [x9], 4
        STR     s3, [x7], 4
        DUP     s2, v2.s[1]
        DUP     s3, v3.s[1]
7:
        TBZ     x1, 1, 8f
        STR     h0, [x6], 2
        STR     h1, [x8], 2
        DUP     h0, v0.h[1]
        DUP     h1, v1.h[1]
        STR     h2, [x9], 2
        STR     h3, [x7], 2
        DUP     h2, v2.h[1]
        DUP     h3, v3.h[1]
8:
        TBZ     x1, 0, 9f
        STR     b0, [x6]
        STR     b1, [x8]
        STR     b2, [x9]
        STR     b3, [x7]
9:
        # Restore d8,d12-d15 from stack
        LDP     d14, d15, [sp, 32]
        LDP     d12, d13, [sp, 16]
        LDR     d8,  [sp], 48
        RET

END_FUNCTION xnn_qu8_gemm_minmax_rndnu_ukernel_4x16c4__aarch64_neondot_ld128

#ifdef __ELF__
.section ".note.GNU-stack","",%progbits
#endif