;++ ; ; Copyright (c) 2020 Intel Corporation. All rights reserved. ; ; Licensed under the MIT License. ; ; Module Name: ; ; QgemvU8S8KernelAvxVnni.asm ; ; Abstract: ; ; This module implements the kernels for the quantized integer matrix/vector ; multiply operation (QGEMV). ; ; This implementation uses AVXVNNI instructions. ; ;-- .xlist INCLUDE mlasi.inc INCLUDE AssembleAvxVnni.inc .list EXTERN MlasMaskMoveAvx:NEAR EXTERN MlasTranspose4x4BytesAvx:NEAR ; ; Stack frame layout for the U8S8 kernel. ; GemvU8S8KernelFrame STRUCT SavedXmm6 OWORD ? SavedXmm7 OWORD ? SavedXmm8 OWORD ? SavedXmm9 OWORD ? SavedXmm10 OWORD ? Padding QWORD ? SavedRdi QWORD ? SavedRsi QWORD ? SavedRbx QWORD ? SavedRbp QWORD ? ReturnAddress QWORD ? PreviousP1Home QWORD ? PreviousP2Home QWORD ? PreviousP3Home QWORD ? PreviousP4Home QWORD ? CountN QWORD ? ldb QWORD ? GemvU8S8KernelFrame ENDS ;++ ; ; Routine Description: ; ; This routine is an inner kernel to compute matrix/vector multiplication. ; ; Arguments: ; ; A (rcx) - Supplies the address of vector A. ; ; B (rdx) - Supplies the address of matrix B. ; ; C (r8) - Supplies the address of matrix C. ; ; CountK (r9) - Supplies the number of columns from vector A and the number ; of rows from matrix B to iterate over. ; ; CountN - Supplies the number of columns from matrix B and matrix C to iterate ; over. ; ; ldb - Supplies the first dimension of matrix B. ; ; Return Value: ; ; None. ; ;-- NESTED_ENTRY MlasGemvU8S8KernelAvxVnni, _TEXT rex_push_reg rbp push_reg rbx push_reg rsi push_reg rdi alloc_stack (GemvU8S8KernelFrame.SavedRdi) save_xmm128 xmm6,GemvU8S8KernelFrame.SavedXmm6 save_xmm128 xmm7,GemvU8S8KernelFrame.SavedXmm7 save_xmm128 xmm8,GemvU8S8KernelFrame.SavedXmm8 save_xmm128 xmm9,GemvU8S8KernelFrame.SavedXmm9 save_xmm128 xmm10,GemvU8S8KernelFrame.SavedXmm10 END_PROLOGUE mov rsi,rdx mov rdi,GemvU8S8KernelFrame.ldb[rsp] mov r10,GemvU8S8KernelFrame.CountN[rsp] mov r11,rsp ; set ZeroMode to any non-zero value ; ; Process 4 rows of matrix B in a loop. ; sub r9,4 jb ProcessRemainingRows ProcessRowLoop4: mov rdx,rsi ; reload matrix B lea rsi,[rsi+rdi*4] ; advance matrix B by 4 rows mov rbx,r8 ; reload matrix C mov rbp,r10 ; reload CountN vpbroadcastd ymm0,DWORD PTR [rcx] add rcx,4 ; advance matrix A by 4 bytes ; ; Process sets of 32 columns from the 4 rows in a loop. ; ; Some permute operations are deferred until the final store of the 4x32 block ; as these permutes are expensive. ; ProcessColumnLoop4By32: cmp rbp,32 jb ProcessColumnLoop4By8 lea rax,[rdx+rdi*2] ; compute matrix B plus 2 rows vmovdqu ymm2,YMMWORD PTR [rdx] vmovdqu ymm3,YMMWORD PTR [rdx+rdi] vmovdqu ymm4,YMMWORD PTR [rax] vmovdqu ymm5,YMMWORD PTR [rax+rdi] vpunpcklbw ymm1,ymm2,ymm3 ; interleave row data bytes vpunpckhbw ymm2,ymm2,ymm3 vpxor ymm7,ymm7,ymm7 vpunpcklbw ymm3,ymm4,ymm5 vpunpckhbw ymm4,ymm4,ymm5 vpxor ymm8,ymm8,ymm8 vpunpcklwd ymm5,ymm1,ymm3 ; interleave row data words vpunpckhwd ymm1,ymm1,ymm3 vpxor ymm9,ymm9,ymm9 vpunpcklwd ymm3,ymm2,ymm4 vpunpckhwd ymm2,ymm2,ymm4 vpxor ymm10,ymm10,ymm10 VpdpbusdsYmmYmmYmm ymm7,ymm0,ymm5 VpdpbusdsYmmYmmYmm ymm8,ymm0,ymm1 VpdpbusdsYmmYmmYmm ymm9,ymm0,ymm3 VpdpbusdsYmmYmmYmm ymm10,ymm0,ymm2 test r11,r11 ; ZeroMode? jnz SkipAccumulateOutput4By32 vpaddd ymm7,ymm7,YMMWORD PTR [rbx] vpaddd ymm8,ymm8,YMMWORD PTR [rbx+32] vpaddd ymm9,ymm9,YMMWORD PTR [rbx+64] vpaddd ymm10,ymm10,YMMWORD PTR [rbx+96] SkipAccumulateOutput4By32: cmp r9,4 ; final 4x32 block? jae StoreOutput4By32 vperm2i128 ymm4,ymm7,ymm8,31h ; interleave vector results vperm2i128 ymm7,ymm7,ymm8,20h vperm2i128 ymm8,ymm9,ymm10,20h vperm2i128 ymm10,ymm9,ymm10,31h vmovaps ymm9,ymm4 StoreOutput4By32: vmovdqu YMMWORD PTR [rbx],ymm7 vmovdqu YMMWORD PTR [rbx+32],ymm8 vmovdqu YMMWORD PTR [rbx+64],ymm9 vmovdqu YMMWORD PTR [rbx+96],ymm10 add rdx,32 ; advance matrix B by 32 bytes add rbx,32*4 ; advance matrix C by 32 columns sub rbp,32 ; decrement CountN jnz ProcessColumnLoop4By32 AdvanceRowLoop4: xor r11,r11 ; clear ZeroMode sub r9,4 ; decrement CountK jae ProcessRowLoop4 ProcessRemainingRows: add r9,4 ; correct for over-subtract above jnz ProcessRemainingSmallK ; ; Restore non-volatile registers and return. ; ExitKernel: vzeroupper movaps xmm6,GemvU8S8KernelFrame.SavedXmm6[rsp] movaps xmm7,GemvU8S8KernelFrame.SavedXmm7[rsp] movaps xmm8,GemvU8S8KernelFrame.SavedXmm8[rsp] movaps xmm9,GemvU8S8KernelFrame.SavedXmm9[rsp] movaps xmm10,GemvU8S8KernelFrame.SavedXmm10[rsp] add rsp,(GemvU8S8KernelFrame.SavedRdi) BEGIN_EPILOGUE pop rdi pop rsi pop rbx pop rbp ret ; ; Process sets of 8 columns from the 4 rows in a loop. ; ProcessColumnLoop4By8: cmp ebp,8 jb ProcessColumn4By4 lea rax,[rdx+rdi*2] ; compute matrix B plus 2 rows vmovq xmm2,QWORD PTR [rdx] vmovq xmm3,QWORD PTR [rdx+rdi] vmovq xmm4,QWORD PTR [rax] vmovq xmm5,QWORD PTR [rax+rdi] vpunpcklbw xmm2,xmm2,xmm3 ; interleave row data bytes vpunpcklbw xmm4,xmm4,xmm5 vpunpcklwd xmm1,xmm2,xmm4 ; interleave row data words vpunpckhwd xmm2,xmm2,xmm4 vinserti128 ymm1,ymm1,xmm2,1 ; concatenate vector vpxor ymm8,ymm8,ymm8 VpdpbusdsYmmYmmYmm ymm8,ymm0,ymm1 test r11,r11 ; ZeroMode? jnz SkipAccumulateOutput4By8 vpaddd ymm8,ymm8,YMMWORD PTR [rbx] SkipAccumulateOutput4By8: vmovdqu YMMWORD PTR [rbx],ymm8 add rdx,8 ; advance matrix B by 8 bytes add rbx,8*4 ; advance matrix C by 8 columns sub ebp,8 ; decrement CountN jnz ProcessColumnLoop4By8 jmp AdvanceRowLoop4 ; ; Process a set of 4 columns from the 4 rows. ; ProcessColumn4By4: test ebp,4 ; (CountN & 4) != 0? jz ProcessColumn4BySmallN lea rax,[rdx+rdi*2] ; compute matrix B plus 2 rows vmovd xmm1,DWORD PTR [rdx] vpinsrd xmm1,xmm1,DWORD PTR [rdx+rdi],1 vpinsrd xmm1,xmm1,DWORD PTR [rax],2 vpinsrd xmm1,xmm1,DWORD PTR [rax+rdi],3 vpshufb xmm1,xmm1,XMMWORD PTR [MlasTranspose4x4BytesAvx] vpxor xmm8,xmm8,xmm8 VpdpbusdsXmmXmmXmm xmm8,xmm0,xmm1 test r11,r11 ; ZeroMode? jnz SkipAccumulateOutput4By4 vpaddd xmm8,xmm8,XMMWORD PTR [rbx] SkipAccumulateOutput4By4: vmovdqu XMMWORD PTR [rbx],xmm8 and ebp,3 ; (CountN & 3) != 0? jz AdvanceRowLoop4 add rdx,4 ; advance matrix B by 4 bytes add rbx,4*4 ; advance matrix C by 4 columns ; ; Process the remaining 1 to 3 columns from the 4 rows. ; ProcessColumn4BySmallN: mov DWORD PTR GemvU8S8KernelFrame.CountN[rsp],ebp vbroadcastss xmm2,DWORD PTR GemvU8S8KernelFrame.CountN[rsp] vpcmpgtd xmm2,xmm2,XMMWORD PTR [MlasMaskMoveAvx] vpxor xmm1,xmm1,xmm1 lea rax,[rdx+rdi*2] ; compute matrix B plus 2 rows cmp ebp,2 ; (CountN & 2) != 0? jb ProcessColumn4By1 vpinsrw xmm1,xmm1,WORD PTR [rdx],0 vpinsrw xmm1,xmm1,WORD PTR [rdx+rdi],2 vpinsrw xmm1,xmm1,WORD PTR [rax],4 vpinsrw xmm1,xmm1,WORD PTR [rax+rdi],6 je ComputeOutput4BySmallN vpinsrb xmm1,xmm1,BYTE PTR [rdx+2],2 vpinsrb xmm1,xmm1,BYTE PTR [rdx+rdi+2],6 vpinsrb xmm1,xmm1,BYTE PTR [rax+2],10 vpinsrb xmm1,xmm1,BYTE PTR [rax+rdi+2],14 jmp ComputeOutput4BySmallN ProcessColumn4By1: vpinsrb xmm1,xmm1,BYTE PTR [rdx],0 vpinsrb xmm1,xmm1,BYTE PTR [rdx+rdi],4 vpinsrb xmm1,xmm1,BYTE PTR [rax],8 vpinsrb xmm1,xmm1,BYTE PTR [rax+rdi],12 ComputeOutput4BySmallN: vpshufb xmm1,xmm1,XMMWORD PTR [MlasTranspose4x4BytesAvx] vpxor xmm8,xmm8,xmm8 VpdpbusdsXmmXmmXmm xmm8,xmm0,xmm1 test r11,r11 ; ZeroMode? jnz StoreOutput4BySmallN vpmaskmovd xmm3,xmm2,XMMWORD PTR [rbx] vpaddd xmm8,xmm8,xmm3 StoreOutput4BySmallN: vpmaskmovd XMMWORD PTR [rbx],xmm2,xmm8 jmp AdvanceRowLoop4 ; ; Broadcast the remaining 1 to 3 values from vector A. ; ProcessRemainingSmallK: vpxor xmm5,xmm5,xmm5 ; keep zero vector for vpinsrb/vpinsrw cmp r9d,2 jb LoadVectorASingleRemainingByte vpinsrw xmm0,xmm5,WORD PTR [rcx],0 je BroadcastVectorARemainingBytes vpinsrb xmm0,xmm0,BYTE PTR [rcx+2],2 jmp BroadcastVectorARemainingBytes LoadVectorASingleRemainingByte: vpinsrb xmm0,xmm5,BYTE PTR [rcx],0 BroadcastVectorARemainingBytes: vpshufd xmm0,xmm0,0 ; broadcast values ; ; Process a set of 4 columns from the remaining rows. ; ProcessColumnLoopSmallKBy4: cmp r10,4 jb ProcessColumnLoopSmallKBySmallN vmovd xmm1,DWORD PTR [rsi] cmp r9d,2 jb ComputeOutputSmallKBy4 vpinsrd xmm1,xmm1,DWORD PTR [rsi+rdi],1 je ComputeOutputSmallKBy4 vpinsrd xmm1,xmm1,DWORD PTR [rsi+rdi*2],2 ComputeOutputSmallKBy4: vpshufb xmm1,xmm1,XMMWORD PTR [MlasTranspose4x4BytesAvx] vpxor xmm8,xmm8,xmm8 VpdpbusdsXmmXmmXmm xmm8,xmm0,xmm1 test r11,r11 ; ZeroMode? jnz SkipAccumulateOutputSmallKBy4 vpaddd xmm8,xmm8,XMMWORD PTR [r8] SkipAccumulateOutputSmallKBy4: vmovdqu XMMWORD PTR [r8],xmm8 add rsi,4 ; advance matrix B by 4 bytes add r8,4*4 ; advance matrix C by 4 columns sub r10,4 ; decrement CountN jnz ProcessColumnLoopSmallKBy4 jmp ExitKernel ; ; Process the remaining 1 to 3 columns from the remaining rows. ; ; Single step through each of the columns to keep code size small for the ; uncommon path (typically the row count is a multiple of 4). ; ProcessColumnLoopSmallKBySmallN: vpinsrb xmm1,xmm5,BYTE PTR [rsi],0 cmp r9d,2 jb ComputeOutputSmallKBySmallN vpinsrb xmm1,xmm1,BYTE PTR [rsi+rdi],1 je ComputeOutputSmallKBySmallN vpinsrb xmm1,xmm1,BYTE PTR [rsi+rdi*2],2 ComputeOutputSmallKBySmallN: vpxor xmm8,xmm8,xmm8 VpdpbusdsXmmXmmXmm xmm8,xmm0,xmm1 test r11,r11 ; ZeroMode? jnz SkipAccumulateOutputSmallKBySmallN vmovd xmm3,DWORD PTR [r8] vpaddd xmm8,xmm8,xmm3 SkipAccumulateOutputSmallKBySmallN: vmovd DWORD PTR [r8],xmm8 inc rsi ; advance matrix B by 1 byte add r8,4 ; advance matrix C by 1 column dec r10 jnz ProcessColumnLoopSmallKBySmallN jmp ExitKernel NESTED_END MlasGemvU8S8KernelAvxVnni, _TEXT END