;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
;   QgemmU8S8KernelAvx2.asm
;
; Abstract:
;
;   This module implements the kernels for the quantized integer matrix/matrix
;   multiply operation (QGEMM).
;
;   This implementation uses AVX2 instructions.
;   Support for AVX-VNNI-INT8 for certain code paths.
;
;--

        .xlist
INCLUDE mlasi.inc
INCLUDE AssembleAvxVnni.inc
        .list

        EXTERN  MlasMaskMoveTableAvx:NEAR

;
; Stack frame layout for the Int8 CopyPackA routine.
;

GemmInt8CopyPackAFrame STRUCT

        PaddedMatrixAData OWORD 4 DUP (?)
        SavedXmm6 OWORD ?
        SavedXmm7 OWORD ?
        SavedXmm8 OWORD ?
        SavedXmm9 OWORD ?
        SavedXmm10 OWORD ?
        Padding QWORD ?
        SavedR13 QWORD ?
        SavedR12 QWORD ?
        SavedRdi QWORD ?
        SavedRsi QWORD ?
        SavedRbx QWORD ?
        SavedRbp QWORD ?
        ReturnAddress QWORD ?
        PreviousP1Home QWORD ?
        PreviousP2Home QWORD ?
        PreviousP3Home QWORD ?
        PreviousP4Home QWORD ?
        CountK QWORD ?
        RowSumBuffer QWORD ?

GemmInt8CopyPackAFrame ENDS

;
; Stack frame layout for the Int8 CopyPackB routine.
;

GemmInt8CopyPackBFrame STRUCT

        PaddedMatrixBData OWORD 4 DUP (?)
        SavedXmm6 OWORD ?
        SavedXmm7 OWORD ?
        SavedXmm8 OWORD ?
        SavedXmm9 OWORD ?
        Padding QWORD ?
        SavedRdi QWORD ?
        SavedRsi QWORD ?
        SavedRbx QWORD ?
        SavedRbp QWORD ?
        ReturnAddress QWORD ?
        PreviousP1Home QWORD ?
        PreviousP2Home QWORD ?
        PreviousP3Home QWORD ?
        PreviousP4Home QWORD ?
        CountK QWORD ?
        ColumnSumBuffer QWORD ?
        BIsSigned QWORD ?

GemmInt8CopyPackBFrame ENDS

;++
;
; Routine Description:
;
;   This routine copies elements from the source matrix to the destination
;   packed buffer.
;
; Arguments:
;
;   D (rcx) - Supplies the address of the destination packed buffer.
;
;   A (rdx) - Supplies the address of the source matrix.
;
;   lda (r8) - Supplies the number of elements per row of the source matrix.
;
;   CountM (r9) - Supplies the number of rows of the source matrix to copy.
;
;   CountK - Supplies the number of columns of the source matrix to copy.
;
;   RowSumBuffer - Supplies the address of the buffer to receive the sums of
;       the elements along each of the rows.
;
; Return Value:
;
;   None.
;
;--

MlasGemmCopyPackAAvx2 MACRO ASigned

        rex_push_reg rbp
        push_reg rbx
        push_reg rsi
        push_reg rdi
        push_reg r12
        push_reg r13
        alloc_stack (GemmInt8CopyPackAFrame.SavedR13)
        save_xmm128 xmm6,GemmInt8CopyPackAFrame.SavedXmm6
        save_xmm128 xmm7,GemmInt8CopyPackAFrame.SavedXmm7
        save_xmm128 xmm8,GemmInt8CopyPackAFrame.SavedXmm8
        save_xmm128 xmm9,GemmInt8CopyPackAFrame.SavedXmm9
        save_xmm128 xmm10,GemmInt8CopyPackAFrame.SavedXmm10

        END_PROLOGUE

        mov     rdi,rcx
        mov     rsi,rdx
        mov     r10,GemmInt8CopyPackAFrame.CountK[rsp]
        lea     r11,[r10+3]
        and     r11,NOT 3                   ; align CountK up to quad count
        mov     r12,GemmInt8CopyPackAFrame.RowSumBuffer[rsp]
        vpcmpeqw ymm8,ymm8,ymm8             ; generate word vector [0xFFFF]
        vpsrlw  ymm8,ymm8,15                ; generate word vector [0x0001]
        vpsllw  ymm9,ymm8,8                 ; generate word vector [0x0100]
        vpor    ymm9,ymm8,ymm9              ; generate word vector [0x0101]

;
; Compute the conditional load/store mask for an unaligned CountK.
;

        mov     eax,r10d
        and     eax,15                      ; isolate unaligned count
        add     eax,3
        shr     eax,2                       ; align unaligned count to quad count
        neg     rax
        lea     rbx,MlasMaskMoveTableAvx+8*4
        vmovdqu xmm10,XMMWORD PTR [rbx+rax*4]

;
; Zero initialize the padded stack buffers.
;

        vpxor   xmm0,xmm0,xmm0
        vmovdqu YMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp],ymm0
        vmovdqu YMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp+32],ymm0

;
; Process 4 rows of matrix A in a loop.
;

        sub     r9,4
        jb      ProcessRemainingRows

ProcessNextRowM4:
        vpxor   xmm0,xmm0,xmm0              ; clear row accumulators
        vpxor   xmm1,xmm1,xmm1
        vpxor   xmm2,xmm2,xmm2
        vpxor   xmm3,xmm3,xmm3
        lea     r13,[r8+r8*2]               ; compute lda * 3
        lea     rax,[r11+r11*2]             ; compute output stride * 3
        mov     rdx,rsi
        mov     rcx,rdi
        lea     rsi,[rsi+r8*4]              ; advance next matrix A by 4 rows
        lea     rdi,[rdi+r11*4]             ; advance next matrix D by 4 rows
        mov     rbx,r10                     ; reload columns remaining
        sub     rbx,32
        jb      ProcessRemainingColumnsM4

ProcessNextColumnLoopM4:
        vmovdqu ymm4,YMMWORD PTR [rdx]
        vmovdqu ymm5,YMMWORD PTR [rdx+r8]
        vmovdqu ymm6,YMMWORD PTR [rdx+r8*2]
        vmovdqu ymm7,YMMWORD PTR [rdx+r13]
        vmovdqu YMMWORD PTR [rcx],ymm4
        vmovdqu YMMWORD PTR [rcx+r11],ymm5
        vmovdqu YMMWORD PTR [rcx+r11*2],ymm6
        vmovdqu YMMWORD PTR [rcx+rax],ymm7
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
        VpdpbssdYmmYmmYmm ymm1,ymm5,ymm9
        VpdpbssdYmmYmmYmm ymm2,ymm6,ymm9
        VpdpbssdYmmYmmYmm ymm3,ymm7,ymm9
ELSE
        vpmaddubsw ymm4,ymm4,ymm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
        vpmaddubsw ymm5,ymm5,ymm9
        vpaddw  ymm1,ymm1,ymm5
        vpmaddubsw ymm6,ymm6,ymm9
        vpaddw  ymm2,ymm2,ymm6
        vpmaddubsw ymm7,ymm7,ymm9
        vpaddw  ymm3,ymm3,ymm7
ENDIF
        add     rdx,32                      ; advance matrix A by 32 bytes
        add     rcx,32                      ; advance matrix D by 32 bytes
        sub     rbx,32                      ; subtract columns remaining
        jae     ProcessNextColumnLoopM4

ProcessRemainingColumnsM4:
        add     rbx,32                      ; correct for over-subtract above
        jz      ReduceRowSumBufferM4
        test    bl,16                       ; (CountK & 16) != 0?
        jz      CopyRemainingCountKLessThan16M4
        vmovdqu xmm4,XMMWORD PTR [rdx]
        vmovdqu xmm5,XMMWORD PTR [rdx+r8]
        vmovdqu xmm6,XMMWORD PTR [rdx+r8*2]
        vmovdqu xmm7,XMMWORD PTR [rdx+r13]
        vmovdqu XMMWORD PTR [rcx],xmm4
        vmovdqu XMMWORD PTR [rcx+r11],xmm5
        vmovdqu XMMWORD PTR [rcx+r11*2],xmm6
        vmovdqu XMMWORD PTR [rcx+rax],xmm7
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
        VpdpbssdYmmYmmYmm ymm1,ymm5,ymm9
        VpdpbssdYmmYmmYmm ymm2,ymm6,ymm9
        VpdpbssdYmmYmmYmm ymm3,ymm7,ymm9
ELSE
        vpmaddubsw xmm4,xmm4,xmm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
        vpmaddubsw xmm5,xmm5,xmm9
        vpaddw  ymm1,ymm1,ymm5
        vpmaddubsw xmm6,xmm6,xmm9
        vpaddw  ymm2,ymm2,ymm6
        vpmaddubsw xmm7,xmm7,xmm9
        vpaddw  ymm3,ymm3,ymm7
ENDIF
        add     rdx,16                      ; advance matrix A by 16 bytes
        add     rcx,16                      ; advance matrix D by 16 bytes
        test    bl,15                       ; test for unaligned columns
        jz      ReduceRowSumBufferM4

;
; Copy the unaligned CountK columns to a zero padded stack buffer.
;

CopyRemainingCountKLessThan16M4:
.errnz  GemmInt8CopyPackAFrame.PaddedMatrixAData
        mov     rbp,rsp                     ; GemmInt8CopyPackAFrame.PaddedMatrixAData
        test    bl,8                        ; (CountK & 8) != 0?
        jz      CopyRemainingCountKLessThan8M4
        mov     rax,QWORD PTR [rdx]
        mov     QWORD PTR [rbp],rax
        mov     rax,QWORD PTR [rdx+r8]
        mov     QWORD PTR [rbp+16],rax
        mov     rax,QWORD PTR [rdx+r8*2]
        mov     QWORD PTR [rbp+32],rax
        mov     rax,QWORD PTR [rdx+r13]
        mov     QWORD PTR [rbp+48],rax
        add     rdx,8
        add     rbp,8                       ; advance padded buffer destination

CopyRemainingCountKLessThan8M4:
        test    bl,4                        ; (CountK & 4) != 0?
        jz      CopyRemainingCountKLessThan4M4
        mov     eax,DWORD PTR [rdx]
        mov     DWORD PTR [rbp],eax
        mov     eax,DWORD PTR [rdx+r8]
        mov     DWORD PTR [rbp+16],eax
        mov     eax,DWORD PTR [rdx+r8*2]
        mov     DWORD PTR [rbp+32],eax
        mov     eax,DWORD PTR [rdx+r13]
        mov     DWORD PTR [rbp+48],eax
        add     rdx,4
        add     rbp,4                       ; advance padded buffer destination

CopyRemainingCountKLessThan4M4:
        test    bl,2                        ; (CountK & 2) != 0?
        jz      CopyRemainingCountKLessThan2M4
        movzx   eax,WORD PTR [rdx]
        mov     WORD PTR [rbp],ax
        movzx   eax,WORD PTR [rdx+r8]
        mov     WORD PTR [rbp+16],ax
        movzx   eax,WORD PTR [rdx+r8*2]
        mov     WORD PTR [rbp+32],ax
        movzx   eax,WORD PTR [rdx+r13]
        mov     WORD PTR [rbp+48],ax
        add     rdx,2
        add     rbp,2                       ; advance padded buffer destination

CopyRemainingCountKLessThan2M4:
        test    bl,1                        ; (CountK & 1) != 0?
        jz      ProcessPaddedMatrixADataM4
        movzx   eax,BYTE PTR [rdx]
        mov     BYTE PTR [rbp],al
        movzx   eax,BYTE PTR [rdx+r8]
        mov     BYTE PTR [rbp+16],al
        movzx   eax,BYTE PTR [rdx+r8*2]
        mov     BYTE PTR [rbp+32],al
        movzx   eax,BYTE PTR [rdx+r13]
        mov     BYTE PTR [rbp+48],al

;
; Process the remaining CountK columns using the zero padded stack buffer.
;

ProcessPaddedMatrixADataM4:
        vmovdqu xmm4,XMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp]
        vmovdqu xmm5,XMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp+16]
        vmovdqu xmm6,XMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp+32]
        vmovdqu xmm7,XMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp+48]
        lea     rax,[rcx+r11*2]             ; compute matrix D plus 2 rows
        vpmaskmovd XMMWORD PTR [rcx],xmm10,xmm4
        vpmaskmovd XMMWORD PTR [rcx+r11],xmm10,xmm5
        vpmaskmovd XMMWORD PTR [rax],xmm10,xmm6
        vpmaskmovd XMMWORD PTR [rax+r11],xmm10,xmm7
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
        VpdpbssdYmmYmmYmm ymm1,ymm5,ymm9
        VpdpbssdYmmYmmYmm ymm2,ymm6,ymm9
        VpdpbssdYmmYmmYmm ymm3,ymm7,ymm9
ELSE
        vpmaddubsw xmm4,xmm4,xmm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
        vpmaddubsw xmm5,xmm5,xmm9
        vpaddw  ymm1,ymm1,ymm5
        vpmaddubsw xmm6,xmm6,xmm9
        vpaddw  ymm2,ymm2,ymm6
        vpmaddubsw xmm7,xmm7,xmm9
        vpaddw  ymm3,ymm3,ymm7
ENDIF

;
; Reduce the sums for the four rows of output.
;

ReduceRowSumBufferM4:
IF ASigned EQ 1
        vphaddd ymm0,ymm0,ymm1              ; reduce and interleave Sum1/Sum0
ELSE
        vpmaddwd ymm0,ymm0,ymm8             ; horizontal word+word=dword per row
        vpmaddwd ymm1,ymm1,ymm8
        vphaddd ymm0,ymm0,ymm1              ; reduce and interleave Sum1/Sum0
        vpmaddwd ymm2,ymm2,ymm8
        vpmaddwd ymm3,ymm3,ymm8
ENDIF
        vphaddd ymm1,ymm2,ymm3              ; reduce and interleave Sum3/Sum2
        vphaddd ymm0,ymm0,ymm1              ; reduce and interleave Sum3/Sum2/Sum1/Sum0
        vextracti128 xmm1,ymm0,1            ; extract high dwords
        vpaddd  xmm0,xmm0,xmm1              ; reduce low/high dwords
        vmovdqu XMMWORD PTR [r12],xmm0
        add     r12,4*4                     ; advance row sum buffer by 4 dwords
        sub     r9,4                        ; subtract rows remaining
        jae     ProcessNextRowM4

ProcessRemainingRows:
        add     r9,4                        ; correct for over-subtract above
        jz      ExitRoutine

;
; Process a single row of matrix A in a loop.
;

ProcessNextRowM1:
        vpxor   xmm0,xmm0,xmm0              ; clear row accumulator
        mov     rdx,rsi
        mov     rcx,rdi
        add     rsi,r8
        add     rdi,r11
        mov     rbx,r10                     ; reload columns remaining
        sub     rbx,32
        jb      ProcessRemainingColumnsM1

ProcessNextColumnLoopM1:
        vmovdqu ymm4,YMMWORD PTR [rdx]
        vmovdqu YMMWORD PTR [rcx],ymm4
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
ELSE
        vpmaddubsw ymm4,ymm4,ymm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
ENDIF
        add     rdx,32                      ; advance matrix A by 32 bytes
        add     rcx,32                      ; advance matrix D by 32 bytes
        sub     rbx,32                      ; subtract columns remaining
        jae     ProcessNextColumnLoopM1

ProcessRemainingColumnsM1:
        add     rbx,32                      ; correct for over-subtract above
        jz      ReduceRowSumBufferM1
        test    bl,16                       ; (CountK & 16) != 0?
        jz      CopyRemainingCountKLessThan16M1
        vmovdqu xmm4,XMMWORD PTR [rdx]
        vmovdqu XMMWORD PTR [rcx],xmm4
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
ELSE
        vpmaddubsw xmm4,xmm4,xmm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
ENDIF
        add     rdx,16                      ; advance matrix A by 16 bytes
        add     rcx,16                      ; advance matrix D by 16 bytes
        test    bl,15                       ; test for unaligned columns
        jz      ReduceRowSumBufferM1

;
; Copy the unaligned CountK columns to a zero padded stack buffer.
;

CopyRemainingCountKLessThan16M1:
.errnz  GemmInt8CopyPackAFrame.PaddedMatrixAData
        mov     rbp,rsp                     ; GemmInt8CopyPackAFrame.PaddedMatrixAData
        test    bl,8                        ; (CountK & 8) != 0?
        jz      CopyRemainingCountKLessThan8M1
        mov     rax,QWORD PTR [rdx]
        mov     QWORD PTR [rbp],rax
        add     rdx,8
        add     rbp,8                       ; advance padded buffer destination

CopyRemainingCountKLessThan8M1:
        test    bl,4                        ; (CountK & 4) != 0?
        jz      CopyRemainingCountKLessThan4M1
        mov     eax,DWORD PTR [rdx]
        mov     DWORD PTR [rbp],eax
        add     rdx,4
        add     rbp,4                       ; advance padded buffer destination

CopyRemainingCountKLessThan4M1:
        test    bl,2                        ; (CountK & 2) != 0?
        jz      CopyRemainingCountKLessThan2M1
        movzx   eax,WORD PTR [rdx]
        mov     WORD PTR [rbp],ax
        add     rdx,2
        add     rbp,2                       ; advance padded buffer destination

CopyRemainingCountKLessThan2M1:
        test    bl,1                        ; (CountK & 1) != 0?
        jz      ProcessPaddedMatrixADataM1
        movzx   eax,BYTE PTR [rdx]
        mov     BYTE PTR [rbp],al

;
; Process the remaining CountK columns using the zero padded stack buffer.
;

ProcessPaddedMatrixADataM1:
        vmovdqu xmm4,XMMWORD PTR GemmInt8CopyPackAFrame.PaddedMatrixAData[rsp]
        vpmaskmovd XMMWORD PTR [rcx],xmm10,xmm4
IF ASigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm9
ELSE
        vpmaddubsw ymm4,ymm4,ymm9           ; horizontal byte+byte=word per row
        vpaddw  ymm0,ymm0,ymm4              ; add words to row accumulators
ENDIF

;
; Reduce the sum for the single row of output.
;

ReduceRowSumBufferM1:
IF ASigned EQ 0
        vpmaddwd ymm0,ymm0,ymm8             ; horizontal word+word=dword per row
ENDIF
        vextracti128 xmm1,ymm0,1            ; extract high dwords
        vpaddd  xmm0,xmm0,xmm1              ; reduction
        vphaddd xmm0,xmm0,xmm0
        vphaddd xmm0,xmm0,xmm0
        vmovd   DWORD PTR [r12],xmm0
        add     r12,4                       ; advance row sum buffer by 1 dword
        dec     r9                          ; decrement rows remaining
        jnz     ProcessNextRowM1

;
; Restore non-volatile registers and return.
;

ExitRoutine:
        vzeroupper
        movaps  xmm6,GemmInt8CopyPackAFrame.SavedXmm6[rsp]
        movaps  xmm7,GemmInt8CopyPackAFrame.SavedXmm7[rsp]
        movaps  xmm8,GemmInt8CopyPackAFrame.SavedXmm8[rsp]
        movaps  xmm9,GemmInt8CopyPackAFrame.SavedXmm9[rsp]
        movaps  xmm10,GemmInt8CopyPackAFrame.SavedXmm10[rsp]
        add     rsp,(GemmInt8CopyPackAFrame.SavedR13)

        BEGIN_EPILOGUE

        pop     r13
        pop     r12
        pop     rdi
        pop     rsi
        pop     rbx
        pop     rbp
        ret

        ENDM

        NESTED_ENTRY MlasGemmU8S8CopyPackAAvx2, _TEXT
        MlasGemmCopyPackAAvx2 0
        NESTED_END MlasGemmU8S8CopyPackAAvx2, _TEXT

        NESTED_ENTRY MlasGemmS8CopyPackAAvx2Vnni, _TEXT
        MlasGemmCopyPackAAvx2 1
        NESTED_END MlasGemmS8CopyPackAAvx2Vnni, _TEXT

;++
;
; Routine Description:
;
;   This routine copies elements from the source matrix to the destination
;   packed buffer.
;
; Arguments:
;
;   D (rcx) - Supplies the address of the destination packed buffer.
;
;   B (rdx) - Supplies the address of the source matrix.
;
;   ldb (r8) - Supplies the number of elements per row of the source matrix.
;
;   CountN (r9) - Supplies the number of columns of the source matrix to copy.
;
;   CountK - Supplies the number of rows of the source matrix to copy.
;
;   ColumnSumBuffer - Supplies the address of the buffer to receive the sums of
;       the elements along each of the columns.
;
;   BIsSigned - Supplies true if the source matrix is signed data, else false
;       if the source matrix is unsigned data.
;
; Return Value:
;
;   None.
;
;--

MlasGemmCopyPackBAvx2 MACRO IsVnni, BSigned

        rex_push_reg rbp
        push_reg rbx
        push_reg rsi
        push_reg rdi
        alloc_stack (GemmInt8CopyPackBFrame.SavedRdi)
        save_xmm128 xmm6,GemmInt8CopyPackBFrame.SavedXmm6
        save_xmm128 xmm7,GemmInt8CopyPackBFrame.SavedXmm7
        save_xmm128 xmm8,GemmInt8CopyPackBFrame.SavedXmm8
        save_xmm128 xmm9,GemmInt8CopyPackBFrame.SavedXmm9

        END_PROLOGUE

        mov     rsi,rdx
        lea     rdi,[r8+r8*2]               ; compute ldb * 3
        mov     r10,GemmInt8CopyPackBFrame.CountK[rsp]
        mov     r11,GemmInt8CopyPackBFrame.ColumnSumBuffer[rsp]
        vpcmpeqw ymm7,ymm7,ymm7             ; generate word vector [0xFFFF]
        vpsrlw  ymm7,ymm7,15                ; generate word vector [0x0001]
        vpsllw  ymm8,ymm7,8                 ; generate word vector [0x0100]
        vpor    ymm8,ymm7,ymm8              ; generate word vector [0x0101]

;
; Compute the bit flip vector to adjust input from U8 to S8.
;

        vpxor   xmm9,xmm9,xmm9              ; generate word vector [0x0000]
IF IsVnni EQ 0
        cmp     BYTE PTR GemmInt8CopyPackBFrame.BIsSigned[rsp],0
        jnz     SkipUnsignedBitFlipVector
        vpsllw  ymm9,ymm8,7                 ; generate word vector [0x8080]
ENDIF
SkipUnsignedBitFlipVector:

;
; Process 16 columns of matrix B in a loop.
;

        sub     r9,16
        jb      ProcessRemainingColumns

ProcessNextColumnN16:
        vpxor   xmm0,xmm0,xmm0              ; clear column accumulators
        vpxor   xmm1,xmm1,xmm1
        mov     rdx,rsi
        add     rsi,16                      ; advance next matrix B by 16 columns
        mov     rbx,r10                     ; reload rows remaining
        sub     rbx,4
        jb      ProcessRemainingRowsN16

ProcessNextRowLoopN16:
        vmovdqu xmm2,XMMWORD PTR [rdx]      ; load 4 rows
        vmovdqu xmm3,XMMWORD PTR [rdx+r8]
        vmovdqu xmm4,XMMWORD PTR [rdx+r8*2]
        vmovdqu xmm5,XMMWORD PTR [rdx+rdi]
        lea     rdx,[rdx+r8*4]              ; advance matrix B by 4 rows

InterleaveRowDataN16:
        vpunpcklbw xmm6,xmm2,xmm3           ; interleave row data
        vpunpckhbw xmm3,xmm2,xmm3
        vpunpcklbw xmm2,xmm4,xmm5
        vpunpckhbw xmm5,xmm4,xmm5
        vpunpcklwd xmm4,xmm6,xmm2
        vpunpckhwd xmm6,xmm6,xmm2
        vpunpcklwd xmm2,xmm3,xmm5
        vpunpckhwd xmm3,xmm3,xmm5
        vinserti128 ymm4,ymm4,xmm6,1
        vinserti128 ymm2,ymm2,xmm3,1
IF IsVnni EQ 0
        vpxor   ymm4,ymm4,ymm9              ; optionally adjust unsigned data
        vpxor   ymm2,ymm2,ymm9
ENDIF
        vmovdqu YMMWORD PTR [rcx],ymm4      ; store interleaved rows
        vmovdqu YMMWORD PTR [rcx+32],ymm2
IF IsVnni EQ 1
    IF BSigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm8
        VpdpbssdYmmYmmYmm ymm1,ymm2,ymm8
    ELSE
        VpdpbuudYmmYmmYmm ymm0,ymm4,ymm8
        VpdpbuudYmmYmmYmm ymm1,ymm2,ymm8
    ENDIF
ELSE
        vpmaddubsw ymm4,ymm8,ymm4           ; horizontal byte+byte=word per row
        vpmaddwd ymm4,ymm4,ymm7             ; horizontal word+word=dword per row
        vpaddd  ymm0,ymm0,ymm4              ; accumulate per column
        vpmaddubsw ymm2,ymm8,ymm2
        vpmaddwd ymm2,ymm2,ymm7
        vpaddd  ymm1,ymm1,ymm2
ENDIF
        add     rcx,64                      ; advance matrix D by 64 bytes
        sub     rbx,4                       ; subtract rows remaining
        jae     ProcessNextRowLoopN16

;
; Process the less than 4 remaining rows where the row has 16 columns.
;

ProcessRemainingRowsN16:
        add     rbx,4                       ; correct for over-subtract above
        jz      StoreColumnSumBufferN16
        vmovdqu xmm2,XMMWORD PTR [rdx]
        vmovaps xmm3,xmm9
        vmovaps xmm4,xmm9
        vmovaps xmm5,xmm9
        xor     ebx,ebx                     ; no more rows remaining
        test    r10b,2                      ; (CountK & 2) != 0?
        jz      InterleaveRowDataN16
        vmovdqu xmm3,XMMWORD PTR [rdx+r8]
        test    r10b,1                      ; (CountK & 1) != 0?
        jz      InterleaveRowDataN16
        vmovdqu xmm4,XMMWORD PTR [rdx+r8*2]
        jmp     InterleaveRowDataN16

StoreColumnSumBufferN16:
        vmovdqu YMMWORD PTR [r11],ymm0
        vmovdqu YMMWORD PTR [r11+32],ymm1
        add     r11,16*4                    ; advance column sum buffer by 16 dwords
        sub     r9,16                       ; subtract columns remaining
        jae     ProcessNextColumnN16

ProcessRemainingColumns:
        add     r9,16                       ; correct for over-subtract above
        jnz     ProcessColumnNUnaligned

;
; Restore non-volatile registers and return.
;

ExitRoutine:
        vzeroupper
        movaps  xmm6,GemmInt8CopyPackBFrame.SavedXmm6[rsp]
        movaps  xmm7,GemmInt8CopyPackBFrame.SavedXmm7[rsp]
        movaps  xmm8,GemmInt8CopyPackBFrame.SavedXmm8[rsp]
        movaps  xmm9,GemmInt8CopyPackBFrame.SavedXmm9[rsp]
        add     rsp,(GemmInt8CopyPackBFrame.SavedRdi)

        BEGIN_EPILOGUE

        pop     rdi
        pop     rsi
        pop     rbx
        pop     rbp
        ret

;
; Process the remaining columns of matrix B.
;

ProcessColumnNUnaligned:
        vpxor   xmm0,xmm0,xmm0              ; clear column accumulators
        vpxor   xmm1,xmm1,xmm1
        vmovdqu YMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp],ymm9
        vmovdqu YMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp+32],ymm9
        sub     r10,4
        jb      ProcessRemainingRowsNUnaligned

ProcessNextRowLoopNUnaligned:
        mov     rdx,rsi
.errnz  GemmInt8CopyPackBFrame.PaddedMatrixBData
        mov     rbp,rsp                     ; GemmInt8CopyPackBFrame.PaddedMatrixBData
        test    r9b,8                       ; (CountN & 8) != 0?
        jz      CopyRemainingCountNLessThan8K4
        mov     rax,QWORD PTR [rdx]
        mov     QWORD PTR [rbp],rax
        mov     rax,QWORD PTR [rdx+r8]
        mov     QWORD PTR [rbp+16],rax
        mov     rax,QWORD PTR [rdx+r8*2]
        mov     QWORD PTR [rbp+32],rax
        mov     rax,QWORD PTR [rdx+rdi]
        mov     QWORD PTR [rbp+48],rax
        add     rdx,8                       ; advance matrix B
        add     rbp,8                       ; advance padded buffer destination

CopyRemainingCountNLessThan8K4:
        test    r9b,4                       ; (CountN & 4) != 0?
        jz      CopyRemainingCountNLessThan4K4
        mov     eax,DWORD PTR [rdx]
        mov     DWORD PTR [rbp],eax
        mov     eax,DWORD PTR [rdx+r8]
        mov     DWORD PTR [rbp+16],eax
        mov     eax,DWORD PTR [rdx+r8*2]
        mov     DWORD PTR [rbp+32],eax
        mov     eax,DWORD PTR [rdx+rdi]
        mov     DWORD PTR [rbp+48],eax
        add     rdx,4                       ; advance matrix B
        add     rbp,4                       ; advance padded buffer destination

CopyRemainingCountNLessThan4K4:
        test    r9b,2                       ; (CountN & 2) != 0?
        jz      CopyRemainingCountNLessThan2K4
        movzx   eax,WORD PTR [rdx]
        mov     WORD PTR [rbp],ax
        movzx   eax,WORD PTR [rdx+r8]
        mov     WORD PTR [rbp+16],ax
        movzx   eax,WORD PTR [rdx+r8*2]
        mov     WORD PTR [rbp+32],ax
        movzx   eax,WORD PTR [rdx+rdi]
        mov     WORD PTR [rbp+48],ax
        add     rdx,2                       ; advance matrix B
        add     rbp,2                       ; advance padded buffer destination

CopyRemainingCountNLessThan2K4:
        test    r9b,1                       ; (CountN & 1) != 0?
        jz      ProcessPaddedMatrixBData
        movzx   eax,BYTE PTR [rdx]
        mov     BYTE PTR [rbp],al
        movzx   eax,BYTE PTR [rdx+r8]
        mov     BYTE PTR [rbp+16],al
        movzx   eax,BYTE PTR [rdx+r8*2]
        mov     BYTE PTR [rbp+32],al
        movzx   eax,BYTE PTR [rdx+rdi]
        mov     BYTE PTR [rbp+48],al

ProcessPaddedMatrixBData:
        vmovdqu xmm2,XMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp]
        vmovdqu xmm3,XMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp+16]
        vmovdqu xmm4,XMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp+32]
        vmovdqu xmm5,XMMWORD PTR GemmInt8CopyPackBFrame.PaddedMatrixBData[rsp+48]
        vpunpcklbw xmm6,xmm2,xmm3           ; interleave row data
        vpunpckhbw xmm3,xmm2,xmm3
        vpunpcklbw xmm2,xmm4,xmm5
        vpunpckhbw xmm5,xmm4,xmm5
        vpunpcklwd xmm4,xmm6,xmm2
        vpunpckhwd xmm6,xmm6,xmm2
        vpunpcklwd xmm2,xmm3,xmm5
        vpunpckhwd xmm3,xmm3,xmm5
        vinserti128 ymm4,ymm4,xmm6,1
        vinserti128 ymm2,ymm2,xmm3,1
IF IsVnni EQ 0
        vpxor   ymm4,ymm4,ymm9              ; optionally adjust unsigned data
        vpxor   ymm2,ymm2,ymm9
ENDIF
        vmovdqu YMMWORD PTR [rcx],ymm4      ; store interleaved rows
        vmovdqu YMMWORD PTR [rcx+32],ymm2
IF IsVnni EQ 1
    IF BSigned EQ 1
        VpdpbssdYmmYmmYmm ymm0,ymm4,ymm8
        VpdpbssdYmmYmmYmm ymm1,ymm2,ymm8
    ELSE
        VpdpbuudYmmYmmYmm ymm0,ymm4,ymm8
        VpdpbuudYmmYmmYmm ymm1,ymm2,ymm8
    ENDIF
ELSE
        vpmaddubsw ymm4,ymm8,ymm4           ; horizontal byte+byte=word per row
        vpmaddwd ymm4,ymm4,ymm7             ; horizontal word+word=dword per row
        vpaddd  ymm0,ymm0,ymm4              ; accumulate per column
        vpmaddubsw ymm2,ymm8,ymm2
        vpmaddwd ymm2,ymm2,ymm7
        vpaddd  ymm1,ymm1,ymm2
ENDIF
        lea     rsi,[rsi+r8*4]              ; advance next matrix B by 4 rows
        add     rcx,64                      ; advance matrix D by 64 bytes
        sub     r10,4                       ; subtract rows remaining
        jae     ProcessNextRowLoopNUnaligned

ProcessRemainingRowsNUnaligned:
        add     r10,4
        jz      StoreColumnSumBufferNUnaligned

;
; Process the less than 4 remaining rows where the row has less than 16 columns.
;

.errnz  GemmInt8CopyPackBFrame.PaddedMatrixBData
        mov     rbp,rsp                     ; GemmInt8CopyPackBFrame.PaddedMatrixBData
        vmovdqu YMMWORD PTR [rbp],ymm9
        vmovdqu YMMWORD PTR [rbp+32],ymm9

CopyUnalignedRowLoop:
        lea     rdi,[rbp+16]                ; advance next padded buffer by 16 bytes
        mov     rdx,rsi
        test    r9b,8                       ; (CountN & 8) != 0?
        jz      CopyRemainingCountNLessThan8KSmall
        mov     rax,QWORD PTR [rdx]
        mov     QWORD PTR [rbp],rax
        add     rdx,8                       ; advance matrix B
        add     rbp,8                       ; advance padded buffer destination

CopyRemainingCountNLessThan8KSmall:
        test    r9b,4                       ; (CountN & 4) != 0?
        jz      CopyRemainingCountNLessThan4KSmall
        mov     eax,DWORD PTR [rdx]
        mov     DWORD PTR [rbp],eax
        add     rdx,4                       ; advance matrix B
        add     rbp,4                       ; advance padded buffer destination

CopyRemainingCountNLessThan4KSmall:
        test    r9b,2                       ; (CountN & 2) != 0?
        jz      CopyRemainingCountNLessThan2KSmall
        movzx   eax,WORD PTR [rdx]
        mov     WORD PTR [rbp],ax
        add     rdx,2                       ; advance matrix B
        add     rbp,2                       ; advance padded buffer destination

CopyRemainingCountNLessThan2KSmall:
        test    r9b,1                       ; (CountN & 1) != 0?
        jz      DoneCopyRemainingCountNKSmall
        movzx   eax,BYTE PTR [rdx]
        mov     BYTE PTR [rbp],al

DoneCopyRemainingCountNKSmall:
        dec     r10
        jz      ProcessPaddedMatrixBData
        add     rsi,r8                      ; advance next matrix B by 1 row
        mov     rbp,rdi
        jmp     CopyUnalignedRowLoop

StoreColumnSumBufferNUnaligned:
        vmovdqu YMMWORD PTR [r11],ymm0
        vmovdqu YMMWORD PTR [r11+32],ymm1
        jmp     ExitRoutine

ENDM

        NESTED_ENTRY MlasGemmU8S8CopyPackBAvx2, _TEXT
        MlasGemmCopyPackBAvx2 0                         ; sign variable not checked if IsVnni = 0
        NESTED_END MlasGemmU8S8CopyPackBAvx2, _TEXT

        NESTED_ENTRY MlasGemmU8CopyPackBAvx2Vnni, _TEXT
        MlasGemmCopyPackBAvx2 1, 0
        NESTED_END MlasGemmU8CopyPackBAvx2Vnni, _TEXT

        NESTED_ENTRY MlasGemmS8CopyPackBAvx2Vnni, _TEXT
        MlasGemmCopyPackBAvx2 1, 1
        NESTED_END MlasGemmS8CopyPackBAvx2Vnni, _TEXT

        END
