/*++ Copyright (c) Microsoft Corporation. All rights reserved. Licensed under the MIT License. Module Name: SymQgemmS8KernelSdot.asm Abstract: This module implements the kernels for the quantized integer matrix/matrix multiply operation (QGEMM), where the right hand side is symmetrically quantized, i.e. zero point being zero. This kernel only requires prepacking of the right hand side, which is usually constant. When the packed right hand side is cached, we achieves higher performance by avoid packing all together. --*/ #include "kxarm64.h" #include "AssembleDotProduct.h" // // Stack frame layout for the S8S8 kernel. // #define GemmS8S8KernelFrame_SavedRegisters (4 * 8) #define GemmS8S8KernelFrame_ColumnSumBuffer (0 + GemmS8S8KernelFrame_SavedRegisters) TEXTAREA /*++ Routine Description: This routine is an inner kernel to compute matrix multiplication for a set of rows. Arguments: A (x0) - Supplies the address of matrix A. B (x1) - Supplies the address of matrix B. The matrix data has been packed using MlasGemmQuantCopyPackB. C (x2) - Supplies the address of matrix C. PackedCountK (x3) - Supplies the number of packed columns from matrix A and the number of packed rows from matrix B to iterate over. Packed K should be 16x CountM (x4) - Supplies the maximum number of rows that can be processed for matrix A and matrix C. The actual number of rows handled for this invocation depends on the kernel implementation. CountN (x5) - Supplies the number of columns from matrix B and matrix C to iterate over. ldc (x6) - Supplies the first dimension of matrix C. lda (x7) - Supplies the first dimension of matrix A. ColumnSumBuffer - Supplies the sum of each column from matrix B multiplied by the zero point offset of matrix A. These values are accumulated into every column of matrix C. Return Value: Returns the number of rows handled. --*/ NESTED_ENTRY MlasSymQgemmS8KernelSdot PROLOG_SAVE_REG_PAIR d8,d9,#-GemmS8S8KernelFrame_SavedRegisters! PROLOG_NOP ldr x8,[sp,#GemmS8S8KernelFrame_ColumnSumBuffer] PROLOG_SAVE_REG_PAIR d10,d11,#16 // compute C pointers: x2, x16, x17, x6 cmp x4,#2 // M < 2 ? add x16,x2,x6,lsl #2 // x16 -> C1 add x17,x2,x6,lsl #3 // x17 -> C2 csel x16,x2,x16,lo // if M < 2 x16/C1 -> C0 csel x17,x16,x17,ls // if M <= 2 x17/C2 -> C1 cmp x4,#4 // M < 4 ? mov x12,#4 // set max M to 4 add x6,x16,x6,lsl #3 // x6 -> C3 mov x9,x0 // save A0 mov x10,x3 // save K csel x6,x17,x6,lo // if M < 4 x6/C3 -> C2 csel x4,x12,x4,hi // if M > 4 M = 4; // Register Usage // B (x1) -> 4x16 // ---------------------------------------------------------------------------- // |v4.b[0]..v4.b[12] v5.b[0]..v5.b[12] v6.b[0]..v6.b[12] v7.b[0]..v7.b[12]| // | ... ... ... ... ... ... ... ... | // |v4.b[3]..v4.b[15] v5.b[3]..v5.b[15] v6.b[3]..v6.b[15] v7.b[3]..v7.b[15]| // A 4x4 ---------------------------------------------------------------------------- // ------------------ ---------------------------------------------------------------------------- // x0 |v0.b[0]..v0.b[3]| |v16.s[0]_v16.s[3] v20.s[0]_v20.s[3] v24.s[0]_v24.s[3] v28.s[0]_v28.s[3]| x2 // x12 |v1.b[0]..v1.b[3]| |v17.s[0]_v17.s[3] v21.s[0]_v21.s[3] v25.s[0]_v25.s[3] v29.s[0]_v29.s[3]| x16 // x13 |v2.b[0]..v2.b[3]| |v18.s[0]_v18.s[3] v22.s[0]_v22.s[3] v26.s[0]_v26.s[3] v30.s[0]_v30.s[3]| x17 // x14 |v3.b[0]..v3.b[3]| |v19.s[0]_v19.s[3] v23.s[0]_v23.s[3] v27.s[0]_v27.s[3] v31.s[0]_v31.s[3]| x6 // ------------------ ---------------------------------------------------------------------------- ProcessNextColumnLoop ldr q16,[x8],#16 // Init accumulators with column sums ldr q20,[x8],#16 ldr q24,[x8],#16 ldr q28,[x8],#16 mov x0,x9 // reload A0 cmp x4,#2 // M < 2 ? add x12,x9,x7 // x12 -> A1 add x13,x0,x7,lsl #1 // x13 -> A2 ldr q4,[x1],#16 // Load B csel x12,x0,x12,lo // if M < 2 A1 -> A0 csel x13,x12,x13,ls // if M <= 2 A2 -> A1 cmp x4,4 // M < 4 ? add x14,x12,x7,lsl #1 // x14 -> A3 ldr q5,[x1],#16 csel x14,x13,x14,lo // if M < 4 A3 -> A2 ldr d0,[x0],#8 // Load A0 1st/2nd block of 4 mov v17.16b,v16.16b mov v18.16b,v16.16b ldr d1,[x12],#8 // Load A1 mov v19.16b,v16.16b mov v21.16b,v20.16b ldr d2,[x13],#8 // Load A2 mov v22.16b,v20.16b mov v23.16b,v20.16b ldr d3,[x14],#8 // Load A3 mov v25.16b,v24.16b mov v26.16b,v24.16b ldr q6,[x1],#16 mov v27.16b,v24.16b mov v29.16b,v28.16b ldr q7,[x1],#16 subs x3,x10,#2 // one loop iteration and epilogue consume k = 32 mov v30.16b,v28.16b mov v31.16b,v28.16b b.lo BlockLoopEpilogue // Need 32 k for main loop BlockLoop sdot v16.4s,v4.16b,v0.4b[0] sdot v17.4s,v4.16b,v1.4b[0] ldr d8,[x0],#8 // Load A0 3rd/4th block of 4 sdot v18.4s,v4.16b,v2.4b[0] sdot v19.4s,v4.16b,v3.4b[0] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v0.4b[0] sdot v21.4s,v5.16b,v1.4b[0] ldr d9,[x12],#8 sdot v22.4s,v5.16b,v2.4b[0] sdot v23.4s,v5.16b,v3.4b[0] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v0.4b[0] sdot v25.4s,v6.16b,v1.4b[0] ldr d10,[x13],#8 sdot v26.4s,v6.16b,v2.4b[0] sdot v27.4s,v6.16b,v3.4b[0] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v0.4b[0] sdot v29.4s,v7.16b,v1.4b[0] ldr d11,[x14],#8 sdot v30.4s,v7.16b,v2.4b[0] sdot v31.4s,v7.16b,v3.4b[0] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v0.4b[1] sdot v17.4s,v4.16b,v1.4b[1] sdot v18.4s,v4.16b,v2.4b[1] sdot v19.4s,v4.16b,v3.4b[1] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v0.4b[1] sdot v21.4s,v5.16b,v1.4b[1] sdot v22.4s,v5.16b,v2.4b[1] sdot v23.4s,v5.16b,v3.4b[1] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v0.4b[1] sdot v25.4s,v6.16b,v1.4b[1] sdot v26.4s,v6.16b,v2.4b[1] sdot v27.4s,v6.16b,v3.4b[1] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v0.4b[1] sdot v29.4s,v7.16b,v1.4b[1] sdot v30.4s,v7.16b,v2.4b[1] sdot v31.4s,v7.16b,v3.4b[1] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v8.4b[0] sdot v17.4s,v4.16b,v9.4b[0] ldr d0,[x0],#8 sdot v18.4s,v4.16b,v10.4b[0] sdot v19.4s,v4.16b,v11.4b[0] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v8.4b[0] sdot v21.4s,v5.16b,v9.4b[0] ldr d1,[x12],#8 sdot v22.4s,v5.16b,v10.4b[0] sdot v23.4s,v5.16b,v11.4b[0] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v8.4b[0] sdot v25.4s,v6.16b,v9.4b[0] ldr d2,[x13],#8 sdot v26.4s,v6.16b,v10.4b[0] sdot v27.4s,v6.16b,v11.4b[0] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v8.4b[0] sdot v29.4s,v7.16b,v9.4b[0] ldr d3,[x14],#8 sdot v30.4s,v7.16b,v10.4b[0] sdot v31.4s,v7.16b,v11.4b[0] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v8.4b[1] sdot v17.4s,v4.16b,v9.4b[1] sdot v18.4s,v4.16b,v10.4b[1] sdot v19.4s,v4.16b,v11.4b[1] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v8.4b[1] sdot v21.4s,v5.16b,v9.4b[1] sdot v22.4s,v5.16b,v10.4b[1] sdot v23.4s,v5.16b,v11.4b[1] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v8.4b[1] sdot v25.4s,v6.16b,v9.4b[1] sdot v26.4s,v6.16b,v10.4b[1] sdot v27.4s,v6.16b,v11.4b[1] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v8.4b[1] sdot v29.4s,v7.16b,v9.4b[1] subs x3,x3,#1 // k -= 16 sdot v30.4s,v7.16b,v10.4b[1] sdot v31.4s,v7.16b,v11.4b[1] ldr q7,[x1],#16 b.hs BlockLoop BlockLoopEpilogue sdot v16.4s,v4.16b,v0.4b[0] sdot v17.4s,v4.16b,v1.4b[0] ldr d8,[x0],#8 sdot v18.4s,v4.16b,v2.4b[0] sdot v19.4s,v4.16b,v3.4b[0] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v0.4b[0] sdot v21.4s,v5.16b,v1.4b[0] ldr d9,[x12],#8 sdot v22.4s,v5.16b,v2.4b[0] sdot v23.4s,v5.16b,v3.4b[0] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v0.4b[0] sdot v25.4s,v6.16b,v1.4b[0] ldr d10,[x13],#8 sdot v26.4s,v6.16b,v2.4b[0] sdot v27.4s,v6.16b,v3.4b[0] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v0.4b[0] sdot v29.4s,v7.16b,v1.4b[0] ldr d11,[x14],#8 sdot v30.4s,v7.16b,v2.4b[0] sdot v31.4s,v7.16b,v3.4b[0] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v0.4b[1] sdot v17.4s,v4.16b,v1.4b[1] sdot v18.4s,v4.16b,v2.4b[1] sdot v19.4s,v4.16b,v3.4b[1] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v0.4b[1] sdot v21.4s,v5.16b,v1.4b[1] sdot v22.4s,v5.16b,v2.4b[1] sdot v23.4s,v5.16b,v3.4b[1] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v0.4b[1] sdot v25.4s,v6.16b,v1.4b[1] sdot v26.4s,v6.16b,v2.4b[1] sdot v27.4s,v6.16b,v3.4b[1] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v0.4b[1] sdot v29.4s,v7.16b,v1.4b[1] sdot v30.4s,v7.16b,v2.4b[1] sdot v31.4s,v7.16b,v3.4b[1] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v8.4b[0] sdot v17.4s,v4.16b,v9.4b[0] sdot v18.4s,v4.16b,v10.4b[0] sdot v19.4s,v4.16b,v11.4b[0] ldr q4,[x1],#16 sdot v20.4s,v5.16b,v8.4b[0] sdot v21.4s,v5.16b,v9.4b[0] sdot v22.4s,v5.16b,v10.4b[0] sdot v23.4s,v5.16b,v11.4b[0] ldr q5,[x1],#16 sdot v24.4s,v6.16b,v8.4b[0] sdot v25.4s,v6.16b,v9.4b[0] sdot v26.4s,v6.16b,v10.4b[0] sdot v27.4s,v6.16b,v11.4b[0] ldr q6,[x1],#16 sdot v28.4s,v7.16b,v8.4b[0] sdot v29.4s,v7.16b,v9.4b[0] sdot v30.4s,v7.16b,v10.4b[0] sdot v31.4s,v7.16b,v11.4b[0] ldr q7,[x1],#16 sdot v16.4s,v4.16b,v8.4b[1] sdot v17.4s,v4.16b,v9.4b[1] sdot v18.4s,v4.16b,v10.4b[1] sdot v19.4s,v4.16b,v11.4b[1] sdot v20.4s,v5.16b,v8.4b[1] sdot v21.4s,v5.16b,v9.4b[1] sdot v22.4s,v5.16b,v10.4b[1] sdot v23.4s,v5.16b,v11.4b[1] sdot v24.4s,v6.16b,v8.4b[1] sdot v25.4s,v6.16b,v9.4b[1] sdot v26.4s,v6.16b,v10.4b[1] sdot v27.4s,v6.16b,v11.4b[1] sdot v28.4s,v7.16b,v8.4b[1] sdot v29.4s,v7.16b,v9.4b[1] subs x5,x5,#16 // adjust CountN remaining sdot v30.4s,v7.16b,v10.4b[1] sdot v31.4s,v7.16b,v11.4b[1] blo StoreOutputPartial stp q16,q20,[x2],#32 stp q24,q28,[x2],#32 stp q17,q21,[x16],#32 stp q25,q29,[x16],#32 stp q18,q22,[x17],#32 stp q26,q30,[x17],#32 stp q19,q23,[x6],#32 stp q27,q31,[x6],#32 cbnz x5,ProcessNextColumnLoop ExitKernel mov x0,x4 // return number of rows handled EPILOG_RESTORE_REG_PAIR d10,d11,#16 EPILOG_RESTORE_REG_PAIR d8,d9,#GemmS8S8KernelFrame_SavedRegisters! EPILOG_RETURN // // Store the partial 1 to 15 columns either overwriting the output matrix or // accumulating into the existing contents of the output matrix. // StoreOutputPartial tbz x5,#3,StoreOutputPartial4 stp q16,q20,[x2],#32 mov v16.16b,v24.16b // shift remaining elements down mov v20.16b,v28.16b stp q17,q21,[x16],#32 mov v17.16b,v25.16b mov v21.16b,v29.16b stp q18,q22,[x17],#32 mov v18.16b,v26.16b mov v22.16b,v30.16b stp q19,q23,[x6],#32 mov v19.16b,v27.16b mov v23.16b,v31.16b StoreOutputPartial4 tbz x5,#2,StoreOutputPartial2 st1 {v16.4s},[x2],#16 mov v16.16b,v20.16b // shift remaining elements down st1 {v17.4s},[x16],#16 mov v17.16b,v21.16b st1 {v18.4s},[x17],#16 mov v18.16b,v22.16b st1 {v19.4s},[x6],#16 mov v19.16b,v23.16b StoreOutputPartial2 tbz x5,#1,StoreOutputPartial1 st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v17.2s},[x16],#8 dup v17.4s,v17.s[2] st1 {v18.2s},[x17],#8 dup v18.4s,v18.s[2] st1 {v19.2s},[x6],#8 dup v19.4s,v19.s[2] StoreOutputPartial1 tbz x5,#0,ExitKernel st1 {v16.s}[0],[x2] st1 {v17.s}[0],[x16] st1 {v18.s}[0],[x17] st1 {v19.s}[0],[x6] b ExitKernel NESTED_END MlasSymQgemmS8KernelSdot END