/*++

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<MLAS_SYMM_GEMM_S8S8_KERNEL_SDOT>.

    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
