saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

MLAS: Implement U8S8 GEMV kernels (#2069)

Browse source 
This implements an optimization for U8S8 MlasGemm when M=1, aka GEMV.
This commit is contained in:
Tracy Sharpe 2019-10-09 11:54:16 -07:00 committed by GitHub
parent eee9c55030
commit 57e0099425
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 1679 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
check_512.s
View file

@ -1,7 +0,0 @@
.intel_syntax noprefix
infiniteLoop:
jmp main
main:
vxorpd zmm0,zmm0,zmm0
jmp infiniteLoop

27
cmake/onnxruntime_mlas.cmake
View file

@ -46,8 +46,11 @@ if(MSVC)
set(mlas_platform_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemvU8S8KernelAvx2.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8S8KernelAvx512BW.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemvU8S8KernelAvx512BW.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8S8KernelAvx512Vnni.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemvU8S8KernelAvx512Vnni.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx512BW.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Vnni.asm
@ -170,6 +173,7 @@ else()
set(mlas_platform_srcs_avx2
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemvU8S8KernelAvx2.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/DgemmKernelFma3.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/SgemmKernelFma3.S
@ -180,9 +184,8 @@ else()
)
set_source_files_properties(${mlas_platform_srcs_avx2} PROPERTIES COMPILE_FLAGS "-mavx2 -mfma")
# Some platforms do not support 512xx flags but still able to compile the source
# Others support the flag and refuse to compile w/o the flag.
# Some platforms do not support AVX512 flags but still able to compile the source
# Others support the flag and refuse to compile without the flag.
# We have to run all 3 checks
check_cxx_compiler_flag("-mavx512f" HAS_AVX512F)
if(HAS_AVX512F)
@ -194,13 +197,13 @@ else()
endif()
set(CMAKE_REQUIRED_FLAGS ${AVX512_NEEDED})
check_cxx_source_compiles(
"int main() {
asm(\"vpxord %zmm0,%zmm0,%zmm0\");
return 0;
}"
AVX512_COMPILES)
check_cxx_source_compiles("
int main() {
asm(\"vpxord %zmm0,%zmm0,%zmm0\");
return 0;
}"
AVX512_COMPILES
)
if(AVX512_COMPILES)
set(mlas_platform_srcs_avx512f
@ -215,7 +218,9 @@ AVX512_COMPILES)
set(mlas_platform_srcs_avx512bw
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8S8KernelAvx512BW.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemvU8S8KernelAvx512BW.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8S8KernelAvx512Vnni.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemvU8S8KernelAvx512Vnni.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512BW.S
${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Vnni.S
)
@ -223,7 +228,7 @@ AVX512_COMPILES)
set_source_files_properties(${mlas_platform_srcs_avx512bw} PROPERTIES COMPILE_FLAGS "-mavx512bw")
endif()
else()
# Do not compile CPP support for 512xx
# Disable platform support for AVX512.
set_source_files_properties(${mlas_common_srcs} PROPERTIES COMPILE_FLAGS "-DMLAS_AVX512_UNSUPPORTED")
endif()

8
onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm
View file

@ -548,8 +548,8 @@ InterleaveRowDataN16:
vpunpckhwd xmm6,xmm6,xmm2
vpunpcklwd xmm2,xmm3,xmm5
vpunpckhwd xmm3,xmm3,xmm5
vinsertf128 ymm4,ymm4,xmm6,1
vinsertf128 ymm2,ymm2,xmm3,1
vinserti128 ymm4,ymm4,xmm6,1
vinserti128 ymm2,ymm2,xmm3,1
vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows
vmovdqu YMMWORD PTR [rcx+32],ymm2
vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row
@ -698,8 +698,8 @@ ProcessPaddedMatrixBData:
vpunpckhwd xmm6,xmm6,xmm2
vpunpcklwd xmm2,xmm3,xmm5
vpunpckhwd xmm3,xmm3,xmm5
vinsertf128 ymm4,ymm4,xmm6,1
vinsertf128 ymm2,ymm2,xmm3,1
vinserti128 ymm4,ymm4,xmm6,1
vinserti128 ymm2,ymm2,xmm3,1
vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows
vmovdqu YMMWORD PTR [rcx+32],ymm2
vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row

374
onnxruntime/core/mlas/lib/amd64/QgemvU8S8KernelAvx2.asm Normal file
View file

@ -0,0 +1,374 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; QgemvU8S8KernelAvx2.asm
;
; Abstract:
;
; This module implements the kernels for the quantized integer matrix/vector
; multiply operation (QGEMV).
;
; This implementation uses AVX2 instructions.
;
;--
.xlist
INCLUDE mlasi.inc
.list
EXTERN MlasMaskMoveAvx:NEAR
EXTERN MlasTranspose4x4BytesAvx:NEAR
;
; Stack frame layout for the U8S8 kernel.
;
GemvU8S8KernelFrame STRUCT
SavedXmm6 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 MlasGemvU8S8KernelAvx2, _TEXT
rex_push_reg rbp
push_reg rbx
push_reg rsi
push_reg rdi
alloc_stack (GemvU8S8KernelFrame.SavedRdi)
save_xmm128 xmm6,GemvU8S8KernelFrame.SavedXmm6
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
vpcmpeqw ymm6,ymm6,ymm6 ; generate word vector [0xFFFF]
vpsrlw ymm6,ymm6,15 ; generate word vector [0x0001]
;
; 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
vpunpcklbw ymm3,ymm4,ymm5
vpunpckhbw ymm4,ymm4,ymm5
vpunpcklwd ymm5,ymm1,ymm3 ; interleave row data words
vpunpckhwd ymm1,ymm1,ymm3
vpunpcklwd ymm3,ymm2,ymm4
vpunpckhwd ymm2,ymm2,ymm4
vpmaddubsw ymm5,ymm0,ymm5 ; multiply and reduce
vpmaddwd ymm5,ymm5,ymm6
vpmaddubsw ymm1,ymm0,ymm1
vpmaddwd ymm1,ymm1,ymm6
vpmaddubsw ymm3,ymm0,ymm3
vpmaddwd ymm3,ymm3,ymm6
vpmaddubsw ymm2,ymm0,ymm2
vpmaddwd ymm2,ymm2,ymm6
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4By32
vpaddd ymm5,ymm5,YMMWORD PTR [rbx]
vpaddd ymm1,ymm1,YMMWORD PTR [rbx+32]
vpaddd ymm3,ymm3,YMMWORD PTR [rbx+64]
vpaddd ymm2,ymm2,YMMWORD PTR [rbx+96]
SkipAccumulateOutput4By32:
cmp r9,4 ; final 4x32 block?
jae StoreOutput4By32
vperm2i128 ymm4,ymm5,ymm1,31h ; interleave vector results
vperm2i128 ymm5,ymm5,ymm1,20h
vperm2i128 ymm1,ymm3,ymm2,20h
vperm2i128 ymm2,ymm3,ymm2,31h
vmovaps ymm3,ymm4
StoreOutput4By32:
vmovdqu YMMWORD PTR [rbx],ymm5
vmovdqu YMMWORD PTR [rbx+32],ymm1
vmovdqu YMMWORD PTR [rbx+64],ymm3
vmovdqu YMMWORD PTR [rbx+96],ymm2
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]
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
vpmaddubsw ymm1,ymm0,ymm1 ; multiply and reduce
vpmaddwd ymm1,ymm1,ymm6
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4By8
vpaddd ymm1,ymm1,YMMWORD PTR [rbx]
SkipAccumulateOutput4By8:
vmovdqu YMMWORD PTR [rbx],ymm1
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]
vpmaddubsw xmm1,xmm0,xmm1 ; multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4By4
vpaddd xmm1,xmm1,XMMWORD PTR [rbx]
SkipAccumulateOutput4By4:
vmovdqu XMMWORD PTR [rbx],xmm1
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]
vpmaddubsw xmm1,xmm0,xmm1 ; multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 ; ZeroMode?
jnz StoreOutput4BySmallN
vpmaskmovd xmm3,xmm2,XMMWORD PTR [rbx]
vpaddd xmm1,xmm1,xmm3
StoreOutput4BySmallN:
vpmaskmovd XMMWORD PTR [rbx],xmm2,xmm1
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]
vpmaddubsw xmm1,xmm0,xmm1 ; multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutputSmallKBy4
vpaddd xmm1,xmm1,XMMWORD PTR [r8]
SkipAccumulateOutputSmallKBy4:
vmovdqu XMMWORD PTR [r8],xmm1
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:
vpmaddubsw xmm1,xmm0,xmm1 ; multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutputSmallKBySmallN
vmovd xmm3,DWORD PTR [r8]
vpaddd xmm1,xmm1,xmm3
SkipAccumulateOutputSmallKBySmallN:
vmovd DWORD PTR [r8],xmm1
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 MlasGemvU8S8KernelAvx2, _TEXT
END

31
onnxruntime/core/mlas/lib/amd64/QgemvU8S8KernelAvx512BW.asm Normal file
View file

@ -0,0 +1,31 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; QgemvU8S8KernelAvx512BW.asm
;
; Abstract:
;
; This module implements the kernels for the quantized integer matrix/vector
; multiply operation (QGEMV).
;
; This implementation uses AVX512BW instructions.
;
;--
.xlist
INCLUDE mlasi.inc
INCLUDE QgemvU8S8KernelAvx512Common.inc
.list
;
; Generate the GEMV kernel.
;
GemvU8S8KernelAvx512Function Avx512BW
END

372
onnxruntime/core/mlas/lib/amd64/QgemvU8S8KernelAvx512Common.inc Normal file
View file

@ -0,0 +1,372 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; QgemvU8S8KernelAvx512Common.inc
;
; Abstract:
;
; This module contains common kernel macros and structures for the quantized
; integer matrix/vector multiply operation (QGEMV) for the AVX512BW and
; AVX512VNNI kernels.
;
;--
GemvU8S8KernelFrame STRUCT
SavedRdi QWORD ?
SavedRsi QWORD ?
SavedRbx QWORD ?
SavedRbp QWORD ?
ReturnAddress QWORD ?
PreviousP1Home QWORD ?
PreviousP2Home QWORD ?
PreviousP3Home QWORD ?
PreviousP4Home QWORD ?
CountN QWORD ?
ldb QWORD ?
GemvU8S8KernelFrame ENDS
;
; Macro Description:
;
; This macro generates the common AVX512 code for the inner kernel to compute
; matrix/vector multiplication.
;
; Arguments:
;
; Isa - Supplies the instruction set architecture string for function tags.
;
GemvU8S8KernelAvx512Function MACRO Isa
;++
;
; 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 MlasGemvU8S8Kernel&Isa&, _TEXT
rex_push_reg rbp
push_reg rbx
push_reg rsi
push_reg rdi
END_PROLOGUE
mov rdi,rcx
mov rsi,rdx
mov r10,GemvU8S8KernelFrame.CountN[rsp]
mov ecx,r10d
and ecx,15 ; isolate unaligned count
mov eax,1
shl eax,cl
dec eax
kmovw k1,eax ; compute vector load/store mask
mov rcx,GemvU8S8KernelFrame.ldb[rsp]
mov r11,rsp ; set ZeroMode to any non-zero value
IFIDNI <Isa>, <Avx512BW>
mov eax,1
vpbroadcastw zmm29,eax
ENDIF
;
; Process 4 rows of matrix B in a loop.
;
sub r9,4
jb ProcessRemainingRows
ProcessRowLoop4:
mov rdx,rsi ; reload matrix B
lea rsi,[rsi+rcx*4] ; advance matrix B by 4 rows
mov rbx,r8 ; reload matrix C
mov rbp,r10 ; reload CountN
vpbroadcastd zmm28,DWORD PTR [rdi]
add rdi,4 ; advance matrix A by 4 bytes
;
; Process sets of 64 columns from the 4 rows in a loop.
;
; Some permute operations are deferred until the final store of the 4x64 block
; as these permutes are expensive.
;
ProcessColumnLoop4By64:
cmp rbp,64
jb ProcessColumnLoop4By16
lea rax,[rdx+rcx*2] ; compute matrix B plus 2 rows
vmovdqu32 zmm16,ZMMWORD PTR [rdx]
vmovdqu32 zmm17,ZMMWORD PTR [rdx+rcx]
vmovdqu32 zmm18,ZMMWORD PTR [rax]
vmovdqu32 zmm19,ZMMWORD PTR [rax+rcx]
vpunpcklbw zmm20,zmm16,zmm17 ; interleave row data bytes
vpunpckhbw zmm21,zmm16,zmm17
vpunpcklbw zmm22,zmm18,zmm19
vpunpckhbw zmm23,zmm18,zmm19
vpunpcklwd zmm16,zmm20,zmm22 ; interleave row data words
vpunpckhwd zmm17,zmm20,zmm22
vpunpcklwd zmm18,zmm21,zmm23
vpunpckhwd zmm19,zmm21,zmm23
IFIDNI <Isa>, <Avx512BW>
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
vpmaddubsw zmm17,zmm28,zmm17
vpmaddwd zmm21,zmm17,zmm29
vpmaddubsw zmm18,zmm28,zmm18
vpmaddwd zmm22,zmm18,zmm29
vpmaddubsw zmm19,zmm28,zmm19
vpmaddwd zmm23,zmm19,zmm29
ELSE
vpxord zmm20,zmm20,zmm20
vpxord zmm21,zmm21,zmm21
vpxord zmm22,zmm22,zmm22
vpxord zmm23,zmm23,zmm23
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
VpdpbusdsZmmZmmZmm zmm21,zmm28,zmm17
VpdpbusdsZmmZmmZmm zmm22,zmm28,zmm18
VpdpbusdsZmmZmmZmm zmm23,zmm28,zmm19
ENDIF
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4By64
vpaddd zmm20,zmm20,ZMMWORD PTR [rbx]
vpaddd zmm21,zmm21,ZMMWORD PTR [rbx+16*4]
vpaddd zmm22,zmm22,ZMMWORD PTR [rbx+32*4]
vpaddd zmm23,zmm23,ZMMWORD PTR [rbx+48*4]
SkipAccumulateOutput4By64:
cmp r9,4 ; final 4x64 block?
jae StoreOutput4By64
vextracti32x4 XMMWORD PTR [rbx],zmm20,0
vextracti32x4 XMMWORD PTR [rbx+4*4],zmm21,0
vextracti32x4 XMMWORD PTR [rbx+8*4],zmm22,0
vextracti32x4 XMMWORD PTR [rbx+12*4],zmm23,0
vextracti32x4 XMMWORD PTR [rbx+16*4],zmm20,1
vextracti32x4 XMMWORD PTR [rbx+20*4],zmm21,1
vextracti32x4 XMMWORD PTR [rbx+24*4],zmm22,1
vextracti32x4 XMMWORD PTR [rbx+28*4],zmm23,1
vextracti32x4 XMMWORD PTR [rbx+32*4],zmm20,2
vextracti32x4 XMMWORD PTR [rbx+36*4],zmm21,2
vextracti32x4 XMMWORD PTR [rbx+40*4],zmm22,2
vextracti32x4 XMMWORD PTR [rbx+44*4],zmm23,2
vextracti32x4 XMMWORD PTR [rbx+48*4],zmm20,3
vextracti32x4 XMMWORD PTR [rbx+52*4],zmm21,3
vextracti32x4 XMMWORD PTR [rbx+56*4],zmm22,3
vextracti32x4 XMMWORD PTR [rbx+60*4],zmm23,3
jmp AdvanceColumnLoop64
StoreOutput4By64:
vmovdqu32 ZMMWORD PTR [rbx],zmm20
vmovdqu32 ZMMWORD PTR [rbx+16*4],zmm21
vmovdqu32 ZMMWORD PTR [rbx+32*4],zmm22
vmovdqu32 ZMMWORD PTR [rbx+48*4],zmm23
AdvanceColumnLoop64:
add rdx,64 ; advance matrix B by 64 bytes
add rbx,64*4 ; advance matrix C by 64 columns
sub rbp,64 ; decrement CountN
jnz ProcessColumnLoop4By64
AdvanceRowLoop4:
xor r11,r11 ; clear ZeroMode
sub r9,4 ; decrement CountK
jae ProcessRowLoop4
ProcessRemainingRows:
add r9,4 ; correct for over-subtract above
jnz ProcessRemainingSmallK
ExitKernel:
vzeroupper
BEGIN_EPILOGUE
pop rdi
pop rsi
pop rbx
pop rbp
ret
;
; Process sets of 16 columns from the 4 rows in a loop or process the remaining
; 1 to 15 columns.
;
ProcessColumnLoop4By16:
lea rax,[rdx+rcx*2] ; compute matrix B plus 2 rows
cmp ebp,16
jb LoadPartialVector4BySmallN
vmovdqu xmm2,XMMWORD PTR [rdx]
vmovdqu xmm3,XMMWORD PTR [rdx+rcx]
vmovdqu xmm4,XMMWORD PTR [rax]
vmovdqu xmm5,XMMWORD PTR [rax+rcx]
jmp ComputeOutput4By16
LoadPartialVector4BySmallN:
vmovdqu8 zmm2{k1}{z},ZMMWORD PTR [rdx]
vmovdqu8 zmm3{k1}{z},ZMMWORD PTR [rdx+rcx]
vmovdqu8 zmm4{k1}{z},ZMMWORD PTR [rax]
vmovdqu8 zmm5{k1}{z},ZMMWORD PTR [rax+rcx]
ComputeOutput4By16:
vpunpcklbw xmm1,xmm2,xmm3 ; interleave row data bytes
vpunpckhbw xmm2,xmm2,xmm3
vpunpcklbw xmm3,xmm4,xmm5
vpunpckhbw xmm4,xmm4,xmm5
vpunpcklwd xmm5,xmm1,xmm3 ; interleave row data words
vpunpckhwd xmm1,xmm1,xmm3
vpunpcklwd xmm3,xmm2,xmm4
vpunpckhwd xmm2,xmm2,xmm4
vinserti128 ymm5,ymm5,xmm1,1 ; concatenate 256-bit vector
vinserti128 ymm3,ymm3,xmm2,1
vshufi32x4 zmm16,zmm5,zmm3,044h ; concatenate 512-bit vector
IFIDNI <Isa>, <Avx512BW>
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
ELSE
vpxord zmm20,zmm20,zmm20
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
ENDIF
cmp ebp,16
jb StorePartialVector4BySmallN
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4By16
vpaddd zmm20,zmm20,ZMMWORD PTR [rbx]
SkipAccumulateOutput4By16:
vmovdqu32 ZMMWORD PTR [rbx],zmm20
add rdx,16 ; advance matrix B by 16 bytes
add rbx,16*4 ; advance matrix C by 16 columns
sub ebp,16 ; decrement CountN
jnz ProcessColumnLoop4By16
jmp AdvanceRowLoop4
StorePartialVector4BySmallN:
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutput4BySmallN
vpaddd zmm20{k1}{z},zmm20,ZMMWORD PTR [rbx]
SkipAccumulateOutput4BySmallN:
vmovdqu32 ZMMWORD PTR [rbx]{k1},zmm20
jmp AdvanceRowLoop4
;
; Broadcast the remaining 1 to 3 values from vector A.
;
ProcessRemainingSmallK:
vpxor xmm0,xmm0,xmm0
cmp r9d,2
jb LoadVectorASingleRemainingByte
vpinsrw xmm0,xmm0,WORD PTR [rdi],0
je BroadcastVectorARemainingBytes
vpinsrb xmm0,xmm0,BYTE PTR [rdi+2],2
jmp BroadcastVectorARemainingBytes
LoadVectorASingleRemainingByte:
vpinsrb xmm0,xmm0,BYTE PTR [rdi],0
BroadcastVectorARemainingBytes:
vpbroadcastd zmm28,xmm0 ; broadcast values
;
; Process sets of 16 columns from the remaining rows in a loop or process the
; remaining 1 to 15 columns.
;
ProcessColumnLoopSmallKBy16:
vpxor xmm3,xmm3,xmm3 ; clear optional row vectors
vpxor xmm4,xmm4,xmm4
vpxor xmm5,xmm5,xmm5
cmp r10d,16
jb LoadPartialVectorSmallKBySmallN
vmovdqu xmm2,XMMWORD PTR [rsi]
cmp r9d,2
jb ComputeOutputSmallKBy16
vmovdqu xmm3,XMMWORD PTR [rsi+rcx]
je ComputeOutputSmallKBy16
vmovdqu xmm4,XMMWORD PTR [rsi+rcx*2]
jmp ComputeOutputSmallKBy16
LoadPartialVectorSmallKBySmallN:
vmovdqu8 zmm2{k1}{z},ZMMWORD PTR [rsi]
cmp r9d,2
jb ComputeOutputSmallKBy16
vmovdqu8 zmm3{k1}{z},ZMMWORD PTR [rsi+rcx]
je ComputeOutputSmallKBy16
vmovdqu8 zmm4{k1}{z},ZMMWORD PTR [rsi+rcx*2]
jmp ComputeOutputSmallKBy16
ComputeOutputSmallKBy16:
vpunpcklbw xmm1,xmm2,xmm3 ; interleave row data bytes
vpunpckhbw xmm2,xmm2,xmm3
vpunpcklbw xmm3,xmm4,xmm5
vpunpckhbw xmm4,xmm4,xmm5
vpunpcklwd xmm5,xmm1,xmm3 ; interleave row data words
vpunpckhwd xmm1,xmm1,xmm3
vpunpcklwd xmm3,xmm2,xmm4
vpunpckhwd xmm2,xmm2,xmm4
vinserti128 ymm5,ymm5,xmm1,1 ; concatenate 256-bit vector
vinserti128 ymm3,ymm3,xmm2,1
vshufi32x4 zmm16,zmm5,zmm3,044h ; concatenate 512-bit vector
IFIDNI <Isa>, <Avx512BW>
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
ELSE
vpxord zmm20,zmm20,zmm20
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
ENDIF
cmp r10d,16
jb StorePartialVectorSmallKBySmallN
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutputSmallKBy16
vpaddd zmm20,zmm20,ZMMWORD PTR [r8]
SkipAccumulateOutputSmallKBy16:
vmovdqu32 ZMMWORD PTR [r8],zmm20
add rsi,16 ; advance matrix B by 16 bytes
add r8,16*4 ; advance matrix C by 16 columns
sub r10d,16 ; decrement CountN
jnz ProcessColumnLoopSmallKBy16
jmp ExitKernel
StorePartialVectorSmallKBySmallN:
test r11,r11 ; ZeroMode?
jnz SkipAccumulateOutputSmallKBySmallN
vpaddd zmm20{k1}{z},zmm20,ZMMWORD PTR [r8]
SkipAccumulateOutputSmallKBySmallN:
vmovdqu32 ZMMWORD PTR [r8]{k1},zmm20
jmp ExitKernel
NESTED_END MlasGemvU8S8Kernel&Isa&, _TEXT
ENDM

32
onnxruntime/core/mlas/lib/amd64/QgemvU8S8KernelAvx512Vnni.asm Normal file
View file

@ -0,0 +1,32 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; QgemvU8S8KernelAvx512Vnni.asm
;
; Abstract:
;
; This module implements the kernels for the quantized integer matrix/vector
; multiply operation (QGEMV).
;
; This implementation uses AVX512VNNI instructions.
;
;--
.xlist
INCLUDE mlasi.inc
INCLUDE QgemvU8S8KernelAvx512Common.inc
INCLUDE AssembleAvx512Vnni.inc
.list
;
; Generate the GEMV kernel.
;
GemvU8S8KernelAvx512Function Avx512Vnni
END

32
onnxruntime/core/mlas/lib/amd64/QgemvU8X8KernelCommon.inc Normal file
View file

@ -0,0 +1,32 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; QgemvU8X8KernelCommon.inc
;
; Abstract:
;
; This module contains common kernel macros and structures for the quantized
; integer matrix/matrix multiply operation (QGEMM) for the AVX2 kernels.
;
;--
GemvU8X8KernelFrame STRUCT
SavedRdi QWORD ?
SavedRsi QWORD ?
SavedRbx QWORD ?
SavedRbp QWORD ?
ReturnAddress QWORD ?
PreviousP1Home QWORD ?
PreviousP2Home QWORD ?
PreviousP3Home QWORD ?
PreviousP4Home QWORD ?
CountN QWORD ?
ldb QWORD ?
GemvU8X8KernelFrame ENDS

17
onnxruntime/core/mlas/lib/mlasi.h
View file

@ -274,6 +274,19 @@ size_t
typedef MLAS_GEMM_U8S8_KERNEL* PMLAS_GEMM_U8S8_KERNEL;
typedef
size_t
(MLASCALL MLAS_GEMV_U8S8_KERNEL)(
const uint8_t* A,
const int8_t* B,
int32_t* C,
size_t CountK,
size_t CountN,
size_t ldb
);
typedef MLAS_GEMV_U8S8_KERNEL* PMLAS_GEMV_U8S8_KERNEL;
typedef
void
(MLASCALL MLAS_GEMM_U8U8_COPY_PACKA_ROUTINE)(
@ -460,8 +473,11 @@ extern "C" {
MLAS_GEMM_U8S8_COPY_PACKA_ROUTINE MlasGemmU8S8CopyPackAAvx2;
MLAS_GEMM_U8S8_COPY_PACKB_ROUTINE MlasGemmU8S8CopyPackBAvx2;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx2;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx2;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512BW;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512BW;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512Vnni;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Vnni;
MLAS_GEMM_U8U8_COPY_PACKA_ROUTINE MlasGemmU8U8CopyPackAAvx2;
MLAS_GEMM_U8U8_COPY_PACKB_ROUTINE MlasGemmU8U8CopyPackBAvx2;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx2;
@ -599,6 +615,7 @@ struct MLAS_PLATFORM {
PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine;
PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine;
PMLAS_GEMM_DOUBLE_KERNEL GemmDoubleKernel;
PMLAS_GEMV_U8S8_KERNEL GemvU8S8Kernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwFloatKernel;
PMLAS_CONV_FLOAT_KERNEL ConvNchwcFloatKernel;
PMLAS_CONV_DEPTHWISE_FLOAT_KERNEL ConvDepthwiseFloatKernel;

3
onnxruntime/core/mlas/lib/platform.cpp
View file

@ -176,6 +176,7 @@ Return Value:
this->GemmU8S8CopyPackARoutine = MlasGemmU8S8CopyPackAAvx2;
this->GemmU8S8CopyPackBRoutine = MlasGemmU8S8CopyPackBAvx2;
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx2;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx2;
this->GemmU8U8CopyPackARoutine = MlasGemmU8U8CopyPackAAvx2;
this->GemmU8U8CopyPackBRoutine = MlasGemmU8U8CopyPackBAvx2;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx2;
@ -218,6 +219,7 @@ Return Value:
if ((Cpuid7[1] & 0x40000000) != 0) {
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512BW;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512BW;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx512BW;
//
@ -227,6 +229,7 @@ Return Value:
if ((Cpuid7[2] & 0x800) != 0) {
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Vnni;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Vnni;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx512Vnni;
}
}

19
onnxruntime/core/mlas/lib/qgemm.cpp
View file

@ -31,6 +31,13 @@ Abstract:
#ifdef MLAS_TARGET_AMD64_IX86
//
// Stores a vector to transpose a 4x4 byte vector using vpshufb.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint8_t MlasTranspose4x4BytesAvx[16], 16) =
{ 0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15 };
//
// U8S8 implementation using SSE2 intrinsics.
//
@ -1122,6 +1129,18 @@ MlasGemm(
MLAS_UNREFERENCED_PARAMETER(ThreadPool);
#if defined(MLAS_TARGET_AMD64)
if (M == 1 && offa == 0 && offb == 0) {
if (MlasPlatform.GemvU8S8Kernel != nullptr) {
MlasPlatform.GemvU8S8Kernel(A, B, C, K, N, ldb);
return;
}
}
#endif
size_t CountK;
for (size_t k = 0; k < K; k += CountK) {

8
onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S
View file

@ -485,8 +485,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2):
vpunpckhwd xmm6,xmm6,xmm2
vpunpcklwd xmm2,xmm3,xmm5
vpunpckhwd xmm3,xmm3,xmm5
vinsertf128 ymm4,ymm4,xmm6,1
vinsertf128 ymm2,ymm2,xmm3,1
vinserti128 ymm4,ymm4,xmm6,1
vinserti128 ymm2,ymm2,xmm3,1
vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows
vmovdqu YMMWORD PTR [rdi+32],ymm2
vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row
@ -626,8 +626,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2):
vpunpckhwd xmm6,xmm6,xmm2
vpunpcklwd xmm2,xmm3,xmm5
vpunpckhwd xmm3,xmm3,xmm5
vinsertf128 ymm4,ymm4,xmm6,1
vinsertf128 ymm2,ymm2,xmm3,1
vinserti128 ymm4,ymm4,xmm6,1
vinserti128 ymm2,ymm2,xmm3,1
vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows
vmovdqu YMMWORD PTR [rdi+32],ymm2
vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row

345
onnxruntime/core/mlas/lib/x86_64/QgemvU8S8KernelAvx2.S Normal file
View file

@ -0,0 +1,345 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
QgemvU8S8KernelAvx2.s
Abstract:
This module implements the kernels for the quantized integer matrix/vector
multiply operation (QGEMV).
This implementation uses AVX2 instructions.
--*/
#include "asmmacro.h"
.intel_syntax noprefix
//
// Stack frame layout for the U8S8 kernel.
//
.equ .LGemvU8S8KernelFrame_mask, -8
.equ .LGemvU8S8KernelFrame_SavedRbx, 0
.equ .LGemvU8S8KernelFrame_SavedRbp, 8
.equ .LGemvU8S8KernelFrame_ReturnAddress, 16
.text
/*++
Routine Description:
This routine is an inner kernel to compute matrix/vector multiplication.
Arguments:
A (rdi) - Supplies the address of vector A.
B (rsi) - Supplies the address of matrix B.
C (rdx) - Supplies the address of matrix C.
CountK (rcx) - Supplies the number of columns from vector A and the number
of rows from matrix B to iterate over.
CountN (r8) - Supplies the number of columns from matrix B and matrix C to
iterate over.
ldb (r9) - Supplies the first dimension of matrix B.
Return Value:
None.
--*/
.globl C_UNDERSCORE(MlasGemvU8S8KernelAvx2)
C_UNDERSCORE(MlasGemvU8S8KernelAvx2):
push rbp
push rbx
mov r10,rdx
mov r11,rsp # set ZeroMode to any non-zero value
vpcmpeqw ymm6,ymm6,ymm6 # generate word vector [0xFFFF]
vpsrlw ymm6,ymm6,15 # generate word vector [0x0001]
//
// Process 4 rows of matrix B in a loop.
//
sub rcx,4
jb .LProcessRemainingRows
.LProcessRowLoop4:
mov rdx,rsi # reload matrix B
lea rsi,[rsi+r9*4] # advance matrix B by 4 rows
mov rbx,r10 # reload matrix C
mov rbp,r8 # reload CountN
vpbroadcastd ymm0,DWORD PTR [rdi]
add rdi,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.
//
.LProcessColumnLoop4By32:
cmp rbp,32
jb .LProcessColumnLoop4By8
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
vmovdqu ymm2,YMMWORD PTR [rdx]
vmovdqu ymm3,YMMWORD PTR [rdx+r9]
vmovdqu ymm4,YMMWORD PTR [rax]
vmovdqu ymm5,YMMWORD PTR [rax+r9]
vpunpcklbw ymm1,ymm2,ymm3 # interleave row data bytes
vpunpckhbw ymm2,ymm2,ymm3
vpunpcklbw ymm3,ymm4,ymm5
vpunpckhbw ymm4,ymm4,ymm5
vpunpcklwd ymm5,ymm1,ymm3 # interleave row data words
vpunpckhwd ymm1,ymm1,ymm3
vpunpcklwd ymm3,ymm2,ymm4
vpunpckhwd ymm2,ymm2,ymm4
vpmaddubsw ymm5,ymm0,ymm5 # multiply and reduce
vpmaddwd ymm5,ymm5,ymm6
vpmaddubsw ymm1,ymm0,ymm1
vpmaddwd ymm1,ymm1,ymm6
vpmaddubsw ymm3,ymm0,ymm3
vpmaddwd ymm3,ymm3,ymm6
vpmaddubsw ymm2,ymm0,ymm2
vpmaddwd ymm2,ymm2,ymm6
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4By32
vpaddd ymm5,ymm5,YMMWORD PTR [rbx]
vpaddd ymm1,ymm1,YMMWORD PTR [rbx+32]
vpaddd ymm3,ymm3,YMMWORD PTR [rbx+64]
vpaddd ymm2,ymm2,YMMWORD PTR [rbx+96]
.LSkipAccumulateOutput4By32:
cmp rcx,4 # final 4x32 block?
jae .LStoreOutput4By32
vperm2i128 ymm4,ymm5,ymm1,0x31 # interleave vector results
vperm2i128 ymm5,ymm5,ymm1,0x20
vperm2i128 ymm1,ymm3,ymm2,0x20
vperm2i128 ymm2,ymm3,ymm2,0x31
vmovaps ymm3,ymm4
.LStoreOutput4By32:
vmovdqu YMMWORD PTR [rbx],ymm5
vmovdqu YMMWORD PTR [rbx+32],ymm1
vmovdqu YMMWORD PTR [rbx+64],ymm3
vmovdqu YMMWORD PTR [rbx+96],ymm2
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 .LProcessColumnLoop4By32
.LAdvanceRowLoop4:
xor r11,r11 # clear ZeroMode
sub rcx,4 # decrement CountK
jae .LProcessRowLoop4
.LProcessRemainingRows:
add rcx,4 # correct for over-subtract above
jnz .LProcessRemainingSmallK
//
// Restore non-volatile registers and return.
//
.LExitKernel:
vzeroupper
pop rbx
pop rbp
ret
//
// Process sets of 8 columns from the 4 rows in a loop.
//
.LProcessColumnLoop4By8:
cmp ebp,8
jb .LProcessColumn4By4
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
vmovq xmm2,QWORD PTR [rdx]
vmovq xmm3,QWORD PTR [rdx+r9]
vmovq xmm4,QWORD PTR [rax]
vmovq xmm5,QWORD PTR [rax+r9]
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
vpmaddubsw ymm1,ymm0,ymm1 # multiply and reduce
vpmaddwd ymm1,ymm1,ymm6
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4By8
vpaddd ymm1,ymm1,YMMWORD PTR [rbx]
.LSkipAccumulateOutput4By8:
vmovdqu YMMWORD PTR [rbx],ymm1
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 .LProcessColumnLoop4By8
jmp .LAdvanceRowLoop4
//
// Process a set of 4 columns from the 4 rows.
//
.LProcessColumn4By4:
test ebp,4 # (CountN & 4) != 0?
jz .LProcessColumn4BySmallN
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
vmovd xmm1,DWORD PTR [rdx]
vpinsrd xmm1,xmm1,DWORD PTR [rdx+r9],1
vpinsrd xmm1,xmm1,DWORD PTR [rax],2
vpinsrd xmm1,xmm1,DWORD PTR [rax+r9],3
vpshufb xmm1,xmm1,XMMWORD PTR C_UNDERSCORE(MlasTranspose4x4BytesAvx)[rip]
vpmaddubsw xmm1,xmm0,xmm1 # multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4By4
vpaddd xmm1,xmm1,XMMWORD PTR [rbx]
.LSkipAccumulateOutput4By4:
vmovdqu XMMWORD PTR [rbx],xmm1
and ebp,3 # (CountN & 3) != 0?
jz .LAdvanceRowLoop4
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.
//
.LProcessColumn4BySmallN:
mov DWORD PTR .LGemvU8S8KernelFrame_mask[rsp],ebp
vbroadcastss xmm2,DWORD PTR .LGemvU8S8KernelFrame_mask[rsp]
vpcmpgtd xmm2,xmm2,XMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip]
vpxor xmm1,xmm1,xmm1
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
cmp ebp,2 # (CountN & 2) != 0?
jb .LProcessColumn4By1
vpinsrw xmm1,xmm1,WORD PTR [rdx],0
vpinsrw xmm1,xmm1,WORD PTR [rdx+r9],2
vpinsrw xmm1,xmm1,WORD PTR [rax],4
vpinsrw xmm1,xmm1,WORD PTR [rax+r9],6
je .LComputeOutput4BySmallN
vpinsrb xmm1,xmm1,BYTE PTR [rdx+2],2
vpinsrb xmm1,xmm1,BYTE PTR [rdx+r9+2],6
vpinsrb xmm1,xmm1,BYTE PTR [rax+2],10
vpinsrb xmm1,xmm1,BYTE PTR [rax+r9+2],14
jmp .LComputeOutput4BySmallN
.LProcessColumn4By1:
vpinsrb xmm1,xmm1,BYTE PTR [rdx],0
vpinsrb xmm1,xmm1,BYTE PTR [rdx+r9],4
vpinsrb xmm1,xmm1,BYTE PTR [rax],8
vpinsrb xmm1,xmm1,BYTE PTR [rax+r9],12
.LComputeOutput4BySmallN:
vpshufb xmm1,xmm1,XMMWORD PTR C_UNDERSCORE(MlasTranspose4x4BytesAvx)[rip]
vpmaddubsw xmm1,xmm0,xmm1 # multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 # ZeroMode?
jnz .LStoreOutput4BySmallN
vpmaskmovd xmm3,xmm2,XMMWORD PTR [rbx]
vpaddd xmm1,xmm1,xmm3
.LStoreOutput4BySmallN:
vpmaskmovd XMMWORD PTR [rbx],xmm2,xmm1
jmp .LAdvanceRowLoop4
//
// Broadcast the remaining 1 to 3 values from vector A.
//
.LProcessRemainingSmallK:
vpxor xmm5,xmm5,xmm5 # keep zero vector for vpinsrb/vpinsrw
cmp ecx,2
jb .LLoadVectorASingleRemainingByte
vpinsrw xmm0,xmm5,WORD PTR [rdi],0
je .LBroadcastVectorARemainingBytes
vpinsrb xmm0,xmm0,BYTE PTR [rdi+2],2
jmp .LBroadcastVectorARemainingBytes
.LLoadVectorASingleRemainingByte:
vpinsrb xmm0,xmm5,BYTE PTR [rdi],0
.LBroadcastVectorARemainingBytes:
vpshufd xmm0,xmm0,0 # broadcast values
//
// Process a set of 4 columns from the remaining rows.
//
.LProcessColumnLoopSmallKBy4:
cmp r8d,4
jb .LProcessColumnLoopSmallKBySmallN
vmovd xmm1,DWORD PTR [rsi]
cmp ecx,2
jb .LComputeOutputSmallKBy4
vpinsrd xmm1,xmm1,DWORD PTR [rsi+r9],1
je .LComputeOutputSmallKBy4
vpinsrd xmm1,xmm1,DWORD PTR [rsi+r9*2],2
.LComputeOutputSmallKBy4:
vpshufb xmm1,xmm1,XMMWORD PTR C_UNDERSCORE(MlasTranspose4x4BytesAvx)[rip]
vpmaddubsw xmm1,xmm0,xmm1 # multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutputSmallKBy4
vpaddd xmm1,xmm1,XMMWORD PTR [r10]
.LSkipAccumulateOutputSmallKBy4:
vmovdqu XMMWORD PTR [r10],xmm1
add rsi,4 # advance matrix B by 4 bytes
add r10,4*4 # advance matrix C by 4 columns
sub r8d,4 # decrement CountN
jnz .LProcessColumnLoopSmallKBy4
jmp .LExitKernel
//
// 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).
//
.LProcessColumnLoopSmallKBySmallN:
vpinsrb xmm1,xmm5,BYTE PTR [rsi],0
cmp ecx,2
jb .LComputeOutputSmallKBySmallN
vpinsrb xmm1,xmm1,BYTE PTR [rsi+r9],1
je .LComputeOutputSmallKBySmallN
vpinsrb xmm1,xmm1,BYTE PTR [rsi+r9*2],2
.LComputeOutputSmallKBySmallN:
vpmaddubsw xmm1,xmm0,xmm1 # multiply and reduce
vpmaddwd xmm1,xmm1,xmm6
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutputSmallKBySmallN
vmovd xmm3,DWORD PTR [r10]
vpaddd xmm1,xmm1,xmm3
.LSkipAccumulateOutputSmallKBySmallN:
vmovd DWORD PTR [r10],xmm1
inc rsi # advance matrix B by 1 byte
add r10,4 # advance matrix C by 1 column
dec r8
jnz .LProcessColumnLoopSmallKBySmallN
jmp .LExitKernel
.end

33
onnxruntime/core/mlas/lib/x86_64/QgemvU8S8KernelAvx512BW.S Normal file
View file

@ -0,0 +1,33 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
QgemvU8S8KernelAvx512BW.s
Abstract:
This module implements the kernels for the quantized integer matrix/vector
multiply operation (QGEMV).
This implementation uses AVX512BW instructions.
--*/
#include "asmmacro.h"
#include "QgemvU8S8KernelAvx512Common.h"
.intel_syntax noprefix
.text
//
// Generate the GEMV kernel.
//
GemvU8S8KernelAvx512Function Avx512BW
.end

356
onnxruntime/core/mlas/lib/x86_64/QgemvU8S8KernelAvx512Common.h Normal file
View file

@ -0,0 +1,356 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
QgemvU8S8KernelAvx512Common.h
Abstract:
This module contains common kernel macros and structures for the quantized
integer matrix/vector multiply operation (QGEMV) for the AVX512BW and
AVX512VNNI kernels.
--*/
//
// Stack frame layout for the U8S8 kernel.
//
.equ .LGemvU8S8KernelFrame_SavedRbx, 0
.equ .LGemvU8S8KernelFrame_SavedRbp, 8
.equ .LGemvU8S8KernelFrame_ReturnAddress, 16
/*++
Macro Description:
This macro generates the common AVX512 code for the inner kernel to compute
matrix/vector multiplication.
Arguments:
Isa - Supplies the instruction set architecture string for function tags.
--*/
.macro GemvU8S8KernelAvx512Function Isa
/*++
Routine Description:
This routine is an inner kernel to compute matrix/vector multiplication.
Arguments:
A (rdi) - Supplies the address of vector A.
B (rsi) - Supplies the address of matrix B.
C (rdx) - Supplies the address of matrix C.
CountK (rcx) - Supplies the number of columns from vector A and the number
of rows from matrix B to iterate over.
CountN (r8) - Supplies the number of columns from matrix B and matrix C to
iterate over.
ldb (r9) - Supplies the first dimension of matrix B.
Return Value:
None.
--*/
.globl C_UNDERSCORE(MlasGemvU8S8Kernel\Isa\())
C_UNDERSCORE(MlasGemvU8S8Kernel\Isa\()):
push rbp
push rbx
mov rbx,rcx
mov ecx,r8d
and ecx,15 # isolate unaligned count
mov eax,1
shl eax,cl
dec eax
kmovw k1,eax # compute vector load/store mask
mov rcx,rbx
mov r10,rdx
mov r11,rsp # set ZeroMode to any non-zero value
.ifeqs "\Isa\()", "Avx512BW"
mov eax,1
vpbroadcastw zmm29,eax
.endif
//
// Process 4 rows of matrix B in a loop.
//
sub rcx,4
jb .LProcessRemainingRows
.LProcessRowLoop4:
mov rdx,rsi # reload matrix B
lea rsi,[rsi+r9*4] # advance matrix B by 4 rows
mov rbx,r10 # reload matrix C
mov rbp,r8 # reload CountN
vpbroadcastd zmm28,DWORD PTR [rdi]
add rdi,4 # advance matrix A by 4 bytes
//
// Process sets of 64 columns from the 4 rows in a loop.
//
// Some permute operations are deferred until the final store of the 4x64 block
// as these permutes are expensive.
//
.LProcessColumnLoop4By64:
cmp rbp,64
jb .LProcessColumnLoop4By16
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
vmovdqu32 zmm16,ZMMWORD PTR [rdx]
vmovdqu32 zmm17,ZMMWORD PTR [rdx+r9]
vmovdqu32 zmm18,ZMMWORD PTR [rax]
vmovdqu32 zmm19,ZMMWORD PTR [rax+r9]
vpunpcklbw zmm20,zmm16,zmm17 # interleave row data bytes
vpunpckhbw zmm21,zmm16,zmm17
vpunpcklbw zmm22,zmm18,zmm19
vpunpckhbw zmm23,zmm18,zmm19
vpunpcklwd zmm16,zmm20,zmm22 # interleave row data words
vpunpckhwd zmm17,zmm20,zmm22
vpunpcklwd zmm18,zmm21,zmm23
vpunpckhwd zmm19,zmm21,zmm23
.ifeqs "\Isa\()", "Avx512BW"
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
vpmaddubsw zmm17,zmm28,zmm17
vpmaddwd zmm21,zmm17,zmm29
vpmaddubsw zmm18,zmm28,zmm18
vpmaddwd zmm22,zmm18,zmm29
vpmaddubsw zmm19,zmm28,zmm19
vpmaddwd zmm23,zmm19,zmm29
.else
vpxord zmm20,zmm20,zmm20
vpxord zmm21,zmm21,zmm21
vpxord zmm22,zmm22,zmm22
vpxord zmm23,zmm23,zmm23
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
VpdpbusdsZmmZmmZmm zmm21,zmm28,zmm17
VpdpbusdsZmmZmmZmm zmm22,zmm28,zmm18
VpdpbusdsZmmZmmZmm zmm23,zmm28,zmm19
.endif
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4By64
vpaddd zmm20,zmm20,ZMMWORD PTR [rbx]
vpaddd zmm21,zmm21,ZMMWORD PTR [rbx+16*4]
vpaddd zmm22,zmm22,ZMMWORD PTR [rbx+32*4]
vpaddd zmm23,zmm23,ZMMWORD PTR [rbx+48*4]
.LSkipAccumulateOutput4By64:
cmp rcx,4 # final 4x64 block?
jae .LStoreOutput4By64
vextracti32x4 XMMWORD PTR [rbx],zmm20,0
vextracti32x4 XMMWORD PTR [rbx+4*4],zmm21,0
vextracti32x4 XMMWORD PTR [rbx+8*4],zmm22,0
vextracti32x4 XMMWORD PTR [rbx+12*4],zmm23,0
vextracti32x4 XMMWORD PTR [rbx+16*4],zmm20,1
vextracti32x4 XMMWORD PTR [rbx+20*4],zmm21,1
vextracti32x4 XMMWORD PTR [rbx+24*4],zmm22,1
vextracti32x4 XMMWORD PTR [rbx+28*4],zmm23,1
vextracti32x4 XMMWORD PTR [rbx+32*4],zmm20,2
vextracti32x4 XMMWORD PTR [rbx+36*4],zmm21,2
vextracti32x4 XMMWORD PTR [rbx+40*4],zmm22,2
vextracti32x4 XMMWORD PTR [rbx+44*4],zmm23,2
vextracti32x4 XMMWORD PTR [rbx+48*4],zmm20,3
vextracti32x4 XMMWORD PTR [rbx+52*4],zmm21,3
vextracti32x4 XMMWORD PTR [rbx+56*4],zmm22,3
vextracti32x4 XMMWORD PTR [rbx+60*4],zmm23,3
jmp .LAdvanceColumnLoop64
.LStoreOutput4By64:
vmovdqu32 ZMMWORD PTR [rbx],zmm20
vmovdqu32 ZMMWORD PTR [rbx+16*4],zmm21
vmovdqu32 ZMMWORD PTR [rbx+32*4],zmm22
vmovdqu32 ZMMWORD PTR [rbx+48*4],zmm23
.LAdvanceColumnLoop64:
add rdx,64 # advance matrix B by 64 bytes
add rbx,64*4 # advance matrix C by 64 columns
sub rbp,64 # decrement CountN
jnz .LProcessColumnLoop4By64
.LAdvanceRowLoop4:
xor r11,r11 # clear ZeroMode
sub rcx,4 # decrement CountK
jae .LProcessRowLoop4
.LProcessRemainingRows:
add rcx,4 # correct for over-subtract above
jnz .LProcessRemainingSmallK
.LExitKernel:
vzeroupper
pop rbx
pop rbp
ret
//
// Process sets of 16 columns from the 4 rows in a loop or process the remaining
// 1 to 15 columns.
//
.LProcessColumnLoop4By16:
lea rax,[rdx+r9*2] # compute matrix B plus 2 rows
cmp ebp,16
jb .LLoadPartialVector4BySmallN
vmovdqu xmm2,XMMWORD PTR [rdx]
vmovdqu xmm3,XMMWORD PTR [rdx+r9]
vmovdqu xmm4,XMMWORD PTR [rax]
vmovdqu xmm5,XMMWORD PTR [rax+r9]
jmp .LComputeOutput4By16
.LLoadPartialVector4BySmallN:
vmovdqu8 zmm2{k1}{z},ZMMWORD PTR [rdx]
vmovdqu8 zmm3{k1}{z},ZMMWORD PTR [rdx+r9]
vmovdqu8 zmm4{k1}{z},ZMMWORD PTR [rax]
vmovdqu8 zmm5{k1}{z},ZMMWORD PTR [rax+r9]
.LComputeOutput4By16:
vpunpcklbw xmm1,xmm2,xmm3 # interleave row data bytes
vpunpckhbw xmm2,xmm2,xmm3
vpunpcklbw xmm3,xmm4,xmm5
vpunpckhbw xmm4,xmm4,xmm5
vpunpcklwd xmm5,xmm1,xmm3 # interleave row data words
vpunpckhwd xmm1,xmm1,xmm3
vpunpcklwd xmm3,xmm2,xmm4
vpunpckhwd xmm2,xmm2,xmm4
vinserti128 ymm5,ymm5,xmm1,1 # concatenate 256-bit vector
vinserti128 ymm3,ymm3,xmm2,1
vshufi32x4 zmm16,zmm5,zmm3,0x44 # concatenate 512-bit vector
.ifeqs "\Isa\()", "Avx512BW"
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
.else
vpxord zmm20,zmm20,zmm20
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
.endif
cmp ebp,16
jb .LStorePartialVector4BySmallN
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4By16
vpaddd zmm20,zmm20,ZMMWORD PTR [rbx]
.LSkipAccumulateOutput4By16:
vmovdqu32 ZMMWORD PTR [rbx],zmm20
add rdx,16 # advance matrix B by 16 bytes
add rbx,16*4 # advance matrix C by 16 columns
sub ebp,16 # decrement CountN
jnz .LProcessColumnLoop4By16
jmp .LAdvanceRowLoop4
.LStorePartialVector4BySmallN:
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutput4BySmallN
vpaddd zmm20{k1}{z},zmm20,ZMMWORD PTR [rbx]
.LSkipAccumulateOutput4BySmallN:
vmovdqu32 ZMMWORD PTR [rbx]{k1},zmm20
jmp .LAdvanceRowLoop4
//
// Broadcast the remaining 1 to 3 values from vector A.
//
.LProcessRemainingSmallK:
vpxor xmm0,xmm0,xmm0
cmp ecx,2
jb .LLoadVectorASingleRemainingByte
vpinsrw xmm0,xmm0,WORD PTR [rdi],0
je .LBroadcastVectorARemainingBytes
vpinsrb xmm0,xmm0,BYTE PTR [rdi+2],2
jmp .LBroadcastVectorARemainingBytes
.LLoadVectorASingleRemainingByte:
vpinsrb xmm0,xmm0,BYTE PTR [rdi],0
.LBroadcastVectorARemainingBytes:
vpbroadcastd zmm28,xmm0 # broadcast values
//
// Process sets of 16 columns from the remaining rows in a loop or process the
// remaining 1 to 15 columns.
//
.LProcessColumnLoopSmallKBy16:
vpxor xmm3,xmm3,xmm3 # clear optional row vectors
vpxor xmm4,xmm4,xmm4
vpxor xmm5,xmm5,xmm5
cmp r8d,16
jb .LLoadPartialVectorSmallKBySmallN
vmovdqu xmm2,XMMWORD PTR [rsi]
cmp ecx,2
jb .LComputeOutputSmallKBy16
vmovdqu xmm3,XMMWORD PTR [rsi+r9]
je .LComputeOutputSmallKBy16
vmovdqu xmm4,XMMWORD PTR [rsi+r9*2]
jmp .LComputeOutputSmallKBy16
.LLoadPartialVectorSmallKBySmallN:
vmovdqu8 zmm2{k1}{z},ZMMWORD PTR [rsi]
cmp ecx,2
jb .LComputeOutputSmallKBy16
vmovdqu8 zmm3{k1}{z},ZMMWORD PTR [rsi+r9]
je .LComputeOutputSmallKBy16
vmovdqu8 zmm4{k1}{z},ZMMWORD PTR [rsi+r9*2]
jmp .LComputeOutputSmallKBy16
.LComputeOutputSmallKBy16:
vpunpcklbw xmm1,xmm2,xmm3 # interleave row data bytes
vpunpckhbw xmm2,xmm2,xmm3
vpunpcklbw xmm3,xmm4,xmm5
vpunpckhbw xmm4,xmm4,xmm5
vpunpcklwd xmm5,xmm1,xmm3 # interleave row data words
vpunpckhwd xmm1,xmm1,xmm3
vpunpcklwd xmm3,xmm2,xmm4
vpunpckhwd xmm2,xmm2,xmm4
vinserti128 ymm5,ymm5,xmm1,1 # concatenate 256-bit vector
vinserti128 ymm3,ymm3,xmm2,1
vshufi32x4 zmm16,zmm5,zmm3,0x44 # concatenate 512-bit vector
.ifeqs "\Isa\()", "Avx512BW"
vpmaddubsw zmm16,zmm28,zmm16
vpmaddwd zmm20,zmm16,zmm29
.else
vpxord zmm20,zmm20,zmm20
VpdpbusdsZmmZmmZmm zmm20,zmm28,zmm16
.endif
cmp r8d,16
jb .LStorePartialVectorSmallKBySmallN
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutputSmallKBy16
vpaddd zmm20,zmm20,ZMMWORD PTR [r10]
.LSkipAccumulateOutputSmallKBy16:
vmovdqu32 ZMMWORD PTR [r10],zmm20
add rsi,16 # advance matrix B by 16 bytes
add r10,16*4 # advance matrix C by 16 columns
sub r8d,16 # decrement CountN
jnz .LProcessColumnLoopSmallKBy16
jmp .LExitKernel
.LStorePartialVectorSmallKBySmallN:
test r11,r11 # ZeroMode?
jnz .LSkipAccumulateOutputSmallKBySmallN
vpaddd zmm20{k1}{z},zmm20,ZMMWORD PTR [r10]
.LSkipAccumulateOutputSmallKBySmallN:
vmovdqu32 ZMMWORD PTR [r10]{k1},zmm20
jmp .LExitKernel
.endm

34
onnxruntime/core/mlas/lib/x86_64/QgemvU8S8KernelAvx512Vnni.S Normal file
View file

@ -0,0 +1,34 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
QgemvU8S8KernelAvx512Vnni.s
Abstract:
This module implements the kernels for the quantized integer matrix/vector
multiply operation (QGEMV).
This implementation uses AVX512VNNI instructions.
--*/
#include "asmmacro.h"
#include "QgemvU8S8KernelAvx512Common.h"
#include "AssembleAvx512Vnni.h"
.intel_syntax noprefix
.text
//
// Generate the GEMV kernel.
//
GemvU8S8KernelAvx512Function Avx512Vnni
.end

9
onnxruntime/test/mlas/unittest.cpp
View file

@ -243,7 +243,7 @@ private:
std::fill_n(CReference, M * N, -0.5f);
MlasGemm(TransA, TransB, M, N, K, T(alpha), A, lda, B, ldb, T(beta), C, ldc, threadpool);
ReferenceSgemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, CReference, ldc);
ReferenceGemm(TransA, TransB, M, N, K, alpha, A, lda, B, ldb, beta, CReference, ldc);
for (size_t f = 0; f < M * N; f++) {
// Sensitive to comparing positive/negative zero.
@ -254,7 +254,7 @@ private:
}
void
ReferenceSgemm(
ReferenceGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
@ -566,6 +566,11 @@ public:
for (size_t b = 256; b < 320; b += 32) {
Test(b, b, b, 85, 173);
}
for (size_t b = 1; b < 96; b++) {
Test(1, b, 32, 0, 0);
Test(1, 32, b, 0, 0);
Test(1, b, b, 0, 0);
}
}
void