diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index 6213f016d9..6cb94816e2 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -55,7 +55,6 @@ if(MSVC) ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemvU8S8KernelAvx512Vnni.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Core.asm - ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Vnni.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/DgemmKernelSse2.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/DgemmKernelAvx.asm ${ONNXRUNTIME_ROOT}/core/mlas/lib/amd64/DgemmKernelFma3.asm @@ -255,7 +254,6 @@ else() ${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/QgemmU8U8KernelAvx512Core.S - ${ONNXRUNTIME_ROOT}/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Vnni.S ) if(HAS_AVX512CORE) set_source_files_properties(${mlas_platform_srcs_avx512core} PROPERTIES COMPILE_FLAGS "-mavx512bw -mavx512dq -mavx512vl") diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 5a399731f8..b342b712fe 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -147,35 +147,19 @@ MlasGemm( MLAS_THREADPOOL* ThreadPool ); +template void MLASCALL MlasGemm( size_t M, size_t N, size_t K, - const uint8_t* A, + const AType* A, size_t lda, - uint8_t offa, - const int8_t* B, + AType offa, + const BType* B, size_t ldb, - int8_t offb, - int32_t* C, - size_t ldc, - MLAS_THREADPOOL* ThreadPool - ); - -void -MLASCALL -MlasGemm( - size_t M, - size_t N, - size_t K, - const uint8_t* A, - size_t lda, - uint8_t offa, - const uint8_t* B, - size_t ldb, - uint8_t offb, + BType offb, int32_t* C, size_t ldc, MLAS_THREADPOOL* ThreadPool diff --git a/onnxruntime/core/mlas/lib/amd64/ErfKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/ErfKernelFma3.asm index 96b3df9513..c372d36e2e 100644 --- a/onnxruntime/core/mlas/lib/amd64/ErfKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/ErfKernelFma3.asm @@ -119,16 +119,16 @@ ErfKernelFrame ENDS alloc_stack (ErfKernelFrame.ReturnAddress) - save_xmm128_avx xmm6,ErfKernelFrame.SavedXmm6 - save_xmm128_avx xmm7,ErfKernelFrame.SavedXmm7 - save_xmm128_avx xmm8,ErfKernelFrame.SavedXmm8 - save_xmm128_avx xmm9,ErfKernelFrame.SavedXmm9 - save_xmm128_avx xmm10,ErfKernelFrame.SavedXmm10 - save_xmm128_avx xmm11,ErfKernelFrame.SavedXmm11 - save_xmm128_avx xmm12,ErfKernelFrame.SavedXmm12 - save_xmm128_avx xmm13,ErfKernelFrame.SavedXmm13 - save_xmm128_avx xmm14,ErfKernelFrame.SavedXmm14 - save_xmm128_avx xmm15,ErfKernelFrame.SavedXmm15 + save_xmm128 xmm6,ErfKernelFrame.SavedXmm6 + save_xmm128 xmm7,ErfKernelFrame.SavedXmm7 + save_xmm128 xmm8,ErfKernelFrame.SavedXmm8 + save_xmm128 xmm9,ErfKernelFrame.SavedXmm9 + save_xmm128 xmm10,ErfKernelFrame.SavedXmm10 + save_xmm128 xmm11,ErfKernelFrame.SavedXmm11 + save_xmm128 xmm12,ErfKernelFrame.SavedXmm12 + save_xmm128 xmm13,ErfKernelFrame.SavedXmm13 + save_xmm128 xmm14,ErfKernelFrame.SavedXmm14 + save_xmm128 xmm15,ErfKernelFrame.SavedXmm15 END_PROLOGUE @@ -548,16 +548,16 @@ LBiggerNumbersRemaining: LErfBatchExp: vzeroupper - vmovaps xmm6,ErfKernelFrame.SavedXmm6[rsp] - vmovaps xmm7,ErfKernelFrame.SavedXmm7[rsp] - vmovaps xmm8,ErfKernelFrame.SavedXmm8[rsp] - vmovaps xmm9,ErfKernelFrame.SavedXmm9[rsp] - vmovaps xmm10,ErfKernelFrame.SavedXmm10[rsp] - vmovaps xmm11,ErfKernelFrame.SavedXmm11[rsp] - vmovaps xmm12,ErfKernelFrame.SavedXmm12[rsp] - vmovaps xmm13,ErfKernelFrame.SavedXmm13[rsp] - vmovaps xmm14,ErfKernelFrame.SavedXmm14[rsp] - vmovaps xmm15,ErfKernelFrame.SavedXmm15[rsp] + movaps xmm6,ErfKernelFrame.SavedXmm6[rsp] + movaps xmm7,ErfKernelFrame.SavedXmm7[rsp] + movaps xmm8,ErfKernelFrame.SavedXmm8[rsp] + movaps xmm9,ErfKernelFrame.SavedXmm9[rsp] + movaps xmm10,ErfKernelFrame.SavedXmm10[rsp] + movaps xmm11,ErfKernelFrame.SavedXmm11[rsp] + movaps xmm12,ErfKernelFrame.SavedXmm12[rsp] + movaps xmm13,ErfKernelFrame.SavedXmm13[rsp] + movaps xmm14,ErfKernelFrame.SavedXmm14[rsp] + movaps xmm15,ErfKernelFrame.SavedXmm15[rsp] add rsp,(ErfKernelFrame.ReturnAddress) BEGIN_EPILOGUE diff --git a/onnxruntime/core/mlas/lib/amd64/FgemmKernelSse2Common.inc b/onnxruntime/core/mlas/lib/amd64/FgemmKernelSse2Common.inc index 276098ebb3..960263a642 100644 --- a/onnxruntime/core/mlas/lib/amd64/FgemmKernelSse2Common.inc +++ b/onnxruntime/core/mlas/lib/amd64/FgemmKernelSse2Common.inc @@ -13,7 +13,7 @@ ; This module implements the kernels for the floating point matrix/matrix ; multiply operation (SGEMM and DGEMM). ; -; This implementation uses AVX fused multiply/add instructions. +; This implementation uses SSE2 instructions. ; ;-- diff --git a/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm index 5d81c23722..a51be4b7e6 100644 --- a/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm @@ -98,16 +98,16 @@ LogisticKernelFrame ENDS alloc_stack (LogisticKernelFrame.ReturnAddress) - save_xmm128_avx xmm6,LogisticKernelFrame.SavedXmm6 - save_xmm128_avx xmm7,LogisticKernelFrame.SavedXmm7 - save_xmm128_avx xmm8,LogisticKernelFrame.SavedXmm8 - save_xmm128_avx xmm9,LogisticKernelFrame.SavedXmm9 - save_xmm128_avx xmm10,LogisticKernelFrame.SavedXmm10 - save_xmm128_avx xmm11,LogisticKernelFrame.SavedXmm11 - save_xmm128_avx xmm12,LogisticKernelFrame.SavedXmm12 - save_xmm128_avx xmm13,LogisticKernelFrame.SavedXmm13 - save_xmm128_avx xmm14,LogisticKernelFrame.SavedXmm14 - save_xmm128_avx xmm15,LogisticKernelFrame.SavedXmm15 + save_xmm128 xmm6,LogisticKernelFrame.SavedXmm6 + save_xmm128 xmm7,LogisticKernelFrame.SavedXmm7 + save_xmm128 xmm8,LogisticKernelFrame.SavedXmm8 + save_xmm128 xmm9,LogisticKernelFrame.SavedXmm9 + save_xmm128 xmm10,LogisticKernelFrame.SavedXmm10 + save_xmm128 xmm11,LogisticKernelFrame.SavedXmm11 + save_xmm128 xmm12,LogisticKernelFrame.SavedXmm12 + save_xmm128 xmm13,LogisticKernelFrame.SavedXmm13 + save_xmm128 xmm14,LogisticKernelFrame.SavedXmm14 + save_xmm128 xmm15,LogisticKernelFrame.SavedXmm15 END_PROLOGUE @@ -185,16 +185,16 @@ ProcessRemainingCount: ExitKernel: vzeroupper - vmovaps xmm6,LogisticKernelFrame.SavedXmm6[rsp] - vmovaps xmm7,LogisticKernelFrame.SavedXmm7[rsp] - vmovaps xmm8,LogisticKernelFrame.SavedXmm8[rsp] - vmovaps xmm9,LogisticKernelFrame.SavedXmm9[rsp] - vmovaps xmm10,LogisticKernelFrame.SavedXmm10[rsp] - vmovaps xmm11,LogisticKernelFrame.SavedXmm11[rsp] - vmovaps xmm12,LogisticKernelFrame.SavedXmm12[rsp] - vmovaps xmm13,LogisticKernelFrame.SavedXmm13[rsp] - vmovaps xmm14,LogisticKernelFrame.SavedXmm14[rsp] - vmovaps xmm15,LogisticKernelFrame.SavedXmm15[rsp] + movaps xmm6,LogisticKernelFrame.SavedXmm6[rsp] + movaps xmm7,LogisticKernelFrame.SavedXmm7[rsp] + movaps xmm8,LogisticKernelFrame.SavedXmm8[rsp] + movaps xmm9,LogisticKernelFrame.SavedXmm9[rsp] + movaps xmm10,LogisticKernelFrame.SavedXmm10[rsp] + movaps xmm11,LogisticKernelFrame.SavedXmm11[rsp] + movaps xmm12,LogisticKernelFrame.SavedXmm12[rsp] + movaps xmm13,LogisticKernelFrame.SavedXmm13[rsp] + movaps xmm14,LogisticKernelFrame.SavedXmm14[rsp] + movaps xmm15,LogisticKernelFrame.SavedXmm15[rsp] add rsp,(LogisticKernelFrame.ReturnAddress) BEGIN_EPILOGUE diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm index a4d9ec8468..b6ac85f756 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm @@ -62,6 +62,7 @@ GemmU8S8CopyPackBFrame STRUCT SavedXmm6 OWORD ? SavedXmm7 OWORD ? SavedXmm8 OWORD ? + SavedXmm9 OWORD ? Padding QWORD ? SavedRdi QWORD ? SavedRsi QWORD ? @@ -75,6 +76,7 @@ GemmU8S8CopyPackBFrame STRUCT CountK QWORD ? ColumnSumVector QWORD ? offa QWORD ? + BTypeIsSigned QWORD ? GemmU8S8CopyPackBFrame ENDS @@ -481,6 +483,9 @@ ExitRoutine: ; offa - Supplies the zero point offset for the other source matrix of the ; matrix multiplication. ; +; BTypeIsSigned - Supplies true if the source matrix is signed data, else +; false if the the source matrix is unsigned data. +; ; Return Value: ; ; None. @@ -497,6 +502,7 @@ ExitRoutine: save_xmm128 xmm6,GemmU8S8CopyPackBFrame.SavedXmm6 save_xmm128 xmm7,GemmU8S8CopyPackBFrame.SavedXmm7 save_xmm128 xmm8,GemmU8S8CopyPackBFrame.SavedXmm8 + save_xmm128 xmm9,GemmU8S8CopyPackBFrame.SavedXmm9 END_PROLOGUE @@ -510,6 +516,17 @@ ExitRoutine: vpsllw ymm0,ymm8,8 ; generate word vector [0x0100] vpor ymm8,ymm8,ymm0 ; generate word vector [0x0101] +; +; Compute the bit flip vector to adjust input from U8 to S8. +; + + vpxor xmm9,xmm9,xmm9 ; generate word vector [0x0000] + cmp BYTE PTR GemmU8S8CopyPackBFrame.BTypeIsSigned[rsp],0 + jnz SkipUnsignedBitFlipVector + vpsllw ymm9,ymm8,7 ; generate word vector [0x8080] + +SkipUnsignedBitFlipVector: + ; ; Process 16 columns of matrix B in a loop. ; @@ -544,6 +561,8 @@ InterleaveRowDataN16: vpunpckhwd xmm3,xmm3,xmm5 vinserti128 ymm4,ymm4,xmm6,1 vinserti128 ymm2,ymm2,xmm3,1 + vpxor ymm4,ymm4,ymm9 ; optionally adjust unsigned data + vpxor ymm2,ymm2,ymm9 vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows vmovdqu YMMWORD PTR [rcx+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row @@ -562,9 +581,9 @@ ProcessRemainingRowsN16: add rbx,4 ; correct for over-subtract above jz ReduceColumnSumVectorN16 vmovdqu xmm2,XMMWORD PTR [rdx] - vpxor xmm3,xmm3,xmm3 - vpxor xmm4,xmm4,xmm4 - vpxor xmm5,xmm5,xmm5 + vmovaps xmm3,xmm9 + vmovaps xmm4,xmm9 + vmovaps xmm5,xmm9 xor ebx,ebx ; no more rows remaining test r10b,2 ; (CountK & 2) != 0? jz InterleaveRowDataN16 @@ -596,6 +615,7 @@ ExitRoutine: movaps xmm6,GemmU8S8CopyPackBFrame.SavedXmm6[rsp] movaps xmm7,GemmU8S8CopyPackBFrame.SavedXmm7[rsp] movaps xmm8,GemmU8S8CopyPackBFrame.SavedXmm8[rsp] + movaps xmm9,GemmU8S8CopyPackBFrame.SavedXmm9[rsp] add rsp,(GemmU8S8CopyPackBFrame.SavedRdi) BEGIN_EPILOGUE @@ -613,8 +633,8 @@ ExitRoutine: ProcessColumnNUnaligned: vpxor xmm0,xmm0,xmm0 ; clear column accumulators vpxor xmm1,xmm1,xmm1 - vmovdqu YMMWORD PTR GemmU8S8CopyPackBFrame.PaddedMatrixBData[rsp],ymm0 - vmovdqu YMMWORD PTR GemmU8S8CopyPackBFrame.PaddedMatrixBData[rsp+32],ymm0 + vmovdqu YMMWORD PTR GemmU8S8CopyPackBFrame.PaddedMatrixBData[rsp],ymm9 + vmovdqu YMMWORD PTR GemmU8S8CopyPackBFrame.PaddedMatrixBData[rsp+32],ymm9 sub r10,4 jb ProcessRemainingRowsNUnaligned @@ -690,6 +710,8 @@ ProcessPaddedMatrixBData: vpunpckhwd xmm3,xmm3,xmm5 vinserti128 ymm4,ymm4,xmm6,1 vinserti128 ymm2,ymm2,xmm3,1 + vpxor ymm4,ymm4,ymm9 ; optionally adjust unsigned data + vpxor ymm2,ymm2,ymm9 vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows vmovdqu YMMWORD PTR [rcx+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row @@ -711,9 +733,8 @@ ProcessRemainingRowsNUnaligned: .errnz GemmU8S8CopyPackBFrame.PaddedMatrixBData mov rbp,rsp ; GemmU8S8CopyPackBFrame.PaddedMatrixBData - vpxor xmm6,xmm6,xmm6 - vmovdqu YMMWORD PTR [rbp],ymm6 - vmovdqu YMMWORD PTR [rbp+32],ymm6 + vmovdqu YMMWORD PTR [rbp],ymm9 + vmovdqu YMMWORD PTR [rbp+32],ymm9 CopyUnalignedRowLoop: lea rdi,[rbp+16] ; advance next padded buffer by 16 bytes diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Common.inc b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Common.inc deleted file mode 100644 index 7afacbbdd4..0000000000 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Common.inc +++ /dev/null @@ -1,64 +0,0 @@ -;++ -; -; Copyright (c) Microsoft Corporation. All rights reserved. -; -; Licensed under the MIT License. -; -; Module Name: -; -; QgemmU8U8KernelAvx512Common.inc -; -; Abstract: -; -; This module contains common kernel macros and structures for the quantized -; integer matrix/matrix multiply operation (QGEMM) for the AVX512 core and -; AVX512VNNI kernels. -; -;-- - -INCLUDE QgemmU8X8KernelAvx512Common.inc - -; -; Macro Description: -; -; This macro generates code to execute the block compute macro multiple -; times and advancing the matrix A and matrix B data pointers. -; -; Arguments: -; -; ColumnCount - Supplies the number of columns to produce. -; -; RowCount - Supplies the number of rows to produce. -; -; Implicit Arguments: -; -; rbx - Supplies the address into the matrix A data plus 3 rows. -; -; rcx - Supplies the address into the matrix A data. -; -; rdx - Supplies the address into the matrix B data. -; -; r9 - Supplies the length in bytes of a row from matrix A. -; -; r14 - Supplies the stride in bytes of between packed blocks of matrix B. -; -; zmm14-zmm31 - Supplies the block accumulators. -; - -ComputeBlockLoop MACRO ColumnCount, RowCount - - LOCAL ComputeBlockBy1Loop - - mov rsi,r9 ; reload row length remaining - -ComputeBlockBy1Loop: - ComputeBlock ColumnCount, RowCount, 0, 0 - add rcx,4 ; advance matrix A by 1 pair -IF RowCount GT 3 - add rbx,4 ; advance matrix A plus 3 rows by 1 pair -ENDIF - add rdx,32 ; advance matrix B - sub rsi,4 - jnz ComputeBlockBy1Loop - - ENDM diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Core.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Core.asm index cd80710300..554aa1c981 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Core.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Core.asm @@ -19,7 +19,7 @@ .xlist INCLUDE mlasi.inc -INCLUDE QgemmU8U8KernelAvx512Common.inc +INCLUDE QgemmU8X8KernelAvx512Common.inc .list ; @@ -118,6 +118,51 @@ ENDIF ENDM +; +; Macro Description: +; +; This macro generates code to execute the block compute macro multiple +; times and advancing the matrix A and matrix B data pointers. +; +; Arguments: +; +; ColumnCount - Supplies the number of columns to produce. +; +; RowCount - Supplies the number of rows to produce. +; +; Implicit Arguments: +; +; rbx - Supplies the address into the matrix A data plus 3 rows. +; +; rcx - Supplies the address into the matrix A data. +; +; rdx - Supplies the address into the matrix B data. +; +; r9 - Supplies the length in bytes of a row from matrix A. +; +; r14 - Supplies the stride in bytes of between packed blocks of matrix B. +; +; zmm14-zmm31 - Supplies the block accumulators. +; + +ComputeBlockLoop MACRO ColumnCount, RowCount + + LOCAL ComputeBlockBy1Loop + + mov rsi,r9 ; reload row length remaining + +ComputeBlockBy1Loop: + ComputeBlock ColumnCount, RowCount, 0, 0 + add rcx,4 ; advance matrix A by 1 pair +IF RowCount GT 3 + add rbx,4 ; advance matrix A plus 3 rows by 1 pair +ENDIF + add rdx,32 ; advance matrix B + sub rsi,4 + jnz ComputeBlockBy1Loop + + ENDM + ; ; Generate the GEMM kernel. ; diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Vnni.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Vnni.asm deleted file mode 100644 index 99c72a4f2c..0000000000 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx512Vnni.asm +++ /dev/null @@ -1,110 +0,0 @@ -;++ -; -; Copyright (c) Microsoft Corporation. All rights reserved. -; -; Licensed under the MIT License. -; -; Module Name: -; -; QgemmU8U8KernelAvx512Vnni.asm -; -; Abstract: -; -; This module implements the kernels for the quantized integer matrix/matrix -; multiply operation (QGEMM). -; -; This implementation uses AVX512VNNI instructions. -; -;-- - - .xlist -INCLUDE mlasi.inc -INCLUDE QgemmU8U8KernelAvx512Common.inc -INCLUDE AssembleAvx512Vnni.inc - .list - -; -; Macro Description: -; -; This macro generates code to multiply and accumulate each row of the output -; block. -; -; Arguments: -; -; ColumnCount - Supplies the number of columns to produce. -; -; RowCount - Supplies the number of rows to produce. -; -; VectorOffset - Supplies the byte offset from matrix B to fetch elements. -; -; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements. -; -; Implicit Arguments: -; -; rbx - Supplies the address into the matrix A data plus 3 rows. -; -; rcx - Supplies the address into the matrix A data. -; -; rdx - Supplies the address into the matrix B data. -; -; r9 - Supplies the length in bytes of a row from matrix A. -; -; r14 - Supplies the stride in bytes of between packed blocks of matrix B. -; -; zmm14-zmm31 - Supplies the block accumulators. -; - -ComputeBlock MACRO ColumnCount, RowCount, VectorOffset, BroadcastOffset - -IF ColumnCount GE 32 -IF ColumnCount GE 48 - vpmovzxbw zmm0,YMMWORD PTR [rdx+VectorOffset] - vpmovzxbw zmm1,YMMWORD PTR [rdx+r14+VectorOffset] - vpmovzxbw zmm2,YMMWORD PTR [rdx+r14*2+VectorOffset] -ELSE - vpmovzxbw zmm1,YMMWORD PTR [rdx+VectorOffset] - vpmovzxbw zmm2,YMMWORD PTR [rdx+r14+VectorOffset] -ENDIF - EmitIfCountGE RowCount, 1, - EmitIfCount2GE RowCount, 1, ColumnCount, 48, - EmitIfCount2GE RowCount, 1, ColumnCount, 32, - EmitIfCount2GE RowCount, 1, ColumnCount, 16, - EmitIfCountGE RowCount, 2, - EmitIfCount2GE RowCount, 2, ColumnCount, 48, - EmitIfCount2GE RowCount, 2, ColumnCount, 32, - EmitIfCount2GE RowCount, 2, ColumnCount, 16, - EmitIfCountGE RowCount, 3, - EmitIfCount2GE RowCount, 3, ColumnCount, 48, - EmitIfCount2GE RowCount, 3, ColumnCount, 32, - EmitIfCount2GE RowCount, 3, ColumnCount, 16, - EmitIfCountGE RowCount, 4, - EmitIfCount2GE RowCount, 4, ColumnCount, 48, - EmitIfCount2GE RowCount, 4, ColumnCount, 32, - EmitIfCount2GE RowCount, 4, ColumnCount, 16, - EmitIfCountGE RowCount, 5, - EmitIfCount2GE RowCount, 5, ColumnCount, 48, - EmitIfCount2GE RowCount, 5, ColumnCount, 32, - EmitIfCount2GE RowCount, 5, ColumnCount, 16, - EmitIfCountGE RowCount, 6, - EmitIfCount2GE RowCount, 6, ColumnCount, 48, - EmitIfCount2GE RowCount, 6, ColumnCount, 32, - EmitIfCount2GE RowCount, 6, ColumnCount, 16, -ELSE - vpmovzxbw zmm2,YMMWORD PTR [rdx+VectorOffset] - EmitIfCountGE RowCount, 1, - EmitIfCountGE RowCount, 2, - EmitIfCountGE RowCount, 3, - EmitIfCountGE RowCount, 4, - EmitIfCountGE RowCount, 5, - EmitIfCountGE RowCount, 6, -ENDIF - - ENDM - -; -; Generate the GEMM kernel. -; - -GemmU8X8KernelAvx512Function U8U8, Avx512Vnni - - END diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc index 86142a75f7..ecfe53d271 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc @@ -128,8 +128,8 @@ IF RowCount GT 3 ENDIF ComputeBlockLoop ColumnCount, RowCount IF RowCount GT 3 - lea rbx,[r8+rax*2] ; compute matrix C plus 3 rows - add rbx,rax + lea rbx,[rax*2+rax] + add rbx,r8 ; compute matrix C plus 3 rows ENDIF ENDM diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelM1Avx.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelM1Avx.asm index ee9debe77f..418c8e332e 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelM1Avx.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelM1Avx.asm @@ -85,9 +85,9 @@ SgemmKernelM1Frame ENDS push_reg rbx push_reg rsi alloc_stack (SgemmKernelM1Frame.SavedRsi) - save_xmm128_avx xmm6,SgemmKernelM1Frame.SavedXmm6 - save_xmm128_avx xmm7,SgemmKernelM1Frame.SavedXmm7 - save_xmm128_avx xmm8,SgemmKernelM1Frame.SavedXmm8 + save_xmm128 xmm6,SgemmKernelM1Frame.SavedXmm6 + save_xmm128 xmm7,SgemmKernelM1Frame.SavedXmm7 + save_xmm128 xmm8,SgemmKernelM1Frame.SavedXmm8 END_PROLOGUE @@ -217,9 +217,9 @@ ProcessRemainingCountK: ExitKernel: vzeroupper - vmovaps xmm6,SgemmKernelM1Frame.SavedXmm6[rsp] - vmovaps xmm7,SgemmKernelM1Frame.SavedXmm7[rsp] - vmovaps xmm8,SgemmKernelM1Frame.SavedXmm8[rsp] + movaps xmm6,SgemmKernelM1Frame.SavedXmm6[rsp] + movaps xmm7,SgemmKernelM1Frame.SavedXmm7[rsp] + movaps xmm8,SgemmKernelM1Frame.SavedXmm8[rsp] add rsp,(SgemmKernelM1Frame.SavedRsi) BEGIN_EPILOGUE @@ -348,8 +348,8 @@ ProcessRemainingCountN1: push_reg rbx push_reg rsi alloc_stack (SgemmKernelM1Frame.SavedRsi) - save_xmm128_avx xmm6,SgemmKernelM1Frame.SavedXmm6 - save_xmm128_avx xmm7,SgemmKernelM1Frame.SavedXmm7 + save_xmm128 xmm6,SgemmKernelM1Frame.SavedXmm6 + save_xmm128 xmm7,SgemmKernelM1Frame.SavedXmm7 END_PROLOGUE @@ -465,8 +465,8 @@ ProcessRemainingCountN: ExitKernel: vzeroupper - vmovaps xmm6,SgemmKernelM1Frame.SavedXmm6[rsp] - vmovaps xmm7,SgemmKernelM1Frame.SavedXmm7[rsp] + movaps xmm6,SgemmKernelM1Frame.SavedXmm6[rsp] + movaps xmm7,SgemmKernelM1Frame.SavedXmm7[rsp] add rsp,(SgemmKernelM1Frame.SavedRsi) BEGIN_EPILOGUE diff --git a/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm index ad5fc4e207..6cb3c55af2 100644 --- a/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm @@ -98,16 +98,16 @@ TanhKernelFrame ENDS alloc_stack (TanhKernelFrame.ReturnAddress) - save_xmm128_avx xmm6,TanhKernelFrame.SavedXmm6 - save_xmm128_avx xmm7,TanhKernelFrame.SavedXmm7 - save_xmm128_avx xmm8,TanhKernelFrame.SavedXmm8 - save_xmm128_avx xmm9,TanhKernelFrame.SavedXmm9 - save_xmm128_avx xmm10,TanhKernelFrame.SavedXmm10 - save_xmm128_avx xmm11,TanhKernelFrame.SavedXmm11 - save_xmm128_avx xmm12,TanhKernelFrame.SavedXmm12 - save_xmm128_avx xmm13,TanhKernelFrame.SavedXmm13 - save_xmm128_avx xmm14,TanhKernelFrame.SavedXmm14 - save_xmm128_avx xmm15,TanhKernelFrame.SavedXmm15 + save_xmm128 xmm6,TanhKernelFrame.SavedXmm6 + save_xmm128 xmm7,TanhKernelFrame.SavedXmm7 + save_xmm128 xmm8,TanhKernelFrame.SavedXmm8 + save_xmm128 xmm9,TanhKernelFrame.SavedXmm9 + save_xmm128 xmm10,TanhKernelFrame.SavedXmm10 + save_xmm128 xmm11,TanhKernelFrame.SavedXmm11 + save_xmm128 xmm12,TanhKernelFrame.SavedXmm12 + save_xmm128 xmm13,TanhKernelFrame.SavedXmm13 + save_xmm128 xmm14,TanhKernelFrame.SavedXmm14 + save_xmm128 xmm15,TanhKernelFrame.SavedXmm15 END_PROLOGUE @@ -177,16 +177,16 @@ ProcessRemainingCount: ExitKernel: vzeroupper - vmovaps xmm6,TanhKernelFrame.SavedXmm6[rsp] - vmovaps xmm7,TanhKernelFrame.SavedXmm7[rsp] - vmovaps xmm8,TanhKernelFrame.SavedXmm8[rsp] - vmovaps xmm9,TanhKernelFrame.SavedXmm9[rsp] - vmovaps xmm10,TanhKernelFrame.SavedXmm10[rsp] - vmovaps xmm11,TanhKernelFrame.SavedXmm11[rsp] - vmovaps xmm12,TanhKernelFrame.SavedXmm12[rsp] - vmovaps xmm13,TanhKernelFrame.SavedXmm13[rsp] - vmovaps xmm14,TanhKernelFrame.SavedXmm14[rsp] - vmovaps xmm15,TanhKernelFrame.SavedXmm15[rsp] + movaps xmm6,TanhKernelFrame.SavedXmm6[rsp] + movaps xmm7,TanhKernelFrame.SavedXmm7[rsp] + movaps xmm8,TanhKernelFrame.SavedXmm8[rsp] + movaps xmm9,TanhKernelFrame.SavedXmm9[rsp] + movaps xmm10,TanhKernelFrame.SavedXmm10[rsp] + movaps xmm11,TanhKernelFrame.SavedXmm11[rsp] + movaps xmm12,TanhKernelFrame.SavedXmm12[rsp] + movaps xmm13,TanhKernelFrame.SavedXmm13[rsp] + movaps xmm14,TanhKernelFrame.SavedXmm14[rsp] + movaps xmm15,TanhKernelFrame.SavedXmm15[rsp] add rsp,(TanhKernelFrame.ReturnAddress) BEGIN_EPILOGUE diff --git a/onnxruntime/core/mlas/lib/amd64/mlasi.inc b/onnxruntime/core/mlas/lib/amd64/mlasi.inc index 39e3ee6b71..d7a8edba9b 100644 --- a/onnxruntime/core/mlas/lib/amd64/mlasi.inc +++ b/onnxruntime/core/mlas/lib/amd64/mlasi.inc @@ -19,27 +19,6 @@ INCLUDE macamd64.inc .list -; -; Macro Description: -; -; This macro uses AVX instructions to save a vector register as part of a -; function prologue as an alternative to save_xmm128. -; -; Arguments: -; -; Register - Supplies the vector register to be saved. -; -; Offset - Supplies the stack frame offset to store the contents of the -; vector register. -; - -save_xmm128_avx MACRO Register, Offset - - vmovaps Offset[rsp], Register - .savexmm128 Register, Offset - - ENDM - ; ; Macro Description: ; diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 919399b38a..6be271309b 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -22,6 +22,7 @@ Abstract: #include #include #include +#include #if defined(_WIN32) #include @@ -246,31 +247,22 @@ typedef MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* PMLAS_SGEMM_TRANSPOSE_PACKB_BL typedef void -(MLASCALL MLAS_GEMM_U8S8_COPY_PACKA_ROUTINE)( - uint8_t* D, +(MLASCALL MLAS_GEMM_U8X8_OPERATION)( + const struct MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, + size_t M, + size_t N, + size_t K, const uint8_t* A, size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumVector, - int16_t offb - ); - -typedef MLAS_GEMM_U8S8_COPY_PACKA_ROUTINE* PMLAS_GEMM_U8S8_COPY_PACKA_ROUTINE; - -typedef -void -(MLASCALL MLAS_GEMM_U8S8_COPY_PACKB_ROUTINE)( - int8_t* D, - const int8_t* B, + int16_t offa, + const uint8_t* B, size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumVector, - int16_t offa + int16_t offb, + int32_t* C, + size_t ldc ); -typedef MLAS_GEMM_U8S8_COPY_PACKB_ROUTINE* PMLAS_GEMM_U8S8_COPY_PACKB_ROUTINE; +typedef MLAS_GEMM_U8X8_OPERATION* PMLAS_GEMM_U8X8_OPERATION; typedef size_t @@ -303,34 +295,6 @@ size_t typedef MLAS_GEMV_U8S8_KERNEL* PMLAS_GEMV_U8S8_KERNEL; -typedef -void -(MLASCALL MLAS_GEMM_U8U8_COPY_PACKA_ROUTINE)( - int16_t* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumVector, - int16_t offb - ); - -typedef MLAS_GEMM_U8U8_COPY_PACKA_ROUTINE* PMLAS_GEMM_U8U8_COPY_PACKA_ROUTINE; - -typedef -void -(MLASCALL MLAS_GEMM_U8U8_COPY_PACKB_ROUTINE)( - uint8_t* D, - const uint8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumVector, - int16_t offa - ); - -typedef MLAS_GEMM_U8U8_COPY_PACKB_ROUTINE* PMLAS_GEMM_U8U8_COPY_PACKB_ROUTINE; - typedef size_t (MLASCALL MLAS_GEMM_U8U8_KERNEL)( @@ -349,24 +313,6 @@ size_t typedef MLAS_GEMM_U8U8_KERNEL* PMLAS_GEMM_U8U8_KERNEL; -typedef -void -(MLASCALL MLAS_GEMM_X8X8_OPERATION)( - size_t M, - size_t N, - size_t K, - const uint8_t* A, - size_t lda, - int16_t offa, - const uint8_t* B, - size_t ldb, - int16_t offb, - int32_t* C, - size_t ldc - ); - -typedef MLAS_GEMM_X8X8_OPERATION* PMLAS_GEMM_X8X8_OPERATION; - typedef void (MLASCALL MLAS_CONV_FLOAT_KERNEL)( @@ -540,26 +486,18 @@ extern "C" { #endif #if defined(MLAS_TARGET_AMD64_IX86) - MLAS_GEMM_U8S8_COPY_PACKA_ROUTINE MlasGemmU8S8CopyPackASse; - MLAS_GEMM_U8S8_COPY_PACKB_ROUTINE MlasGemmU8S8CopyPackBSse; - MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelSse; - MLAS_GEMM_U8U8_COPY_PACKA_ROUTINE MlasGemmU8U8CopyPackASse; - MLAS_GEMM_U8U8_COPY_PACKB_ROUTINE MlasGemmU8U8CopyPackBSse; - MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelSse; + MLAS_GEMM_U8X8_OPERATION MlasGemmU8X8OperationSse; + MLAS_GEMM_U8X8_OPERATION MlasGemmU8S8OperationAvx2; + MLAS_GEMM_U8X8_OPERATION MlasGemmU8U8OperationAvx2; #if defined(MLAS_TARGET_AMD64) - 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 MlasGemmU8S8KernelAvx512Core; MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Core; 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; MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Core; - MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Vnni; #endif #endif @@ -682,12 +620,8 @@ struct MLAS_PLATFORM { #if defined(MLAS_TARGET_AMD64_IX86) PMLAS_GEMM_FLOAT_KERNEL GemmFloatKernel; - PMLAS_GEMM_U8S8_COPY_PACKA_ROUTINE GemmU8S8CopyPackARoutine; - PMLAS_GEMM_U8S8_COPY_PACKB_ROUTINE GemmU8S8CopyPackBRoutine; - PMLAS_GEMM_U8S8_KERNEL GemmU8S8Kernel; - PMLAS_GEMM_U8U8_COPY_PACKA_ROUTINE GemmU8U8CopyPackARoutine; - PMLAS_GEMM_U8U8_COPY_PACKB_ROUTINE GemmU8U8CopyPackBRoutine; - PMLAS_GEMM_U8U8_KERNEL GemmU8U8Kernel; + PMLAS_GEMM_U8X8_OPERATION GemmU8S8Operation; + PMLAS_GEMM_U8X8_OPERATION GemmU8U8Operation; #endif #if defined(MLAS_TARGET_AMD64) @@ -695,7 +629,9 @@ struct MLAS_PLATFORM { PMLAS_SGEMM_KERNEL_M1_ROUTINE KernelM1TransposeBRoutine; PMLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE TransposePackB16x4Routine; PMLAS_GEMM_DOUBLE_KERNEL GemmDoubleKernel; + PMLAS_GEMM_U8S8_KERNEL GemmU8S8Kernel; PMLAS_GEMV_U8S8_KERNEL GemvU8S8Kernel; + PMLAS_GEMM_U8U8_KERNEL GemmU8U8Kernel; PMLAS_CONV_FLOAT_KERNEL ConvNchwFloatKernel; PMLAS_CONV_FLOAT_KERNEL ConvNchwcFloatKernel; PMLAS_CONV_DEPTHWISE_FLOAT_KERNEL ConvDepthwiseFloatKernel; @@ -831,6 +767,9 @@ typedef __m128i MLAS_INT32X4; typedef __vector float MLAS_FLOAT32X4; typedef __vector int MLAS_INT32X4; typedef __vector unsigned MLAS_UINT32X4; +#else +typedef float MLAS_FLOAT32X4 __attribute__ ((vector_size(16))); +typedef int32_t MLAS_INT32X4 __attribute__ ((vector_size(16))); #endif MLAS_FORCEINLINE @@ -937,19 +876,8 @@ MlasStoreAlignedFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector) MLAS_UNREFERENCED_PARAMETER(Buffer); MLAS_UNREFERENCED_PARAMETER(Vector); vec_st(Vector, 0, Buffer); -#endif -} - -MLAS_FORCEINLINE -void -MlasStoreLowHalfFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector) -{ -#if defined(MLAS_NEON_INTRINSICS) - vst1_f32(Buffer, vget_low_f32(Vector)); -#elif defined(MLAS_SSE2_INTRINSICS) - _mm_storel_pi((__m64*)Buffer, Vector); -#elif defined(MLAS_VSX_INTRINSICS) - *((int64_t*)Buffer) = ((__vector int64_t)Vector)[0]; +#else + MlasStoreFloat32x4(Buffer, Vector); #endif } @@ -969,6 +897,22 @@ MlasStoreLaneFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector) #endif } +MLAS_FORCEINLINE +void +MlasStoreLowHalfFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector) +{ +#if defined(MLAS_NEON_INTRINSICS) + vst1_f32(Buffer, vget_low_f32(Vector)); +#elif defined(MLAS_SSE2_INTRINSICS) + _mm_storel_pi((__m64*)Buffer, Vector); +#elif defined(MLAS_VSX_INTRINSICS) + *((int64_t*)Buffer) = ((__vector int64_t)Vector)[0]; +#else + MlasStoreLaneFloat32x4<0>(&Buffer[0], Vector); + MlasStoreLaneFloat32x4<1>(&Buffer[1], Vector); +#endif +} + template MLAS_FORCEINLINE float @@ -1100,8 +1044,10 @@ MlasMaximumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) return vmaxq_f32(Vector1, Vector2); #elif defined(MLAS_SSE2_INTRINSICS) return _mm_max_ps(Vector1, Vector2); -#else +#elif defined(MLAS_VSX_INTRINSICS) return vec_sel(Vector2, Vector1, vec_cmpgt(Vector1, Vector2)); +#else +#error Unsupported architecture. #endif } @@ -1113,8 +1059,10 @@ MlasMinimumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2) return vminq_f32(Vector1, Vector2); #elif defined(MLAS_SSE2_INTRINSICS) return _mm_min_ps(Vector1, Vector2); -#else +#elif defined(MLAS_VSX_INTRINSICS) return vec_sel(Vector2, Vector1, vec_cmpgt(Vector2, Vector1)); +#else +#error Unsupported architecture. #endif } @@ -1298,10 +1246,8 @@ MlasShiftLeftInt32x4(MLAS_INT32X4 Vector) return vshlq_n_s32(Vector, ShiftCount); #elif defined(MLAS_SSE2_INTRINSICS) return _mm_slli_epi32(Vector, ShiftCount); -#elif defined(MLAS_VSX_INTRINSICS) - return vec_sl(Vector, MLAS_UINT32X4(MlasBroadcastInt32x4(ShiftCount))); #else -#error Unsupported architecture. + return Vector << ShiftCount; #endif } diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 6d507fd1c7..2ee83779b6 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -115,12 +115,8 @@ Return Value: // this->GemmFloatKernel = MlasGemmFloatKernelSse; - this->GemmU8S8CopyPackARoutine = MlasGemmU8S8CopyPackASse; - this->GemmU8S8CopyPackBRoutine = MlasGemmU8S8CopyPackBSse; - this->GemmU8S8Kernel = MlasGemmU8S8KernelSse; - this->GemmU8U8CopyPackARoutine = MlasGemmU8U8CopyPackASse; - this->GemmU8U8CopyPackBRoutine = MlasGemmU8U8CopyPackBSse; - this->GemmU8U8Kernel = MlasGemmU8U8KernelSse; + this->GemmU8S8Operation = MlasGemmU8X8OperationSse; + this->GemmU8U8Operation = MlasGemmU8X8OperationSse; #if defined(MLAS_TARGET_AMD64) @@ -199,12 +195,10 @@ Return Value: if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) { - this->GemmU8S8CopyPackARoutine = MlasGemmU8S8CopyPackAAvx2; - this->GemmU8S8CopyPackBRoutine = MlasGemmU8S8CopyPackBAvx2; + this->GemmU8S8Operation = MlasGemmU8S8OperationAvx2; this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx2; this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx2; - this->GemmU8U8CopyPackARoutine = MlasGemmU8U8CopyPackAAvx2; - this->GemmU8U8CopyPackBRoutine = MlasGemmU8U8CopyPackBAvx2; + this->GemmU8U8Operation = MlasGemmU8U8OperationAvx2; this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx2; this->GemmFloatKernel = MlasGemmFloatKernelFma3; @@ -261,9 +255,9 @@ Return Value: if ((Cpuid7[2] & 0x800) != 0) { + this->GemmU8U8Operation = MlasGemmU8S8OperationAvx2; this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Vnni; this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Vnni; - this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx512Vnni; } } diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index 377628c356..9bede5864b 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -17,21 +17,15 @@ Abstract: #include "mlasi.h" -// -// Define the default strides to step through slices of the input matrices. -// - -#define MLAS_GEMM_X8X8_STRIDEM 24 -#define MLAS_GEMM_X8X8_STRIDEN 256 -#define MLAS_GEMM_X8X8_STRIDEK 128 - // // Define the parameters to execute segments of a QGEMM operation on worker // threads. // -struct MLAS_GEMM_X8X8_WORK_BLOCK { - PMLAS_GEMM_X8X8_OPERATION GemmX8X8Operation; +struct MLAS_GEMM_U8X8_WORK_BLOCK +{ + int32_t ThreadCountM; + int32_t ThreadCountN; size_t M; size_t N; size_t K; @@ -41,587 +35,23 @@ struct MLAS_GEMM_X8X8_WORK_BLOCK { size_t ldb; int32_t* C; size_t ldc; - int32_t ThreadCountM; - int32_t ThreadCountN; - size_t StrideM; - size_t StrideN; int16_t offa; int16_t offb; + bool BTypeIsSigned; }; #ifdef MLAS_TARGET_AMD64_IX86 // -// Stores a vector to transpose a 4x4 byte vector using vpshufb. +// Define the default strides to step through slices of the input matrices. // -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. -// +#define MLAS_GEMM_U8X8_STRIDEM_SSE 12 +#define MLAS_GEMM_U8X8_STRIDEN_SSE 128 +#define MLAS_GEMM_U8X8_STRIDEK_SSE 128 void -MLASCALL -MlasGemmU8S8CopyPackASse( - uint8_t* D, - const uint8_t* A, - size_t lda, - size_t CountM, - size_t CountK, - int32_t* RowSumVector, - int16_t offb - ) -/*++ - -Routine Description: - - This routine copies elements from the source matrix to the destination - packed buffer. - -Arguments: - - D - Supplies the address of the destination packed buffer. - - A - Supplies the address of the source matrix. - - lda - Supplies the number of elements per row of the source matrix. - - CountM - Supplies the number of rows of the source matrix to copy. - - CountK - Supplies the number of columns of the source matrix to copy. - - RowSumVector - Supplies the address of the buffer to receive the sums of - the elements from each of the rows. Each sum has also been multiplied - by the zero point offset. - - offb - Supplies the zero point offset for the other source matrix of the - matrix multiplication. - -Return Value: - - None. - ---*/ -{ - const __m128i ZeroVector = _mm_setzero_si128(); - const __m128i OffsetBroadcast = _mm_set1_epi16(offb); - uint8_t PaddedMatrixAData[8] = { 0 }; - - // - // Process a single row of matrix A in a loop. - // - - while (CountM > 0) { - - const uint8_t* a = A; - size_t k = CountK; - __m128i RowSum = ZeroVector; - - // - // Copy the source bytes to the packed buffer. - // - // The packed buffer has the same data ordering as the source bytes, - // but CountK is aligned up to a multiple of 4 to maintain 32-bit - // alignment. All extra bytes are zero-padded. - // - // These values are also zero-extended and accumulated into an - // intermediate per-row accumulator. CountK cannot be greater than 128 - // to avoid overflowing these signed 16-bit accumulators. - // - - while (k >= 8) { - - __m128i Bytes = _mm_loadl_epi64((__m128i*)&a[0]); - _mm_storel_epi64((__m128i*)&D[0], Bytes); - - RowSum = _mm_add_epi16(RowSum, _mm_unpacklo_epi8(Bytes, ZeroVector)); - - D += 8; - a += 8; - k -= 8; - } - - if (k > 0) { - - // - // Copy the remaining bytes to the zero padded stack buffer. - // - - uint8_t* padded = PaddedMatrixAData; - uint8_t* padded_end = padded + k; - - do { - padded[0] = a[0]; - padded++; - a++; - } while (padded < padded_end); - - __m128i Bytes = _mm_loadl_epi64((__m128i*)PaddedMatrixAData); - _mm_storel_epi64((__m128i*)&D[0], Bytes); - - RowSum = _mm_add_epi16(RowSum, _mm_unpacklo_epi8(Bytes, ZeroVector)); - - // - // Copy quads of 8-bit values from the vector to the packed - // buffer and rotate the vector for the next iteration. - // - - for (size_t quads = (k + 3) / 4; quads > 0; quads--) { - *((int32_t*)D) = _mm_cvtsi128_si32(Bytes); - D += 4; - Bytes = _mm_shuffle_epi32(Bytes, _MM_SHUFFLE(0, 3, 2, 1)); - } - } - - // - // Reduce the sum for the single row of output and multiply by the - // zero point offset of the other source matrix. - // - - RowSum = _mm_madd_epi16(RowSum, OffsetBroadcast); - RowSum = _mm_add_epi32(RowSum, _mm_shuffle_epi32(RowSum, _MM_SHUFFLE(3, 2, 3, 2))); - RowSum = _mm_add_epi32(RowSum, _mm_shuffle_epi32(RowSum, _MM_SHUFFLE(0, 1, 0, 1))); - - *RowSumVector++ = _mm_cvtsi128_si32(RowSum); - - A += lda; - CountM -= 1; - } -} - -void -MLASCALL -MlasGemmU8S8CopyPackBSse( - int8_t* D, - const int8_t* B, - size_t ldb, - size_t CountN, - size_t CountK, - int32_t* ColumnSumVector, - int16_t offa - ) -/*++ - -Routine Description: - - This routine copies elements from the source matrix to the destination - packed buffer. - -Arguments: - - D - Supplies the address of the destination packed buffer. - - B - Supplies the address of the source matrix. - - ldb - Supplies the number of elements per row of the source matrix. - - CountN - Supplies the number of columns of the source matrix to copy. - - CountK - Supplies the number of rows of the source matrix to copy. - - ColumnSumVector - Supplies the address of the buffer to receive the sums of - the elements from each of the columns. Each sum has also been multiplied - by the zero point offset. - - offa - Supplies the zero point offset for the other source matrix of the - matrix multiplication. - -Return Value: - - None. - ---*/ -{ - const __m128i ZeroVector = _mm_setzero_si128(); - const __m128i OffsetBroadcast = _mm_set1_epi16(offa); - int8_t PaddedMatrixBData[16] = { 0 }; - - // - // Process 8 columns of matrix B in a loop. - // - - while (CountN >= 8) { - - const int8_t* b = B; - size_t k = CountK; - __m128i ColumnSum0 = ZeroVector; - __m128i ColumnSum1 = ZeroVector; - - // - // Interleave 2 rows of matrix B and write to the packed buffer. - // - // These values are also sign-extended and accumulated into an - // intermediate per-column accumulator. CountK cannot be greater than - // 128 to avoid overflowing these signed 16-bit accumulators. - // - - while (k >= 2) { - - __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]); - __m128i BytesRow1 = _mm_loadl_epi64((__m128i*)&b[ldb]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); - - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - __m128i WordsLow = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum0 = _mm_add_epi16(ColumnSum0, WordsLow); - __m128i WordsHigh = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum1 = _mm_add_epi16(ColumnSum1, WordsHigh); - - b += ldb * 2; - D += 16; - k -= 2; - } - - if (k > 0) { - - // - // Process the remaining row of matrix B. - // - - __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, ZeroVector); - - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - __m128i WordsLow = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum0 = _mm_add_epi16(ColumnSum0, WordsLow); - __m128i WordsHigh = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum1 = _mm_add_epi16(ColumnSum1, WordsHigh); - - D += 16; - } - - // - // The number of rows written to the packed buffer should be a multiple - // of 4. Zero pad the packed buffer if the block is not complete. - // - - if (((CountK - 1) & 2) == 0) { - - _mm_storeu_si128((__m128i*)&D[0], ZeroVector); - - D += 16; - } - - ColumnSum0 = _mm_madd_epi16(ColumnSum0, OffsetBroadcast); - ColumnSum1 = _mm_madd_epi16(ColumnSum1, OffsetBroadcast); - - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSum0); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSum1); - - ColumnSumVector += 8; - - B += 8; - CountN -= 8; - } - - // - // Process the remaining columns of matrix B. - // - - if (CountN > 0) { - - const int8_t* b = B; - size_t k = CountK; - __m128i ColumnSum0 = ZeroVector; - __m128i ColumnSum1 = ZeroVector; - - // - // Interleave 2 rows of matrix B and write to the packed buffer. - // - // These values are also sign-extended and accumulated into an - // intermediate per-column accumulator. CountK cannot be greater than - // 128 to avoid overflowing these signed 16-bit accumulators. - // - - while (k >= 2) { - - // - // Copy the remaining columns to the zero padded stack buffer. - // - - const int8_t* bcopy = b; - int8_t* padded = PaddedMatrixBData; - int8_t* padded_end = padded + CountN; - - do { - padded[0] = bcopy[0]; - padded[8] = bcopy[ldb]; - padded++; - bcopy++; - } while (padded < padded_end); - - __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[0]); - __m128i BytesRow1 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[8]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); - - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - __m128i WordsLow = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum0 = _mm_add_epi16(ColumnSum0, WordsLow); - __m128i WordsHigh = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum1 = _mm_add_epi16(ColumnSum1, WordsHigh); - - b += ldb * 2; - D += 16; - k -= 2; - } - - if (k > 0) { - - // - // Copy the remaining columns to the zero padded stack buffer. - // - - const int8_t* bcopy = b; - int8_t* padded = PaddedMatrixBData; - int8_t* padded_end = padded + CountN; - - do { - padded[0] = bcopy[0]; - padded++; - bcopy++; - } while (padded < padded_end); - - __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[0]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, ZeroVector); - - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - __m128i WordsLow = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum0 = _mm_add_epi16(ColumnSum0, WordsLow); - __m128i WordsHigh = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BytesInterleaved), 8); - ColumnSum1 = _mm_add_epi16(ColumnSum1, WordsHigh); - - D += 16; - } - - // - // The number of rows written to the packed buffer should be a multiple - // of 4. Zero pad the packed buffer if the block is not complete. - // - - if (((CountK - 1) & 2) == 0) { - - _mm_storeu_si128((__m128i*)&D[0], ZeroVector); - - D += 16; - } - - ColumnSum0 = _mm_madd_epi16(ColumnSum0, OffsetBroadcast); - ColumnSum1 = _mm_madd_epi16(ColumnSum1, OffsetBroadcast); - - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSum0); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSum1); - } -} - -size_t -MLASCALL -MlasGemmU8S8KernelSse( - const uint8_t* A, - const int8_t* B, - int32_t* C, - size_t PairCountK, - size_t CountM, - size_t CountN, - size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, - int32_t DepthValue, - bool ZeroMode - ) -/*++ - -Routine Description: - - This routine is an inner kernel to compute matrix multiplication for a - set of rows. - -Arguments: - - A - Supplies the address of matrix A. The matrix data has been packed - using MlasGemmU8S8CopyPackASse. - - B - Supplies the address of matrix B. The matrix data has been packed - using MlasGemmU8S8CopyPackBSse. - - C - Supplies the address of matrix C. - - PairCountK - Supplies the number of paired columns from matrix A and the - number of paired rows from matrix B to iterate over. - - CountM - Supplies the maximum number of rows that can be processed for - matrix A and matrix C. The actual number of rows handled for this - invocation depends on the kernel implementation. - - CountN - Supplies the number of columns from matrix B and matrix C to iterate - over. - - ldc - Supplies the first dimension of matrix C. - - RowSumVector - Supplies the sum of each row from matrix A multiplied by the - zero point offset of matrix B. These values are accumulated into every - row of matrix C. - - ColumnSumVector - Supplies the sum of each column from matrix B multiplied - by the zero point offset of matrix A. These values are accumulated into - every column of matrix C. - - DepthValue - Supplies the value CountK multiplied by the zero point offset - of matrixA multplied by the zero point offset of matrix B. This value is - accumulated into every element of matrix C. - - ZeroMode - Supplies true if the output matrix must be zero initialized, - else false if the output matrix is accumulated into. - -Return Value: - - Returns the number of rows handled. - ---*/ -{ - const __m128i ZeroVector = _mm_setzero_si128(); - - MLAS_UNREFERENCED_PARAMETER(CountM); - MLAS_UNREFERENCED_PARAMETER(ldc); - - while (CountN > 0) { - - // - // Initialize the accumulators with the sum of the global depth value - // constant, the column sums, and the row sums. - // - - __m128i Accumulator0 = _mm_set1_epi32(DepthValue); - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_set1_epi32(RowSumVector[0])); - __m128i Accumulator1 = Accumulator0; - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&ColumnSumVector[0])); - Accumulator1 = _mm_add_epi32(Accumulator1, _mm_loadu_si128((__m128i*)&ColumnSumVector[4])); - ColumnSumVector += 8; - - // - // Broadcast each pair of 16-bit values from the matrix A and multiply - // with the zero-extended pair of 16-bit values from matrix B, and add - // the 32-bit intermediate into the accumulator registers. - // - - const uint8_t* a = A; - size_t k = PairCountK; - - while (k > 0) { - - __m128i AElements = _mm_unpacklo_epi8(_mm_cvtsi32_si128(*((int32_t*)a)), ZeroVector); - - __m128i BElements; - __m128i Intermediate0; - __m128i Intermediate1; - - BElements = _mm_loadu_si128((__m128i*)&B[0]); - Intermediate0 = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BElements), 8); - Intermediate1 = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BElements), 8); - - __m128i AElements0 = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(0, 0, 0, 0)); - - Intermediate0 = _mm_madd_epi16(Intermediate0, AElements0); - Intermediate1 = _mm_madd_epi16(Intermediate1, AElements0); - - Accumulator0 = _mm_add_epi32(Accumulator0, Intermediate0); - Accumulator1 = _mm_add_epi32(Accumulator1, Intermediate1); - - BElements = _mm_loadu_si128((__m128i*)&B[16]); - Intermediate0 = _mm_srai_epi16(_mm_unpacklo_epi8(ZeroVector, BElements), 8); - Intermediate1 = _mm_srai_epi16(_mm_unpackhi_epi8(ZeroVector, BElements), 8); - - __m128i AElements1 = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(1, 1, 1, 1)); - - Intermediate0 = _mm_madd_epi16(Intermediate0, AElements1); - Intermediate1 = _mm_madd_epi16(Intermediate1, AElements1); - - Accumulator0 = _mm_add_epi32(Accumulator0, Intermediate0); - Accumulator1 = _mm_add_epi32(Accumulator1, Intermediate1); - - a += 4; - B += 32; - k -= 1; - } - - // - // Output the accumulator block after optionally accumulating the values - // from matrix C. - // - - if (CountN >= 8) { - - if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&C[0])); - Accumulator1 = _mm_add_epi32(Accumulator1, _mm_loadu_si128((__m128i*)&C[4])); - } - - _mm_storeu_si128((__m128i*)&C[0], Accumulator0); - _mm_storeu_si128((__m128i*)&C[4], Accumulator1); - - C += 8; - CountN -= 8; - - } else { - - // - // Output the remaining partial output block. - // - - if ((CountN & 4) != 0) { - - if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&C[0])); - } - - _mm_storeu_si128((__m128i*)&C[0], Accumulator0); - C += 4; - - Accumulator0 = Accumulator1; - } - - if ((CountN & 2) != 0) { - - if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadl_epi64((__m128i*)&C[0])); - } - - _mm_storel_epi64((__m128i*)&C[0], Accumulator0); - C += 2; - - Accumulator0 = _mm_shuffle_epi32(Accumulator0, _MM_SHUFFLE(1, 0, 3, 2)); - } - - if ((CountN & 1) != 0) { - - int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulator0); - - if (!ZeroMode) { - AccumulatorValue += C[0]; - } - - C[0] = AccumulatorValue; - } - - CountN = 0; - } - } - - return 1; -} - -// -// U8U8 implementation using SSE2 intrinsics. -// - -void -MLASCALL -MlasGemmU8U8CopyPackASse( +MlasGemmU8X8CopyPackASse( int16_t* D, const uint8_t* A, size_t lda, @@ -751,11 +181,43 @@ Return Value: } } +template void -MLASCALL -MlasGemmU8U8CopyPackBSse( - uint8_t* D, - const uint8_t* B, +MlasGemmU8X8CopyPackBProcessSse( + int16_t* D, + __m128i BytesRow0, + __m128i BytesRow1, + __m128i ZeroVector, + __m128i ColumnSums[2] + ) +{ + __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); + __m128i WordsInterleaved[2]; + + // + // Zero or sign extend the bytes to words. + // + + if (std::is_same::value) { + WordsInterleaved[0] = _mm_unpacklo_epi8(BytesInterleaved, ZeroVector); + WordsInterleaved[1] = _mm_unpackhi_epi8(BytesInterleaved, ZeroVector); + } else { + WordsInterleaved[0] = _mm_srai_epi16(_mm_unpacklo_epi8(BytesInterleaved, BytesInterleaved), 8); + WordsInterleaved[1] = _mm_srai_epi16(_mm_unpackhi_epi8(BytesInterleaved, BytesInterleaved), 8); + } + + ColumnSums[0] = _mm_add_epi16(ColumnSums[0], WordsInterleaved[0]); + ColumnSums[1] = _mm_add_epi16(ColumnSums[1], WordsInterleaved[1]); + + _mm_storeu_si128((__m128i*)&D[0], WordsInterleaved[0]); + _mm_storeu_si128((__m128i*)&D[8], WordsInterleaved[1]); +} + +template +void +MlasGemmU8X8CopyPackBSse( + int16_t* D, + const BType* B, size_t ldb, size_t CountN, size_t CountK, @@ -796,7 +258,7 @@ Return Value: { const __m128i ZeroVector = _mm_setzero_si128(); const __m128i OffsetBroadcast = _mm_set1_epi16(offa); - uint8_t PaddedMatrixBData[16] = { 0 }; + BType PaddedMatrixBData[16] = { 0 }; // // Process 8 columns of matrix B in a loop. @@ -804,13 +266,15 @@ Return Value: while (CountN >= 8) { - const uint8_t* b = B; + const BType* b = B; size_t k = CountK; - __m128i ColumnSum0 = ZeroVector; - __m128i ColumnSum1 = ZeroVector; + __m128i ColumnSums[2]; + + ColumnSums[0] = ZeroVector; + ColumnSums[1] = ZeroVector; // - // Interleave 2 rows of matrix B and write to the packed buffer. + // Interleave rows of matrix B and write to the packed buffer. // // These values are also zero-extended and accumulated into an // intermediate per-column accumulator. CountK cannot be greater than @@ -821,12 +285,8 @@ Return Value: __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]); __m128i BytesRow1 = _mm_loadl_epi64((__m128i*)&b[ldb]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - ColumnSum0 = _mm_add_epi16(ColumnSum0, _mm_unpacklo_epi8(BytesInterleaved, ZeroVector)); - ColumnSum1 = _mm_add_epi16(ColumnSum1, _mm_unpackhi_epi8(BytesInterleaved, ZeroVector)); + MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, BytesRow1, ZeroVector, ColumnSums); b += ldb * 2; D += 16; @@ -835,26 +295,23 @@ Return Value: if (k > 0) { - // - // Process the remaining row of matrix B. - // - __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&b[0]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, ZeroVector); - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - ColumnSum0 = _mm_add_epi16(ColumnSum0, _mm_unpacklo_epi8(BytesInterleaved, ZeroVector)); - ColumnSum1 = _mm_add_epi16(ColumnSum1, _mm_unpackhi_epi8(BytesInterleaved, ZeroVector)); + MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, ZeroVector, ZeroVector, ColumnSums); D += 16; } - ColumnSum0 = _mm_madd_epi16(ColumnSum0, OffsetBroadcast); - ColumnSum1 = _mm_madd_epi16(ColumnSum1, OffsetBroadcast); + // + // Reduce the sum for the packed columns and multiply by the zero point + // offset of the other source matrix. + // - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSum0); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSum1); + ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OffsetBroadcast); + ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OffsetBroadcast); + + _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSums[0]); + _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSums[1]); ColumnSumVector += 8; @@ -868,20 +325,23 @@ Return Value: if (CountN > 0) { - const uint8_t* b = B; + const BType* b = B; size_t k = CountK; - __m128i ColumnSum0 = ZeroVector; - __m128i ColumnSum1 = ZeroVector; + __m128i ColumnSums[2]; + + ColumnSums[0] = ZeroVector; + ColumnSums[1] = ZeroVector; + + // + // Interleave rows of matrix B using an intermediate zero padded stack + // buffer and write to the packed buffer. + // while (k >= 2) { - // - // Copy the remaining columns to the zero padded stack buffer. - // - - const uint8_t* bcopy = b; - uint8_t* padded = PaddedMatrixBData; - uint8_t* padded_end = padded + CountN; + const BType* bcopy = b; + BType* padded = PaddedMatrixBData; + BType* padded_end = padded + CountN; do { padded[0] = bcopy[0]; @@ -892,12 +352,8 @@ Return Value: __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[0]); __m128i BytesRow1 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[8]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, BytesRow1); - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - ColumnSum0 = _mm_add_epi16(ColumnSum0, _mm_unpacklo_epi8(BytesInterleaved, ZeroVector)); - ColumnSum1 = _mm_add_epi16(ColumnSum1, _mm_unpackhi_epi8(BytesInterleaved, ZeroVector)); + MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, BytesRow1, ZeroVector, ColumnSums); b += ldb * 2; D += 16; @@ -906,13 +362,9 @@ Return Value: if (k > 0) { - // - // Copy the remaining columns to the zero padded stack buffer. - // - - const uint8_t* bcopy = b; - uint8_t* padded = PaddedMatrixBData; - uint8_t* padded_end = padded + CountN; + const BType* bcopy = b; + BType* padded = PaddedMatrixBData; + BType* padded_end = padded + CountN; do { padded[0] = bcopy[0]; @@ -921,12 +373,8 @@ Return Value: } while (padded < padded_end); __m128i BytesRow0 = _mm_loadl_epi64((__m128i*)&PaddedMatrixBData[0]); - __m128i BytesInterleaved = _mm_unpacklo_epi8(BytesRow0, ZeroVector); - _mm_storeu_si128((__m128i*)&D[0], BytesInterleaved); - - ColumnSum0 = _mm_add_epi16(ColumnSum0, _mm_unpacklo_epi8(BytesInterleaved, ZeroVector)); - ColumnSum1 = _mm_add_epi16(ColumnSum1, _mm_unpackhi_epi8(BytesInterleaved, ZeroVector)); + MlasGemmU8X8CopyPackBProcessSse(D, BytesRow0, ZeroVector, ZeroVector, ColumnSums); } // @@ -934,24 +382,36 @@ Return Value: // offset of the other source matrix. // - ColumnSum0 = _mm_madd_epi16(ColumnSum0, OffsetBroadcast); - ColumnSum1 = _mm_madd_epi16(ColumnSum1, OffsetBroadcast); + ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OffsetBroadcast); + ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OffsetBroadcast); - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSum0); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSum1); + _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSums[0]); + _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSums[1]); } } -size_t -MLASCALL -MlasGemmU8U8KernelSse( +MLAS_FORCEINLINE +void +MlasGemmU8X8MultiplyAccumulateRowSse( + __m128i ABroadcast, + const int16_t* B, + __m128i Accumulators[2] + ) +{ + __m128i BElements0 = _mm_load_si128((__m128i*)&B[0]); + __m128i BElements1 = _mm_load_si128((__m128i*)&B[8]); + + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_madd_epi16(BElements0, ABroadcast)); + Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_madd_epi16(BElements1, ABroadcast)); +} + +void +MlasGemmU8X8KernelSse( const int16_t* A, - const uint8_t* B, + const int16_t* B, int32_t* C, size_t PairCountK, - size_t CountM, size_t CountN, - size_t ldc, const int32_t* RowSumVector, const int32_t* ColumnSumVector, int32_t DepthValue, @@ -962,30 +422,24 @@ MlasGemmU8U8KernelSse( Routine Description: This routine is an inner kernel to compute matrix multiplication for a - set of rows. + single row. Arguments: A - Supplies the address of matrix A. The matrix data has been packed - using MlasGemmU8U8CopyPackASse. + using MlasGemmU8X8CopyPackASse. B - Supplies the address of matrix B. The matrix data has been packed - using MlasGemmU8U8CopyPackBSse. + using MlasGemmU8X8CopyPackBSse. C - Supplies the address of matrix C. PairCountK - Supplies the number of paired columns from matrix A and the number of paired rows from matrix B to iterate over. - CountM - Supplies the maximum number of rows that can be processed for - matrix A and matrix C. The actual number of rows handled for this - invocation depends on the kernel implementation. - CountN - Supplies the number of columns from matrix B and matrix C to iterate over. - ldc - Supplies the first dimension of matrix C. - RowSumVector - Supplies the sum of each row from matrix A multiplied by the zero point offset of matrix B. These values are accumulated into every row of matrix C. @@ -1003,51 +457,61 @@ Arguments: Return Value: - Returns the number of rows handled. + None. --*/ { - const __m128i ZeroVector = _mm_setzero_si128(); - - MLAS_UNREFERENCED_PARAMETER(CountM); - MLAS_UNREFERENCED_PARAMETER(ldc); - while (CountN > 0) { + __m128i Accumulators[2]; + // // Initialize the accumulators with the sum of the global depth value // constant, the column sums, and the row sums. // - __m128i Accumulator0 = _mm_set1_epi32(DepthValue); - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_set1_epi32(RowSumVector[0])); - __m128i Accumulator1 = Accumulator0; - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&ColumnSumVector[0])); - Accumulator1 = _mm_add_epi32(Accumulator1, _mm_loadu_si128((__m128i*)&ColumnSumVector[4])); + Accumulators[0] = _mm_set1_epi32(DepthValue); + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_set1_epi32(RowSumVector[0])); + Accumulators[1] = Accumulators[0]; + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&ColumnSumVector[0])); + Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&ColumnSumVector[4])); ColumnSumVector += 8; // // Broadcast each pair of 16-bit values from the matrix A and multiply - // with the zero-extended pair of 16-bit values from matrix B, and add - // the 32-bit intermediate into the accumulator registers. + // with the pair of 16-bit values from matrix B, and add the 32-bit + // intermediate into the accumulator registers. // const int16_t* a = A; size_t k = PairCountK; + while (k >= 4) { + + __m128i AElements = _mm_loadu_si128((__m128i*)a); + __m128i ABroadcast; + + ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(0, 0, 0, 0)); + MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[0], Accumulators); + + ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(1, 1, 1, 1)); + MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[16], Accumulators); + + ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(2, 2, 2, 2)); + MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[32], Accumulators); + + ABroadcast = _mm_shuffle_epi32(AElements, _MM_SHUFFLE(3, 3, 3, 3)); + MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[48], Accumulators); + + a += 4 * 2; + B += 4 * 16; + k -= 4; + } + while (k > 0) { - __m128i AElements = _mm_set1_epi32(*((int32_t*)a)); - __m128i BElements0 = _mm_loadu_si128((__m128i*)&B[0]); - - __m128i Intermediate0 = _mm_unpacklo_epi8(BElements0, ZeroVector); - __m128i Intermediate1 = _mm_unpackhi_epi8(BElements0, ZeroVector); - - Intermediate0 = _mm_madd_epi16(Intermediate0, AElements); - Intermediate1 = _mm_madd_epi16(Intermediate1, AElements); - - Accumulator0 = _mm_add_epi32(Accumulator0, Intermediate0); - Accumulator1 = _mm_add_epi32(Accumulator1, Intermediate1); + __m128i ABroadcast = _mm_set1_epi32(*((int32_t*)a)); + MlasGemmU8X8MultiplyAccumulateRowSse(ABroadcast, &B[0], Accumulators); a += 2; B += 16; @@ -1062,12 +526,12 @@ Return Value: if (CountN >= 8) { if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&C[0])); - Accumulator1 = _mm_add_epi32(Accumulator1, _mm_loadu_si128((__m128i*)&C[4])); + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); + Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&C[4])); } - _mm_storeu_si128((__m128i*)&C[0], Accumulator0); - _mm_storeu_si128((__m128i*)&C[4], Accumulator1); + _mm_storeu_si128((__m128i*)&C[0], Accumulators[0]); + _mm_storeu_si128((__m128i*)&C[4], Accumulators[1]); C += 8; CountN -= 8; @@ -1081,30 +545,30 @@ Return Value: if ((CountN & 4) != 0) { if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadu_si128((__m128i*)&C[0])); + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&C[0])); } - _mm_storeu_si128((__m128i*)&C[0], Accumulator0); + _mm_storeu_si128((__m128i*)&C[0], Accumulators[0]); C += 4; - Accumulator0 = Accumulator1; + Accumulators[0] = Accumulators[1]; } if ((CountN & 2) != 0) { if (!ZeroMode) { - Accumulator0 = _mm_add_epi32(Accumulator0, _mm_loadl_epi64((__m128i*)&C[0])); + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadl_epi64((__m128i*)&C[0])); } - _mm_storel_epi64((__m128i*)&C[0], Accumulator0); + _mm_storel_epi64((__m128i*)&C[0], Accumulators[0]); C += 2; - Accumulator0 = _mm_shuffle_epi32(Accumulator0, _MM_SHUFFLE(1, 0, 3, 2)); + Accumulators[0] = _mm_shuffle_epi32(Accumulators[0], _MM_SHUFFLE(3, 2, 3, 2)); } if ((CountN & 1) != 0) { - int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulator0); + int32_t AccumulatorValue = _mm_cvtsi128_si32(Accumulators[0]); if (!ZeroMode) { AccumulatorValue += C[0]; @@ -1116,13 +580,12 @@ Return Value: CountN = 0; } } - - return 1; } void MLASCALL -MlasGemmU8S8Operation( +MlasGemmU8X8OperationSse( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, size_t M, size_t N, size_t K, @@ -1144,6 +607,8 @@ Routine Description: Arguments: + WorkBlock - Supplies the structure containing the GEMM parameters. + M - Supplies the number of rows of matrix A and matrix C. N - Supplies the number of columns of matrix B and matrix C. @@ -1173,27 +638,15 @@ Return Value: --*/ { - MLAS_DECLSPEC_ALIGN(uint8_t PanelA[MLAS_GEMM_X8X8_STRIDEM * MLAS_GEMM_X8X8_STRIDEK], 64); - MLAS_DECLSPEC_ALIGN(int8_t PanelB[MLAS_GEMM_X8X8_STRIDEN * MLAS_GEMM_X8X8_STRIDEK], 64); + MLAS_DECLSPEC_ALIGN(int16_t PanelA[MLAS_GEMM_U8X8_STRIDEM_SSE * MLAS_GEMM_U8X8_STRIDEK_SSE], 16); + MLAS_DECLSPEC_ALIGN(int16_t PanelB[MLAS_GEMM_U8X8_STRIDEN_SSE * MLAS_GEMM_U8X8_STRIDEK_SSE], 16); - MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_X8X8_STRIDEM], 16); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_X8X8_STRIDEN], 16); + MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_U8X8_STRIDEM_SSE], 16); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_U8X8_STRIDEN_SSE], 16); - size_t StrideM = MLAS_GEMM_X8X8_STRIDEM; - size_t StrideN = MLAS_GEMM_X8X8_STRIDEN; - size_t StrideK = MLAS_GEMM_X8X8_STRIDEK; - -#if defined(MLAS_TARGET_AMD64) - - if (M == 1 && offa == 0 && offb == 0) { - - if (MlasPlatform.GemvU8S8Kernel != nullptr) { - MlasPlatform.GemvU8S8Kernel(A, (const int8_t*)B, C, K, N, ldb); - return; - } - } - -#endif + size_t StrideM = MLAS_GEMM_U8X8_STRIDEM_SSE; + size_t StrideN = MLAS_GEMM_U8X8_STRIDEN_SSE; + size_t StrideK = MLAS_GEMM_U8X8_STRIDEK_SSE; // // Step through each slice of matrix B along the K dimension. @@ -1203,11 +656,7 @@ Return Value: for (size_t k = 0; k < K; k += CountK) { - CountK = StrideK; - - if (CountK > (K - k)) { - CountK = K - k; - } + CountK = (std::min)(K - k, StrideK); // // Step through each slice of matrix B along the N dimension. @@ -1217,45 +666,289 @@ Return Value: for (size_t n = 0; n < N; n += CountN) { - CountN = StrideN; + CountN = (std::min)(N - n, StrideN); - if (CountN > (N - n)) { - CountN = N - n; + // + // Copy a panel of matrix B to a local packed buffer. + // + + const uint8_t* b = B + n + k * ldb; + + if (WorkBlock->BTypeIsSigned) { + MlasGemmU8X8CopyPackBSse(PanelB, (const int8_t*)b, ldb, CountN, + CountK, ColumnSumVector, -int16_t(offa)); + } else { + MlasGemmU8X8CopyPackBSse(PanelB, (const uint8_t*)b, ldb, CountN, + CountK, ColumnSumVector, -int16_t(offa)); } - const int8_t* b = (const int8_t*)B + n + k * ldb; + // + // Step through each slice of matrix A along the M dimension. + // - MlasPlatform.GemmU8S8CopyPackBRoutine(PanelB, b, ldb, CountN, - CountK, ColumnSumVector, -int16_t(offa)); + const int32_t DepthValue = int32_t(CountK) * offa * offb; + const size_t PairCountK = (CountK + 1) / 2; + int32_t* c = C + n; size_t CountM; for (size_t m = 0; m < M; m += CountM) { - CountM = StrideM; + CountM = (std::min)(M - m, StrideM); - if (CountM > (M - m)) { - CountM = M - m; - } + // + // Copy a panel of matrix A to a local packed buffer. + // - MlasPlatform.GemmU8S8CopyPackARoutine(PanelA, A + k + m * lda, - lda, CountM, CountK, RowSumVector, -int16_t(offb)); + MlasGemmU8X8CopyPackASse(PanelA, A + k + m * lda, lda, CountM, + CountK, RowSumVector, -int16_t(offb)); - uint8_t* pa = PanelA; - int32_t* c = C + n + m * ldc; + // + // Step through the rows of the local packed buffer. + // + int16_t* pa = PanelA; int32_t* RowSums = RowSumVector; - size_t RowsRemaining = CountM; - size_t RowsHandled; - - size_t QuadCountK = (CountK + 3) / 4; while (RowsRemaining > 0) { + MlasGemmU8X8KernelSse(pa, PanelB, c, PairCountK, CountN, + RowSums, ColumnSumVector, DepthValue, k == 0); + + c += ldc; + pa += 2 * PairCountK; + RowSums += 1; + RowsRemaining -= 1; + } + } + } + } +} + +#endif + +#ifdef MLAS_TARGET_AMD64 + +// +// Define the default strides to step through slices of the input matrices. +// + +#define MLAS_GEMM_U8X8_STRIDEM_AVX2 24 +#define MLAS_GEMM_U8X8_STRIDEN_AVX2 256 +#define MLAS_GEMM_U8X8_STRIDEK_AVX2 128 + +// +// 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 }; + +// +// Define the prototypes of the AVX2/AVX512 routines written in assembly. +// + +extern "C" { + + void + MLASCALL + MlasGemmU8S8CopyPackAAvx2( + uint8_t* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumVector, + int32_t offb + ); + + void + MLASCALL + MlasGemmU8S8CopyPackBAvx2( + int8_t* D, + const int8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumVector, + int32_t offa, + bool BTypeIsSigned + ); + + void + MLASCALL + MlasGemmU8U8CopyPackAAvx2( + int16_t* D, + const uint8_t* A, + size_t lda, + size_t CountM, + size_t CountK, + int32_t* RowSumVector, + int32_t offb + ); + + void + MLASCALL + MlasGemmU8U8CopyPackBAvx2( + uint8_t* D, + const uint8_t* B, + size_t ldb, + size_t CountN, + size_t CountK, + int32_t* ColumnSumVector, + int32_t offa + ); +} + +void +MLASCALL +MlasGemmU8S8OperationAvx2( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + int16_t offa, + const uint8_t* B, + size_t ldb, + int16_t offb, + int32_t* C, + size_t ldc + ) +/*++ + +Routine Description: + + This module implements the quantized integer matrix/matrix multiply + operation (QGEMM). + + This implementation supports AVX2/AVX512 U8S8 and AVX512VNNI U8S8/U8U8. + +Arguments: + + WorkBlock - Supplies the structure containing the GEMM parameters. + + M - Supplies the number of rows of matrix A and matrix C. + + N - Supplies the number of columns of matrix B and matrix C. + + K - Supplies the number of columns of matrix A and the number of rows of + matrix B. + + A - Supplies the address of matrix A. + + lda - Supplies the first dimension of matrix A. + + offa - Supplies the zero point offset of matrix A. + + B - Supplies the address of matrix B. + + ldb - Supplies the first dimension of matrix B. + + offb - Supplies the zero point offset of matrix B. + + C - Supplies the address of matrix C. + + ldc - Supplies the first dimension of matrix C. + +Return Value: + + None. + +--*/ +{ + MLAS_DECLSPEC_ALIGN(uint8_t PanelA[MLAS_GEMM_U8X8_STRIDEM_AVX2 * MLAS_GEMM_U8X8_STRIDEK_AVX2], 64); + MLAS_DECLSPEC_ALIGN(int8_t PanelB[MLAS_GEMM_U8X8_STRIDEN_AVX2 * MLAS_GEMM_U8X8_STRIDEK_AVX2], 64); + + MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_U8X8_STRIDEM_AVX2], 16); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_U8X8_STRIDEN_AVX2], 16); + + size_t StrideM = MLAS_GEMM_U8X8_STRIDEM_AVX2; + size_t StrideN = MLAS_GEMM_U8X8_STRIDEN_AVX2; + size_t StrideK = MLAS_GEMM_U8X8_STRIDEK_AVX2; + + if (WorkBlock->BTypeIsSigned) { + + if (M == 1 && offa == 0 && offb == 0) { + + if (MlasPlatform.GemvU8S8Kernel != nullptr) { + MlasPlatform.GemvU8S8Kernel(A, (const int8_t*)B, C, K, N, ldb); + return; + } + } + + } else { + + offb = int8_t(offb ^ 0x80); + } + + // + // Step through each slice of matrix B along the K dimension. + // + + size_t CountK; + + for (size_t k = 0; k < K; k += CountK) { + + CountK = (std::min)(K - k, StrideK); + + // + // Step through each slice of matrix B along the N dimension. + // + + size_t CountN; + + for (size_t n = 0; n < N; n += CountN) { + + CountN = (std::min)(N - n, StrideN); + + // + // Copy a panel of matrix B to a local packed buffer. + // + + const int8_t* b = (const int8_t*)B + n + k * ldb; + + MlasGemmU8S8CopyPackBAvx2(PanelB, b, ldb, CountN, + CountK, ColumnSumVector, -int16_t(offa), WorkBlock->BTypeIsSigned); + + // + // Step through each slice of matrix A along the M dimension. + // + + const int32_t DepthValue = int32_t(CountK) * offa * offb; + size_t QuadCountK = (CountK + 3) / 4; + + int32_t* c = C + n; + size_t CountM; + + for (size_t m = 0; m < M; m += CountM) { + + CountM = (std::min)(M - m, StrideM); + + // + // Copy a panel of matrix A to a local packed buffer. + // + + MlasGemmU8S8CopyPackAAvx2(PanelA, A + k + m * lda, + lda, CountM, CountK, RowSumVector, -int16_t(offb)); + + // + // Step through the rows of the local packed buffer. + // + + uint8_t* pa = PanelA; + int32_t* RowSums = RowSumVector; + size_t RowsRemaining = CountM; + + while (RowsRemaining > 0) { + + size_t RowsHandled; + RowsHandled = MlasPlatform.GemmU8S8Kernel(pa, PanelB, c, QuadCountK, RowsRemaining, CountN, ldc, RowSums, - ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); + ColumnSumVector, DepthValue, k == 0); RowsRemaining -= RowsHandled; c += ldc * RowsHandled; @@ -1269,7 +962,8 @@ Return Value: void MLASCALL -MlasGemmU8U8Operation( +MlasGemmU8U8OperationAvx2( + const MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, size_t M, size_t N, size_t K, @@ -1289,8 +983,12 @@ Routine Description: This module implements the quantized integer matrix/matrix multiply operation (QGEMM). + This implementation supports AVX2/AVX512 U8U8. + Arguments: + WorkBlock - Supplies the structure containing the GEMM parameters. + M - Supplies the number of rows of matrix A and matrix C. N - Supplies the number of columns of matrix B and matrix C. @@ -1320,15 +1018,17 @@ Return Value: --*/ { - MLAS_DECLSPEC_ALIGN(int16_t PanelA[MLAS_GEMM_X8X8_STRIDEM * MLAS_GEMM_X8X8_STRIDEK], 64); - MLAS_DECLSPEC_ALIGN(uint8_t PanelB[MLAS_GEMM_X8X8_STRIDEN * MLAS_GEMM_X8X8_STRIDEK], 64); + MLAS_DECLSPEC_ALIGN(int16_t PanelA[MLAS_GEMM_U8X8_STRIDEM_AVX2 * MLAS_GEMM_U8X8_STRIDEK_AVX2], 64); + MLAS_DECLSPEC_ALIGN(uint8_t PanelB[MLAS_GEMM_U8X8_STRIDEN_AVX2 * MLAS_GEMM_U8X8_STRIDEK_AVX2], 64); - MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_X8X8_STRIDEM], 16); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_X8X8_STRIDEN], 16); + MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[MLAS_GEMM_U8X8_STRIDEM_AVX2], 16); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[MLAS_GEMM_U8X8_STRIDEN_AVX2], 16); - size_t StrideM = MLAS_GEMM_X8X8_STRIDEM; - size_t StrideN = MLAS_GEMM_X8X8_STRIDEN; - size_t StrideK = MLAS_GEMM_X8X8_STRIDEK; + size_t StrideM = MLAS_GEMM_U8X8_STRIDEM_AVX2; + size_t StrideN = MLAS_GEMM_U8X8_STRIDEN_AVX2; + size_t StrideK = MLAS_GEMM_U8X8_STRIDEK_AVX2; + + MLAS_UNREFERENCED_PARAMETER(WorkBlock); // // Step through each slice of matrix B along the K dimension. @@ -1338,11 +1038,7 @@ Return Value: for (size_t k = 0; k < K; k += CountK) { - CountK = StrideK; - - if (CountK > (K - k)) { - CountK = K - k; - } + CountK = (std::min)(K - k, StrideK); // // Step through each slice of matrix B along the N dimension. @@ -1352,58 +1048,70 @@ Return Value: for (size_t n = 0; n < N; n += CountN) { - CountN = StrideN; + CountN = (std::min)(N - n, StrideN); - if (CountN > (N - n)) { - CountN = N - n; - } + // + // Copy a panel of matrix B to a local packed buffer. + // const uint8_t* b = (const uint8_t*)B + n + k * ldb; - MlasPlatform.GemmU8U8CopyPackBRoutine(PanelB, b, ldb, CountN, - CountK, ColumnSumVector, -int16_t(offa)); + MlasGemmU8U8CopyPackBAvx2(PanelB, b, ldb, CountN, CountK, ColumnSumVector, + -int16_t(offa)); + // + // Step through each slice of matrix A along the M dimension. + // + + const int32_t DepthValue = int32_t(CountK) * offa * offb; + size_t PairCountK = (CountK + 1) / 2; + + int32_t* c = C + n; size_t CountM; for (size_t m = 0; m < M; m += CountM) { - CountM = StrideM; + CountM = (std::min)(M - m, StrideM); - if (CountM > (M - m)) { - CountM = M - m; - } + // + // Copy a panel of matrix A to a local packed buffer. + // - MlasPlatform.GemmU8U8CopyPackARoutine(PanelA, A + k + m * lda, - lda, CountM, CountK, RowSumVector, -int16_t(offb)); + MlasGemmU8U8CopyPackAAvx2(PanelA, A + k + m * lda, lda, CountM, + CountK, RowSumVector, -int16_t(offb)); + + // + // Step through the rows of the local packed buffer. + // int16_t* pa = PanelA; - int32_t* c = C + n + m * ldc; - int32_t* RowSums = RowSumVector; - size_t RowsRemaining = CountM; - size_t RowsHandled; - - size_t PairCountK = (CountK + 1) / 2; while (RowsRemaining > 0) { + size_t RowsHandled; + RowsHandled = MlasPlatform.GemmU8U8Kernel(pa, PanelB, c, PairCountK, RowsRemaining, CountN, ldc, RowSums, - ColumnSumVector, int32_t(CountK) * offa * offb, k == 0); + ColumnSumVector, DepthValue, k == 0); - RowsRemaining -= RowsHandled; c += ldc * RowsHandled; pa += 2 * PairCountK * RowsHandled; RowSums += RowsHandled; + RowsRemaining -= RowsHandled; } } } } } +#endif + +#ifdef MLAS_TARGET_AMD64_IX86 + void -MlasGemmX8X8Threaded( +MlasGemmU8X8Threaded( void* Context, int32_t ThreadId ) @@ -1426,7 +1134,7 @@ Return Value: --*/ { - const auto* WorkBlock = (MLAS_GEMM_X8X8_WORK_BLOCK*)Context; + const auto* WorkBlock = (MLAS_GEMM_U8X8_WORK_BLOCK*)Context; const int32_t ThreadCountM = WorkBlock->ThreadCountM; const int32_t ThreadCountN = WorkBlock->ThreadCountN; @@ -1476,13 +1184,25 @@ Return Value: const uint8_t* b = WorkBlock->B + n; int32_t* c = WorkBlock->C + n + m * ldc; - WorkBlock->GemmX8X8Operation(CountM, CountN, WorkBlock->K, a, lda, + PMLAS_GEMM_U8X8_OPERATION GemmU8X8Operation; + +#if defined(MLAS_TARGET_AMD64) + if (WorkBlock->BTypeIsSigned) { + GemmU8X8Operation = MlasPlatform.GemmU8S8Operation; + } else { + GemmU8X8Operation = MlasPlatform.GemmU8U8Operation; + } +#else + GemmU8X8Operation = MlasGemmU8X8OperationSse; +#endif + + GemmU8X8Operation(WorkBlock, CountM, CountN, WorkBlock->K, a, lda, WorkBlock->offa, b, ldb, WorkBlock->offb, c, ldc); } void -MlasGemmX8X8Schedule( - MLAS_GEMM_X8X8_WORK_BLOCK* WorkBlock, +MlasGemmU8X8Schedule( + MLAS_GEMM_U8X8_WORK_BLOCK* WorkBlock, MLAS_THREADPOOL* ThreadPool ) /*++ @@ -1559,21 +1279,22 @@ Return Value: WorkBlock->ThreadCountN = 1; } - MlasExecuteThreaded(MlasGemmX8X8Threaded, WorkBlock, TargetThreadCount, ThreadPool); + MlasExecuteThreaded(MlasGemmU8X8Threaded, WorkBlock, TargetThreadCount, ThreadPool); } +template void MLASCALL MlasGemm( size_t M, size_t N, size_t K, - const uint8_t* A, + const AType* A, size_t lda, - uint8_t offa, - const int8_t* B, + AType offa, + const BType* B, size_t ldb, - int8_t offb, + BType offb, int32_t* C, size_t ldc, MLAS_THREADPOOL* ThreadPool @@ -1619,7 +1340,7 @@ Return Value: --*/ { - MLAS_GEMM_X8X8_WORK_BLOCK WorkBlock; + MLAS_GEMM_U8X8_WORK_BLOCK WorkBlock; // // Capture the GEMM parameters to the work block. @@ -1636,15 +1357,34 @@ Return Value: WorkBlock.ldc = ldc; WorkBlock.offa = int16_t(offa); WorkBlock.offb = int16_t(offb); - WorkBlock.GemmX8X8Operation = MlasGemmU8S8Operation; + WorkBlock.BTypeIsSigned = std::is_signed::value; // // Schedule the operation across a set of worker threads. // - MlasGemmX8X8Schedule(&WorkBlock, ThreadPool); + MlasGemmU8X8Schedule(&WorkBlock, ThreadPool); } +template +void +MLASCALL +MlasGemm( + size_t M, + size_t N, + size_t K, + const uint8_t* A, + size_t lda, + uint8_t offa, + const int8_t* B, + size_t ldb, + int8_t offb, + int32_t* C, + size_t ldc, + MLAS_THREADPOOL* ThreadPool + ); + +template void MLASCALL MlasGemm( @@ -1660,72 +1400,6 @@ MlasGemm( int32_t* C, size_t ldc, MLAS_THREADPOOL* ThreadPool - ) -/*++ - -Routine Description: - - This module implements the quantized integer matrix/matrix multiply - operation (QGEMM). - -Arguments: - - M - Supplies the number of rows of matrix A and matrix C. - - N - Supplies the number of columns of matrix B and matrix C. - - K - Supplies the number of columns of matrix A and the number of rows of - matrix B. - - A - Supplies the address of matrix A. - - lda - Supplies the first dimension of matrix A. - - offa - Supplies the zero point offset of matrix A. - - B - Supplies the address of matrix B. - - ldb - Supplies the first dimension of matrix B. - - offb - Supplies the zero point offset of matrix B. - - C - Supplies the address of matrix C. - - ldc - Supplies the first dimension of matrix C. - - ThreadPool - Supplies the thread pool object to use, else nullptr if the - base library threading support should be used. - -Return Value: - - None. - ---*/ -{ - MLAS_GEMM_X8X8_WORK_BLOCK WorkBlock; - - // - // Capture the GEMM parameters to the work block. - // - - WorkBlock.M = M; - WorkBlock.N = N; - WorkBlock.K = K; - WorkBlock.A = A; - WorkBlock.lda = lda; - WorkBlock.B = B; - WorkBlock.ldb = ldb; - WorkBlock.C = C; - WorkBlock.ldc = ldc; - WorkBlock.offa = int16_t(offa); - WorkBlock.offb = int16_t(offb); - WorkBlock.GemmX8X8Operation = MlasGemmU8U8Operation; - - // - // Schedule the operation across a set of worker threads. - // - - MlasGemmX8X8Schedule(&WorkBlock, ThreadPool); -} + ); #endif diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S index 6a6edd7e80..028cc9ceae 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S @@ -45,6 +45,7 @@ Abstract: .equ .LGemmU8S8CopyPackBFrame_SavedRbp, 8 .equ .LGemmU8S8CopyPackBFrame_ReturnAddress, 16 .equ .LGemmU8S8CopyPackBFrame_offa, 24 + .equ .LGemmU8S8CopyPackBFrame_BTypeIsSigned, 32 .text @@ -427,6 +428,9 @@ Arguments: offa - Supplies the zero point offset for the other source matrix of the matrix multiplication. + BTypeIsSigned - Supplies true if the source matrix is signed data, else + false if the the source matrix is unsigned data. + Return Value: None. @@ -447,6 +451,17 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vpsllw ymm0,ymm8,8 # generate word vector [0x0100] vpor ymm8,ymm8,ymm0 # generate word vector [0x0101] +// +// Compute the bit flip vector to adjust input from U8 to S8. +// + + vpxor xmm9,xmm9,xmm9 # generate word vector [0x0000] + cmp BYTE PTR .LGemmU8S8CopyPackBFrame_BTypeIsSigned[rsp],0 + jnz .LCopyPackB.SkipUnsignedBitFlipVector + vpsllw ymm9,ymm8,7 # generate word vector [0x8080] + +.LCopyPackB.SkipUnsignedBitFlipVector: + // // Process 16 columns of matrix B in a loop. // @@ -481,6 +496,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vpunpckhwd xmm3,xmm3,xmm5 vinserti128 ymm4,ymm4,xmm6,1 vinserti128 ymm2,ymm2,xmm3,1 + vpxor ymm4,ymm4,ymm9 # optionally adjust unsigned data + vpxor ymm2,ymm2,ymm9 vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows vmovdqu YMMWORD PTR [rdi+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row @@ -499,9 +516,9 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): add rbx,4 # correct for over-subtract above jz .LCopyPackB.ReduceColumnSumVectorN16 vmovdqu xmm2,XMMWORD PTR [rdx] - vpxor xmm3,xmm3,xmm3 - vpxor xmm4,xmm4,xmm4 - vpxor xmm5,xmm5,xmm5 + vmovaps xmm3,xmm9 + vmovaps xmm4,xmm9 + vmovaps xmm5,xmm9 xor ebx,ebx # no more rows remaining test r8b,2 # (CountK & 2) != 0? jz .LCopyPackB.InterleaveRowDataN16 @@ -542,8 +559,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.ProcessColumnNUnaligned: vpxor xmm0,xmm0,xmm0 # clear column accumulators vpxor xmm1,xmm1,xmm1 - vmovdqu YMMWORD PTR .LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp],ymm0 - vmovdqu YMMWORD PTR .LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp+32],ymm0 + vmovdqu YMMWORD PTR .LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp],ymm9 + vmovdqu YMMWORD PTR .LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp+32],ymm9 sub r8,4 jb .LCopyPackB.ProcessRemainingRowsNUnaligned @@ -618,6 +635,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vpunpckhwd xmm3,xmm3,xmm5 vinserti128 ymm4,ymm4,xmm6,1 vinserti128 ymm2,ymm2,xmm3,1 + vpxor ymm4,ymm4,ymm9 # optionally adjust unsigned data + vpxor ymm2,ymm2,ymm9 vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows vmovdqu YMMWORD PTR [rdi+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row @@ -638,9 +657,8 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): // lea rbp,.LGemmU8S8CopyPackBFrame_PaddedMatrixBData[rsp] - vpxor xmm6,xmm6,xmm6 - vmovdqu YMMWORD PTR [rbp],ymm6 - vmovdqu YMMWORD PTR [rbp+32],ymm6 + vmovdqu YMMWORD PTR [rbp],ymm9 + vmovdqu YMMWORD PTR [rbp+32],ymm9 .LCopyPackB.CopyUnalignedRowLoop: lea r11,[rbp+16] # advance next padded buffer by 16 bytes diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Common.h b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Common.h deleted file mode 100644 index 9872817b3c..0000000000 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Common.h +++ /dev/null @@ -1,64 +0,0 @@ -/*++ - -Copyright (c) Microsoft Corporation. All rights reserved. - -Licensed under the MIT License. - -Module Name: - - QgemmU8U8KernelAvx512Common.h - -Abstract: - - This module contains common kernel macros and structures for the quantized - integer matrix/matrix multiply operation (QGEMM) for the AVX512 core and - AVX512VNNI kernels. - ---*/ - -#include "QgemmU8X8KernelAvx512Common.h" - -/*++ - - Macro Description: - - This macro generates code to execute the block compute macro multiple - times and advancing the matrix A and matrix B data pointers. - -Arguments: - - ColumnCount - Supplies the number of columns to produce. - - RowCount - Supplies the number of rows to produce. - -Implicit Arguments: - - rbx - Supplies the address into the matrix A data plus 3 rows. - - rdi - Supplies the address into the matrix A data. - - rsi - Supplies the address into the matrix B data. - - rcx - Supplies the length in bytes of a row from matrix A. - - r14 - Supplies the stride in bytes of between packed blocks of matrix B. - - zmm14-zmm31 - Supplies the block accumulators. - ---*/ - - .macro ComputeBlockLoop ColumnCount, RowCount - - mov rbp,rcx # reload row length remaining - -.LComputeBlockBy1Loop\@: - ComputeBlock \ColumnCount\(), \RowCount\(), 0, 0 - add rdi,4 # advance matrix A by 1 pair -.if \RowCount\() > 3 - add rbx,4 # advance matrix A plus 3 rows by 1 pair -.endif - add rsi,32 # advance matrix B - sub rbp,4 - jnz .LComputeBlockBy1Loop\@ - - .endm diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Core.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Core.S index d00021616a..8fca218e3b 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Core.S +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Core.S @@ -18,7 +18,7 @@ Abstract: --*/ #include "asmmacro.h" -#include "QgemmU8U8KernelAvx512Common.h" +#include "QgemmU8X8KernelAvx512Common.h" .intel_syntax noprefix @@ -124,6 +124,51 @@ Implicit Arguments: .endm +/*++ + + Macro Description: + + This macro generates code to execute the block compute macro multiple + times and advancing the matrix A and matrix B data pointers. + +Arguments: + + ColumnCount - Supplies the number of columns to produce. + + RowCount - Supplies the number of rows to produce. + +Implicit Arguments: + + rbx - Supplies the address into the matrix A data plus 3 rows. + + rdi - Supplies the address into the matrix A data. + + rsi - Supplies the address into the matrix B data. + + rcx - Supplies the length in bytes of a row from matrix A. + + r14 - Supplies the stride in bytes of between packed blocks of matrix B. + + zmm14-zmm31 - Supplies the block accumulators. + +--*/ + + .macro ComputeBlockLoop ColumnCount, RowCount + + mov rbp,rcx # reload row length remaining + +.LComputeBlockBy1Loop\@: + ComputeBlock \ColumnCount\(), \RowCount\(), 0, 0 + add rdi,4 # advance matrix A by 1 pair +.if \RowCount\() > 3 + add rbx,4 # advance matrix A plus 3 rows by 1 pair +.endif + add rsi,32 # advance matrix B + sub rbp,4 + jnz .LComputeBlockBy1Loop\@ + + .endm + // // Generate the GEMM kernel. // diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Vnni.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Vnni.S deleted file mode 100644 index c946f50bee..0000000000 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx512Vnni.S +++ /dev/null @@ -1,114 +0,0 @@ -/*++ - -Copyright (c) Microsoft Corporation. All rights reserved. - -Licensed under the MIT License. - -Module Name: - - QgemmU8U8KernelAvx512Vnni.s - -Abstract: - - This module implements the kernels for the quantized integer matrix/matrix - multiply operation (QGEMM). - - This implementation uses AVX512VNNI instructions. - ---*/ - -#include "asmmacro.h" -#include "QgemmU8U8KernelAvx512Common.h" -#include "AssembleAvx512Vnni.h" - - .intel_syntax noprefix - - .text - -/*++ - -Macro Description: - - This macro generates code to multiply and accumulate each row of the output - block. - -Arguments: - - ColumnCount - Supplies the number of columns to produce. - - RowCount - Supplies the number of rows to produce. - - VectorOffset - Supplies the byte offset from matrix B to fetch elements. - - BroadcastOffset - Supplies the byte offset from matrix A to fetch elements. - -Implicit Arguments: - - rbx - Supplies the address into the matrix A data plus 3 rows. - - rdi - Supplies the address into the matrix A data. - - rsi - Supplies the address into the matrix B data. - - rcx - Supplies the length in bytes of a row from matrix A. - - r14 - Supplies the stride in bytes of between packed blocks of matrix B. - - zmm14-zmm31 - Supplies the block accumulators. - ---*/ - - .macro ComputeBlock ColumnCount, RowCount, VectorOffset, BroadcastOffset - -.if \ColumnCount\() >= 32 -.if \ColumnCount\() >= 48 - vpmovzxbw zmm0,YMMWORD PTR [rsi+\VectorOffset\()] - vpmovzxbw zmm1,YMMWORD PTR [rsi+r14+\VectorOffset\()] - vpmovzxbw zmm2,YMMWORD PTR [rsi+r14*2+\VectorOffset\()] -.else - vpmovzxbw zmm1,YMMWORD PTR [rsi+\VectorOffset\()] - vpmovzxbw zmm2,YMMWORD PTR [rsi+r14+\VectorOffset\()] -.endif - EmitIfCountGE \RowCount\(), 1, "vpbroadcastd zmm3,DWORD PTR [rdi+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 1, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm26,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 1, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm20,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 1, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm14,zmm3,zmm2" - EmitIfCountGE \RowCount\(), 2, "vpbroadcastd zmm3,DWORD PTR [rdi+rcx+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 2, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm27,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 2, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm21,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 2, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm15,zmm3,zmm2" - EmitIfCountGE \RowCount\(), 3, "vpbroadcastd zmm3,DWORD PTR [rdi+rcx*2+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 3, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm28,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 3, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm22,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 3, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm16,zmm3,zmm2" - EmitIfCountGE \RowCount\(), 4, "vpbroadcastd zmm3,DWORD PTR [rbx+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 4, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm29,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 4, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm23,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 4, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm17,zmm3,zmm2" - EmitIfCountGE \RowCount\(), 5, "vpbroadcastd zmm3,DWORD PTR [rbx+rcx+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 5, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm30,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 5, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm24,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 5, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm18,zmm3,zmm2" - EmitIfCountGE \RowCount\(), 6, "vpbroadcastd zmm3,DWORD PTR [rbx+rcx*2+\BroadcastOffset\()]" - EmitIfCount2GE \RowCount\(), 6, \ColumnCount\(), 48, "VpdpwssdZmmZmmZmm zmm31,zmm3,zmm0" - EmitIfCount2GE \RowCount\(), 6, \ColumnCount\(), 32, "VpdpwssdZmmZmmZmm zmm25,zmm3,zmm1" - EmitIfCount2GE \RowCount\(), 6, \ColumnCount\(), 16, "VpdpwssdZmmZmmZmm zmm19,zmm3,zmm2" -.else - vpmovzxbw zmm2,YMMWORD PTR [rsi+\VectorOffset\()] - EmitIfCountGE \RowCount\(), 1, "VpdpwssdZmmZmmBroadcast zmm14,zmm2,rdi,\BroadcastOffset\()" - EmitIfCountGE \RowCount\(), 2, "VpdpwssdZmmZmmBroadcast zmm15,zmm2,rdi,\BroadcastOffset\(),rcx,1" - EmitIfCountGE \RowCount\(), 3, "VpdpwssdZmmZmmBroadcast zmm16,zmm2,rdi,\BroadcastOffset\(),rcx,2" - EmitIfCountGE \RowCount\(), 4, "VpdpwssdZmmZmmBroadcast zmm17,zmm2,rbx,\BroadcastOffset\()" - EmitIfCountGE \RowCount\(), 5, "VpdpwssdZmmZmmBroadcast zmm18,zmm2,rbx,\BroadcastOffset\(),rcx,1" - EmitIfCountGE \RowCount\(), 6, "VpdpwssdZmmZmmBroadcast zmm19,zmm2,rbx,\BroadcastOffset\(),rcx,2" -.endif - - .endm - -// -// Generate the GEMM kernel. -// - -GemmU8X8KernelAvx512Function U8U8, Avx512Vnni - - .end diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h index 172bd9fab1..7dbff45389 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h @@ -106,8 +106,8 @@ Implicit Arguments: .endif ComputeBlockLoop \ColumnCount\(), \RowCount\() .if \RowCount\() > 3 - lea rbx,[rdx+rax*2] # compute matrix C plus 3 rows - add rbx,rax + lea rbx,[rax*2+rax] + add rbx,rdx # compute matrix C plus 3 rows .endif .endm