/*++

Copyright (c) Microsoft Corporation. All rights reserved.

Licensed under the MIT License.

Module Name:

    ConvSymS8KernelDot.asm

Abstract:

    This module implements the kernels for the symmetric quantized integer
    convolution operation.

--*/

#include "kxarm64.h"
#include "AssembleDotProduct.h"

#define     MLAS_CONV_SYM_FLAG_PER_CHANNEL_SCALE 2

//
// Stack frame layout for the symmetric convolution kernel.
// d8-d15, x19-x30 need to be preserved if used
//
#define     ConvSymFrame_SavedRegisters       (6 * 8)
#define     ConvSymFrame_PostProcessParams    0 + ConvSymFrame_SavedRegisters
#define     ConvSymFrame_KernelFlags          8 + ConvSymFrame_SavedRegisters

#define     ConvSymPostProcessParams_Bias      0
#define     ConvSymPostProcessParams_Scale     8
#define     ConvSymPostProcessParams_Min       16
#define     ConvSymPostProcessParams_Max       20
#define     ConvSymPostProcessParams_ZeroPoint 24

        TEXTAREA

/*++

Routine Description:

    This routine is the inner kernel to compute a convolution for the elements
    of an output row for a set of filter rows.

Arguments:

    Input (x0) - Points to the indirection buffer. Every pointer in the indirection
        buffer points at a InputChannels length vector (either from the input tensor
        or a vector of padding values). These are grouped in batches of length
        KernelSize.  These batches are then repeated OutputCount times.

    Filter (x1) - Points to the filter buffer.

    Output (x2) - Points the output buffer.

    KernelSize (x3/x9) - Size of the kernel (most commonly. 3x3=9, 5x5=25).
                         Must be > 1

    InputChannels (x4/x7) - Number of input channels.

    OutputChannels (x5) - Number of output channels.

    ChannelCount (x6) - Number of output channels this iteration produces.

    OutputCount (x7) - Number of output elements this iteration produces.

        This implementation requires the count to be no larger than 4.

    PostProcessParams (x8) - Points to the post process parameter block.

    KernelFlags - (w10) Additional flags controlling the operation.

Return Value:

    None.

--*/
        NESTED_ENTRY MlasConvSymS8KernelDot

        PROLOG_SAVE_REG_PAIR  d8,d9,#-ConvSymFrame_SavedRegisters!
        PROLOG_NOP    ldr     x8,[sp,#ConvSymFrame_PostProcessParams]
        PROLOG_SAVE_REG       d10,#16
        PROLOG_NOP    cmp     x7,2                    // OutputCount < 2 ?
        PROLOG_SAVE_REG       d11,#24
        PROLOG_NOP    add     x16,x2,x5               // x16 -> C1
        PROLOG_SAVE_REG       x19,#32
        lsl     x3,x3,#3                // KernelSize * sizeof(int8_t*)
        csel    x16,x2,x16,lo           // if OutputCount < 2  x16/C1 -> C0
        add     x4,x4,3                 // InputChannels align to 4
        add     x17,x16,x5              // x17 -> C2
        ldr     x11,[x8,#ConvSymPostProcessParams_Bias]
        csel    x17,x16,x17,ls          // if OutputCount <= 2  x17/C2 -> C1
        bic     x4,x4,3
        cmp     x7,4                    // OutputCount < 4 ?
        ldr     w10,[sp,#ConvSymFrame_KernelFlags]
        add     x5,x17,x5               // x5 -> C3
        ldr     x19,[x8,#ConvSymPostProcessParams_Scale]
        csel    x5,x17,x5,lo            // if OutputCount < 4  x5/C3 -> C2

        // TODO!! tiptoe around loading biases if we need to support
        // output channels none divisible by 16
OutputChannelLoop
        ldp     q16,q20,[x11],32        // Init accumulators with biases
        mov     v17.16b,v16.16b
        mov     v18.16b,v16.16b
        ldp     q24,q28,[x11],32
        mov     v19.16b,v16.16b
        mov     v21.16b,v20.16b
        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
        mov     x9,x3                   // restore KernelSize * sizeof(int8_t*)

KernelSizeLoop
        ldr     x12,[x0]                // x12 -> A0
        cmp     x16,x2
        b.eq    SkipLoadA1              // C1==C0 -> A0=A1=A2=A3
        cmp     x17,x16
        lsl     x14,x3,#1
        ldr     x13,[x0,x3]             // x13 -> A1
        b.eq    SkipLoadA2              // C2==C1 -> A1=A2=A3
        cmp     x5,x17
        add     x15,x3,x3,lsl#1
        ldr     x14,[x0,x14]            // x14 -> A2
        b.eq    SkipLoadA3              // C3==C2 -> A2=A3
        ldr     x15,[x0,x15]            // x15 -> A3
        b       FinishLoadAPtr
SkipLoadA1
        mov     x13,x12
SkipLoadA2
        mov     x14,x13
SkipLoadA3
        mov     x15,x14

// 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       ----------------------------------------------------------------------------
//     ------------------ ----------------------------------------------------------------------------
// x12 |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
// x13 |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
// x14 |v2.b[0]..v2.b[3]| |v18.s[0]_v18.s[3] v22.s[0]_v23.s[3]  v26.s[0]_v26.s[3]   v30.s[0]_v31.s[3]| x17
// x15 |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]|  x5
//     ------------------ ----------------------------------------------------------------------------

FinishLoadAPtr
        subs    x7,x4,16                // Need 16 input channels for loop
        add     x0,x0,8                 // indirect A advance to next pointer, prepare for kernel size loop
        b.lo    InChannels8

        ldr     d0,[x12],8
        ldr     q4,[x1],16
        ldr     d1,[x13],8
        subs    x7,x7,16
        ldr     d2,[x14],8
        ldr     d3,[x15],8
        ldr     q5,[x1],16
        ldr     q6,[x1],16
        ldr     q7,[x1],16
        b.lo    InChLoopEpilogue        // Need 32 input channels for main loop

InputChannelLoop
        SdotByElement    16, 4, 0,0
        SdotByElement    17, 4, 1,0
        ldr     d8,[x12],8
        SdotByElement    18, 4, 2,0
        SdotByElement    19, 4, 3,0
        ldr     q4,[x1],16
        SdotByElement    20, 5, 0,0
        SdotByElement    21, 5, 1,0
        ldr     d9,[x13],8
        SdotByElement    22, 5, 2,0
        SdotByElement    23, 5, 3,0
        ldr     q5,[x1],16
        SdotByElement    24, 6, 0,0
        SdotByElement    25, 6, 1,0
        ldr     d10,[x14],8
        SdotByElement    26, 6, 2,0
        SdotByElement    27, 6, 3,0
        ldr     q6,[x1],16
        SdotByElement    28, 7, 0,0
        SdotByElement    29, 7, 1,0
        ldr     d11,[x15],8
        SdotByElement    30, 7, 2,0
        SdotByElement    31, 7, 3,0
        ldr     q7,[x1],16
        SdotByElement    16, 4, 0,1
        SdotByElement    17, 4, 1,1
        SdotByElement    18, 4, 2,1
        SdotByElement    19, 4, 3,1
        ldr     q4,[x1],16
        SdotByElement    20, 5, 0,1
        SdotByElement    21, 5, 1,1
        SdotByElement    22, 5, 2,1
        SdotByElement    23, 5, 3,1
        ldr     q5,[x1],16
        SdotByElement    24, 6, 0,1
        SdotByElement    25, 6, 1,1
        SdotByElement    26, 6, 2,1
        SdotByElement    27, 6, 3,1
        ldr     q6,[x1],16
        SdotByElement    28, 7, 0,1
        SdotByElement    29, 7, 1,1
        SdotByElement    30, 7, 2,1
        SdotByElement    31, 7, 3,1
        ldr     q7,[x1],16
        SdotByElement    16, 4, 8,0
        SdotByElement    17, 4, 9,0
        ldr     d0,[x12],8
        SdotByElement    18, 4,10,0
        SdotByElement    19, 4,11,0
        ldr     q4,[x1],16
        SdotByElement    20, 5, 8,0
        SdotByElement    21, 5, 9,0
        ldr     d1,[x13],8
        SdotByElement    22, 5,10,0
        SdotByElement    23, 5,11,0
        ldr     q5,[x1],16
        SdotByElement    24, 6, 8,0
        SdotByElement    25, 6, 9,0
        ldr     d2,[x14],8
        SdotByElement    26, 6,10,0
        SdotByElement    27, 6,11,0
        ldr     q6,[x1],16
        SdotByElement    28, 7, 8,0
        SdotByElement    29, 7, 9,0
        ldr     d3,[x15],8
        SdotByElement    30, 7,10,0
        SdotByElement    31, 7,11,0
        ldr     q7,[x1],16
        SdotByElement    16, 4, 8,1
        SdotByElement    17, 4, 9,1
        SdotByElement    18, 4,10,1
        SdotByElement    19, 4,11,1
        ldr     q4,[x1],16
        SdotByElement    20, 5, 8,1
        SdotByElement    21, 5, 9,1
        SdotByElement    22, 5,10,1
        SdotByElement    23, 5,11,1
        ldr     q5,[x1],16
        SdotByElement    24, 6, 8,1
        SdotByElement    25, 6, 9,1
        SdotByElement    26, 6,10,1
        SdotByElement    27, 6,11,1
        ldr     q6,[x1],16
        SdotByElement    28, 7, 8,1
        SdotByElement    29, 7, 9,1
        subs    x7,x7,16                // InputChannels -= 16
        SdotByElement    30, 7,10,1
        SdotByElement    31, 7,11,1
        ldr     q7,[x1],16
        b.hs    InputChannelLoop

InChLoopEpilogue
        SdotByElement    16, 4, 0,0
        SdotByElement    17, 4, 1,0
        ldr     d8,[x12],8
        SdotByElement    18, 4, 2,0
        SdotByElement    19, 4, 3,0
        ldr     q4,[x1],16
        SdotByElement    20, 5, 0,0
        SdotByElement    21, 5, 1,0
        ldr     d9,[x13],8
        SdotByElement    22, 5, 2,0
        SdotByElement    23, 5, 3,0
        ldr     q5,[x1],16
        SdotByElement    24, 6, 0,0
        SdotByElement    25, 6, 1,0
        ldr     d10,[x14],8
        SdotByElement    26, 6, 2,0
        SdotByElement    27, 6, 3,0
        ldr     q6,[x1],16
        SdotByElement    28, 7, 0,0
        SdotByElement    29, 7, 1,0
        ldr     d11,[x15],8
        SdotByElement    30, 7, 2,0
        SdotByElement    31, 7, 3,0
        ldr     q7,[x1],16
        SdotByElement    16, 4, 0,1
        SdotByElement    17, 4, 1,1
        SdotByElement    18, 4, 2,1
        SdotByElement    19, 4, 3,1
        ldr     q4,[x1],16
        SdotByElement    20, 5, 0,1
        SdotByElement    21, 5, 1,1
        SdotByElement    22, 5, 2,1
        SdotByElement    23, 5, 3,1
        ldr     q5,[x1],16
        SdotByElement    24, 6, 0,1
        SdotByElement    25, 6, 1,1
        SdotByElement    26, 6, 2,1
        SdotByElement    27, 6, 3,1
        ldr     q6,[x1],16
        SdotByElement    28, 7, 0,1
        SdotByElement    29, 7, 1,1
        SdotByElement    30, 7, 2,1
        SdotByElement    31, 7, 3,1
        ldr     q7,[x1],16
        SdotByElement    16, 4, 8,0
        SdotByElement    17, 4, 9,0
        SdotByElement    18, 4,10,0
        SdotByElement    19, 4,11,0
        ldr     q4,[x1],16
        SdotByElement    20, 5, 8,0
        SdotByElement    21, 5, 9,0
        SdotByElement    22, 5,10,0
        SdotByElement    23, 5,11,0
        ldr     q5,[x1],16
        SdotByElement    24, 6, 8,0
        SdotByElement    25, 6, 9,0
        SdotByElement    26, 6,10,0
        SdotByElement    27, 6,11,0
        ldr     q6,[x1],16
        SdotByElement    28, 7, 8,0
        SdotByElement    29, 7, 9,0
        SdotByElement    30, 7,10,0
        SdotByElement    31, 7,11,0
        ldr     q7,[x1],16
        SdotByElement    16, 4, 8,1
        SdotByElement    17, 4, 9,1
        SdotByElement    18, 4,10,1
        SdotByElement    19, 4,11,1
        SdotByElement    20, 5, 8,1
        SdotByElement    21, 5, 9,1
        SdotByElement    22, 5,10,1
        SdotByElement    23, 5,11,1
        SdotByElement    24, 6, 8,1
        SdotByElement    25, 6, 9,1
        SdotByElement    26, 6,10,1
        SdotByElement    27, 6,11,1
        SdotByElement    28, 7, 8,1
        SdotByElement    29, 7, 9,1
        tst     x7,15
        SdotByElement    30, 7,10,1
        SdotByElement    31, 7,11,1
        b.ne    InChannels8             // 4 ~ 12 InputChannels
        subs    x9,x9,8                 // KernelSize-=1
        b.hi    KernelSizeLoop

Requantize
        tst     w10,#MLAS_CONV_SYM_FLAG_PER_CHANNEL_SCALE
        ldr     w13,[x8,#ConvSymPostProcessParams_ZeroPoint]
        beq     BroadcastScaleValue
        ldp     q0,q1,[x19],32          // load scale vector
        ldp     q2,q3,[x19],32
        b       AccumulatorsToFloat

BroadcastScaleValue
        ld1r    {v0.4s},[x19]           // load scale Value
        mov     v1.16b, v0.16b
        mov     v2.16b, v0.16b
        mov     v3.16b, v0.16b

AccumulatorsToFloat
        scvtf   v16.4s,v16.4s           // convert to float
        scvtf   v17.4s,v17.4s
        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
        scvtf   v24.4s,v24.4s
        scvtf   v25.4s,v25.4s
        scvtf   v26.4s,v26.4s
        scvtf   v27.4s,v27.4s
        scvtf   v28.4s,v28.4s
        scvtf   v29.4s,v29.4s
        scvtf   v30.4s,v30.4s
        scvtf   v31.4s,v31.4s
        fmul    v16.4s,v16.4s,v0.4s     // multiply by scale
        fmul    v17.4s,v17.4s,v0.4s
        fmul    v18.4s,v18.4s,v0.4s
        fmul    v19.4s,v19.4s,v0.4s
        fmul    v20.4s,v20.4s,v1.4s
        fmul    v21.4s,v21.4s,v1.4s
        fmul    v22.4s,v22.4s,v1.4s
        fmul    v23.4s,v23.4s,v1.4s
        fmul    v24.4s,v24.4s,v2.4s
        fmul    v25.4s,v25.4s,v2.4s
        fmul    v26.4s,v26.4s,v2.4s
        fmul    v27.4s,v27.4s,v2.4s
        fmul    v28.4s,v28.4s,v3.4s
        fmul    v29.4s,v29.4s,v3.4s
        fmul    v30.4s,v30.4s,v3.4s
        fmul    v31.4s,v31.4s,v3.4s
        fcvtns  v16.4s,v16.4s           // convert to int
        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
        fcvtns  v24.4s,v24.4s
        fcvtns  v25.4s,v25.4s
        fcvtns  v26.4s,v26.4s
        fcvtns  v27.4s,v27.4s
        fcvtns  v28.4s,v28.4s
        fcvtns  v29.4s,v29.4s
        fcvtns  v30.4s,v30.4s
        fcvtns  v31.4s,v31.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
        dup     v4.8h,w13               // zero point
        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,v4.8h
        sqadd   v17.8h,v17.8h,v4.8h
        sqadd   v18.8h,v18.8h,v4.8h
        sqadd   v19.8h,v19.8h,v4.8h
        sqadd   v24.8h,v24.8h,v4.8h
        sqadd   v25.8h,v25.8h,v4.8h
        sqadd   v26.8h,v26.8h,v4.8h
        sqadd   v27.8h,v27.8h,v4.8h
        sqxtn   v0.8b,v16.8h
        sqxtn   v1.8b,v17.8h
        sqxtn   v2.8b,v18.8h
        sqxtn   v3.8b,v19.8h
        sqxtn2  v0.16b,v24.8h
        sqxtn2  v1.16b,v25.8h
        subs    x6,x6,16            // processed 16 output channels
        sqxtn2  v2.16b,v26.8h
        sqxtn2  v3.16b,v27.8h
        b.lo    PartialStore

        st1     {v3.16b},[x5],16    // Store full 4 x 16
        st1     {v2.16b},[x17],16
        sub     x0,x0,x3            // Restore pointer to A: a -= ks
        st1     {v1.16b},[x16],16
        st1     {v0.16b},[x2],16
        b.hi    OutputChannelLoop

ExitKernel
        EPILOG_RESTORE_REG       x19,#32
        EPILOG_RESTORE_REG_PAIR  d10,d11,#16
        EPILOG_RESTORE_REG_PAIR  d8,d9,#ConvSymFrame_SavedRegisters!
        EPILOG_RETURN

InChannels8
        tbz     x7,3,InChannels4
        ldr     d0,[x12],8
        ldr     q4,[x1],16
        ldr     d1,[x13],8
        ldr     d2,[x14],8
        ldr     d3,[x15],8
        ldr     q5,[x1],16
        SdotByElement    16, 4, 0,0
        SdotByElement    17, 4, 1,0
        ldp     q6, q7, [x1], 32
        SdotByElement    18, 4, 2,0
        SdotByElement    19, 4, 3,0
        SdotByElement    20, 5, 0,0
        SdotByElement    21, 5, 1,0
        SdotByElement    22, 5, 2,0
        SdotByElement    23, 5, 3,0
        SdotByElement    24, 6, 0,0
        SdotByElement    25, 6, 1,0
        ldp     q4, q5, [x1], 32
        SdotByElement    26, 6, 2,0
        SdotByElement    27, 6, 3,0
        SdotByElement    28, 7, 0,0
        SdotByElement    29, 7, 1,0
        SdotByElement    30, 7, 2,0
        SdotByElement    31, 7, 3,0
        SdotByElement    16, 4, 0,1
        SdotByElement    17, 4, 1,1
        ldp     q6, q7, [x1], 32
        SdotByElement    18, 4, 2,1
        SdotByElement    19, 4, 3,1
        SdotByElement    20, 5, 0,1
        SdotByElement    21, 5, 1,1
        SdotByElement    22, 5, 2,1
        SdotByElement    23, 5, 3,1
        SdotByElement    24, 6, 0,1
        SdotByElement    25, 6, 1,1
        SdotByElement    26, 6, 2,1
        SdotByElement    27, 6, 3,1
        SdotByElement    28, 7, 0,1
        SdotByElement    29, 7, 1,1
        SdotByElement    30, 7, 2,1
        SdotByElement    31, 7, 3,1
        tbz     x7,2,SkipInCh4

InChannels4
        ldr     s0,[x12],4
        ldr     q4,[x1],16
        ldr     s1,[x13],4
        ldr     s2,[x14],4
        ldr     s3,[x15],4
        ldr     q5,[x1],16
        SdotByElement    16, 4, 0,0
        SdotByElement    17, 4, 1,0
        ldp     q6, q7, [x1], 32
        SdotByElement    18, 4, 2,0
        SdotByElement    19, 4, 3,0
        SdotByElement    20, 5, 0,0
        SdotByElement    21, 5, 1,0
        SdotByElement    22, 5, 2,0
        SdotByElement    23, 5, 3,0
        SdotByElement    24, 6, 0,0
        SdotByElement    25, 6, 1,0
        SdotByElement    26, 6, 2,0
        SdotByElement    27, 6, 3,0
        SdotByElement    28, 7, 0,0
        SdotByElement    29, 7, 1,0
        SdotByElement    30, 7, 2,0
        SdotByElement    31, 7, 3,0

SkipInCh4
        subs    x9,x9,8             // ks -= 1
        b.hi    KernelSizeLoop
        b       Requantize

PartialStore
        tbz     x6,3,LT8Store
        str     d3,[x5],8           // no less than 8 channels
        str     d2,[x17],8
        dup     d3,v3.d[1]
        dup     d2,v2.d[1]
        str     d1,[x16],8
        str     d0,[x2],8
        dup     d1,v1.d[1]
        dup     d0,v0.d[1]
LT8Store
        tbz     x6,2,LT4Store
        str     s3,[x5],4
        str     s2,[x17],4
        dup     s3,v3.s[1]
        dup     s2,v2.s[1]
        str     s1,[x16],4
        str     s0,[x2],4
        dup     s1,v1.s[1]
        dup     s0,v0.s[1]
LT4Store
        tbz     x6,1, LT2Store
        str     h3,[x5],2
        str     h2,[x17],2
        dup     h3,v3.h[1]
        dup     h2,v2.h[1]
        str     h1,[x16],2
        str     h0,[x2],2
        dup     h1,v1.h[1]
        dup     h0,v0.h[1]
LT2Store
        tbz     x6,0,ExitKernel
        str     b3,[x5]
        str     b2,[x17]
        str     b1,[x16]
        str     b0,[x2]
        b       ExitKernel

        NESTED_END MlasConvSymS8KernelDot

        END
