diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm index 751ce80553..9aec9a195e 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8S8KernelAvx2.asm @@ -47,8 +47,7 @@ GemmU8S8CopyPackAFrame STRUCT PreviousP3Home QWORD ? PreviousP4Home QWORD ? CountK QWORD ? - RowSumVector QWORD ? - offb QWORD ? + RowSumBuffer QWORD ? GemmU8S8CopyPackAFrame ENDS @@ -74,9 +73,8 @@ GemmU8S8CopyPackBFrame STRUCT PreviousP3Home QWORD ? PreviousP4Home QWORD ? CountK QWORD ? - ColumnSumVector QWORD ? - offa QWORD ? - BTypeIsSigned QWORD ? + ColumnSumBuffer QWORD ? + BIsSigned QWORD ? GemmU8S8CopyPackBFrame ENDS @@ -99,11 +97,8 @@ GemmU8S8CopyPackBFrame ENDS ; ; 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. -; -; offb - Supplies the zero point offset for the other source matrix of the -; matrix multiplication. +; RowSumBuffer - Supplies the address of the buffer to receive the sums of +; the elements along each of the rows. ; ; Return Value: ; @@ -133,12 +128,11 @@ GemmU8S8CopyPackBFrame ENDS mov r10,GemmU8S8CopyPackAFrame.CountK[rsp] lea r11,[r10+3] and r11,NOT 3 ; align CountK up to quad count - mov r12,GemmU8S8CopyPackAFrame.RowSumVector[rsp] - vpbroadcastw xmm8,WORD PTR GemmU8S8CopyPackAFrame.offb[rsp] - vpcmpeqw ymm9,ymm9,ymm9 ; generate word vector [0xFFFF] - vpsrlw ymm9,ymm9,15 ; generate word vector [0x0001] - vpsllw ymm0,ymm9,8 ; generate word vector [0x0100] - vpor ymm9,ymm9,ymm0 ; generate word vector [0x0101] + mov r12,GemmU8S8CopyPackAFrame.RowSumBuffer[rsp] + vpcmpeqw ymm8,ymm8,ymm8 ; generate word vector [0xFFFF] + vpsrlw ymm8,ymm8,15 ; generate word vector [0x0001] + vpsllw ymm9,ymm8,8 ; generate word vector [0x0100] + vpor ymm9,ymm8,ymm9 ; generate word vector [0x0101] ; ; Compute the conditional load/store mask for an unaligned CountK. @@ -206,7 +200,7 @@ ProcessNextColumnLoopM4: ProcessRemainingColumnsM4: add rbx,32 ; correct for over-subtract above - jz ReduceRowSumVectorM4 + jz ReduceRowSumBufferM4 test bl,16 ; (CountK & 16) != 0? jz CopyRemainingCountKLessThan16M4 vmovdqu xmm4,XMMWORD PTR [rdx] @@ -228,7 +222,7 @@ ProcessRemainingColumnsM4: add rdx,16 ; advance matrix A by 16 bytes add rcx,16 ; advance matrix D by 16 bytes test bl,15 ; test for unaligned columns - jz ReduceRowSumVectorM4 + jz ReduceRowSumBufferM4 ; ; Copy the unaligned CountK columns to a zero padded stack buffer. @@ -317,15 +311,18 @@ ProcessPaddedMatrixADataM4: ; Reduce the sums for the four rows of output. ; -ReduceRowSumVectorM4: - vphaddw ymm0,ymm0,ymm1 ; reduce and interleave Sum1/Sum0 - vphaddw ymm1,ymm2,ymm3 ; reduce and interleave Sum3/Sum2 - vphaddw ymm0,ymm0,ymm1 ; reduce and interleave Sum3/Sum2/Sum1/Sum0 - vextracti128 xmm1,ymm0,1 ; extract high pairs - vpaddw xmm0,xmm0,xmm1 ; reduce low/high pairs - vpmaddwd xmm0,xmm0,xmm8 ; multiply by offset and reduce 32-bit sum +ReduceRowSumBufferM4: + vpmaddwd ymm0,ymm0,ymm8 ; horizontal word+word=dword per row + vpmaddwd ymm1,ymm1,ymm8 + vphaddd ymm0,ymm0,ymm1 ; reduce and interleave Sum1/Sum0 + vpmaddwd ymm2,ymm2,ymm8 + vpmaddwd ymm3,ymm3,ymm8 + vphaddd ymm1,ymm2,ymm3 ; reduce and interleave Sum3/Sum2 + vphaddd ymm0,ymm0,ymm1 ; reduce and interleave Sum3/Sum2/Sum1/Sum0 + vextracti128 xmm1,ymm0,1 ; extract high dwords + vpaddd xmm0,xmm0,xmm1 ; reduce low/high dwords vmovdqu XMMWORD PTR [r12],xmm0 - add r12,4*4 ; advance row sum vector by 4 DWORDs + add r12,4*4 ; advance row sum buffer by 4 dwords sub r9,4 ; subtract rows remaining jae ProcessNextRowM4 @@ -359,7 +356,7 @@ ProcessNextColumnLoopM1: ProcessRemainingColumnsM1: add rbx,32 ; correct for over-subtract above - jz ReduceRowSumVectorM1 + jz ReduceRowSumBufferM1 test bl,16 ; (CountK & 16) != 0? jz CopyRemainingCountKLessThan16M1 vmovdqu xmm4,XMMWORD PTR [rdx] @@ -369,7 +366,7 @@ ProcessRemainingColumnsM1: add rdx,16 ; advance matrix A by 16 bytes add rcx,16 ; advance matrix D by 16 bytes test bl,15 ; test for unaligned columns - jz ReduceRowSumVectorM1 + jz ReduceRowSumBufferM1 ; ; Copy the unaligned CountK columns to a zero padded stack buffer. @@ -415,20 +412,20 @@ ProcessPaddedMatrixADataM1: vmovdqu xmm4,XMMWORD PTR GemmU8S8CopyPackAFrame.PaddedMatrixAData[rsp] vpmaskmovd XMMWORD PTR [rcx],xmm10,xmm4 vpmaddubsw ymm4,ymm4,ymm9 ; horizontal byte+byte=word per row - vpaddw ymm0,ymm0,ymm4 ; accumulate per row along columns + vpaddw ymm0,ymm0,ymm4 ; add words to row accumulators ; ; Reduce the sum for the single row of output. ; -ReduceRowSumVectorM1: - vextracti128 xmm1,ymm0,1 ; extract high pairs - vpaddw xmm0,xmm0,xmm1 ; reduction - vphaddw xmm0,xmm0,xmm0 - vphaddw xmm0,xmm0,xmm0 - vpmaddwd xmm0,xmm0,xmm8 ; multiply by offset and reduce +ReduceRowSumBufferM1: + vpmaddwd ymm0,ymm0,ymm8 ; horizontal word+word=dword per row + vextracti128 xmm1,ymm0,1 ; extract high dwords + vpaddd xmm0,xmm0,xmm1 ; reduction + vphaddd xmm0,xmm0,xmm0 + vphaddd xmm0,xmm0,xmm0 vmovd DWORD PTR [r12],xmm0 - add r12,4 ; advance row sum vector by 1 DWORD + add r12,4 ; advance row sum buffer by 1 dword dec r9 ; decrement rows remaining jnz ProcessNextRowM1 @@ -476,15 +473,11 @@ ExitRoutine: ; ; 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. +; ColumnSumBuffer - Supplies the address of the buffer to receive the sums of +; the elements along each of the columns. ; -; 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. +; BIsSigned - Supplies true if the source matrix is signed data, else false +; if the source matrix is unsigned data. ; ; Return Value: ; @@ -509,19 +502,18 @@ ExitRoutine: mov rsi,rdx lea rdi,[r8+r8*2] ; compute ldb * 3 mov r10,GemmU8S8CopyPackBFrame.CountK[rsp] - mov r11,GemmU8S8CopyPackBFrame.ColumnSumVector[rsp] - vpbroadcastw ymm7,WORD PTR GemmU8S8CopyPackBFrame.offa[rsp] - vpcmpeqw ymm8,ymm8,ymm8 ; generate word vector [0xFFFF] - vpsrlw ymm8,ymm8,15 ; generate word vector [0x0001] - vpsllw ymm0,ymm8,8 ; generate word vector [0x0100] - vpor ymm8,ymm8,ymm0 ; generate word vector [0x0101] + mov r11,GemmU8S8CopyPackBFrame.ColumnSumBuffer[rsp] + vpcmpeqw ymm7,ymm7,ymm7 ; generate word vector [0xFFFF] + vpsrlw ymm7,ymm7,15 ; generate word vector [0x0001] + vpsllw ymm8,ymm7,8 ; generate word vector [0x0100] + vpor ymm8,ymm7,ymm8 ; generate word vector [0x0101] ; ; Compute the bit flip vector to adjust input from U8 to S8. ; vpxor xmm9,xmm9,xmm9 ; generate word vector [0x0000] - cmp BYTE PTR GemmU8S8CopyPackBFrame.BTypeIsSigned[rsp],0 + cmp BYTE PTR GemmU8S8CopyPackBFrame.BIsSigned[rsp],0 jnz SkipUnsignedBitFlipVector vpsllw ymm9,ymm8,7 ; generate word vector [0x8080] @@ -566,9 +558,11 @@ InterleaveRowDataN16: vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows vmovdqu YMMWORD PTR [rcx+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row - vpaddw ymm0,ymm0,ymm4 ; add words to row accumulators + vpmaddwd ymm4,ymm4,ymm7 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column vpmaddubsw ymm2,ymm8,ymm2 - vpaddw ymm1,ymm1,ymm2 + vpmaddwd ymm2,ymm2,ymm7 + vpaddd ymm1,ymm1,ymm2 add rcx,64 ; advance matrix D by 64 bytes sub rbx,4 ; subtract rows remaining jae ProcessNextRowLoopN16 @@ -579,7 +573,7 @@ InterleaveRowDataN16: ProcessRemainingRowsN16: add rbx,4 ; correct for over-subtract above - jz ReduceColumnSumVectorN16 + jz StoreColumnSumBufferN16 vmovdqu xmm2,XMMWORD PTR [rdx] vmovaps xmm3,xmm9 vmovaps xmm4,xmm9 @@ -593,12 +587,10 @@ ProcessRemainingRowsN16: vmovdqu xmm4,XMMWORD PTR [rdx+r8*2] jmp InterleaveRowDataN16 -ReduceColumnSumVectorN16: - vpmaddwd ymm0,ymm0,ymm7 ; multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm7 ; multiply by offset and reduce +StoreColumnSumBufferN16: vmovdqu YMMWORD PTR [r11],ymm0 vmovdqu YMMWORD PTR [r11+32],ymm1 - add r11,16*4 ; advance column sum vector by 16 DWORDs + add r11,16*4 ; advance column sum buffer by 16 dwords sub r9,16 ; subtract columns remaining jae ProcessNextColumnN16 @@ -715,9 +707,11 @@ ProcessPaddedMatrixBData: vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows vmovdqu YMMWORD PTR [rcx+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 ; horizontal byte+byte=word per row - vpaddw ymm0,ymm0,ymm4 ; add words to row accumulators + vpmaddwd ymm4,ymm4,ymm7 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column vpmaddubsw ymm2,ymm8,ymm2 - vpaddw ymm1,ymm1,ymm2 + vpmaddwd ymm2,ymm2,ymm7 + vpaddd ymm1,ymm1,ymm2 lea rsi,[rsi+r8*4] ; advance next matrix B by 4 rows add rcx,64 ; advance matrix D by 64 bytes sub r10,4 ; subtract rows remaining @@ -725,7 +719,7 @@ ProcessPaddedMatrixBData: ProcessRemainingRowsNUnaligned: add r10,4 - jz ReduceColumnSumVectorNUnaligned + jz StoreColumnSumBufferNUnaligned ; ; Process the less than 4 remaining rows where the row has less than 16 columns. @@ -775,9 +769,7 @@ DoneCopyRemainingCountNKSmall: mov rbp,rdi jmp CopyUnalignedRowLoop -ReduceColumnSumVectorNUnaligned: - vpmaddwd ymm0,ymm0,ymm7 ; multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm7 ; multiply by offset and reduce +StoreColumnSumBufferNUnaligned: vmovdqu YMMWORD PTR [r11],ymm0 vmovdqu YMMWORD PTR [r11+32],ymm1 jmp ExitRoutine @@ -958,11 +950,11 @@ ENDIF ; ; ldc - Supplies the first dimension of matrix C. ; -; RowSumVector - Supplies the sum of each row from matrix A multiplied by the +; RowSumBuffer - 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 +; ColumnSumBuffer - 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. ; @@ -1005,8 +997,8 @@ ENDIF shl r9,2 ; convert to row length movzx r10,BYTE PTR GemmU8X8KernelFrame.ZeroMode[rsp] mov r11,GemmU8X8KernelFrame.CountM[rsp] - mov r12,GemmU8X8KernelFrame.RowSumVector[rsp] - mov r13,GemmU8X8KernelFrame.ColumnSumVector[rsp] + mov r12,GemmU8X8KernelFrame.RowSumBuffer[rsp] + mov r13,GemmU8X8KernelFrame.ColumnSumBuffer[rsp] vpcmpeqw ymm12,ymm12,ymm12 ; generate 256-bit word vector [0xFFFF] vpsrlw ymm12,ymm12,15 ; generate 256-bit word vector [0x0001] diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm index 6a573f3ef6..d30823756f 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm @@ -46,8 +46,7 @@ GemmU8U8CopyPackAFrame STRUCT PreviousP3Home QWORD ? PreviousP4Home QWORD ? CountK QWORD ? - RowSumVector QWORD ? - offb QWORD ? + RowSumBuffer QWORD ? GemmU8U8CopyPackAFrame ENDS @@ -67,8 +66,7 @@ GemmU8U8CopyPackBFrame STRUCT PreviousP3Home QWORD ? PreviousP4Home QWORD ? CountK QWORD ? - ColumnSumVector QWORD ? - offa QWORD ? + ColumnSumBuffer QWORD ? GemmU8U8CopyPackBFrame ENDS @@ -96,12 +94,8 @@ GemmU8U8CopyPackBFrame ENDS ; ; 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. +; RowSumBuffer - Supplies the address of the buffer to receive the sums of +; the elements along each of the rows. ; ; Return Value: ; @@ -130,8 +124,9 @@ GemmU8U8CopyPackBFrame ENDS mov r10,GemmU8U8CopyPackAFrame.CountK[rsp] lea r11,[r10+1] and r11,NOT 1 ; align CountK up to pair count - mov r12,GemmU8U8CopyPackAFrame.RowSumVector[rsp] - vpbroadcastw xmm8,WORD PTR GemmU8U8CopyPackAFrame.offb[rsp] + mov r12,GemmU8U8CopyPackAFrame.RowSumBuffer[rsp] + vpcmpeqw ymm8,ymm8,ymm8 ; generate word vector [0xFFFF] + vpsrlw ymm8,ymm8,15 ; generate word vector [0x0001] ; ; Compute the conditional load/store mask for an unaligned CountK. @@ -205,7 +200,7 @@ ProcessNextColumnLoopM4: ProcessRemainingColumnsM4: add rbx,16 ; correct for over-subtract above - jz ReduceRowSumVectorM4 + jz ReduceRowSumBufferM4 ; ; Copy the unaligned CountK columns to a zero padded stack buffer. @@ -290,21 +285,21 @@ ProcessPaddedMatrixADataM4: vpaddw ymm3,ymm3,ymm7 ; -; Reduce the sums for the four rows of output. Transpose the intermediate -; accumulators by treating the registers as 32-bit elements containing a pair -; of 16-bit sums. Continue reducing the transposed accumulators to produce the -; final 32-bit vector output. +; Reduce the sums for the four rows of output. ; -ReduceRowSumVectorM4: - vphaddw ymm0,ymm0,ymm1 ; reduce and interleave Sum1/Sum0 - vphaddw ymm1,ymm2,ymm3 ; reduce and interleave Sum3/Sum2 - vphaddw ymm0,ymm0,ymm1 ; reduce and interleave Sum3/Sum2/Sum1/Sum0 - vextracti128 xmm1,ymm0,1 ; extract high pairs - vpaddw xmm0,xmm0,xmm1 ; reduce low/high pairs - vpmaddwd xmm0,xmm0,xmm8 ; multiply by offset and reduce 32-bit sum +ReduceRowSumBufferM4: + vpmaddwd ymm0,ymm0,ymm8 ; horizontal word+word=dword per row + vpmaddwd ymm1,ymm1,ymm8 + vphaddd ymm0,ymm0,ymm1 ; reduce and interleave Sum1/Sum0 + vpmaddwd ymm2,ymm2,ymm8 + vpmaddwd ymm3,ymm3,ymm8 + vphaddd ymm1,ymm2,ymm3 ; reduce and interleave Sum3/Sum2 + vphaddd ymm0,ymm0,ymm1 ; reduce and interleave Sum3/Sum2/Sum1/Sum0 + vextracti128 xmm1,ymm0,1 ; extract high dwords + vpaddd xmm0,xmm0,xmm1 ; reduce low/high dwords vmovdqu XMMWORD PTR [r12],xmm0 - add r12,4*4 ; advance row sum vector by 4 DWORDs + add r12,4*4 ; advance row sum buffer by 4 dwords sub r9,4 ; subtract rows remaining jae ProcessNextRowM4 @@ -337,7 +332,7 @@ ProcessNextColumnLoopM1: ProcessRemainingColumnsM1: add rbx,16 ; correct for over-subtract above - jz ReduceRowSumVectorM1 + jz ReduceRowSumBufferM1 ; ; Copy the unaligned CountK columns to a zero padded stack buffer. @@ -387,14 +382,14 @@ ProcessPaddedMatrixADataM1: ; Reduce the sum for the single row of output. ; -ReduceRowSumVectorM1: - vextracti128 xmm1,ymm0,1 ; extract high pairs - vpaddw xmm0,xmm0,xmm1 ; reduction - vphaddw xmm0,xmm0,xmm0 - vphaddw xmm0,xmm0,xmm0 - vpmaddwd xmm0,xmm0,xmm8 ; multiply by offset and reduce +ReduceRowSumBufferM1: + vpmaddwd ymm0,ymm0,ymm8 ; horizontal word+word=dword per row + vextracti128 xmm1,ymm0,1 ; extract high dwords + vpaddd xmm0,xmm0,xmm1 ; reduction + vphaddd xmm0,xmm0,xmm0 + vphaddd xmm0,xmm0,xmm0 vmovd DWORD PTR [r12],xmm0 - add r12,4 ; advance row sum vector by 1 DWORD + add r12,4 ; advance row sum buffer by 1 dword dec r9 ; decrement rows remaining jnz ProcessNextRowM1 @@ -441,12 +436,8 @@ ExitRoutine: ; ; 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. +; ColumnSumBuffer - Supplies the address of the buffer to receive the sums of +; the elements along each of the columns. ; ; Return Value: ; @@ -465,8 +456,9 @@ ExitRoutine: mov rsi,rdx mov r10,GemmU8U8CopyPackBFrame.CountK[rsp] - mov r11,GemmU8U8CopyPackBFrame.ColumnSumVector[rsp] - vpbroadcastw ymm5,WORD PTR GemmU8U8CopyPackBFrame.offa[rsp] + mov r11,GemmU8U8CopyPackBFrame.ColumnSumBuffer[rsp] + vpcmpeqw ymm5,ymm5,ymm5 ; generate word vector [0xFFFF] + vpsrlw ymm5,ymm5,15 ; generate word vector [0x0001] ; ; Zero initialize the padded stack buffers. @@ -502,29 +494,31 @@ ProcessNextRowLoopN16: vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 add rcx,32 ; advance matrix D by 32 bytes - vpaddw ymm0,ymm0,ymm4 ; accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 sub rbx,2 ; subtract rows remaining jae ProcessNextRowLoopN16 ProcessRemainingRowsN16: add rbx,2 ; correct for over-subtract above - jz ReduceColumnSumVectorN16 + jz StoreColumnSumBufferN16 vpmovzxbw ymm4,XMMWORD PTR [rdx] vmovdqu YMMWORD PTR [rcx],ymm4 ; store interleaved rows vextracti128 xmm3,ymm4,1 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 ; accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 add rcx,32 ; advance matrix D by 32 bytes -ReduceColumnSumVectorN16: - vpmaddwd ymm0,ymm0,ymm5 ; multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm5 ; multiply by offset and reduce +StoreColumnSumBufferN16: vmovdqu YMMWORD PTR [r11],ymm0 vmovdqu YMMWORD PTR [r11+32],ymm1 - add r11,16*4 ; advance column sum vector by 16 DWORDs + add r11,16*4 ; advance column sum buffer by 16 dwords sub r9,16 ; subtract columns remaining jae ProcessNextColumnN16 @@ -607,8 +601,10 @@ ProcessPaddedMatrixBDataK2: vmovdqu XMMWORD PTR [rcx+16],xmm3 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 ; accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 lea rsi,[rsi+r8*2] ; advance next matrix B by 2 rows add rcx,32 ; advance matrix D by 32 bytes sub r10,2 ; subtract columns remaining @@ -616,7 +612,7 @@ ProcessPaddedMatrixBDataK2: ProcessRemainingRowsNUnaligned: add r10,2 - jz ReduceColumnSumVectorNUnaligned + jz StoreColumnSumBufferNUnaligned mov rdx,rsi .errnz GemmU8U8CopyPackBFrame.PaddedMatrixBData mov rbp,rsp ; GemmU8U8CopyPackBFrame.PaddedMatrixBData @@ -655,12 +651,12 @@ ProcessPaddedMatrixBDataK1: vextracti128 xmm3,ymm4,1 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 ; accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 ; horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 ; accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 -ReduceColumnSumVectorNUnaligned: - vpmaddwd ymm0,ymm0,ymm5 ; multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm5 ; multiply by offset and reduce +StoreColumnSumBufferNUnaligned: vmovdqu YMMWORD PTR [r11],ymm0 vmovdqu YMMWORD PTR [r11+32],ymm1 jmp ExitRoutine @@ -852,11 +848,11 @@ ComputeBlockLoopExit: ; ; ldc - Supplies the first dimension of matrix C. ; -; RowSumVector - Supplies the sum of each row from matrix A multiplied by the +; RowSumBuffer - 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 +; ColumnSumBuffer - 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. ; @@ -902,8 +898,8 @@ ComputeBlockLoopExit: shl r9,2 ; convert to row length movzx r10,BYTE PTR GemmU8X8KernelFrame.ZeroMode[rsp] mov r11,GemmU8X8KernelFrame.CountM[rsp] - mov r12,GemmU8X8KernelFrame.RowSumVector[rsp] - mov r13,GemmU8X8KernelFrame.ColumnSumVector[rsp] + mov r12,GemmU8X8KernelFrame.RowSumBuffer[rsp] + mov r13,GemmU8X8KernelFrame.ColumnSumBuffer[rsp] ; ; Process CountM rows of the matrices. diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc index ecfe53d271..9d45b4fa25 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc @@ -48,8 +48,8 @@ GemmU8X8KernelFrame STRUCT CountM QWORD ? CountN QWORD ? ldc QWORD ? - RowSumVector QWORD ? - ColumnSumVector QWORD ? + RowSumBuffer QWORD ? + ColumnSumBuffer QWORD ? DepthValue QWORD ? ZeroMode QWORD ? @@ -77,9 +77,9 @@ GemmU8X8KernelFrame ENDS ; ; r9 - Supplies the length in bytes of a row from matrix A. ; -; r12 - Supplies the address of the row sum vector. +; r12 - Supplies the address of the row sum buffer. ; -; r13 - Supplies the address of the column sum vector. +; r13 - Supplies the address of the column sum buffer. ; ; ymm4-ymm15 - Supplies the block accumulators. ; @@ -95,7 +95,7 @@ ProduceOutputBlock MACRO ColumnCount, RowCount IF ColumnCount EQ 16 vpaddd ymm0,ymm1,YMMWORD PTR [r13] vpaddd ymm1,ymm1,YMMWORD PTR [r13+32] - add r13,16*4 ; advance ColumnSumVector by 16 columns + add r13,16*4 ; advance ColumnSumBuffer by 16 columns ELSE vpaddd ymm1,ymm1,YMMWORD PTR [r13] ENDIF @@ -166,9 +166,9 @@ ENDIF ; ; r10b - Supplies the zero mode flag. ; -; r12 - Supplies the address of the row sum vector. +; r12 - Supplies the address of the row sum buffer. ; -; r13 - Supplies the address of the column sum vector. +; r13 - Supplies the address of the column sum buffer. ; ProcessCountM MACRO RowCount, Fallthrough diff --git a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx512Common.inc b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx512Common.inc index c2825e1eb0..ae02a70a05 100644 --- a/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx512Common.inc +++ b/onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx512Common.inc @@ -39,8 +39,8 @@ GemmU8X8KernelFrame STRUCT CountM QWORD ? CountN QWORD ? ldc QWORD ? - RowSumVector QWORD ? - ColumnSumVector QWORD ? + RowSumBuffer QWORD ? + ColumnSumBuffer QWORD ? DepthValue QWORD ? ZeroMode QWORD ? @@ -68,9 +68,9 @@ GemmU8X8KernelFrame ENDS ; ; r9 - Supplies the length in bytes of a row from matrix A. ; -; r12 - Supplies the address of the row sum vector. +; r12 - Supplies the address of the row sum buffer. ; -; r13 - Supplies the address of the column sum vector. +; r13 - Supplies the address of the column sum buffer. ; ProduceOutputBlock MACRO ColumnCount, RowCount @@ -90,7 +90,7 @@ ELSE vpaddd zmm1,zmm3,ZMMWORD PTR [r13] vpaddd zmm0,zmm3,ZMMWORD PTR [r13+64] ENDIF - add_immed r13,ColumnCount*4 ; advance ColumnSumVector by N columns + add_immed r13,ColumnCount*4 ; advance ColumnSumBuffer by N columns ELSE vpaddd zmm0,zmm3,ZMMWORD PTR [r13] ENDIF @@ -158,9 +158,9 @@ ENDIF ; ; r10b - Supplies the zero mode flag. ; -; r12 - Supplies the address of the row sum vector. +; r12 - Supplies the address of the row sum buffer. ; -; r13 - Supplies the address of the column sum vector. +; r13 - Supplies the address of the column sum buffer. ; ; r14 - Supplies the stride in bytes of between packed blocks of matrix B. ; @@ -321,11 +321,11 @@ GemmU8X8KernelAvx512Function MACRO Type, Isa ; ; ldc - Supplies the first dimension of matrix C. ; -; RowSumVector - Supplies the sum of each row from matrix A multiplied by the +; RowSumBuffer - 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 +; ColumnSumBuffer - 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. ; @@ -364,8 +364,8 @@ GemmU8X8KernelAvx512Function MACRO Type, Isa shl r9,2 ; convert to row length movzx r10,BYTE PTR GemmU8X8KernelFrame.ZeroMode[rsp] mov r11,GemmU8X8KernelFrame.CountM[rsp] - mov r12,GemmU8X8KernelFrame.RowSumVector[rsp] - mov r13,GemmU8X8KernelFrame.ColumnSumVector[rsp] + mov r12,GemmU8X8KernelFrame.RowSumBuffer[rsp] + mov r13,GemmU8X8KernelFrame.ColumnSumBuffer[rsp] mov esi,-1 kmovw k1,esi ; update mask to write all columns IFIDNI , diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index b254b33993..21106a7dff 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -596,18 +596,16 @@ extern "C" { #endif MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernel; - MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8Kernel; - MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8Kernel; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32Kernel; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasLogisticKernel; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasTanhKernel; MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32Kernel; MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32Kernel; MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32Kernel; + MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8Kernel; + MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8Kernel; #if defined(MLAS_TARGET_AMD64) MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernelFma3; - MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8KernelAvx2; - MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8KernelAvx2; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelFma3; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelAvx512F; MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasLogisticKernelFma3; @@ -616,15 +614,14 @@ extern "C" { MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelAvx512F; MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32KernelAvx; MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelAvx; + MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8KernelAvx2; + MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8KernelAvx2; #endif MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32Kernel; -#if defined(MLAS_TARGET_AMD64) - MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx; -#endif - MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32Kernel; #if defined(MLAS_TARGET_AMD64) + MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx; MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32KernelAvx; #endif diff --git a/onnxruntime/core/mlas/lib/qgemm.cpp b/onnxruntime/core/mlas/lib/qgemm.cpp index 8c553728bb..f1c7f89acf 100644 --- a/onnxruntime/core/mlas/lib/qgemm.cpp +++ b/onnxruntime/core/mlas/lib/qgemm.cpp @@ -38,10 +38,34 @@ struct MLAS_GEMM_U8X8_WORK_BLOCK { const float* BiasFloat; uint8_t offa; uint8_t offb; - bool BTypeIsSigned; - bool CTypeIsFloat; + bool BIsSigned; + bool CIsFloat; }; +void +MlasGemmU8X8ScaleSumBuffer( + int32_t* D, + const int32_t* S, + size_t N, + int32_t Scale + ) +{ + for (size_t n = 0; n < N; n++) { + D[n] = S[n] * Scale; + } +} + +MLAS_FORCEINLINE +void +MlasGemmU8X8ScaleSumBuffer( + int32_t* SumBuffer, + size_t N, + int32_t Scale + ) +{ + return MlasGemmU8X8ScaleSumBuffer(SumBuffer, SumBuffer, N, Scale); +} + template MLAS_FORCEINLINE void @@ -68,8 +92,8 @@ Return Value: MLAS_DECLSPEC_ALIGN(typename KernelType::PackedAType PanelA[KernelType::StrideM * KernelType::StrideK], 64); MLAS_DECLSPEC_ALIGN(typename KernelType::PackedBType PanelB[KernelType::StrideN * KernelType::StrideK], 64); - MLAS_DECLSPEC_ALIGN(int32_t RowSumVector[KernelType::StrideM], 64); - MLAS_DECLSPEC_ALIGN(int32_t ColumnSumVector[KernelType::StrideN], 64); + MLAS_DECLSPEC_ALIGN(int32_t RowSumBuffer[KernelType::StrideM], 64); + MLAS_DECLSPEC_ALIGN(int32_t ColumnSumBuffer[KernelType::StrideN], 64); const uint8_t* A = WorkBlock->A; const uint8_t* B = WorkBlock->B; @@ -79,16 +103,18 @@ Return Value: const size_t ldb = WorkBlock->ldb; const size_t ldc = WorkBlock->ldc; + int32_t offa = WorkBlock->offa; + int32_t offb = typename KernelType::OffsetBType(WorkBlock->offb); + // // Flip the sign bit of the zero point offset of matrix B if the kernel uses // signed types and the matrix B data is unsigned. // - int16_t offa = WorkBlock->offa; - int16_t offb = typename KernelType::OffsetBType(WorkBlock->offb); - - if (std::is_signed::value && !WorkBlock->BTypeIsSigned) { - offb = typename KernelType::OffsetBType(offb ^ 0x80); + if (std::is_signed::value) { + if (!WorkBlock->BIsSigned) { + offb = typename KernelType::OffsetBType(offb ^ 0x80); + } } // @@ -119,7 +145,9 @@ Return Value: // KernelType::CopyPackB(PanelB, B + n, ldb, CountN, CountK, - ColumnSumVector, -offa, WorkBlock->BTypeIsSigned); + ColumnSumBuffer, WorkBlock->BIsSigned); + + MlasGemmU8X8ScaleSumBuffer(ColumnSumBuffer, CountN, -offa); // // Step through each slice of matrix A along the M dimension. @@ -141,14 +169,16 @@ Return Value: // KernelType::CopyPackA(PanelA, A + m * lda, lda, CountM, CountK, - RowSumVector, -offb); + RowSumBuffer); + + MlasGemmU8X8ScaleSumBuffer(RowSumBuffer, CountM, -offb); // // Step through the rows of the local packed buffer. // typename KernelType::PackedAType* pa = PanelA; - int32_t* RowSums = RowSumVector; + int32_t* RowSums = RowSumBuffer; size_t RowsRemaining = CountM; bool ZeroMode = (k == 0); @@ -159,10 +189,10 @@ Return Value: size_t RowsHandled; RowsHandled = KernelType::Kernel(pa, PanelB, c, PackedCountK, - RowsRemaining, CountN, ldc, RowSums, ColumnSumVector, + RowsRemaining, CountN, ldc, RowSums, ColumnSumBuffer, DepthValue, ZeroMode); - if (PostProcess && WorkBlock->CTypeIsFloat) { + if (PostProcess && WorkBlock->CIsFloat) { KernelType::OutputFloat(WorkBlock, c, n, RowsHandled, CountN); } @@ -188,8 +218,7 @@ MlasGemmU8X8CopyPackASse( size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ) /*++ @@ -210,12 +239,8 @@ Arguments: 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. + RowSumBuffer - Supplies the address of the buffer to receive the sums of + the elements along each of the rows. Return Value: @@ -224,7 +249,7 @@ Return Value: --*/ { const __m128i ZeroVector = _mm_setzero_si128(); - const __m128i OffsetBroadcast = _mm_set1_epi16(offb); + const __m128i OnesWordBroadcast = _mm_set1_epi16(1); uint8_t PaddedMatrixAData[8] = { 0 }; // @@ -235,7 +260,7 @@ Return Value: const uint8_t* a = A; size_t k = CountK; - __m128i RowSum = ZeroVector; + __m128i ReductionVector = ZeroVector; // // Zero extend the source bytes to 16-bits and write to the packed @@ -255,7 +280,7 @@ Return Value: __m128i Bytes = _mm_loadl_epi64((__m128i*)&a[0]); __m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector); - RowSum = _mm_add_epi16(RowSum, Words); + ReductionVector = _mm_add_epi16(ReductionVector, Words); _mm_storeu_si128((__m128i*)&D[0], Words); @@ -282,7 +307,7 @@ Return Value: __m128i Bytes = _mm_loadl_epi64((__m128i*)PaddedMatrixAData); __m128i Words = _mm_unpacklo_epi8(Bytes, ZeroVector); - RowSum = _mm_add_epi16(RowSum, Words); + ReductionVector = _mm_add_epi16(ReductionVector, Words); // // Copy pairs of 16-bit values from the vector to the packed @@ -297,15 +322,16 @@ Return Value: } // - // Reduce the sum for the single row of output and multiply by the - // zero point offset of the other source matrix. + // Reduce the partial accumulators. // - 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))); + ReductionVector = _mm_madd_epi16(ReductionVector, OnesWordBroadcast); + ReductionVector = _mm_add_epi32(ReductionVector, + _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(3, 2, 3, 2))); + ReductionVector = _mm_add_epi32(ReductionVector, + _mm_shuffle_epi32(ReductionVector, _MM_SHUFFLE(0, 1, 0, 1))); - *RowSumVector++ = _mm_cvtsi128_si32(RowSum); + *RowSumBuffer++ = _mm_cvtsi128_si32(ReductionVector); A += lda; CountM -= 1; @@ -342,9 +368,8 @@ MlasGemmU8X8CopyPackBSse( size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa, - bool BTypeIsSigned + int32_t* ColumnSumBuffer, + bool BIsSigned ) /*++ @@ -365,12 +390,11 @@ Arguments: 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. + ColumnSumBuffer - Supplies the address of the buffer to receive the sums of + the elements along each of the columns. - offa - Supplies the zero point offset for the other source matrix of the - matrix multiplication. + BIsSigned - Supplies true if the source matrix is signed data, else false + if the source matrix is unsigned data. Return Value: @@ -378,8 +402,8 @@ Return Value: --*/ { - const __m128i OffsetBroadcast = _mm_set1_epi16(offa); - const __m128i BitFlipVector = _mm_set1_epi32(BTypeIsSigned ? 0 : 0x80808080); + const __m128i OnesWordBroadcast = _mm_set1_epi16(1); + const __m128i BitFlipVector = _mm_set1_epi32(BIsSigned ? 0 : 0x80808080); // // Process 8 columns of matrix B in a loop. @@ -424,17 +448,16 @@ Return Value: } // - // Reduce the sum for the packed columns and multiply by the zero point - // offset of the other source matrix. + // Reduce the partial accumulators. // - ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OffsetBroadcast); - ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OffsetBroadcast); + ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OnesWordBroadcast); + ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OnesWordBroadcast); - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSums[0]); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSums[1]); + _mm_storeu_si128((__m128i*)&ColumnSumBuffer[0], ColumnSums[0]); + _mm_storeu_si128((__m128i*)&ColumnSumBuffer[4], ColumnSums[1]); - ColumnSumVector += 8; + ColumnSumBuffer += 8; B += 8; CountN -= 8; @@ -506,11 +529,11 @@ Return Value: // offset of the other source matrix. // - ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OffsetBroadcast); - ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OffsetBroadcast); + ColumnSums[0] = _mm_madd_epi16(ColumnSums[0], OnesWordBroadcast); + ColumnSums[1] = _mm_madd_epi16(ColumnSums[1], OnesWordBroadcast); - _mm_storeu_si128((__m128i*)&ColumnSumVector[0], ColumnSums[0]); - _mm_storeu_si128((__m128i*)&ColumnSumVector[4], ColumnSums[1]); + _mm_storeu_si128((__m128i*)&ColumnSumBuffer[0], ColumnSums[0]); + _mm_storeu_si128((__m128i*)&ColumnSumBuffer[4], ColumnSums[1]); } } @@ -536,8 +559,8 @@ MlasGemmU8X8KernelSse( int32_t* C, size_t PackedCountK, size_t CountN, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, int32_t DepthValue, bool ZeroMode ) @@ -564,11 +587,11 @@ Arguments: CountN - Supplies the number of columns from matrix B and matrix C to iterate over. - RowSumVector - Supplies the sum of each row from matrix A multiplied by the + RowSumBuffer - 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 + ColumnSumBuffer - 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. @@ -595,11 +618,11 @@ Return Value: // Accumulators[0] = _mm_set1_epi32(DepthValue); - Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_set1_epi32(RowSumVector[0])); + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_set1_epi32(RowSumBuffer[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; + Accumulators[0] = _mm_add_epi32(Accumulators[0], _mm_loadu_si128((__m128i*)&ColumnSumBuffer[0])); + Accumulators[1] = _mm_add_epi32(Accumulators[1], _mm_loadu_si128((__m128i*)&ColumnSumBuffer[4])); + ColumnSumBuffer += 8; // // Broadcast each pair of 16-bit values from the matrix A and multiply @@ -832,11 +855,10 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ) { - MlasGemmU8X8CopyPackASse(D, A, lda, CountM, CountK, RowSumVector, offb); + MlasGemmU8X8CopyPackASse(D, A, lda, CountM, CountK, RowSumBuffer); } MLAS_FORCEINLINE @@ -848,13 +870,12 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa, - bool BTypeIsSigned + int32_t* ColumnSumBuffer, + bool BIsSigned ) { - MlasGemmU8X8CopyPackBSse(D, B, ldb, CountN, CountK, ColumnSumVector, offa, - BTypeIsSigned); + MlasGemmU8X8CopyPackBSse(D, B, ldb, CountN, CountK, ColumnSumBuffer, + BIsSigned); } MLAS_FORCEINLINE @@ -868,8 +889,8 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE size_t CountM, size_t CountN, size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, int32_t DepthValue, bool ZeroMode ) @@ -877,8 +898,8 @@ struct MLAS_GEMM_U8X8_KERNEL_SSE MLAS_UNREFERENCED_PARAMETER(CountM); MLAS_UNREFERENCED_PARAMETER(ldc); - MlasGemmU8X8KernelSse(A, B, C, PackedCountK, CountN, RowSumVector, - ColumnSumVector, DepthValue, ZeroMode); + MlasGemmU8X8KernelSse(A, B, C, PackedCountK, CountN, RowSumBuffer, + ColumnSumBuffer, DepthValue, ZeroMode); return 1; } @@ -955,8 +976,7 @@ extern "C" { size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ); void @@ -967,9 +987,8 @@ extern "C" { size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa, - bool BTypeIsSigned + int32_t* ColumnSumBuffer, + bool BIsSigned ); void @@ -980,8 +999,7 @@ extern "C" { size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ); void @@ -992,8 +1010,7 @@ extern "C" { size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa + int32_t* ColumnSumBuffer ); } @@ -1017,11 +1034,10 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ) { - MlasGemmU8S8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumVector, offb); + MlasGemmU8S8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumBuffer); } MLAS_FORCEINLINE @@ -1033,13 +1049,12 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa, - bool BTypeIsSigned + int32_t* ColumnSumBuffer, + bool BIsSigned ) { - MlasGemmU8S8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumVector, offa, - BTypeIsSigned); + MlasGemmU8S8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumBuffer, + BIsSigned); } MLAS_FORCEINLINE @@ -1053,14 +1068,14 @@ struct MLAS_GEMM_U8S8_KERNEL_AVX2 size_t CountM, size_t CountN, size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, int32_t DepthValue, bool ZeroMode ) { return MlasPlatform.GemmU8S8Kernel(A, B, C, PackedCountK, CountM, CountN, - ldc, RowSumVector, ColumnSumVector, DepthValue, ZeroMode); + ldc, RowSumBuffer, ColumnSumBuffer, DepthValue, ZeroMode); } MLAS_FORCEINLINE @@ -1107,7 +1122,7 @@ Return Value: --*/ { - if ((WorkBlock->M == 1) && WorkBlock->BTypeIsSigned && !WorkBlock->CTypeIsFloat && + if ((WorkBlock->M == 1) && WorkBlock->BIsSigned && !WorkBlock->CIsFloat && (WorkBlock->offa == 0) && (WorkBlock->offb == 0)) { if (MlasPlatform.GemvU8S8Kernel != nullptr) { @@ -1140,11 +1155,10 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 size_t lda, size_t CountM, size_t CountK, - int32_t* RowSumVector, - int16_t offb + int32_t* RowSumBuffer ) { - MlasGemmU8U8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumVector, offb); + MlasGemmU8U8CopyPackAAvx2(D, A, lda, CountM, CountK, RowSumBuffer); } MLAS_FORCEINLINE @@ -1156,14 +1170,13 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 size_t ldb, size_t CountN, size_t CountK, - int32_t* ColumnSumVector, - int16_t offa, - bool BTypeIsSigned + int32_t* ColumnSumBuffer, + bool BIsSigned ) { - MLAS_UNREFERENCED_PARAMETER(BTypeIsSigned); + MLAS_UNREFERENCED_PARAMETER(BIsSigned); - MlasGemmU8U8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumVector, offa); + MlasGemmU8U8CopyPackBAvx2(D, B, ldb, CountN, CountK, ColumnSumBuffer); } MLAS_FORCEINLINE @@ -1177,14 +1190,14 @@ struct MLAS_GEMM_U8U8_KERNEL_AVX2 size_t CountM, size_t CountN, size_t ldc, - const int32_t* RowSumVector, - const int32_t* ColumnSumVector, + const int32_t* RowSumBuffer, + const int32_t* ColumnSumBuffer, int32_t DepthValue, bool ZeroMode ) { return MlasPlatform.GemmU8U8Kernel(A, B, C, PackedCountK, CountM, CountN, - ldc, RowSumVector, ColumnSumVector, DepthValue, ZeroMode); + ldc, RowSumBuffer, ColumnSumBuffer, DepthValue, ZeroMode); } MLAS_FORCEINLINE @@ -1319,7 +1332,7 @@ Return Value: } #if defined(MLAS_TARGET_AMD64) - if (WorkBlock->BTypeIsSigned) { + if (LocalWorkBlock.BIsSigned) { MlasPlatform.GemmU8S8Operation(&LocalWorkBlock); } else { MlasPlatform.GemmU8U8Operation(&LocalWorkBlock); @@ -1488,7 +1501,7 @@ Return Value: WorkBlock.ldc = ldc; WorkBlock.offa = offa; WorkBlock.offb = offb; - WorkBlock.BTypeIsSigned = std::is_signed::value; + WorkBlock.BIsSigned = std::is_signed::value; // // Schedule the operation across a set of worker threads. @@ -1614,8 +1627,8 @@ Return Value: WorkBlock.BiasFloat = Bias; WorkBlock.offa = offa; WorkBlock.offb = offb; - WorkBlock.BTypeIsSigned = std::is_signed::value; - WorkBlock.CTypeIsFloat = true; + WorkBlock.BIsSigned = std::is_signed::value; + WorkBlock.CIsFloat = true; // // Schedule the operation across a set of worker threads. diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S index 9b8d7a9096..6245aa97df 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8S8KernelAvx2.S @@ -33,7 +33,6 @@ Abstract: .equ .LGemmU8S8CopyPackAFrame_SavedRbx, 16 .equ .LGemmU8S8CopyPackAFrame_SavedRbp, 24 .equ .LGemmU8S8CopyPackAFrame_ReturnAddress, 32 - .equ .LGemmU8S8CopyPackAFrame_offb, 40 // // Stack frame layout for the U8S8 CopyPackB routine. @@ -44,8 +43,7 @@ Abstract: .equ .LGemmU8S8CopyPackBFrame_SavedRbx, 0 .equ .LGemmU8S8CopyPackBFrame_SavedRbp, 8 .equ .LGemmU8S8CopyPackBFrame_ReturnAddress, 16 - .equ .LGemmU8S8CopyPackBFrame_offa, 24 - .equ .LGemmU8S8CopyPackBFrame_BTypeIsSigned, 32 + .equ .LGemmU8S8CopyPackBFrame_BIsSigned, 24 .text @@ -68,13 +66,10 @@ Arguments: CountK (r8) - Supplies the number of columns of the source matrix to copy. - RowSumVector (r9) - Supplies the address of the buffer to receive the sums - of the elements from each of the rows. Each sum has also been multiplied + RowSumBuffer (r9) - Supplies the address of the buffer to receive the sums + of the elements along each of the rows. by the zero point offset. - offb - Supplies the zero point offset for the other source matrix of the - matrix multiplication. - Return Value: None. @@ -93,11 +88,10 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): mov r11,rcx lea r12,[r8+3] and r12,NOT 3 # align CountK up to quad count - vpbroadcastw xmm8,WORD PTR .LGemmU8S8CopyPackAFrame_offb[rsp] - vpcmpeqw ymm9,ymm9,ymm9 # generate word vector [0xFFFF] - vpsrlw ymm9,ymm9,15 # generate word vector [0x0001] - vpsllw ymm0,ymm9,8 # generate word vector [0x0100] - vpor ymm9,ymm9,ymm0 # generate word vector [0x0101] + vpcmpeqw ymm8,ymm8,ymm8 # generate word vector [0xFFFF] + vpsrlw ymm8,ymm8,15 # generate word vector [0x0001] + vpsllw ymm9,ymm8,8 # generate word vector [0x0100] + vpor ymm9,ymm8,ymm9 # generate word vector [0x0101] // // Compute the conditional load/store mask for an unaligned CountK. @@ -165,7 +159,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): .LCopyPackA.ProcessRemainingColumnsM4: add rbx,32 # correct for over-subtract above - jz .LCopyPackA.ReduceRowSumVectorM4 + jz .LCopyPackA.ReduceRowSumBufferM4 test bl,16 # (CountK & 16) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan16M4 vmovdqu xmm4,XMMWORD PTR [rdx] @@ -187,7 +181,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): add rdx,16 # advance matrix A by 16 bytes add rcx,16 # advance matrix D by 16 bytes test bl,15 # test for unaligned columns - jz .LCopyPackA.ReduceRowSumVectorM4 + jz .LCopyPackA.ReduceRowSumBufferM4 // // Copy the unaligned CountK columns to a zero padded stack buffer. @@ -275,15 +269,18 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): // Reduce the sums for the four rows of output. // -.LCopyPackA.ReduceRowSumVectorM4: - vphaddw ymm0,ymm0,ymm1 # reduce and interleave Sum1/Sum0 - vphaddw ymm1,ymm2,ymm3 # reduce and interleave Sum3/Sum2 - vphaddw ymm0,ymm0,ymm1 # reduce and interleave Sum3/Sum2/Sum1/Sum0 - vextracti128 xmm1,ymm0,1 # extract high pairs - vpaddw xmm0,xmm0,xmm1 # reduce low/high pairs - vpmaddwd xmm0,xmm0,xmm8 # multiply by offset and reduce 32-bit sum +.LCopyPackA.ReduceRowSumBufferM4: + vpmaddwd ymm0,ymm0,ymm8 # horizontal word+word=dword per row + vpmaddwd ymm1,ymm1,ymm8 + vphaddd ymm0,ymm0,ymm1 # reduce and interleave Sum1/Sum0 + vpmaddwd ymm2,ymm2,ymm8 + vpmaddwd ymm3,ymm3,ymm8 + vphaddd ymm1,ymm2,ymm3 # reduce and interleave Sum3/Sum2 + vphaddd ymm0,ymm0,ymm1 # reduce and interleave Sum3/Sum2/Sum1/Sum0 + vextracti128 xmm1,ymm0,1 # extract high dwords + vpaddd xmm0,xmm0,xmm1 # reduce low/high dwords vmovdqu XMMWORD PTR [r9],xmm0 - add r9,4*4 # advance row sum vector by 4 DWORDs + add r9,4*4 # advance row sum buffer by 4 dwords sub r11,4 # subtract rows remaining jae .LCopyPackA.ProcessNextRowM4 @@ -317,7 +314,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): .LCopyPackA.ProcessRemainingColumnsM1: add rbx,32 # correct for over-subtract above - jz .LCopyPackA.ReduceRowSumVectorM1 + jz .LCopyPackA.ReduceRowSumBufferM1 test bl,16 # (CountK & 16) != 0? jz .LCopyPackA.CopyRemainingCountKLessThan16M1 vmovdqu xmm4,XMMWORD PTR [rdx] @@ -327,7 +324,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): add rdx,16 # advance matrix A by 16 bytes add rcx,16 # advance matrix D by 16 bytes test bl,15 # test for unaligned columns - jz .LCopyPackA.ReduceRowSumVectorM1 + jz .LCopyPackA.ReduceRowSumBufferM1 // // Copy the unaligned CountK columns to a zero padded stack buffer. @@ -378,14 +375,14 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackAAvx2): // Reduce the sum for the single row of output. // -.LCopyPackA.ReduceRowSumVectorM1: - vextracti128 xmm1,ymm0,1 # extract high pairs - vpaddw xmm0,xmm0,xmm1 # reduction - vphaddw xmm0,xmm0,xmm0 - vphaddw xmm0,xmm0,xmm0 - vpmaddwd xmm0,xmm0,xmm8 # multiply by offset and reduce +.LCopyPackA.ReduceRowSumBufferM1: + vpmaddwd ymm0,ymm0,ymm8 # horizontal word+word=dword per row + vextracti128 xmm1,ymm0,1 # extract high dwords + vpaddd xmm0,xmm0,xmm1 # reduction + vphaddd xmm0,xmm0,xmm0 + vphaddd xmm0,xmm0,xmm0 vmovd DWORD PTR [r9],xmm0 - add r9,4 # advance row sum vector by 1 DWORD + add r9,4 # advance row sum buffer by 1 dword dec r11 # decrement rows remaining jnz .LCopyPackA.ProcessNextRowM1 @@ -421,15 +418,11 @@ Arguments: CountK (r8) - Supplies the number of rows of the source matrix to copy. - ColumnSumVector (r9) - 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. + ColumnSumBuffer (r9) - Supplies the address of the buffer to receive the sums + of the elements along each of the columns. - 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. + BIsSigned - Supplies true if the source matrix is signed data, else false if + the source matrix is unsigned data. Return Value: @@ -445,18 +438,17 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): mov r10,rdx lea r11,[r10+r10*2] # compute ldb * 3 - vpbroadcastw ymm7,WORD PTR .LGemmU8S8CopyPackBFrame_offa[rsp] - vpcmpeqw ymm8,ymm8,ymm8 # generate word vector [0xFFFF] - vpsrlw ymm8,ymm8,15 # generate word vector [0x0001] - vpsllw ymm0,ymm8,8 # generate word vector [0x0100] - vpor ymm8,ymm8,ymm0 # generate word vector [0x0101] + vpcmpeqw ymm7,ymm7,ymm7 # generate word vector [0xFFFF] + vpsrlw ymm7,ymm7,15 # generate word vector [0x0001] + vpsllw ymm8,ymm7,8 # generate word vector [0x0100] + vpor ymm8,ymm7,ymm8 # generate word vector [0x0101] // // Compute the bit flip vector to adjust input from U8 to S8. // vpxor xmm9,xmm9,xmm9 # generate word vector [0x0000] - cmp BYTE PTR .LGemmU8S8CopyPackBFrame_BTypeIsSigned[rsp],0 + cmp BYTE PTR .LGemmU8S8CopyPackBFrame_BIsSigned[rsp],0 jnz .LCopyPackB.SkipUnsignedBitFlipVector vpsllw ymm9,ymm8,7 # generate word vector [0x8080] @@ -501,9 +493,11 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows vmovdqu YMMWORD PTR [rdi+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row - vpaddw ymm0,ymm0,ymm4 # add words to row accumulators + vpmaddwd ymm4,ymm4,ymm7 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column vpmaddubsw ymm2,ymm8,ymm2 - vpaddw ymm1,ymm1,ymm2 + vpmaddwd ymm2,ymm2,ymm7 + vpaddd ymm1,ymm1,ymm2 add rdi,64 # advance matrix D by 64 bytes sub rbx,4 # subtract rows remaining jae .LCopyPackB.ProcessNextRowLoopN16 @@ -514,7 +508,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.ProcessRemainingRowsN16: add rbx,4 # correct for over-subtract above - jz .LCopyPackB.ReduceColumnSumVectorN16 + jz .LCopyPackB.StoreColumnSumBufferN16 vmovdqu xmm2,XMMWORD PTR [rdx] vmovaps xmm3,xmm9 vmovaps xmm4,xmm9 @@ -528,12 +522,10 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vmovdqu xmm4,XMMWORD PTR [rdx+r10*2] jmp .LCopyPackB.InterleaveRowDataN16 -.LCopyPackB.ReduceColumnSumVectorN16: - vpmaddwd ymm0,ymm0,ymm7 # multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm7 # multiply by offset and reduce +.LCopyPackB.StoreColumnSumBufferN16: vmovdqu YMMWORD PTR [r9],ymm0 vmovdqu YMMWORD PTR [r9+32],ymm1 - add r9,16*4 # advance column sum vector by 16 DWORDs + add r9,16*4 # advance column sum buffer by 16 dwords sub rcx,16 # subtract columns remaining jae .LCopyPackB.ProcessNextColumnN16 @@ -640,9 +632,11 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows vmovdqu YMMWORD PTR [rdi+32],ymm2 vpmaddubsw ymm4,ymm8,ymm4 # horizontal byte+byte=word per row - vpaddw ymm0,ymm0,ymm4 # add words to row accumulators + vpmaddwd ymm4,ymm4,ymm7 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column vpmaddubsw ymm2,ymm8,ymm2 - vpaddw ymm1,ymm1,ymm2 + vpmaddwd ymm2,ymm2,ymm7 + vpaddd ymm1,ymm1,ymm2 lea rsi,[rsi+r10*4] # advance next matrix B by 4 rows add rdi,64 # advance matrix D by 64 bytes sub r8,4 # subtract rows remaining @@ -650,7 +644,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): .LCopyPackB.ProcessRemainingRowsNUnaligned: add r8,4 - jz .LCopyPackB.ReduceColumnSumVectorNUnaligned + jz .LCopyPackB.StoreColumnSumBufferNUnaligned // // Process the less than 4 remaining rows where the row has less than 16 columns. @@ -699,9 +693,7 @@ C_UNDERSCORE(MlasGemmU8S8CopyPackBAvx2): mov rbp,r11 jmp .LCopyPackB.CopyUnalignedRowLoop -.LCopyPackB.ReduceColumnSumVectorNUnaligned: - vpmaddwd ymm0,ymm0,ymm7 # multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm7 # multiply by offset and reduce +.LCopyPackB.StoreColumnSumBufferNUnaligned: vmovdqu YMMWORD PTR [r9],ymm0 vmovdqu YMMWORD PTR [r9+32],ymm1 jmp .LCopyPackB.ExitRoutine @@ -884,11 +876,11 @@ Arguments: ldc - Supplies the first dimension of matrix C. - RowSumVector - Supplies the sum of each row from matrix A multiplied by the + RowSumBuffer - 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 + ColumnSumBuffer - 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. @@ -918,8 +910,8 @@ C_UNDERSCORE(MlasGemmU8S8KernelAvx2): shl rcx,2 # convert to row length movzx r10,BYTE PTR .LGemmU8X8KernelFrame_ZeroMode[rsp] mov r11,rdi - mov r12,.LGemmU8X8KernelFrame_RowSumVector[rsp] - mov r13,.LGemmU8X8KernelFrame_ColumnSumVector[rsp] + mov r12,.LGemmU8X8KernelFrame_RowSumBuffer[rsp] + mov r13,.LGemmU8X8KernelFrame_ColumnSumBuffer[rsp] vpcmpeqw ymm12,ymm12,ymm12 # generate 256-bit word vector [0xFFFF] vpsrlw ymm12,ymm12,15 # generate 256-bit word vector [0x0001] diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S index 4eab5027be..1952578d72 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S @@ -33,7 +33,6 @@ Abstract: .equ .LGemmU8U8CopyPackAFrame_SavedRbx, 16 .equ .LGemmU8U8CopyPackAFrame_SavedRbp, 24 .equ .LGemmU8U8CopyPackAFrame_ReturnAddress, 32 - .equ .LGemmU8U8CopyPackAFrame_offb, 40 // // Stack frame layout for the U8U8 CopyPackB routine. @@ -44,7 +43,6 @@ Abstract: .equ .LGemmU8U8CopyPackBFrame_SavedRbx, 0 .equ .LGemmU8U8CopyPackBFrame_SavedRbp, 8 .equ .LGemmU8U8CopyPackBFrame_ReturnAddress, 16 - .equ .LGemmU8U8CopyPackBFrame_offa, 24 .text @@ -72,12 +70,8 @@ Arguments: CountK (r8) - Supplies the number of columns of the source matrix to copy. - RowSumVector (r9) - 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. + RowSumBuffer (r9) - Supplies the address of the buffer to receive the sums + of the elements along each of the rows. Return Value: @@ -97,7 +91,8 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): mov r11,rcx lea r12,[r8+1] and r12,NOT 1 # align CountK up to pair count - vpbroadcastw xmm8,WORD PTR .LGemmU8U8CopyPackAFrame_offb[rsp] + vpcmpeqw ymm8,ymm8,ymm8 # generate word vector [0xFFFF] + vpsrlw ymm8,ymm8,15 # generate word vector [0x0001] // // Compute the conditional load/store mask for an unaligned CountK. @@ -127,10 +122,6 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): // The packed buffer has the same data ordering as the source bytes, but CountK // is aligned up to a multiple of 2 to maintain 32-bit alignment. All padding // bytes are zero filled. -// -// These 16-bit values are also accumulated into an intermediate per-row -// accumulator. CountK cannot be greater than 128 to avoid overflowing these -// signed 16-bit accumulators. // sub r11,4 @@ -171,7 +162,7 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): .LCopyPackA.ProcessRemainingColumnsM4: add rbx,16 # correct for over-subtract above - jz .LCopyPackA.ReduceRowSumVectorM4 + jz .LCopyPackA.ReduceRowSumBufferM4 // // Copy the unaligned CountK columns to a zero padded stack buffer. @@ -255,21 +246,21 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): vpaddw ymm3,ymm3,ymm7 // -// Reduce the sums for the four rows of output. Transpose the intermediate -// accumulators by treating the registers as 32-bit elements containing a pair -// of 16-bit sums. Continue reducing the transposed accumulators to produce the -// final 32-bit vector output. +// Reduce the sums for the four rows of output. // -.LCopyPackA.ReduceRowSumVectorM4: - vphaddw ymm0,ymm0,ymm1 # reduce and interleave Sum1/Sum0 - vphaddw ymm1,ymm2,ymm3 # reduce and interleave Sum3/Sum2 - vphaddw ymm0,ymm0,ymm1 # reduce and interleave Sum3/Sum2/Sum1/Sum0 - vextracti128 xmm1,ymm0,1 # extract high pairs - vpaddw xmm0,xmm0,xmm1 # reduce low/high pairs - vpmaddwd xmm0,xmm0,xmm8 # multiply by offset and reduce 32-bit sum +.LCopyPackA.ReduceRowSumBufferM4: + vpmaddwd ymm0,ymm0,ymm8 # horizontal word+word=dword per row + vpmaddwd ymm1,ymm1,ymm8 + vphaddd ymm0,ymm0,ymm1 # reduce and interleave Sum1/Sum0 + vpmaddwd ymm2,ymm2,ymm8 + vpmaddwd ymm3,ymm3,ymm8 + vphaddd ymm1,ymm2,ymm3 # reduce and interleave Sum3/Sum2 + vphaddd ymm0,ymm0,ymm1 # reduce and interleave Sum3/Sum2/Sum1/Sum0 + vextracti128 xmm1,ymm0,1 # extract high dwords + vpaddd xmm0,xmm0,xmm1 # reduce low/high dwords vmovdqu XMMWORD PTR [r9],xmm0 - add r9,4*4 # advance row sum vector by 4 dwords + add r9,4*4 # advance row sum buffer by 4 dwords sub r11,4 # subtract rows remaining jae .LCopyPackA.ProcessNextRowM4 @@ -302,7 +293,7 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): .LCopyPackA.ProcessRemainingColumnsM1: add rbx,16 # correct for over-subtract above - jz .LCopyPackA.ReduceRowSumVectorM1 + jz .LCopyPackA.ReduceRowSumBufferM1 // // Copy the unaligned CountK columns to a zero padded stack buffer. @@ -351,14 +342,14 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackAAvx2): // Reduce the sum for the single row of output. // -.LCopyPackA.ReduceRowSumVectorM1: - vextracti128 xmm1,ymm0,1 # extract high pairs - vpaddw xmm0,xmm0,xmm1 # reduction - vphaddw xmm0,xmm0,xmm0 - vphaddw xmm0,xmm0,xmm0 - vpmaddwd xmm0,xmm0,xmm8 # multiply by offset and reduce +.LCopyPackA.ReduceRowSumBufferM1: + vpmaddwd ymm0,ymm0,ymm8 # horizontal word+word=dword per row + vextracti128 xmm1,ymm0,1 # extract high dwords + vpaddd xmm0,xmm0,xmm1 # reduction + vphaddd xmm0,xmm0,xmm0 + vphaddd xmm0,xmm0,xmm0 vmovd DWORD PTR [r9],xmm0 - add r9,4 # advance row sum vector by 1 DWORD + add r9,4 # advance row sum buffer by 1 dword dec r11 # decrement rows remaining jnz .LCopyPackA.ProcessNextRowM1 @@ -394,12 +385,8 @@ Arguments: CountK (r8) - Supplies the number of rows of the source matrix to copy. - ColumnSumVector (r9) - 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. + ColumnSumBuffer (r9) - Supplies the address of the buffer to receive the sums + of the elements along each of the columns. Return Value: @@ -414,7 +401,8 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackBAvx2): push rbx mov r10,rdx - vpbroadcastw ymm5,WORD PTR .LGemmU8U8CopyPackBFrame_offa[rsp] + vpcmpeqw ymm5,ymm5,ymm5 # generate word vector [0xFFFF] + vpsrlw ymm5,ymm5,15 # generate word vector [0x0001] // // Zero initialize the padded stack buffers. @@ -450,29 +438,31 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackBAvx2): vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 add rdi,32 # advance matrix D by 32 bytes - vpaddw ymm0,ymm0,ymm4 # accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 sub rbx,2 # subtract rows remaining jae .LCopyPackB.ProcessNextRowLoopN16 .LCopyPackB.ProcessRemainingRowsN16: add rbx,2 # correct for over-subtract above - jz .LCopyPackB.ReduceColumnSumVectorN16 + jz .LCopyPackB.StoreColumnSumBufferN16 vpmovzxbw ymm4,XMMWORD PTR [rdx] vmovdqu YMMWORD PTR [rdi],ymm4 # store interleaved rows vextracti128 xmm3,ymm4,1 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 # accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 add rdi,32 # advance matrix D by 32 bytes -.LCopyPackB.ReduceColumnSumVectorN16: - vpmaddwd ymm0,ymm0,ymm5 # multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm5 # multiply by offset and reduce +.LCopyPackB.StoreColumnSumBufferN16: vmovdqu YMMWORD PTR [r9],ymm0 vmovdqu YMMWORD PTR [r9+32],ymm1 - add r9,64 # advance column sum vector by 16 DWORDs + add r9,64 # advance column sum buffer by 16 dwords sub rcx,16 # subtract columns remaining jae .LCopyPackB.ProcessNextColumnN16 @@ -550,8 +540,10 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackBAvx2): vmovdqu XMMWORD PTR [rdi+16],xmm3 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 # accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 lea rsi,[rsi+r10*2] # advance next matrix B by 2 rows add rdi,32 # advance matrix D by 32 bytes sub r8,2 # subtract rows remaining @@ -559,7 +551,7 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackBAvx2): .LCopyPackB.ProcessRemainingRowsNUnaligned: add r8,2 - jz .LCopyPackB.ReduceColumnSumVectorNUnaligned + jz .LCopyPackB.StoreColumnSumBufferNUnaligned mov rdx,rsi lea rbp,.LGemmU8U8CopyPackBFrame_PaddedMatrixBData[rsp] test cl,8 # (CountN & 8) != 0? @@ -597,12 +589,12 @@ C_UNDERSCORE(MlasGemmU8U8CopyPackBAvx2): vextracti128 xmm3,ymm4,1 vpmovzxbw ymm4,xmm4 vpmovzxbw ymm3,xmm3 - vpaddw ymm0,ymm0,ymm4 # accumulate per column - vpaddw ymm1,ymm1,ymm3 + vpmaddwd ymm4,ymm4,ymm5 # horizontal word+word=dword per row + vpaddd ymm0,ymm0,ymm4 # accumulate per column + vpmaddwd ymm3,ymm3,ymm5 + vpaddd ymm1,ymm1,ymm3 -.LCopyPackB.ReduceColumnSumVectorNUnaligned: - vpmaddwd ymm0,ymm0,ymm5 # multiply by offset and reduce - vpmaddwd ymm1,ymm1,ymm5 # multiply by offset and reduce +.LCopyPackB.StoreColumnSumBufferNUnaligned: vmovdqu YMMWORD PTR [r9],ymm0 vmovdqu YMMWORD PTR [r9+32],ymm1 jmp .LCopyPackB.ExitRoutine @@ -793,11 +785,11 @@ Arguments: ldc - Supplies the first dimension of matrix C. - RowSumVector - Supplies the sum of each row from matrix A multiplied by the + RowSumBuffer - 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 + ColumnSumBuffer - 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. @@ -827,8 +819,8 @@ C_UNDERSCORE(MlasGemmU8U8KernelAvx2): shl rcx,2 # convert to row length movzx r10,BYTE PTR .LGemmU8X8KernelFrame_ZeroMode[rsp] mov r11,rdi - mov r12,.LGemmU8X8KernelFrame_RowSumVector[rsp] - mov r13,.LGemmU8X8KernelFrame_ColumnSumVector[rsp] + mov r12,.LGemmU8X8KernelFrame_RowSumBuffer[rsp] + mov r13,.LGemmU8X8KernelFrame_ColumnSumBuffer[rsp] // // Process CountM rows of the matrices. diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h index 7dbff45389..8a2bf10db8 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h @@ -26,8 +26,8 @@ Abstract: .equ .LGemmU8X8KernelFrame_SavedRbp, 24 .equ .LGemmU8X8KernelFrame_ReturnAddress, 32 .equ .LGemmU8X8KernelFrame_ldc, 40 - .equ .LGemmU8X8KernelFrame_RowSumVector, 48 - .equ .LGemmU8X8KernelFrame_ColumnSumVector, 56 + .equ .LGemmU8X8KernelFrame_RowSumBuffer, 48 + .equ .LGemmU8X8KernelFrame_ColumnSumBuffer, 56 .equ .LGemmU8X8KernelFrame_DepthValue, 64 .equ .LGemmU8X8KernelFrame_ZeroMode, 72 @@ -54,9 +54,9 @@ Implicit Arguments: rcx - Supplies the length in bytes of a row from matrix A. - r12 - Supplies the address of the row sum vector. + r12 - Supplies the address of the row sum buffer. - r13 - Supplies the address of the column sum vector. + r13 - Supplies the address of the column sum buffer. ymm4-ymm15 - Supplies the block accumulators. @@ -73,7 +73,7 @@ Implicit Arguments: .if \ColumnCount\() == 16 vpaddd ymm0,ymm1,YMMWORD PTR [r13] vpaddd ymm1,ymm1,YMMWORD PTR [r13+32] - add r13,16*4 # advance ColumnSumVector by 16 columns + add r13,16*4 # advance ColumnSumBuffer by 16 columns .else vpaddd ymm1,ymm1,YMMWORD PTR [r13] .endif @@ -145,9 +145,9 @@ Implicit Arguments: r10b - Supplies the zero mode flag. - r12 - Supplies the address of the row sum vector. + r12 - Supplies the address of the row sum buffer. - r13 - Supplies the address of the column sum vector. + r13 - Supplies the address of the column sum buffer. --*/ diff --git a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx512Common.h b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx512Common.h index 173f42593d..7db55ba661 100644 --- a/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx512Common.h +++ b/onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx512Common.h @@ -27,8 +27,8 @@ Abstract: .equ .LGemmU8X8KernelFrame_SavedRbp, 32 .equ .LGemmU8X8KernelFrame_ReturnAddress, 40 .equ .LGemmU8X8KernelFrame_ldc, 48 - .equ .LGemmU8X8KernelFrame_RowSumVector, 56 - .equ .LGemmU8X8KernelFrame_ColumnSumVector, 64 + .equ .LGemmU8X8KernelFrame_RowSumBuffer, 56 + .equ .LGemmU8X8KernelFrame_ColumnSumBuffer, 64 .equ .LGemmU8X8KernelFrame_DepthValue, 72 .equ .LGemmU8X8KernelFrame_ZeroMode, 80 @@ -55,9 +55,9 @@ Implicit Arguments: rcx - Supplies the length in bytes of a row from matrix A. - r12 - Supplies the address of the row sum vector. + r12 - Supplies the address of the row sum buffer. - r13 - Supplies the address of the column sum vector. + r13 - Supplies the address of the column sum buffer. --*/ @@ -78,7 +78,7 @@ Implicit Arguments: vpaddd zmm1,zmm3,ZMMWORD PTR [r13] vpaddd zmm0,zmm3,ZMMWORD PTR [r13+64] .endif - add_immed r13,\ColumnCount\()*4 # advance ColumnSumVector by N columns + add_immed r13,\ColumnCount\()*4 # advance ColumnSumBuffer by N columns .else vpaddd zmm0,zmm3,ZMMWORD PTR [r13] .endif @@ -147,9 +147,9 @@ Implicit Arguments: r10b - Supplies the zero mode flag. - r12 - Supplies the address of the row sum vector. + r12 - Supplies the address of the row sum buffer. - r13 - Supplies the address of the column sum vector. + r13 - Supplies the address of the column sum buffer. r14 - Supplies the stride in bytes of between packed blocks of matrix B. @@ -303,11 +303,11 @@ Arguments: ldc - Supplies the first dimension of matrix C. - RowSumVector - Supplies the sum of each row from matrix A multiplied by the + RowSumBuffer - 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 + ColumnSumBuffer - 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. @@ -338,8 +338,8 @@ C_UNDERSCORE(MlasGemm\Type\()Kernel\Isa\()): shl rcx,2 # convert to row length movzx r10,BYTE PTR .LGemmU8X8KernelFrame_ZeroMode[rsp] mov r11,rdi - mov r12,.LGemmU8X8KernelFrame_RowSumVector[rsp] - mov r13,.LGemmU8X8KernelFrame_ColumnSumVector[rsp] + mov r12,.LGemmU8X8KernelFrame_RowSumBuffer[rsp] + mov r13,.LGemmU8X8KernelFrame_ColumnSumBuffer[rsp] mov ebp,-1 kmovw k1,ebp # update mask to write all columns .ifeqs "\Type\()", "U8S8"