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

MLAS: more prepacking kernel changes (#4397)

Browse source 
Kernel changes to support StrideK>128
This commit is contained in:
Tracy Sharpe 2020-07-01 22:09:42 -07:00 committed by GitHub
parent 2d54c89d77
commit 5c23b17196
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 390 additions and 408 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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]

120
onnxruntime/core/mlas/lib/amd64/QgemmU8U8KernelAvx2.asm
View file

@ -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.

14
onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx2Common.inc
View file

@ -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

22
onnxruntime/core/mlas/lib/amd64/QgemmU8X8KernelAvx512Common.inc
View file

@ -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 <Type>, <U8S8>

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

@ -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

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

@ -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<typename KernelType>
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<typename KernelType::OffsetBType>::value && !WorkBlock->BTypeIsSigned) {
offb = typename KernelType::OffsetBType(offb ^ 0x80);
if (std::is_signed<typename KernelType::OffsetBType>::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<BType>::value;
WorkBlock.BIsSigned = std::is_signed<BType>::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<BType>::value;
WorkBlock.CTypeIsFloat = true;
WorkBlock.BIsSigned = std::is_signed<BType>::value;
WorkBlock.CIsFloat = true;
//
// Schedule the operation across a set of worker threads.

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

@ -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]

116
onnxruntime/core/mlas/lib/x86_64/QgemmU8U8KernelAvx2.S
View file

@ -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.

14
onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx2Common.h
View file

@ -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.
--*/

22
onnxruntime/core/mlas/lib/x86_64/QgemmU8X8KernelAvx512Common.h
View file

@ -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"