// 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. $assert REQUANTIZATION in ["FP32", "RNDNU"] #include $REWIND_DECREMENT = {"RNDNU": 15, "FP32": 7}[REQUANTIZATION] # void xnn_qu8_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__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 # 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 # C0 x6 v16 v20 # C1 x8 v17 v21 # C2 x9 v18 v22 # C3 x7 v19 v23 # zero_point v7 v24 v25 v26 v27 # unused v8 v9 v10 v11 v13 v14 v15 v28 v29 v30 v31 BEGIN_FUNCTION xnn_qu8_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__aarch64_neondot_ld128 # Clamp A and C pointers CMP x0, 2 // if mr < 2 ADD x2, x2, 3 // kc = (kc + 3) & ~3 ADD x15, x3, x4 // a1 = a0 + a_stride ADD x8, x6, x7 // c1 = c0 + cm_stride CSEL x15, x3, x15, LO // a1 = a0 CSEL x8, x6, x8, LO // c1 = c0 BIC x2, x2, 3 LDP x12, x11, [sp] // cn_stride, params 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 LD1R {v7.4s}, [x11], 4 // kernel_zero_point CMP x0, 4 // if mr < 4 ADD x4, x13, x4 // a3 = a2 + a_stride ADD x7, x9, x7 // c3 = c2 + cm_stride 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 SUBS x0, x2, 16 // k = kc - 16 MOV v17.16b, v16.16b MOV v18.16b, v16.16b MOV v19.16b, v16.16b MOV v21.16b, v20.16b MOV v22.16b, v20.16b MOV v23.16b, v20.16b MOVI v24.16b, 0 MOVI v25.16b, 0 MOVI v26.16b, 0 MOVI v27.16b, 0 # Is there at least 16 bytes? B.LO 30f # 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 v24.4s, v7.16b, v0.16b // update zero point UDOT v25.4s, v7.16b, v1.16b UDOT v26.4s, v7.16b, v2.16b UDOT v27.4s, v7.16b, v3.16b UDOT v16.4s, v4.16b, v0.4b[0] UDOT v17.4s, v4.16b, v1.4b[0] LDR q6, [x5], 16 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] LDR q4, [x5], 16 UDOT v22.4s, v5.16b, v2.4b[0] UDOT v23.4s, v5.16b, v3.4b[0] UDOT v16.4s, v6.16b, v0.4b[1] UDOT v17.4s, v6.16b, v1.4b[1] LDR q5, [x5], 16 UDOT v18.4s, v6.16b, v2.4b[1] UDOT v19.4s, v6.16b, v3.4b[1] UDOT v20.4s, v4.16b, v0.4b[1] UDOT v21.4s, v4.16b, v1.4b[1] LDR q6, [x5], 16 UDOT v22.4s, v4.16b, v2.4b[1] UDOT v23.4s, v4.16b, v3.4b[1] UDOT v16.4s, v5.16b, v0.4b[2] UDOT v17.4s, v5.16b, v1.4b[2] LDR q4, [x5], 16 UDOT v18.4s, v5.16b, v2.4b[2] UDOT v19.4s, v5.16b, v3.4b[2] UDOT v20.4s, v6.16b, v0.4b[2] UDOT v21.4s, v6.16b, v1.4b[2] LDR q5, [x5], 16 UDOT v22.4s, v6.16b, v2.4b[2] UDOT v23.4s, v6.16b, v3.4b[2] UDOT v16.4s, v4.16b, v0.4b[3] UDOT v17.4s, v4.16b, v1.4b[3] UDOT v18.4s, v4.16b, v2.4b[3] UDOT v19.4s, v4.16b, v3.4b[3] SUBS x0, x0, 16 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] B.HS 1b # Is there a remainder?- 8 bytes of A TBNZ x0, 3, 3f # Is there a remainder?- 4 bytes of A TBNZ x0, 2, 4f 2: ADDP v0.4s, v24.4s, v24.4s ADDP v1.4s, v25.4s, v25.4s ADDP v2.4s, v26.4s, v26.4s ADDP v3.4s, v27.4s, v27.4s ADDP v24.4s, v0.4s, v0.4s ADDP v25.4s, v1.4s, v1.4s ADDP v26.4s, v2.4s, v2.4s ADDP v27.4s, v3.4s, v3.4s # Subtract zero point from accumulators SUB v16.4s, v16.4s, v24.4s SUB v17.4s, v17.4s, v25.4s SUB v18.4s, v18.4s, v26.4s SUB v19.4s, v19.4s, v27.4s SUB v20.4s, v20.4s, v24.4s SUB v21.4s, v21.4s, v25.4s SUB v22.4s, v22.4s, v26.4s SUB v23.4s, v23.4s, v27.4s $if REQUANTIZATION == "RNDNU": # 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 LD1R {v5.4s}, [x11], 4 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 {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 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 $elif REQUANTIZATION == "FP32": # Apply params - scale, bias and clamp SCVTF v16.4s, v16.4s SCVTF v17.4s, v17.4s LD1R {v4.4s}, [x11], 4 SCVTF v18.4s, v18.4s SCVTF v19.4s, v19.4s SCVTF v20.4s, v20.4s SCVTF v21.4s, v21.4s SCVTF v22.4s, v22.4s SCVTF v23.4s, v23.4s FMUL v16.4s, v16.4s, v4.4s FMUL v17.4s, v17.4s, v4.4s FMUL v18.4s, v18.4s, v4.4s FMUL v19.4s, v19.4s, v4.4s FMUL v20.4s, v20.4s, v4.4s FMUL v21.4s, v21.4s, v4.4s FMUL v22.4s, v22.4s, v4.4s FMUL v23.4s, v23.4s, v4.4s FCVTNS v16.4s, v16.4s FCVTNS v17.4s, v17.4s FCVTNS v18.4s, v18.4s FCVTNS v19.4s, v19.4s FCVTNS v20.4s, v20.4s FCVTNS v21.4s, v21.4s FCVTNS v22.4s, v22.4s FCVTNS v23.4s, v23.4s SQXTN v16.4h, v16.4s SQXTN v17.4h, v17.4s SQXTN v18.4h, v18.4s SQXTN v19.4h, v19.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 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 LD1R {v4.16b}, [x11], 1 // clamp min value SQXTUN v0.8b, v16.8h SQXTUN v1.8b, v18.8h LD1R {v5.16b}, [x11] // clamp max value SQXTUN2 v0.16b, v17.8h SQXTUN2 v1.16b, v19.8h SUB x11, x11, ${REWIND_DECREMENT} // rewind params pointer UMAX v0.16b, v0.16b, v4.16b UMAX v1.16b, v1.16b, v4.16b SUBS x1, x1, 8 UMIN v0.16b, v0.16b, v5.16b UMIN v1.16b, v1.16b, v5.16b B.LO 5f # Store full 4 x 8 ST1 {v0.8b}, [x6], x12 SUB x3, x3, x2 // a0 -= kc ST1 {v0.d}[1], [x8], x12 SUB x15, x15, x2 // a1 -= kc ST1 {v1.8b}, [x9], x12 SUB x13, x13, x2 // a2 -= kc ST1 {v1.d}[1], [x7], x12 SUB x4, x4, x2 // a3 -= kc B.NE 0b RET # Remainder- 4-12 bytes of A .p2align 3 30: TBZ x0, 3, 4f 3: LDR d0, [x3], 8 LDR q4, [x5] LDR d1, [x15], 8 LDR d2, [x13], 8 LDR d3, [x4], 8 LDR q5, [x5, 16] UDOT v24.4s, v7.16b, v0.16b // update zero point UDOT v25.4s, v7.16b, v1.16b UDOT v26.4s, v7.16b, v2.16b UDOT v27.4s, v7.16b, v3.16b UDOT v16.4s, v4.16b, v0.4b[0] UDOT v17.4s, v4.16b, v1.4b[0] LDR q6, [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] LDR q4, [x5, 48] UDOT v22.4s, v5.16b, v2.4b[0] UDOT v23.4s, v5.16b, v3.4b[0] UDOT v16.4s, v6.16b, v0.4b[1] UDOT v17.4s, v6.16b, v1.4b[1] UDOT v18.4s, v6.16b, v2.4b[1] UDOT v19.4s, v6.16b, v3.4b[1] ADD x5, x5, 64 UDOT v20.4s, v4.16b, v0.4b[1] UDOT v21.4s, v4.16b, v1.4b[1] UDOT v22.4s, v4.16b, v2.4b[1] UDOT v23.4s, v4.16b, v3.4b[1] TBZ x0, 2, 2b 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 v24.4s, v7.16b, v0.16b // update zero point UDOT v25.4s, v7.16b, v1.16b UDOT v26.4s, v7.16b, v2.16b UDOT v27.4s, v7.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] 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] B 2b # Store odd width .p2align 3 5: TBZ x1, 2, 6f STR s0, [x6], 4 ST1 {v0.s}[2], [x8], 4 STR s1, [x9], 4 ST1 {v1.s}[2], [x7], 4 EXT v0.16b, v0.16b, v0.16b, 4 EXT v1.16b, v1.16b, v1.16b, 4 6: TBZ x1, 1, 7f STR h0, [x6], 2 ST1 {v0.h}[4], [x8], 2 STR h1, [x9], 2 ST1 {v1.h}[4], [x7], 2 EXT v0.16b, v0.16b, v0.16b, 2 EXT v1.16b, v1.16b, v1.16b, 2 7: TBZ x1, 0, 8f STR b0, [x6] ST1 {v0.b}[8], [x8] STR b1, [x9] ST1 {v1.b}[8], [x7] 8: RET END_FUNCTION xnn_qu8_gemm_minmax_${REQUANTIZATION.lower()}_ukernel_4x8c4__aarch64_neondot_ld128 #ifdef __ELF__ .section ".note.GNU-stack","",%progbits #endif