/*++ Copyright (c) Microsoft Corporation. All rights reserved. Licensed under the MIT License. Module Name: QgemmU8X8KernelUdot.asm Abstract: This module implements the kernels for the quantized integer matrix/matrix multiply operation (QGEMM). This implementation uses ARM v8.4 dot product instructions. --*/ #include "kxarm64.h" #include "AssembleDotProduct.h" // // Stack frame layout for the U8X8 kernel. // Defining spaces for saving 2 vector registers, and pointers to parameters // on the stack // #define GemmU8X8KernelFrame_SavedNeonRegisters (2 * 8) #define GemmU8X8KernelFrame_SavedRegisters GemmU8X8KernelFrame_SavedNeonRegisters #define GemmU8X8KernelFrame_ColumnSumBuffer (0 + GemmU8X8KernelFrame_SavedRegisters) #define GemmU8X8KernelFrame_ZeroPointB (8 + GemmU8X8KernelFrame_SavedRegisters) #define GemmU8X8KernelFrame_ZeroMode (16 + GemmU8X8KernelFrame_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. The matrix data has been packed using MlasGemmQuantCopyPackA. 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. 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. RowSumBuffer (x7) - Supplies the sum of each row from matrix A. These values have been pre-scaled by the zero point offset of matrix B if the offset is per-tensor (ZeroPointB is nullptr). Otherwise, these values must be scaled by the per-column zero point offsets of matrix B. These values are accumulated into every row of matrix C. 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. ZeroPointB - Optionally supplies the per-column zero point offsets of matrix B, else nullptr if the matrix B is using per-tensor quantization. ZeroMode - Supplies true if the output matrix must be zero initialized, else false if the output matrix is accumulated into. Return Value: Returns the number of rows handled. --*/ NESTED_ENTRY MlasGemmU8X8KernelUdot PROLOG_SAVE_REG_PAIR d8,d9,#-16! ldr x8,[sp,#GemmU8X8KernelFrame_ColumnSumBuffer] ldr x9,[sp,#GemmU8X8KernelFrame_ZeroPointB] ldrb w13,[sp,#GemmU8X8KernelFrame_ZeroMode] mov x14,x0 ld1 {v8.4s},[x7],#16 // load row sum 0 ~ 4 mov x15,x3 cmp x4,#1 // CountM == 1? beq ProcessLoopM1 cmp x4,#4 // CountM < 4? blo ProcessLoopM2 cmp x4,#8 // CountM < 8? blo ProcessNextColumnLoopM4 ld1 {v9.4s},[x7] // load row sum 5 ~ 8 // // Process 8 rows of the matrices. // Row Sums: v8 ~ v9 // A 4x8 block // /-----------------------------------------| // |v0.b[0] ... v0.b[12] v1.b[0] ... v1.b[12]| // | ... ... | // |v0.b[3] ... v0.b[15] v1.b[3] ... v1.b[15]| // \-----------------------------------------/ // B 8x4 block // /---------------------\ /-----------------------------------------| // |v4.b[0] ... v4.b[3] | |v16.s[0] .. v16.s[3] v17.s[0] .. v17.s[3]| // |v4.b[4] ... v4.b[7] | |v18.s[0] .. v18.s[3] v19.s[0] .. v19.s[3]| // |v4.b[8] ... v4.b[11]| |v20.s[0] .. v20.s[3] v21.s[0] .. v21.s[3]| // |v4.b[12] ... v4.b[15]| |v22.s[0] .. v22.s[3] v23.s[0] .. v23.s[3]| // |v5.b[0] ... v5.b[3] | |v24.s[0] .. v24.s[3] v25.s[0] .. v25.s[3]| // |v5.b[4] ... v5.b[7] | |v26.s[0] .. v26.s[3] v27.s[0] .. v27.s[3]| // |v5.b[8] ... v5.b[11]| |v28.s[0] .. v28.s[3] v29.s[0] .. v29.s[3]| // |v5.b[12] ... v5.b[15]| |v30.s[0] .. v30.s[3] v31.s[0] .. v31.s[3]| // \---------------------/ \-----------------------------------------/ // // unroll for the next 4 in k dimension // /-----------------------------------------| // |v2.b[0] ... v2.b[12] v3.b[0] ... v3.b[12]| // | ... ... | // |v2.b[3] ... v2.b[15] v3.b[3] ... v3.b[15]| // \-----------------------------------------/ // /---------------------\ /-----------------------------------------\ // |v6.b[0] ... v6.b[3] | |v16.s[0] .. v16.s[3] v17.s[0] .. v17.s[3]| // |v6.b[4] ... v6.b[7] | |v18.s[0] .. v18.s[3] v19.s[0] .. v19.s[3]| // |v6.b[8] ... v6.b[11]| |v20.s[0] .. v20.s[3] v21.s[0] .. v21.s[3]| // |v6.b[12] ... v6.b[15]| |v22.s[0] .. v22.s[3] v23.s[0] .. v23.s[3]| // |v7.b[0] ... v7.b[3] | |v24.s[0] .. v24.s[3] v25.s[0] .. v25.s[3]| // |v7.b[4] ... v7.b[7] | |v26.s[0] .. v26.s[3] v27.s[0] .. v27.s[3]| // |v7.b[8] ... v7.b[11]| |v28.s[0] .. v28.s[3] v29.s[0] .. v29.s[3]| // |v7.b[12] ... v7.b[15]| |v30.s[0] .. v30.s[3] v31.s[0] .. v31.s[3]| // \---------------------/ \-----------------------------------------/ // Starting the loop: initialize accumulators with scaled combination // of row and column sums dup v17.4s,v8.s[0] // broadcast row sums dup v19.4s,v8.s[1] dup v21.4s,v8.s[2] dup v23.4s,v8.s[3] dup v25.4s,v9.s[0] dup v27.4s,v9.s[1] dup v29.4s,v9.s[2] dup v31.4s,v9.s[3] ProcessNextColumnLoopM8 mov x0,x14 // reload matrix A ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer[0] mov x3,x15 // reload PackedCountK ld1 {v7.4s},[x8],#16 // load ColumnSumBuffer[4] cbz x9,SkipScaleByZeroPointBM8 // accumulator = zero point B * row sum A + column sum B ld1 {v0.4s},[x9],#16 // load ZeroPointB[0] mul v16.4s,v0.4s,v17.4s mul v18.4s,v0.4s,v19.4s ld1 {v1.4s},[x9],#16 // load ZeroPointB[4] mul v20.4s,v0.4s,v21.4s mul v22.4s,v0.4s,v23.4s mul v24.4s,v0.4s,v25.4s mul v26.4s,v0.4s,v27.4s mul v28.4s,v0.4s,v29.4s mul v30.4s,v0.4s,v31.4s mul v17.4s,v1.4s,v17.4s mul v19.4s,v1.4s,v19.4s mul v21.4s,v1.4s,v21.4s mul v23.4s,v1.4s,v23.4s mul v25.4s,v1.4s,v25.4s mul v27.4s,v1.4s,v27.4s mul v29.4s,v1.4s,v29.4s mul v31.4s,v1.4s,v31.4s // preloading mixed with accumulator inits ld1 {v0.16b},[x1],#16 // load packed B0 add v16.4s,v3.4s,v16.4s add v18.4s,v3.4s,v18.4s ldr q4,[x0],#16 // load packed A0 add v20.4s,v3.4s,v20.4s add v22.4s,v3.4s,v22.4s ldr q5,[x0],#16 // load packed A1 add v24.4s,v3.4s,v24.4s add v26.4s,v3.4s,v26.4s ld1 {v1.16b},[x1],#16 // load packed B1 add v28.4s,v3.4s,v28.4s add v30.4s,v3.4s,v30.4s ldr q6,[x0],#16 // load packed A2 add v17.4s,v7.4s,v17.4s add v19.4s,v7.4s,v19.4s ld1 {v2.16b},[x1],#16 // load packed B0_next4k add v21.4s,v7.4s,v21.4s add v23.4s,v7.4s,v23.4s add v25.4s,v7.4s,v25.4s add v27.4s,v7.4s,v27.4s add v29.4s,v7.4s,v29.4s add v31.4s,v7.4s,v31.4s b ComputeBlockLoopM8 SkipScaleByZeroPointBM8 // accumulator = row sum A + column sum B ld1 {v0.16b},[x1],#16 // load packed B0 add v16.4s,v3.4s,v17.4s add v18.4s,v3.4s,v19.4s ldr q4,[x0],#16 // load packed A0 add v20.4s,v3.4s,v21.4s add v22.4s,v3.4s,v23.4s ldr q5,[x0],#16 // load packed A1 add v24.4s,v3.4s,v25.4s add v26.4s,v3.4s,v27.4s ld1 {v1.16b},[x1],#16 // load packed B1 add v28.4s,v3.4s,v29.4s add v30.4s,v3.4s,v31.4s ldr q6,[x0],#16 // load packed A2 add v17.4s,v7.4s,v17.4s add v19.4s,v7.4s,v19.4s ld1 {v2.16b},[x1],#16 // load packed B0_next4k add v21.4s,v7.4s,v21.4s add v23.4s,v7.4s,v23.4s add v25.4s,v7.4s,v25.4s add v27.4s,v7.4s,v27.4s add v29.4s,v7.4s,v29.4s add v31.4s,v7.4s,v31.4s ComputeBlockLoopM8 sub x3,x3,#1 ld1 {v3.16b},[x1],#16 // load packed B1_next4k UdotByElement 16, 0, 4, 0 UdotByElement 18, 0, 4, 1 ldr q7,[x0],#16 // load packed A3 UdotByElement 20, 0, 4, 2 UdotByElement 22, 0, 4, 3 cbz x3,ComputeBlockLoopFinishM8 UdotByElement 17, 1, 4, 0 UdotByElement 19, 1, 4, 1 UdotByElement 21, 1, 4, 2 UdotByElement 23, 1, 4, 3 ldr q4,[x0],#16 // load packed A0 for next iteration UdotByElement 24, 0, 5, 0 UdotByElement 26, 0, 5, 1 UdotByElement 28, 0, 5, 2 UdotByElement 30, 0, 5, 3 ld1 {v0.16b},[x1],#16 // load packed B0 for next iteration UdotByElement 25, 1, 5, 0 UdotByElement 27, 1, 5, 1 UdotByElement 29, 1, 5, 2 UdotByElement 31, 1, 5, 3 ld1 {v1.16b},[x1],#16 // load packed B1 for next iteration UdotByElement 16, 2, 6, 0 UdotByElement 18, 2, 6, 1 ldr q5,[x0],#16 // load packed A1 for next iteration UdotByElement 20, 2, 6, 2 UdotByElement 22, 2, 6, 3 UdotByElement 17, 3, 6, 0 UdotByElement 19, 3, 6, 1 UdotByElement 21, 3, 6, 2 UdotByElement 23, 3, 6, 3 ldr q6,[x0],#16 // load packed A2 for next iteration UdotByElement 24, 2, 7, 0 UdotByElement 26, 2, 7, 1 UdotByElement 28, 2, 7, 2 UdotByElement 30, 2, 7, 3 ld1 {v2.16b},[x1],#16 // load packed B0_next4k for next iteration UdotByElement 25, 3, 7, 0 UdotByElement 27, 3, 7, 1 UdotByElement 29, 3, 7, 2 UdotByElement 31, 3, 7, 3 b ComputeBlockLoopM8 ComputeBlockLoopFinishM8 // postfix, compute tail values and prepare to write results // We are either about to go to ProcessNextColumnLoopM8 // where x0 and x3 are about to be restored, or exit // when x0 and x3 will not be used. // x4 x7 has finished their task // so we can use x0 x3 x4 x7 as output row pointers UdotByElement 17, 1, 4, 0 UdotByElement 19, 1, 4, 1 add x10,x2,x6,lsl #2 // compute output row 2 add x11,x10,x6,lsl #2 // compute output row 3 UdotByElement 21, 1, 4, 2 UdotByElement 23, 1, 4, 3 add x12,x11,x6,lsl #2 // compute output row 4 add x0,x12,x6,lsl #2 // compute output row 5 UdotByElement 24, 0, 5, 0 UdotByElement 26, 0, 5, 1 add x3,x0,x6,lsl #2 // compute output row 6 add x4,x3,x6,lsl #2 // compute output row 7 UdotByElement 28, 0, 5, 2 UdotByElement 30, 0, 5, 3 add x7,x4,x6,lsl #2 // compute output row 8 subs x5,x5,#8 // adjust CountN remaining UdotByElement 25, 1, 5, 0 UdotByElement 27, 1, 5, 1 UdotByElement 29, 1, 5, 2 UdotByElement 31, 1, 5, 3 UdotByElement 16, 2, 6, 0 UdotByElement 18, 2, 6, 1 UdotByElement 20, 2, 6, 2 UdotByElement 22, 2, 6, 3 UdotByElement 17, 3, 6, 0 UdotByElement 19, 3, 6, 1 UdotByElement 21, 3, 6, 2 UdotByElement 23, 3, 6, 3 UdotByElement 24, 2, 7, 0 UdotByElement 26, 2, 7, 1 UdotByElement 28, 2, 7, 2 UdotByElement 30, 2, 7, 3 UdotByElement 25, 3, 7, 0 UdotByElement 27, 3, 7, 1 UdotByElement 29, 3, 7, 2 UdotByElement 31, 3, 7, 3 blo StoreOutputPartialM8 cbnz x13,SkipAccumulateOutputM8 ldp q0,q1,[x2] ldp q2,q3,[x10] add v16.4s,v16.4s,v0.4s add v17.4s,v17.4s,v1.4s ldp q4,q5,[x11] add v18.4s,v18.4s,v2.4s add v19.4s,v19.4s,v3.4s ldp q6,q7,[x12] add v20.4s,v20.4s,v4.4s add v21.4s,v21.4s,v5.4s ldp q0, q1, [x0] add v22.4s,v22.4s,v6.4s add v23.4s,v23.4s,v7.4s ldp q2, q3, [x3] add v24.4s,v24.4s,v0.4s add v25.4s,v25.4s,v1.4s ldp q4, q5, [x4] add v26.4s,v26.4s,v2.4s add v27.4s,v27.4s,v3.4s ldp q6, q7, [x7] add v28.4s,v28.4s,v4.4s add v29.4s,v29.4s,v5.4s add v30.4s,v30.4s,v6.4s add v31.4s,v31.4s,v7.4s SkipAccumulateOutputM8 stp q16,q17,[x2],#32 dup v17.4s,v8.s[0] // broadcast row sums stp q18,q19,[x10] dup v19.4s,v8.s[1] stp q20,q21,[x11] dup v21.4s,v8.s[2] stp q22,q23,[x12] dup v23.4s,v8.s[3] stp q24,q25,[x0] dup v25.4s,v9.s[0] stp q26,q27,[x3] dup v27.4s,v9.s[1] stp q28,q29,[x4] dup v29.4s,v9.s[2] stp q30,q31,[x7] dup v31.4s,v9.s[3] cbnz x5,ProcessNextColumnLoopM8 ExitKernelM8 mov x0,#8 // return number of rows handled EPILOG_RESTORE_REG_PAIR d8,d9,#16! EPILOG_RETURN // // Store the partial 1 to 7 columns either overwriting the output matrix or // accumulating into the existing contents of the output matrix. // StoreOutputPartialM8 cbz x13,StoreOutputPartialAddModeM8 StoreOutputPartialZeroModeM8 tbz x5,#2,StoreOutputPartial2ZeroModeM8 st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b st1 {v20.4s},[x11],#16 mov v20.16b,v21.16b st1 {v22.4s},[x12],#16 mov v22.16b,v23.16b st1 {v24.4s},[x0],#16 mov v24.16b,v25.16b st1 {v26.4s},[x3],#16 mov v26.16b,v27.16b st1 {v28.4s},[x4],#16 mov v28.16b,v29.16b st1 {v30.4s},[x7],#16 mov v30.16b,v31.16b StoreOutputPartial2ZeroModeM8 tbz x5,#1,StoreOutputPartial1ZeroModeM8 st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] st1 {v20.2s},[x11],#8 dup v20.4s,v20.s[2] st1 {v22.2s},[x12],#8 dup v22.4s,v22.s[2] st1 {v24.2s},[x0],#8 dup v24.4s,v24.s[2] st1 {v26.2s},[x3],#8 dup v26.4s,v26.s[2] st1 {v28.2s},[x4],#8 dup v28.4s,v28.s[2] st1 {v30.2s},[x7],#8 dup v30.4s,v30.s[2] StoreOutputPartial1ZeroModeM8 tbz x5,#0,ExitKernelM8 st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] st1 {v20.s}[0],[x11] st1 {v22.s}[0],[x12] st1 {v24.s}[0],[x0] st1 {v26.s}[0],[x3] st1 {v28.s}[0],[x4] st1 {v30.s}[0],[x7] b ExitKernelM8 StoreOutputPartialAddModeM8 tbz x5,#2,StoreOutputPartial2AddModeM8 ld1 {v0.4s},[x2] ld1 {v1.4s},[x10] ld1 {v2.4s},[x11] ld1 {v3.4s},[x12] ld1 {v4.4s},[x0] ld1 {v5.4s},[x3] ld1 {v6.4s},[x4] ld1 {v7.4s},[x7] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b add v20.4s,v20.4s,v2.4s add v22.4s,v22.4s,v3.4s st1 {v20.4s},[x11],#16 mov v20.16b,v21.16b st1 {v22.4s},[x12],#16 mov v22.16b,v23.16b add v24.4s,v24.4s,v4.4s add v26.4s,v26.4s,v5.4s st1 {v24.4s},[x0],#16 mov v24.16b,v25.16b st1 {v26.4s},[x3],#16 mov v26.16b,v27.16b add v28.4s,v28.4s,v6.4s add v30.4s,v30.4s,v7.4s st1 {v28.4s},[x4],#16 mov v28.16b,v29.16b st1 {v30.4s},[x7],#16 mov v30.16b,v31.16b StoreOutputPartial2AddModeM8 tbz x5,#1,StoreOutputPartial1AddModeM8 ld1 {v0.2s},[x2] ld1 {v1.2s},[x10] ld1 {v2.2s},[x11] ld1 {v3.2s},[x12] ld1 {v4.2s},[x0] ld1 {v5.2s},[x3] ld1 {v6.2s},[x4] ld1 {v7.2s},[x7] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] add v20.4s,v20.4s,v2.4s add v22.4s,v22.4s,v3.4s st1 {v20.2s},[x11],#8 dup v20.4s,v20.s[2] st1 {v22.2s},[x12],#8 dup v22.4s,v22.s[2] add v24.4s,v24.4s,v4.4s add v26.4s,v26.4s,v5.4s st1 {v24.2s},[x0],#8 dup v24.4s,v24.s[2] st1 {v26.2s},[x3],#8 dup v26.4s,v26.s[2] add v28.4s,v28.4s,v6.4s add v30.4s,v30.4s,v7.4s st1 {v28.2s},[x4],#8 dup v28.4s,v28.s[2] st1 {v30.2s},[x7],#8 dup v30.4s,v30.s[2] StoreOutputPartial1AddModeM8 tbz x5,#0,ExitKernelM8 ld1 {v0.s}[0],[x2] ld1 {v1.s}[0],[x10] add v16.4s,v16.4s,v0.4s ld1 {v2.s}[0],[x11] add v18.4s,v18.4s,v1.4s ld1 {v3.s}[0],[x12] add v20.4s,v20.4s,v2.4s st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] add v22.4s,v22.4s,v3.4s st1 {v20.s}[0],[x11] st1 {v22.s}[0],[x12] ld1 {v4.s}[0],[x0] ld1 {v5.s}[0],[x3] ld1 {v6.s}[0],[x4] ld1 {v7.s}[0],[x7] add v24.4s,v24.4s,v4.4s st1 {v24.s}[0],[x0] add v26.4s,v26.4s,v5.4s st1 {v26.s}[0],[x3] add v28.4s,v28.4s,v6.4s st1 {v28.s}[0],[x4] add v30.4s,v30.4s,v7.4s st1 {v30.s}[0],[x7] b ExitKernelM8 // // Process 4 rows of the matrices. // // // The packing layout is setup to have a pair of four quad vectors from // packed matrix A and a pair of eight quad vectors from packed matrix B. // With this scheme, alternating loads from the packed matrices can be // interleaved with the dot product instructions. // // One negative consequence of using four rows here is that the accumulator // register tile is too small for processors with high out of order execution // windows (such as the Apple M1). The dot product instructions for a given // cell are too close to each other to avoid dependencies. To workaround this, // the below loop uses a pair of accumulator registers that are then added // together when the loop finishes. // // A55-based cores are optimized for 64-bit loads, so use 64-bit loads for // packed matrix A. At the time of this implementation, using a wider 128-bit // load didn't affect performance for higher end cores. // // B 4x8 block // /-----------------------------------------| // |v0.b[0] ... v0.b[12] v1.b[0] ... v1.b[12]| // | ... ... | // |v0.b[3] ... v0.b[15] v1.b[3] ... v1.b[15]| // \-----------------------------------------/ // A 4x4 block // /---------------------\ /-----------------------------------------| // |d4.b[0] ... d4.b[3] | |v16.s[0] .. v16.s[3] v17.s[0] .. v17.s[3]| // |d4.b[4] ... d4.b[7] | |v18.s[0] .. v18.s[3] v19.s[0] .. v19.s[3]| // |d5.b[0] ... d5.b[3] | |v20.s[0] .. v20.s[3] v21.s[0] .. v21.s[3]| // |d5.b[4] ... d5.b[7] | |v22.s[0] .. v22.s[3] v23.s[0] .. v23.s[3]| // \---------------------/ \-----------------------------------------/ // unroll for the next 4 in k dimension // /-----------------------------------------| // |v0.b[0] ... v0.b[12] v1.b[0] ... v1.b[12]| // | ... ... | // |v0.b[3] ... v0.b[15] v1.b[3] ... v1.b[15]| // \-----------------------------------------/ // /---------------------\ /-----------------------------------------\ // |d6.b[0] ... d6.b[3] | |v24.s[0] .. v24.s[3] v25.s[0] .. v25.s[3]| // |d6.b[4] ... d6.b[7] | |v26.s[0] .. v26.s[3] v27.s[0] .. v27.s[3]| // |d7.b[0] ... d7.b[3] | |v28.s[0] .. v24.s[3] v29.s[0] .. v29.s[3]| // |d7.b[4] ... d7.b[7] | |v30.s[0] .. v24.s[3] v31.s[0] .. v31.s[3]| // \---------------------/ \-----------------------------------------/ ProcessNextColumnLoopM4 ld1 {v0.16b},[x1],#16 // load packed B0 mov x0,x14 // reload matrix A ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer[0] mov x3,x15 // reload PackedCountK ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer[4] dup v17.4s,v8.s[0] // broadcast row sums dup v19.4s,v8.s[1] dup v21.4s,v8.s[2] dup v23.4s,v8.s[3] cbz x9,SkipScaleByZeroPointBM4 ld1 {v30.4s},[x9],#16 // load ZeroPointB[0] mul v16.4s,v30.4s,v17.4s mul v18.4s,v30.4s,v19.4s ld1 {v31.4s},[x9],#16 // load ZeroPointB[4] mul v20.4s,v30.4s,v21.4s mul v22.4s,v30.4s,v23.4s mul v17.4s,v31.4s,v17.4s mul v19.4s,v31.4s,v19.4s mul v21.4s,v31.4s,v21.4s mul v23.4s,v31.4s,v23.4s add v16.4s,v2.4s,v16.4s add v18.4s,v2.4s,v18.4s add v20.4s,v2.4s,v20.4s add v22.4s,v2.4s,v22.4s add v17.4s,v3.4s,v17.4s add v19.4s,v3.4s,v19.4s add v21.4s,v3.4s,v21.4s add v23.4s,v3.4s,v23.4s b ComputeBlockLoopStartM4 SkipScaleByZeroPointBM4 add v16.4s,v2.4s,v17.4s add v18.4s,v2.4s,v19.4s add v20.4s,v2.4s,v21.4s add v22.4s,v2.4s,v23.4s add v17.4s,v3.4s,v17.4s add v19.4s,v3.4s,v19.4s add v21.4s,v3.4s,v21.4s add v23.4s,v3.4s,v23.4s ComputeBlockLoopStartM4 ldr d4,[x0],#32 // load packed A0.l movi v24.4s,#0 movi v25.4s,#0 ldur d5,[x0,#-24] // load packed A0.h movi v26.4s,#0 movi v27.4s,#0 ldur d6,[x0,#-16] // load packed A1.l movi v28.4s,#0 movi v29.4s,#0 movi v30.4s,#0 movi v31.4s,#0 ComputeBlockLoopM4 ld1 {v1.16b},[x1],#16 // load packed B1 UdotByElement 16, 0, 4, 0 UdotByElement 18, 0, 4, 1 ldur d7,[x0,#-8] // load packed A1.h UdotByElement 20, 0, 5, 0 UdotByElement 22, 0, 5, 1 ld1 {v0.16b},[x1],#16 // load packed B0_next4k UdotByElement 17, 1, 4, 0 UdotByElement 19, 1, 4, 1 sub x3,x3,#1 cbz x3,ComputeBlockLoopFinishM4 ldr d4,[x0],#32 // load packed A0.l for next iteration UdotByElement 21, 1, 5, 0 UdotByElement 23, 1, 5, 1 ld1 {v1.16b},[x1],#16 // load packed B1_next4k UdotByElement 24, 0, 6, 0 UdotByElement 26, 0, 6, 1 ldur d5,[x0,#-24] // load packed A0.h for next iteration UdotByElement 28, 0, 7, 0 UdotByElement 30, 0, 7, 1 ld1 {v0.16b},[x1],#16 // load packed B0 for next iteration UdotByElement 25, 1, 6, 0 UdotByElement 27, 1, 6, 1 ldur d6,[x0,#-16] // load packed A1.l for next iteration UdotByElement 29, 1, 7, 0 UdotByElement 31, 1, 7, 1 b ComputeBlockLoopM4 ComputeBlockLoopFinishM4 UdotByElement 21, 1, 5, 0 UdotByElement 23, 1, 5, 1 ld1 {v1.16b},[x1],#16 // load packed B1_next4k UdotByElement 24, 0, 6, 0 UdotByElement 26, 0, 6, 1 UdotByElement 28, 0, 7, 0 UdotByElement 30, 0, 7, 1 UdotByElement 25, 1, 6, 0 UdotByElement 27, 1, 6, 1 UdotByElement 29, 1, 7, 0 UdotByElement 31, 1, 7, 1 add x10,x2,x6,lsl #2 // compute output row 2 add v16.4s,v16.4s,v24.4s // fold high results into low results add v18.4s,v18.4s,v26.4s add v20.4s,v20.4s,v28.4s add v22.4s,v22.4s,v30.4s add x11,x10,x6,lsl #2 // compute output row 3 add v17.4s,v17.4s,v25.4s add v19.4s,v19.4s,v27.4s add v21.4s,v21.4s,v29.4s add v23.4s,v23.4s,v31.4s add x12,x11,x6,lsl #2 // compute output row 4 subs x5,x5,#8 // adjust CountN remaining blo StoreOutputPartialM4 cbnz x13,SkipAccumulateOutputM4 ldp q0,q1,[x2] ldp q2,q3,[x10] add v16.4s,v16.4s,v0.4s add v17.4s,v17.4s,v1.4s ldp q4,q5,[x11] add v18.4s,v18.4s,v2.4s add v19.4s,v19.4s,v3.4s ldp q6,q7,[x12] add v20.4s,v20.4s,v4.4s add v21.4s,v21.4s,v5.4s add v22.4s,v22.4s,v6.4s add v23.4s,v23.4s,v7.4s SkipAccumulateOutputM4 stp q16,q17,[x2],#32 stp q18,q19,[x10] stp q20,q21,[x11] stp q22,q23,[x12] cbnz x5,ProcessNextColumnLoopM4 ExitKernelM4 mov x0,#4 // return number of rows handled EPILOG_RESTORE_REG_PAIR d8,d9,#16! EPILOG_RETURN // // Store the partial 1 to 7 columns either overwriting the output matrix or // accumulating into the existing contents of the output matrix. // StoreOutputPartialM4 cbz x13,StoreOutputPartialAddModeM4 StoreOutputPartialZeroModeM4 tbz x5,#2,StoreOutputPartial2ZeroModeM4 st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b st1 {v20.4s},[x11],#16 mov v20.16b,v21.16b st1 {v22.4s},[x12],#16 mov v22.16b,v23.16b StoreOutputPartial2ZeroModeM4 tbz x5,#1,StoreOutputPartial1ZeroModeM4 st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] st1 {v20.2s},[x11],#8 dup v20.4s,v20.s[2] st1 {v22.2s},[x12],#8 dup v22.4s,v22.s[2] StoreOutputPartial1ZeroModeM4 tbz x5,#0,ExitKernelM4 st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] st1 {v20.s}[0],[x11] st1 {v22.s}[0],[x12] b ExitKernelM4 StoreOutputPartialAddModeM4 tbz x5,#2,StoreOutputPartial2AddModeM4 ld1 {v0.4s},[x2] ld1 {v1.4s},[x10] ld1 {v2.4s},[x11] ld1 {v3.4s},[x12] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b add v20.4s,v20.4s,v2.4s add v22.4s,v22.4s,v3.4s st1 {v20.4s},[x11],#16 mov v20.16b,v21.16b st1 {v22.4s},[x12],#16 mov v22.16b,v23.16b StoreOutputPartial2AddModeM4 tbz x5,#1,StoreOutputPartial1AddModeM4 ld1 {v0.2s},[x2] ld1 {v1.2s},[x10] ld1 {v2.2s},[x11] ld1 {v3.2s},[x12] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] add v20.4s,v20.4s,v2.4s add v22.4s,v22.4s,v3.4s st1 {v20.2s},[x11],#8 dup v20.4s,v20.s[2] st1 {v22.2s},[x12],#8 dup v22.4s,v22.s[2] StoreOutputPartial1AddModeM4 tbz x5,#0,ExitKernelM4 ld1 {v0.s}[0],[x2] ld1 {v1.s}[0],[x10] add v16.4s,v16.4s,v0.4s ld1 {v2.s}[0],[x11] add v18.4s,v18.4s,v1.4s ld1 {v3.s}[0],[x12] add v20.4s,v20.4s,v2.4s st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] add v22.4s,v22.4s,v3.4s st1 {v20.s}[0],[x11] st1 {v22.s}[0],[x12] b ExitKernelM4 // // Process 2 rows of the matrices. // ProcessLoopM2 dup v9.4s, v8.s[1] dup v8.4s, v8.s[0] ProcessNextColumnLoopM2 ld1 {v0.16b},[x1],#16 // load packed B0 ld1 {v1.16b},[x1],#16 // load packed B1 mov x0,x14 // reload matrix A ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer[0] mov x3,x15 // reload PackedCountK ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer[4] cbz x9,SkipScaleByZeroPointBM2 ld1 {v30.4s},[x9],#16 // load ZeroPointB[0] ld1 {v31.4s},[x9],#16 // load ZeroPointB[4] mul v16.4s,v30.4s,v8.4s mul v18.4s,v30.4s,v9.4s mul v17.4s,v31.4s,v8.4s mul v19.4s,v31.4s,v9.4s ldr d4,[x0],#8 // load packed A0.l add v16.4s,v2.4s,v16.4s add v18.4s,v2.4s,v18.4s ldr d5,[x0],#8 // load packed A0.h add v17.4s,v3.4s,v17.4s add v19.4s,v3.4s,v19.4s b ComputeBlockLoopM2 SkipScaleByZeroPointBM2 ldr d4,[x0],#8 // load packed A0.l add v16.4s,v2.4s,v8.4s add v18.4s,v2.4s,v9.4s ldr d5,[x0],#8 // load packed A0.h add v17.4s,v3.4s,v8.4s add v19.4s,v3.4s,v9.4s ComputeBlockLoopM2 sub x3,x3,#1 ld1 {v6.16b},[x1],#16 // load packed B0 next 4 k ld1 {v7.16b},[x1],#16 // load packed B1 next 4 k UdotByElement 16, 0, 4, 0 UdotByElement 17, 1, 4, 0 UdotByElement 18, 0, 4, 1 UdotByElement 19, 1, 4, 1 cbz x3,ComputeBlockLoopFinishM2 ldr d4,[x0],#8 // load packed A0.l for next iter ld1 {v0.16b},[x1],#16 // load packed B0 for next iter ld1 {v1.16b},[x1],#16 // load packed B1 for next iter UdotByElement 16, 6, 5, 0 UdotByElement 17, 7, 5, 0 UdotByElement 18, 6, 5, 1 UdotByElement 19, 7, 5, 1 ldr d5,[x0],#8 // load packed A0.h for next iter b ComputeBlockLoopM2 ComputeBlockLoopFinishM2 add x10,x2,x6,lsl #2 // compute output row 2 subs x5,x5,#8 // adjust CountN remaining UdotByElement 16, 6, 5, 0 UdotByElement 17, 7, 5, 0 UdotByElement 18, 6, 5, 1 UdotByElement 19, 7, 5, 1 blo StoreOutputPartialM2 cbnz x13,SkipAccumulateOutputM2 ldp q0,q1,[x2] ldp q2,q3,[x10] add v16.4s,v16.4s,v0.4s add v17.4s,v17.4s,v1.4s add v18.4s,v18.4s,v2.4s add v19.4s,v19.4s,v3.4s SkipAccumulateOutputM2 stp q16,q17,[x2],#32 stp q18,q19,[x10] cbnz x5,ProcessNextColumnLoopM2 ExitKernelM2 mov x0,#2 // return number of rows handled EPILOG_RESTORE_REG_PAIR d8,d9,#16! EPILOG_RETURN // // Store the partial 1 to 7 columns either overwriting the output matrix or // accumulating into the existing contents of the output matrix. // StoreOutputPartialM2 cbz x13,StoreOutputPartialAddModeM2 StoreOutputPartialZeroModeM2 tbz x5,#2,StoreOutputPartial2ZeroModeM2 st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b StoreOutputPartial2ZeroModeM2 tbz x5,#1,StoreOutputPartial1ZeroModeM2 st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] StoreOutputPartial1ZeroModeM2 tbz x5,#0,ExitKernelM2 st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] b ExitKernelM2 StoreOutputPartialAddModeM2 tbz x5,#2,StoreOutputPartial2AddModeM2 ld1 {v0.4s},[x2] ld1 {v1.4s},[x10] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down st1 {v18.4s},[x10],#16 mov v18.16b,v19.16b StoreOutputPartial2AddModeM2 tbz x5,#1,StoreOutputPartial1AddModeM2 ld1 {v0.2s},[x2] ld1 {v1.2s},[x10] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down st1 {v18.2s},[x10],#8 dup v18.4s,v18.s[2] StoreOutputPartial1AddModeM2 tbz x5,#0,ExitKernelM2 ld1 {v0.s}[0],[x2] ld1 {v1.s}[0],[x10] add v16.4s,v16.4s,v0.4s add v18.4s,v18.4s,v1.4s st1 {v16.s}[0],[x2] st1 {v18.s}[0],[x10] b ExitKernelM2 // // Process 1 row of the matrices. // ProcessLoopM1 dup v8.4s,v8.s[0] ProcessNextColumnLoopM1 ld1 {v0.16b},[x1],#16 // load packed B0 ld1 {v1.16b},[x1],#16 // load packed B1 mov x0,x14 // reload matrix A ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer0 mov x3,x15 // reload PackedCountK ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer1 cbz x9,SkipScaleByZeroPointBM1 ld1 {v30.4s},[x9],#16 // load ZeroPointB0 ld1 {v31.4s},[x9],#16 // load ZeroPointB1 mul v16.4s,v30.4s,v8.4s mul v17.4s,v31.4s,v8.4s ldr d4,[x0],#8 // load packed A0 ld1 {v6.16b},[x1],#16 // load packed B0 next 4 k ld1 {v7.16b},[x1],#16 // load packed B1 next 4 k add v16.4s,v2.4s,v16.4s add v17.4s,v3.4s,v17.4s b ComputeBlockLoopM1 SkipScaleByZeroPointBM1 ldr d4,[x0],#8 // load packed A0 ld1 {v6.16b},[x1],#16 // load packed B0 next 4 k ld1 {v7.16b},[x1],#16 // load packed B1 next 4 k add v16.4s,v2.4s,v8.4s add v17.4s,v3.4s,v8.4s ComputeBlockLoopM1 sub x3,x3,#1 UdotByElement 16, 0, 4, 0 UdotByElement 17, 1, 4, 0 cbz x3,ComputeBlockLoopFinishM1 ld1 {v0.16b},[x1],#16 // load packed B0 for next iter ld1 {v1.16b},[x1],#16 // load packed B1 for next iter UdotByElement 16, 6, 4, 1 UdotByElement 17, 7, 4, 1 ldr d4,[x0],#8 // load packed A0 for next iter ld1 {v6.16b},[x1],#16 // load packed B0 next 4 k for next iter ld1 {v7.16b},[x1],#16 // load packed B1 next 4 k for next iter b ComputeBlockLoopM1 ComputeBlockLoopFinishM1 subs x5,x5,#8 // adjust CountN remaining UdotByElement 16, 6, 4, 1 UdotByElement 17, 7, 4, 1 blo StoreOutputPartialM1 cbnz x13,SkipAccumulateOutputM1 ldp q0,q1,[x2] add v16.4s,v16.4s,v0.4s add v17.4s,v17.4s,v1.4s SkipAccumulateOutputM1 stp q16,q17,[x2],#32 cbnz x5,ProcessNextColumnLoopM1 ExitKernelM1 mov x0,#1 // return number of rows handled EPILOG_RESTORE_REG_PAIR d8,d9,#16! EPILOG_RETURN // // Store the partial 1 to 7 columns either overwriting the output matrix or // accumulating into the existing contents of the output matrix. // StoreOutputPartialM1 cbz x13,StoreOutputPartialAddModeM1 StoreOutputPartialZeroModeM1 tbz x5,#2,StoreOutputPartial2ZeroModeM1 st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down StoreOutputPartial2ZeroModeM1 tbz x5,#1,StoreOutputPartial1ZeroModeM1 st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down StoreOutputPartial1ZeroModeM1 tbz x5,#0,ExitKernelM1 st1 {v16.s}[0],[x2] b ExitKernelM1 StoreOutputPartialAddModeM1 tbz x5,#2,StoreOutputPartial2AddModeM1 ld1 {v0.4s},[x2] add v16.4s,v16.4s,v0.4s st1 {v16.4s},[x2],#16 mov v16.16b,v17.16b // shift remaining elements down StoreOutputPartial2AddModeM1 tbz x5,#1,StoreOutputPartial1AddModeM1 ld1 {v0.2s},[x2] add v16.4s,v16.4s,v0.4s st1 {v16.2s},[x2],#8 dup v16.4s,v16.s[2] // shift remaining elements down StoreOutputPartial1AddModeM1 tbz x5,#0,ExitKernelM1 ld1 {v0.s}[0],[x2] add v16.4s,v16.4s,v0.4s st1 {v16.s}[0],[x2] b ExitKernelM1 NESTED_END MlasGemmU8X8KernelUdot END