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

Change GEMM kernels to natively handle broader range of row counts (#406)

Browse source 
The outer GEMM loop repeatedly calls the inner GEMM kernel with a row count (the M parameter to GEMM) and the inner kernel decides how many rows it will actually handle. The FMA3 kernel only handled row counts of 1,3,6 to keep code size down. To be competitive however, the FMA3 kernel needs to handle any row count from 1-6.

One example model was issuing a GEMM with M=11 and this had been broken up into 6,3,1,1, but can now be handled as 6,5.

The kernels have been templatized MASM style to avoid the cut/paste code from the original implementation. The Linux variants will be updated after doing some additional work on the MASM variants first.
This commit is contained in:
Tracy Sharpe 2019-01-29 17:20:32 -08:00 committed by GitHub
parent d75bdc5194
commit 20ef8b43a6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 866 additions and 1143 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm
View file

@ -19,7 +19,6 @@
.xlist
INCLUDE mlasi.inc
INCLUDE SgemmKernelCommon.inc
.list
EXTERN MlasMaskMoveAvx:NEAR

515
onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm
View file

@ -40,6 +40,9 @@ INCLUDE SgemmKernelCommon.inc
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 2 rows.
@ -53,7 +56,7 @@ INCLUDE SgemmKernelCommon.inc
; ymm8-ymm15 - Supplies the block accumulators.
;
ComputeBlockAvxBy16 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockAvxBy16 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IF Count EQ 1
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
@ -64,30 +67,26 @@ IF Count EQ 1
ELSE
vmovaps ymm0,YMMWORD PTR [rdx+VectorOffset]
vmovaps ymm1,YMMWORD PTR [rdx+VectorOffset+32]
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vmulps ymm4,ymm3,ymm0
vaddps ymm8,ymm8,ymm4
vmulps ymm5,ymm3,ymm1
vaddps ymm9,ymm9,ymm5
IF Count GE 2
vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]
vmulps ymm6,ymm3,ymm0
vaddps ymm10,ymm10,ymm6
vmulps ymm7,ymm3,ymm1
vaddps ymm11,ymm11,ymm7
ENDIF
IF Count GE 4
vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]
vmulps ymm4,ymm3,ymm0
vaddps ymm12,ymm12,ymm4
vmulps ymm5,ymm3,ymm1
vaddps ymm13,ymm13,ymm5
vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]
vmulps ymm6,ymm3,ymm0
vaddps ymm14,ymm14,ymm6
vmulps ymm7,ymm3,ymm1
vaddps ymm15,ymm15,ymm7
ENDIF
EmitIfCountGE Count, 1, <vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]>
EmitIfCountGE Count, 1, <vmulps ymm4,ymm3,ymm0>
EmitIfCountGE Count, 1, <vaddps ymm8,ymm8,ymm4>
EmitIfCountGE Count, 1, <vmulps ymm5,ymm3,ymm1>
EmitIfCountGE Count, 1, <vaddps ymm9,ymm9,ymm5>
EmitIfCountGE Count, 2, <vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 2, <vmulps ymm6,ymm3,ymm0>
EmitIfCountGE Count, 2, <vaddps ymm10,ymm10,ymm6>
EmitIfCountGE Count, 2, <vmulps ymm7,ymm3,ymm1>
EmitIfCountGE Count, 2, <vaddps ymm11,ymm11,ymm7>
EmitIfCountGE Count, 3, <vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]>
EmitIfCountGE Count, 3, <vmulps ymm4,ymm3,ymm0>
EmitIfCountGE Count, 3, <vaddps ymm12,ymm12,ymm4>
EmitIfCountGE Count, 3, <vmulps ymm5,ymm3,ymm1>
EmitIfCountGE Count, 3, <vaddps ymm13,ymm13,ymm5>
EmitIfCountGE Count, 4, <vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 4, <vmulps ymm6,ymm3,ymm0>
EmitIfCountGE Count, 4, <vaddps ymm14,ymm14,ymm6>
EmitIfCountGE Count, 4, <vmulps ymm7,ymm3,ymm1>
EmitIfCountGE Count, 4, <vaddps ymm15,ymm15,ymm7>
ENDIF
ENDM
@ -108,6 +107,9 @@ ENDIF
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 2 rows.
@ -121,7 +123,7 @@ ENDIF
; ymm8-ymm15 - Supplies the block accumulators.
;
ComputeBlockAvxBy8 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockAvxBy8 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IF Count EQ 1
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
@ -129,22 +131,211 @@ IF Count EQ 1
vaddps ymm9,ymm9,ymm5
ELSE
vmovaps ymm0,YMMWORD PTR [rdx+VectorOffset]
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vmulps ymm5,ymm3,ymm0
vaddps ymm9,ymm9,ymm5
IF Count GE 2
vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]
vmulps ymm7,ymm3,ymm0
vaddps ymm11,ymm11,ymm7
EmitIfCountGE Count, 1, <vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]>
EmitIfCountGE Count, 1, <vmulps ymm5,ymm3,ymm0>
EmitIfCountGE Count, 1, <vaddps ymm9,ymm9,ymm5>
EmitIfCountGE Count, 2, <vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 2, <vmulps ymm7,ymm3,ymm0>
EmitIfCountGE Count, 2, <vaddps ymm11,ymm11,ymm7>
EmitIfCountGE Count, 3, <vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]>
EmitIfCountGE Count, 3, <vmulps ymm5,ymm3,ymm0>
EmitIfCountGE Count, 3, <vaddps ymm13,ymm13,ymm5>
EmitIfCountGE Count, 4, <vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 4, <vmulps ymm7,ymm3,ymm0>
EmitIfCountGE Count, 4, <vaddps ymm15,ymm15,ymm7>
ENDIF
IF Count GE 4
vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]
vmulps ymm5,ymm3,ymm0
vaddps ymm13,ymm13,ymm5
vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]
vmulps ymm7,ymm3,ymm0
vaddps ymm15,ymm15,ymm7
ENDM
;
; ComputeBlockAvxLoop
;
; This macro generates code to execute the block compute macro multiple
; times and advancing the matrix A and matrix B data pointers.
;
; Arguments:
;
; ComputeBlock - Supplies the macro to compute a single block.
;
; Count - Supplies the number of rows to access from matrix A.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus N rows.
;
; rcx - Supplies the address into the matrix A data.
;
; rdx - Supplies the address into the matrix B data.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockAvxLoop MACRO ComputeBlock, Count
IF Count GT 2
lea rbx,[rcx+r10*2] ; compute matrix A plus 2 rows
ENDIF
ComputeBlockLoop ComputeBlock, Count, <Count GT 2>
IF Count GT 2
lea rbx,[r8+rax*2] ; compute matrix C plus 2 rows
ENDIF
ENDM
;
; ProcessCountMAvx
;
; Macro Description:
;
; This macro generates code to compute matrix multiplication for a fixed set
; of rows.
;
; Arguments:
;
; Mode - Supplies the mode of operation for updating the contents of matrix C.
;
; Count - Supplies the number of rows to process.
;
; Implicit Arguments:
;
; rcx - Supplies the address of matrix A.
;
; rdx - Supplies the address of matrix B.
;
; rsi - Supplies the address of matrix A.
;
; rbp - Supplies the number of columns from matrix B and matrix C to iterate
; over.
;
; r8 - Supplies the address of matrix C.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; r10 - Supplies the length in bytes of a row from matrix A.
;
ProcessCountMAvx MACRO Mode, Count, Fallthrough
LOCAL ProcessNextColumnLoop16xN
LOCAL ProcessRemainingCountN
LOCAL OutputMasked8xNBlock
LOCAL OutputMasked16xNBlock
cmp rbp,8
jbe ProcessRemainingCountN
ProcessNextColumnLoop16xN:
EmitIfCountGE Count, 1, <vxorps xmm8,xmm8,xmm8>
EmitIfCountGE Count, 1, <vxorps xmm9,xmm9,xmm9>
EmitIfCountGE Count, 2, <vxorps xmm10,xmm10,xmm10>
EmitIfCountGE Count, 2, <vxorps xmm11,xmm11,xmm11>
EmitIfCountGE Count, 3, <vxorps xmm12,xmm12,xmm12>
EmitIfCountGE Count, 3, <vxorps xmm13,xmm13,xmm13>
EmitIfCountGE Count, 4, <vxorps xmm14,xmm14,xmm14>
EmitIfCountGE Count, 4, <vxorps xmm15,xmm15,xmm15>
ComputeBlockAvxLoop ComputeBlockAvxBy16, Count
EmitIfCountGE Count, 1, <vmulps ymm8,ymm8,ymm2>
EmitIfCountGE Count, 1, <vmulps ymm9,ymm9,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm10,ymm10,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm11,ymm11,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm12,ymm12,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm13,ymm13,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm14,ymm14,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm15,ymm15,ymm2>
sub rbp,16
jb OutputMasked16xNBlock
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vaddps ymm8,ymm8,YMMWORD PTR [r8]>
EmitIfCountGE Count, 1, <vaddps ymm9,ymm9,YMMWORD PTR [r8+32]>
EmitIfCountGE Count, 2, <vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 2, <vaddps ymm11,ymm11,YMMWORD PTR [r8+rax+32]>
EmitIfCountGE Count, 3, <vaddps ymm12,ymm12,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 3, <vaddps ymm13,ymm13,YMMWORD PTR [rbx+32]>
EmitIfCountGE Count, 4, <vaddps ymm14,ymm14,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 4, <vaddps ymm15,ymm15,YMMWORD PTR [rbx+rax+32]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm8>
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8+32],ymm9>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm10>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax+32],ymm11>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [rbx],ymm12>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [rbx+32],ymm13>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx+rax],ymm14>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx+rax+32],ymm15>
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
cmp rbp,8
ja ProcessNextColumnLoop16xN
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN:
EmitIfCountGE Count, 1, <vxorps xmm9,xmm9,xmm9>
EmitIfCountGE Count, 2, <vxorps xmm11,xmm11,xmm11>
EmitIfCountGE Count, 3, <vxorps xmm13,xmm13,xmm13>
EmitIfCountGE Count, 4, <vxorps xmm15,xmm15,xmm15>
ComputeBlockAvxLoop ComputeBlockAvxBy8, Count
EmitIfCountGE Count, 1, <vmulps ymm9,ymm9,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm11,ymm11,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm13,ymm13,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm15,ymm15,ymm2>
cmp rbp,8
jb OutputMasked8xNBlock
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vaddps ymm9,ymm9,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vaddps ymm11,ymm11,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vaddps ymm13,ymm13,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 4, <vaddps ymm15,ymm15,YMMWORD PTR [rbx+rax]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm9>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm11>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [rbx],ymm13>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx+rax],ymm15>
jmp ExitKernel
OutputMasked16xNBlock:
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vaddps ymm8,ymm8,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vaddps ymm12,ymm12,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 4, <vaddps ymm14,ymm14,YMMWORD PTR [rbx+rax]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm8>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm10>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [rbx],ymm12>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx+rax],ymm14>
add r8,8*4 ; advance matrix C by 8 columns
IF Count GT 2
add rbx,8*4 ; advance matrix C plus 2 rows by 8 columns
ENDIF
add rbp,8 ; correct for over-subtract above
OutputMasked8xNBlock:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss xmm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd xmm1,xmm0,XMMWORD PTR [MlasMaskMoveAvx+16]
vpcmpgtd xmm0,xmm0,XMMWORD PTR [MlasMaskMoveAvx]
vinsertf128 ymm0,ymm0,xmm1,1
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmaskmovps ymm8,ymm0,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vmaskmovps ymm10,ymm0,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vmaskmovps ymm12,ymm0,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 4, <vmaskmovps ymm14,ymm0,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 1, <vaddps ymm9,ymm9,ymm8>
EmitIfCountGE Count, 2, <vaddps ymm11,ymm11,ymm10>
EmitIfCountGE Count, 3, <vaddps ymm13,ymm13,ymm12>
EmitIfCountGE Count, 4, <vaddps ymm15,ymm15,ymm14>
ENDIF
EmitIfCountGE Count, 1, <vmaskmovps YMMWORD PTR [r8],ymm0,ymm9>
EmitIfCountGE Count, 2, <vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm11>
EmitIfCountGE Count, 3, <vmaskmovps YMMWORD PTR [rbx],ymm0,ymm13>
EmitIfCountGE Count, 4, <vmaskmovps YMMWORD PTR [rbx+rax],ymm0,ymm15>
IFB <Fallthrough>
jmp ExitKernel
ENDIF
ENDM
@ -202,117 +393,7 @@ SgemmKernelAvxFunction MACRO Mode
cmp r11,4
jb ProcessCountMLessThan4
mov r11d,4 ; return 4 rows handled
cmp rbp,8
jbe ProcessRemainingCountN4
ProcessNextColumnLoop16x4:
vxorps xmm8,xmm8,xmm8 ; clear block accumulators
vxorps xmm9,xmm9,xmm9
vxorps xmm10,xmm10,xmm10
vxorps xmm11,xmm11,xmm11
vxorps xmm12,xmm12,xmm12
vxorps xmm13,xmm13,xmm13
vxorps xmm14,xmm14,xmm14
vxorps xmm15,xmm15,xmm15
lea rbx,[rcx+r10*2] ; compute matrix A plus 2 rows
ComputeBlockAvxLoop ComputeBlockAvxBy16, 4
vmulps ymm8,ymm8,ymm2 ; multiply by alpha
vmulps ymm9,ymm9,ymm2
vmulps ymm10,ymm10,ymm2
vmulps ymm11,ymm11,ymm2
vmulps ymm12,ymm12,ymm2
vmulps ymm13,ymm13,ymm2
vmulps ymm14,ymm14,ymm2
vmulps ymm15,ymm15,ymm2
lea rcx,[r8+rax*2] ; compute matrix C plus 2 rows
sub rbp,16
jb OutputMasked16x4Block
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
vaddps ymm9,ymm9,YMMWORD PTR [r8+32]
vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]
vaddps ymm11,ymm11,YMMWORD PTR [r8+rax+32]
vaddps ymm12,ymm12,YMMWORD PTR [rcx]
vaddps ymm13,ymm13,YMMWORD PTR [rcx+32]
vaddps ymm14,ymm14,YMMWORD PTR [rcx+rax]
vaddps ymm15,ymm15,YMMWORD PTR [rcx+rax+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
vmovups YMMWORD PTR [r8+32],ymm9
vmovups YMMWORD PTR [r8+rax],ymm10
vmovups YMMWORD PTR [r8+rax+32],ymm11
vmovups YMMWORD PTR [rcx],ymm12
vmovups YMMWORD PTR [rcx+32],ymm13
vmovups YMMWORD PTR [rcx+rax],ymm14
vmovups YMMWORD PTR [rcx+rax+32],ymm15
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
cmp rbp,8
ja ProcessNextColumnLoop16x4
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN4:
vxorps xmm9,xmm9,xmm9 ; clear block accumulators
vxorps xmm11,xmm11,xmm11
vxorps xmm13,xmm13,xmm13
vxorps xmm15,xmm15,xmm15
lea rbx,[rcx+r10*2] ; compute matrix A plus 2 rows
ComputeBlockAvxLoop ComputeBlockAvxBy8, 4
vmulps ymm9,ymm9,ymm2 ; multiply by alpha
vmulps ymm11,ymm11,ymm2
vmulps ymm13,ymm13,ymm2
vmulps ymm15,ymm15,ymm2
lea rcx,[r8+rax*2] ; compute matrix C plus 2 rows
cmp rbp,8
jb OutputMasked8x4Block
IFIDNI <Mode>, <Add>
vaddps ymm9,ymm9,YMMWORD PTR [r8]
vaddps ymm11,ymm11,YMMWORD PTR [r8+rax]
vaddps ymm13,ymm13,YMMWORD PTR [rcx]
vaddps ymm15,ymm15,YMMWORD PTR [rcx+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm9
vmovups YMMWORD PTR [r8+rax],ymm11
vmovups YMMWORD PTR [rcx],ymm13
vmovups YMMWORD PTR [rcx+rax],ymm15
jmp ExitKernel
OutputMasked16x4Block:
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]
vaddps ymm12,ymm12,YMMWORD PTR [rcx]
vaddps ymm14,ymm14,YMMWORD PTR [rcx+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
vmovups YMMWORD PTR [r8+rax],ymm10
vmovups YMMWORD PTR [rcx],ymm12
vmovups YMMWORD PTR [rcx+rax],ymm14
add r8,8*4 ; advance matrix C by 8 columns
add rcx,8*4 ; advance matrix C plus 2 rows by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x4Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss xmm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd xmm1,xmm0,XMMWORD PTR [MlasMaskMoveAvx+16]
vpcmpgtd xmm0,xmm0,XMMWORD PTR [MlasMaskMoveAvx]
vinsertf128 ymm0,ymm0,xmm1,1
IFIDNI <Mode>, <Add>
vmaskmovps ymm8,ymm0,YMMWORD PTR [r8]
vmaskmovps ymm10,ymm0,YMMWORD PTR [r8+rax]
vmaskmovps ymm12,ymm0,YMMWORD PTR [rcx]
vmaskmovps ymm14,ymm0,YMMWORD PTR [rcx+rax]
vaddps ymm9,ymm9,ymm8
vaddps ymm11,ymm11,ymm10
vaddps ymm13,ymm13,ymm12
vaddps ymm15,ymm15,ymm14
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm9
vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm11
vmaskmovps YMMWORD PTR [rcx],ymm0,ymm13
vmaskmovps YMMWORD PTR [rcx+rax],ymm0,ymm15
ProcessCountMAvx Mode, 4, Fallthrough
;
; Restore non-volatile registers and return.
@ -330,142 +411,14 @@ ProcessCountMLessThan4:
cmp r11,2
jb ProcessCountMLessThan2
mov r11d,2 ; return 2 rows handled
cmp rbp,8
jbe ProcessRemainingCountN2
ProcessNextColumnLoop16x2:
vxorps xmm8,xmm8,xmm8 ; clear block accumulators
vxorps xmm9,xmm9,xmm9
vxorps xmm10,xmm10,xmm10
vxorps xmm11,xmm11,xmm11
ComputeBlockAvxLoop ComputeBlockAvxBy16, 2
vmulps ymm8,ymm8,ymm2 ; multiply by alpha
vmulps ymm9,ymm9,ymm2
vmulps ymm10,ymm10,ymm2
vmulps ymm11,ymm11,ymm2
sub rbp,16
jb OutputMasked16x2Block
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
vaddps ymm9,ymm9,YMMWORD PTR [r8+32]
vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]
vaddps ymm11,ymm11,YMMWORD PTR [r8+rax+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
vmovups YMMWORD PTR [r8+32],ymm9
vmovups YMMWORD PTR [r8+rax],ymm10
vmovups YMMWORD PTR [r8+rax+32],ymm11
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
cmp rbp,8
ja ProcessNextColumnLoop16x2
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN2:
vxorps xmm9,xmm9,xmm9 ; clear block accumulators
vxorps xmm11,xmm11,xmm11
ComputeBlockAvxLoop ComputeBlockAvxBy8, 2
vmulps ymm9,ymm9,ymm2 ; multiply by alpha
vmulps ymm11,ymm11,ymm2
cmp rbp,8
jb OutputMasked8x2Block
IFIDNI <Mode>, <Add>
vaddps ymm9,ymm9,YMMWORD PTR [r8]
vaddps ymm11,ymm11,YMMWORD PTR [r8+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm9
vmovups YMMWORD PTR [r8+rax],ymm11
jmp ExitKernel
OutputMasked16x2Block:
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
vaddps ymm10,ymm10,YMMWORD PTR [r8+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
vmovups YMMWORD PTR [r8+rax],ymm10
add r8,8*4 ; advance matrix C by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x2Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss xmm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd xmm1,xmm0,XMMWORD PTR [MlasMaskMoveAvx+16]
vpcmpgtd xmm0,xmm0,XMMWORD PTR [MlasMaskMoveAvx]
vinsertf128 ymm0,ymm0,xmm1,1
IFIDNI <Mode>, <Add>
vmaskmovps ymm8,ymm0,YMMWORD PTR [r8]
vmaskmovps ymm10,ymm0,YMMWORD PTR [r8+rax]
vaddps ymm9,ymm9,ymm8
vaddps ymm11,ymm11,ymm10
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm9
vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm11
jmp ExitKernel
ProcessCountMAvx Mode, 2
;
; Process 1 row of the matrices.
;
ProcessCountMLessThan2:
mov r11d,1 ; return 1 row handled
cmp rbp,8
jbe ProcessRemainingCountN1
ProcessNextColumnLoop16x1:
vxorps xmm8,xmm8,xmm8 ; clear block accumulators
vxorps xmm9,xmm9,xmm9
ComputeBlockAvxLoop ComputeBlockAvxBy16, 1
vmulps ymm8,ymm8,ymm2 ; multiply by alpha
vmulps ymm9,ymm9,ymm2
sub rbp,16
jb OutputMasked16x1Block
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
vaddps ymm9,ymm9,YMMWORD PTR [r8+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
vmovups YMMWORD PTR [r8+32],ymm9
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
cmp rbp,8
ja ProcessNextColumnLoop16x1
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN1:
vxorps xmm9,xmm9,xmm9 ; clear block accumulators
ComputeBlockAvxLoop ComputeBlockAvxBy8, 1
vmulps ymm9,ymm9,ymm2 ; multiply by alpha
cmp rbp,8
jb OutputMasked8x1Block
IFIDNI <Mode>, <Add>
vaddps ymm9,ymm9,YMMWORD PTR [r8]
ENDIF
vmovups YMMWORD PTR [r8],ymm9
jmp ExitKernel
OutputMasked16x1Block:
IFIDNI <Mode>, <Add>
vaddps ymm8,ymm8,YMMWORD PTR [r8]
ENDIF
vmovups YMMWORD PTR [r8],ymm8
add r8,8*4 ; advance matrix C by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x1Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss xmm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd xmm1,xmm0,XMMWORD PTR [MlasMaskMoveAvx+16]
vpcmpgtd xmm0,xmm0,XMMWORD PTR [MlasMaskMoveAvx]
vinsertf128 ymm0,ymm0,xmm1,1
IFIDNI <Mode>, <Add>
vmaskmovps ymm8,ymm0,YMMWORD PTR [r8]
vaddps ymm9,ymm9,ymm8
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm9
jmp ExitKernel
ProcessCountMAvx Mode, 1
NESTED_END MlasSgemmKernel&Mode&Avx, _TEXT

714
onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm
View file

@ -25,8 +25,7 @@ INCLUDE SgemmKernelCommon.inc
;
; ComputeBlockAvx512FBy32
;
; This macro multiplies and accumulates for a 32xN block (where N is 1,3,6,12)
; of the output matrix.
; This macro multiplies and accumulates for a 32xN block of the output matrix.
;
; Arguments:
;
@ -36,6 +35,9 @@ INCLUDE SgemmKernelCommon.inc
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 3 rows.
@ -53,8 +55,12 @@ INCLUDE SgemmKernelCommon.inc
; zmm4-zmm27 - Supplies the block accumulators.
;
ComputeBlockAvx512FBy32 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockAvx512FBy32 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IFNB <PrefetchOffset>
prefetcht0 [rdx+VectorOffset+PrefetchOffset]
prefetcht0 [rdx+r12+VectorOffset+PrefetchOffset]
ENDIF
IF Count EQ 1
vbroadcastss zmm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps zmm4,zmm3,ZMMWORD PTR [rdx+VectorOffset]
@ -62,48 +68,42 @@ IF Count EQ 1
ELSE
vmovaps zmm0,ZMMWORD PTR [rdx+VectorOffset]
vmovaps zmm1,ZMMWORD PTR [rdx+r12+VectorOffset]
vbroadcastss zmm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps zmm4,zmm3,zmm0
vfmadd231ps zmm5,zmm3,zmm1
IF Count GE 3
vbroadcastss zmm3,DWORD PTR [rcx+r10+BroadcastOffset]
vfmadd231ps zmm6,zmm3,zmm0
vfmadd231ps zmm7,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [rcx+r10*2+BroadcastOffset]
vfmadd231ps zmm8,zmm3,zmm0
vfmadd231ps zmm9,zmm3,zmm1
ENDIF
IF Count GE 6
vbroadcastss zmm3,DWORD PTR [rbx+BroadcastOffset]
vfmadd231ps zmm10,zmm3,zmm0
vfmadd231ps zmm11,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [rbx+r10+BroadcastOffset]
vfmadd231ps zmm12,zmm3,zmm0
vfmadd231ps zmm13,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [rbx+r10*2+BroadcastOffset]
vfmadd231ps zmm14,zmm3,zmm0
vfmadd231ps zmm15,zmm3,zmm1
ENDIF
IF Count GE 12
vbroadcastss zmm3,DWORD PTR [r13+BroadcastOffset]
vfmadd231ps zmm16,zmm3,zmm0
vfmadd231ps zmm17,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [r13+r10+BroadcastOffset]
vfmadd231ps zmm18,zmm3,zmm0
vfmadd231ps zmm19,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [r13+r10*2+BroadcastOffset]
vfmadd231ps zmm20,zmm3,zmm0
vfmadd231ps zmm21,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [r14+BroadcastOffset]
vfmadd231ps zmm22,zmm3,zmm0
vfmadd231ps zmm23,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [r14+r10+BroadcastOffset]
vfmadd231ps zmm24,zmm3,zmm0
vfmadd231ps zmm25,zmm3,zmm1
vbroadcastss zmm3,DWORD PTR [r14+r10*2+BroadcastOffset]
vfmadd231ps zmm26,zmm3,zmm0
vfmadd231ps zmm27,zmm3,zmm1
ENDIF
EmitIfCountGE Count, 1, <vbroadcastss zmm3,DWORD PTR [rcx+BroadcastOffset]>
EmitIfCountGE Count, 1, <vfmadd231ps zmm4,zmm3,zmm0>
EmitIfCountGE Count, 1, <vfmadd231ps zmm5,zmm3,zmm1>
EmitIfCountGE Count, 2, <vbroadcastss zmm3,DWORD PTR [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 2, <vfmadd231ps zmm6,zmm3,zmm0>
EmitIfCountGE Count, 2, <vfmadd231ps zmm7,zmm3,zmm1>
EmitIfCountGE Count, 3, <vbroadcastss zmm3,DWORD PTR [rcx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 3, <vfmadd231ps zmm8,zmm3,zmm0>
EmitIfCountGE Count, 3, <vfmadd231ps zmm9,zmm3,zmm1>
EmitIfCountGE Count, 4, <vbroadcastss zmm3,DWORD PTR [rbx+BroadcastOffset]>
EmitIfCountGE Count, 4, <vfmadd231ps zmm10,zmm3,zmm0>
EmitIfCountGE Count, 4, <vfmadd231ps zmm11,zmm3,zmm1>
EmitIfCountGE Count, 5, <vbroadcastss zmm3,DWORD PTR [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 5, <vfmadd231ps zmm12,zmm3,zmm0>
EmitIfCountGE Count, 5, <vfmadd231ps zmm13,zmm3,zmm1>
EmitIfCountGE Count, 6, <vbroadcastss zmm3,DWORD PTR [rbx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 6, <vfmadd231ps zmm14,zmm3,zmm0>
EmitIfCountGE Count, 6, <vfmadd231ps zmm15,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r13+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm16,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm17,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r13+r10+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm18,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm19,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r13+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm20,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm21,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r14+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm22,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm23,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r14+r10+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm24,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm25,zmm3,zmm1>
EmitIfCountGE Count, 12, <vbroadcastss zmm3,DWORD PTR [r14+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm26,zmm3,zmm0>
EmitIfCountGE Count, 12, <vfmadd231ps zmm27,zmm3,zmm1>
ENDIF
ENDM
@ -111,8 +111,7 @@ ENDIF
;
; ComputeBlockAvx512FBy16
;
; This macro multiplies and accumulates for a 16xN block (where N is 1,3,6,12)
; of the output matrix.
; This macro multiplies and accumulates for a 16xN block of the output matrix.
;
; Arguments:
;
@ -122,6 +121,9 @@ ENDIF
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 3 rows.
@ -139,27 +141,248 @@ ENDIF
; zmm4-zmm27 - Supplies the block accumulators.
;
ComputeBlockAvx512FBy16 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockAvx512FBy16 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IFNB <PrefetchOffset>
prefetcht0 [rdx+VectorOffset+PrefetchOffset]
ENDIF
vmovaps zmm0,ZMMWORD PTR [rdx+VectorOffset]
vfmadd231ps zmm5,zmm0,DWORD BCST [rcx+BroadcastOffset]
IF Count GE 3
vfmadd231ps zmm7,zmm0,DWORD BCST [rcx+r10+BroadcastOffset]
vfmadd231ps zmm9,zmm0,DWORD BCST [rcx+r10*2+BroadcastOffset]
EmitIfCountGE Count, 1, <vfmadd231ps zmm5,zmm0,DWORD BCST [rcx+BroadcastOffset]>
EmitIfCountGE Count, 2, <vfmadd231ps zmm7,zmm0,DWORD BCST [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 3, <vfmadd231ps zmm9,zmm0,DWORD BCST [rcx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 4, <vfmadd231ps zmm11,zmm0,DWORD BCST [rbx+BroadcastOffset]>
EmitIfCountGE Count, 5, <vfmadd231ps zmm13,zmm0,DWORD BCST [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 6, <vfmadd231ps zmm15,zmm0,DWORD BCST [rbx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm17,zmm0,DWORD BCST [r13+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm19,zmm0,DWORD BCST [r13+r10+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm21,zmm0,DWORD BCST [r13+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm23,zmm0,DWORD BCST [r14+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm25,zmm0,DWORD BCST [r14+r10+BroadcastOffset]>
EmitIfCountGE Count, 12, <vfmadd231ps zmm27,zmm0,DWORD BCST [r14+r10*2+BroadcastOffset]>
ENDM
;
; ComputeBlockAvx512FLoop
;
; This macro generates code to execute the block compute macro multiple
; times and advancing the matrix A and matrix B data pointers.
;
; Arguments:
;
; ComputeBlock - Supplies the macro to compute a single block.
;
; Count - Supplies the number of rows to access from matrix A.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus N rows.
;
; rcx - Supplies the address into the matrix A data.
;
; rdx - Supplies the address into the matrix B data.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockAvx512FLoop MACRO ComputeBlock, Count
IF Count GT 3
lea rbx,[r10*2+r10]
IF Count EQ 12
lea r13,[rcx+rbx*2] ; compute matrix A plus 6 rows
lea r14,[r13+rbx] ; compute matrix A plus 9 rows
ENDIF
IF Count GE 6
vfmadd231ps zmm11,zmm0,DWORD BCST [rbx+BroadcastOffset]
vfmadd231ps zmm13,zmm0,DWORD BCST [rbx+r10+BroadcastOffset]
vfmadd231ps zmm15,zmm0,DWORD BCST [rbx+r10*2+BroadcastOffset]
add rbx,rcx ; compute matrix A plus 3 rows
ENDIF
IF Count GE 12
vfmadd231ps zmm17,zmm0,DWORD BCST [r13+BroadcastOffset]
vfmadd231ps zmm19,zmm0,DWORD BCST [r13+r10+BroadcastOffset]
vfmadd231ps zmm21,zmm0,DWORD BCST [r13+r10*2+BroadcastOffset]
vfmadd231ps zmm23,zmm0,DWORD BCST [r14+BroadcastOffset]
vfmadd231ps zmm25,zmm0,DWORD BCST [r14+r10+BroadcastOffset]
vfmadd231ps zmm27,zmm0,DWORD BCST [r14+r10*2+BroadcastOffset]
ComputeBlockLoop ComputeBlock, Count, <Count GT 3>
IF Count GT 3
lea rbx,[rax*2+rax]
IF Count EQ 12
lea r13,[r8+rbx*2] ; compute matrix C plus 6 rows
lea r14,[r13+rbx] ; compute matrix C plus 9 rows
ENDIF
add rbx,r8 ; compute matrix C plus 3 rows
ENDIF
ENDM
;
; ProcessCountMAvx512F
;
; Macro Description:
;
; This macro generates code to compute matrix multiplication for a fixed set
; of rows.
;
; Arguments:
;
; Mode - Supplies the mode of operation for updating the contents of matrix C.
;
; Count - Supplies the number of rows to process.
;
; Implicit Arguments:
;
; rcx - Supplies the address of matrix A.
;
; rdx - Supplies the address of matrix B.
;
; rsi - Supplies the address of matrix A.
;
; rbp - Supplies the number of columns from matrix B and matrix C to iterate
; over.
;
; r8 - Supplies the address of matrix C.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; r10 - Supplies the length in bytes of a row from matrix A.
;
ProcessCountMAvx512F MACRO Mode, Count
LOCAL ProcessNextColumnLoop32xN
LOCAL Output16xNBlock
LOCAL Output16xNBlockWithMask
LOCAL ProcessRemainingCountN
cmp rbp,16
jbe ProcessRemainingCountN
ProcessNextColumnLoop32xN:
EmitIfCountGE Count, 12, <vmovaps zmm16,zmm4>
; clear upper block accumulators
EmitIfCountGE Count, 12, <vmovaps zmm17,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm18,zmm4>
EmitIfCountGE Count, 12, <vmovaps zmm19,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm20,zmm4>
EmitIfCountGE Count, 12, <vmovaps zmm21,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm22,zmm4>
EmitIfCountGE Count, 12, <vmovaps zmm23,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm24,zmm4>
EmitIfCountGE Count, 12, <vmovaps zmm25,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm26,zmm4>
EmitIfCountGE Count, 12, <vmovaps zmm27,zmm5>
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, Count
add rdx,r12 ; advance matrix B by 16*CountK floats
IFDIFI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmulps zmm4,zmm4,zmm31>
EmitIfCountGE Count, 2, <vmulps zmm6,zmm6,zmm31>
EmitIfCountGE Count, 3, <vmulps zmm8,zmm8,zmm31>
EmitIfCountGE Count, 4, <vmulps zmm10,zmm10,zmm31>
EmitIfCountGE Count, 5, <vmulps zmm12,zmm12,zmm31>
EmitIfCountGE Count, 6, <vmulps zmm14,zmm14,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm16,zmm16,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm18,zmm18,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm20,zmm20,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm22,zmm22,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm24,zmm24,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm26,zmm26,zmm31>
ELSE
EmitIfCountGE Count, 1, <vfmadd213ps zmm4,zmm31,ZMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vfmadd213ps zmm6,zmm31,ZMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vfmadd213ps zmm8,zmm31,ZMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 4, <vfmadd213ps zmm10,zmm31,ZMMWORD PTR [rbx]>
EmitIfCountGE Count, 5, <vfmadd213ps zmm12,zmm31,ZMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 6, <vfmadd213ps zmm14,zmm31,ZMMWORD PTR [rbx+rax*2]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm16,zmm31,ZMMWORD PTR [r13]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm18,zmm31,ZMMWORD PTR [r13+rax]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm20,zmm31,ZMMWORD PTR [r13+rax*2]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm22,zmm31,ZMMWORD PTR [r14]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm24,zmm31,ZMMWORD PTR [r14+rax]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm26,zmm31,ZMMWORD PTR [r14+rax*2]>
ENDIF
EmitIfCountGE Count, 1, <vmovups ZMMWORD PTR [r8],zmm4>
EmitIfCountGE Count, 2, <vmovups ZMMWORD PTR [r8+rax],zmm6>
EmitIfCountGE Count, 3, <vmovups ZMMWORD PTR [r8+rax*2],zmm8>
EmitIfCountGE Count, 4, <vmovups ZMMWORD PTR [rbx],zmm10>
EmitIfCountGE Count, 5, <vmovups ZMMWORD PTR [rbx+rax],zmm12>
EmitIfCountGE Count, 6, <vmovups ZMMWORD PTR [rbx+rax*2],zmm14>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13],zmm16>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13+rax],zmm18>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13+rax*2],zmm20>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14],zmm22>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14+rax],zmm24>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14+rax*2],zmm26>
add r8,16*4 ; advance matrix C by 16 columns
IF Count GT 3
add rbx,16*4 ; advance matrix C plus 3 rows by 16 columns
add r13,16*4 ; advance matrix C plus 6 rows by 16 columns
add r14,16*4 ; advance matrix C plus 9 rows by 16 columns
ENDIF
sub rbp,16
Output16xNBlock:
sub rbp,16
jae Output16xNBlockWithMask
lea ecx,[ebp+16] ; correct for over-subtract above
mov edi,1
shl edi,cl
dec edi
kmovw k1,edi ; update mask for remaining columns
xor ebp,ebp ; no more columns remaining
Output16xNBlockWithMask:
IFDIFI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmulps zmm5,zmm5,zmm31>
EmitIfCountGE Count, 2, <vmulps zmm7,zmm7,zmm31>
EmitIfCountGE Count, 3, <vmulps zmm9,zmm9,zmm31>
EmitIfCountGE Count, 4, <vmulps zmm11,zmm11,zmm31>
EmitIfCountGE Count, 5, <vmulps zmm13,zmm13,zmm31>
EmitIfCountGE Count, 6, <vmulps zmm15,zmm15,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm17,zmm17,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm19,zmm19,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm21,zmm21,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm23,zmm23,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm25,zmm25,zmm31>
EmitIfCountGE Count, 12, <vmulps zmm27,zmm27,zmm31>
ELSE
EmitIfCountGE Count, 1, <vfmadd213ps zmm5{k1},zmm31,ZMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vfmadd213ps zmm7{k1},zmm31,ZMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vfmadd213ps zmm9{k1},zmm31,ZMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 4, <vfmadd213ps zmm11{k1},zmm31,ZMMWORD PTR [rbx]>
EmitIfCountGE Count, 5, <vfmadd213ps zmm13{k1},zmm31,ZMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 6, <vfmadd213ps zmm15{k1},zmm31,ZMMWORD PTR [rbx+rax*2]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm17{k1},zmm31,ZMMWORD PTR [r13]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm19{k1},zmm31,ZMMWORD PTR [r13+rax]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm21{k1},zmm31,ZMMWORD PTR [r13+rax*2]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm23{k1},zmm31,ZMMWORD PTR [r14]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm25{k1},zmm31,ZMMWORD PTR [r14+rax]>
EmitIfCountGE Count, 12, <vfmadd213ps zmm27{k1},zmm31,ZMMWORD PTR [r14+rax*2]>
ENDIF
EmitIfCountGE Count, 1, <vmovups ZMMWORD PTR [r8]{k1},zmm5>
EmitIfCountGE Count, 2, <vmovups ZMMWORD PTR [r8+rax]{k1},zmm7>
EmitIfCountGE Count, 3, <vmovups ZMMWORD PTR [r8+rax*2]{k1},zmm9>
EmitIfCountGE Count, 4, <vmovups ZMMWORD PTR [rbx]{k1},zmm11>
EmitIfCountGE Count, 5, <vmovups ZMMWORD PTR [rbx+rax]{k1},zmm13>
EmitIfCountGE Count, 6, <vmovups ZMMWORD PTR [rbx+rax*2]{k1},zmm15>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13]{k1},zmm17>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13+rax]{k1},zmm19>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r13+rax*2]{k1},zmm21>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14]{k1},zmm23>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14+rax]{k1},zmm25>
EmitIfCountGE Count, 12, <vmovups ZMMWORD PTR [r14+rax*2]{k1},zmm27>
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,16
ja ProcessNextColumnLoop32xN
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN:
EmitIfCountGE Count, 12, <vmovaps zmm17,zmm5>
; clear upper block accumulators
EmitIfCountGE Count, 12, <vmovaps zmm19,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm21,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm23,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm25,zmm5>
EmitIfCountGE Count, 12, <vmovaps zmm27,zmm5>
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy16, Count
jmp Output16xNBlock
ENDM
@ -214,155 +437,33 @@ SgemmKernelAvx512FFunction MACRO Mode
vbroadcastss zmm31,DWORD PTR SgemmKernelFrame.Alpha[rsp]
;
; Process 12 rows of the matrices.
; Process N rows of the matrices.
;
cmp r11,12
jb ProcessCountMLessThan12
mov r11d,12 ; return 12 rows handled
cmp rbp,16
jbe ProcessRemainingCountN12
ProcessCountMAvx512F Mode, 12
ProcessNextColumnLoop32x12:
vmovaps zmm16,zmm4 ; clear upper block accumulators
vmovaps zmm17,zmm5
vmovaps zmm18,zmm4
vmovaps zmm19,zmm5
vmovaps zmm20,zmm4
vmovaps zmm21,zmm5
vmovaps zmm22,zmm4
vmovaps zmm23,zmm5
vmovaps zmm24,zmm4
vmovaps zmm25,zmm5
vmovaps zmm26,zmm4
vmovaps zmm27,zmm5
lea rbx,[r10*2+r10]
lea r13,[rcx+rbx*2] ; compute matrix A plus 6 rows
lea r14,[r13+rbx] ; compute matrix A plus 9 rows
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, 12
add rdx,r12 ; advance matrix B by 16*CountK floats
lea rbx,[rax*2+rax]
lea r13,[r8+rbx*2] ; compute matrix C plus 6 rows
lea r14,[r13+rbx] ; compute matrix C plus 9 rows
add rbx,r8 ; compute matrix C plus 3 rows
IFDIFI <Mode>, <Add>
vmulps zmm4,zmm4,zmm31 ; multiply by alpha
vmulps zmm6,zmm6,zmm31
vmulps zmm8,zmm8,zmm31
vmulps zmm10,zmm10,zmm31
vmulps zmm12,zmm12,zmm31
vmulps zmm14,zmm14,zmm31
vmulps zmm16,zmm16,zmm31
vmulps zmm18,zmm18,zmm31
vmulps zmm20,zmm20,zmm31
vmulps zmm22,zmm22,zmm31
vmulps zmm24,zmm24,zmm31
vmulps zmm26,zmm26,zmm31
ELSE
vfmadd213ps zmm4,zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm6,zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm8,zmm31,ZMMWORD PTR [r8+rax*2]
vfmadd213ps zmm10,zmm31,ZMMWORD PTR [rbx]
vfmadd213ps zmm12,zmm31,ZMMWORD PTR [rbx+rax]
vfmadd213ps zmm14,zmm31,ZMMWORD PTR [rbx+rax*2]
vfmadd213ps zmm16,zmm31,ZMMWORD PTR [r13]
vfmadd213ps zmm18,zmm31,ZMMWORD PTR [r13+rax]
vfmadd213ps zmm20,zmm31,ZMMWORD PTR [r13+rax*2]
vfmadd213ps zmm22,zmm31,ZMMWORD PTR [r14]
vfmadd213ps zmm24,zmm31,ZMMWORD PTR [r14+rax]
vfmadd213ps zmm26,zmm31,ZMMWORD PTR [r14+rax*2]
ENDIF
vmovups ZMMWORD PTR [r8],zmm4
vmovups ZMMWORD PTR [r8+rax],zmm6
vmovups ZMMWORD PTR [r8+rax*2],zmm8
vmovups ZMMWORD PTR [rbx],zmm10
vmovups ZMMWORD PTR [rbx+rax],zmm12
vmovups ZMMWORD PTR [rbx+rax*2],zmm14
vmovups ZMMWORD PTR [r13],zmm16
vmovups ZMMWORD PTR [r13+rax],zmm18
vmovups ZMMWORD PTR [r13+rax*2],zmm20
vmovups ZMMWORD PTR [r14],zmm22
vmovups ZMMWORD PTR [r14+rax],zmm24
vmovups ZMMWORD PTR [r14+rax*2],zmm26
add r8,16*4 ; advance matrix C by 16 columns
sub rbp,16
ProcessCountMLessThan12:
cmp r11,5
ja ProcessCountM6
je ProcessCountM5
cmp r11,3
ja ProcessCountM4
je ProcessCountM3
cmp r11,1
je ProcessCountM1
Output16x12Block:
sub rbp,16
jae Output16x12BlockWithMask
lea ecx,[ebp+16] ; correct for over-subtract above
mov edi,1
shl edi,cl
dec edi
kmovw k1,edi ; update mask for remaining columns
xor ebp,ebp ; no more columns remaining
ProcessCountM2:
ProcessCountMAvx512F Mode, 2
Output16x12BlockWithMask:
lea rbx,[rax*2+rax]
lea r13,[r8+rbx*2] ; compute matrix C plus 6 rows
lea r14,[r13+rbx] ; compute matrix C plus 9 rows
add rbx,r8 ; compute matrix C plus 3 rows
IFDIFI <Mode>, <Add>
vmulps zmm5,zmm5,zmm31 ; multiply by alpha
vmulps zmm7,zmm7,zmm31
vmulps zmm9,zmm9,zmm31
vmulps zmm11,zmm11,zmm31
vmulps zmm13,zmm13,zmm31
vmulps zmm15,zmm15,zmm31
vmulps zmm17,zmm17,zmm31
vmulps zmm19,zmm19,zmm31
vmulps zmm21,zmm21,zmm31
vmulps zmm23,zmm23,zmm31
vmulps zmm25,zmm25,zmm31
vmulps zmm27,zmm27,zmm31
ELSE
vfmadd213ps zmm5{k1},zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm7{k1},zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm9{k1},zmm31,ZMMWORD PTR [r8+rax*2]
vfmadd213ps zmm11{k1},zmm31,ZMMWORD PTR [rbx]
vfmadd213ps zmm13{k1},zmm31,ZMMWORD PTR [rbx+rax]
vfmadd213ps zmm15{k1},zmm31,ZMMWORD PTR [rbx+rax*2]
vfmadd213ps zmm17{k1},zmm31,ZMMWORD PTR [r13]
vfmadd213ps zmm19{k1},zmm31,ZMMWORD PTR [r13+rax]
vfmadd213ps zmm21{k1},zmm31,ZMMWORD PTR [r13+rax*2]
vfmadd213ps zmm23{k1},zmm31,ZMMWORD PTR [r14]
vfmadd213ps zmm25{k1},zmm31,ZMMWORD PTR [r14+rax]
vfmadd213ps zmm27{k1},zmm31,ZMMWORD PTR [r14+rax*2]
ENDIF
vmovups ZMMWORD PTR [r8]{k1},zmm5
vmovups ZMMWORD PTR [r8+rax]{k1},zmm7
vmovups ZMMWORD PTR [r8+rax*2]{k1},zmm9
vmovups ZMMWORD PTR [rbx]{k1},zmm11
vmovups ZMMWORD PTR [rbx+rax]{k1},zmm13
vmovups ZMMWORD PTR [rbx+rax*2]{k1},zmm15
vmovups ZMMWORD PTR [r13]{k1},zmm17
vmovups ZMMWORD PTR [r13+rax]{k1},zmm19
vmovups ZMMWORD PTR [r13+rax*2]{k1},zmm21
vmovups ZMMWORD PTR [r14]{k1},zmm23
vmovups ZMMWORD PTR [r14+rax]{k1},zmm25
vmovups ZMMWORD PTR [r14+rax*2]{k1},zmm27
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,16
ja ProcessNextColumnLoop32x12
test rbp,rbp
jz ExitKernel
ProcessCountM4:
ProcessCountMAvx512F Mode, 4
ProcessRemainingCountN12:
vmovaps zmm17,zmm5 ; clear upper block accumulators
vmovaps zmm19,zmm5
vmovaps zmm21,zmm5
vmovaps zmm23,zmm5
vmovaps zmm25,zmm5
vmovaps zmm27,zmm5
lea rbx,[r10*2+r10]
lea r13,[rcx+rbx*2] ; compute matrix A plus 6 rows
lea r14,[r13+rbx] ; compute matrix A plus 9 rows
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy16, 12
jmp Output16x12Block
ProcessCountM6:
mov r11d,6 ; return 6 rows handled
ProcessCountMAvx512F Mode, 6
;
; Restore non-volatile registers and return.
@ -371,209 +472,14 @@ ProcessRemainingCountN12:
ExitKernel:
SgemmKernelAvxExit RestoreExtra
;
; Process 6 rows of the matrices.
;
ProcessCountM1:
ProcessCountMAvx512F Mode, 1
ProcessCountMLessThan12:
cmp r11,6
jb ProcessCountMLessThan6
mov r11d,6 ; return 6 rows handled
cmp rbp,16
jbe ProcessRemainingCountN6
ProcessCountM3:
ProcessCountMAvx512F Mode, 3
ProcessNextColumnLoop32x6:
lea rbx,[r10*2+r10]
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, 6
add rdx,r12 ; advance matrix B by 16*CountK floats
lea rdi,[r8+rax*2] ; compute matrix C plus 2 rows
lea rbx,[r8+rax*4] ; compute matrix C plus 4 rows
IFDIFI <Mode>, <Add>
vmulps zmm4,zmm4,zmm31 ; multiply by alpha
vmulps zmm6,zmm6,zmm31
vmulps zmm8,zmm8,zmm31
vmulps zmm10,zmm10,zmm31
vmulps zmm12,zmm12,zmm31
vmulps zmm14,zmm14,zmm31
ELSE
vfmadd213ps zmm4,zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm6,zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm8,zmm31,ZMMWORD PTR [rdi]
vfmadd213ps zmm10,zmm31,ZMMWORD PTR [rdi+rax]
vfmadd213ps zmm12,zmm31,ZMMWORD PTR [rbx]
vfmadd213ps zmm14,zmm31,ZMMWORD PTR [rbx+rax]
ENDIF
vmovups ZMMWORD PTR [r8],zmm4
vmovups ZMMWORD PTR [r8+rax],zmm6
vmovups ZMMWORD PTR [rdi],zmm8
vmovups ZMMWORD PTR [rdi+rax],zmm10
vmovups ZMMWORD PTR [rbx],zmm12
vmovups ZMMWORD PTR [rbx+rax],zmm14
add r8,16*4 ; advance matrix C by 16 columns
sub rbp,16
Output16x6Block:
sub rbp,16
jae Output16x6BlockWithMask
lea ecx,[ebp+16] ; correct for over-subtract above
mov edi,1
shl edi,cl
dec edi
kmovw k1,edi ; update mask for remaining columns
xor ebp,ebp ; no more columns remaining
Output16x6BlockWithMask:
lea rdi,[r8+rax*2] ; compute matrix C plus 2 rows
lea rbx,[r8+rax*4] ; compute matrix C plus 4 rows
IFDIFI <Mode>, <Add>
vmulps zmm5,zmm5,zmm31 ; multiply by alpha
vmulps zmm7,zmm7,zmm31
vmulps zmm9,zmm9,zmm31
vmulps zmm11,zmm11,zmm31
vmulps zmm13,zmm13,zmm31
vmulps zmm15,zmm15,zmm31
ELSE
vfmadd213ps zmm5{k1},zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm7{k1},zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm9{k1},zmm31,ZMMWORD PTR [rdi]
vfmadd213ps zmm11{k1},zmm31,ZMMWORD PTR [rdi+rax]
vfmadd213ps zmm13{k1},zmm31,ZMMWORD PTR [rbx]
vfmadd213ps zmm15{k1},zmm31,ZMMWORD PTR [rbx+rax]
ENDIF
vmovups ZMMWORD PTR [r8]{k1},zmm5
vmovups ZMMWORD PTR [r8+rax]{k1},zmm7
vmovups ZMMWORD PTR [rdi]{k1},zmm9
vmovups ZMMWORD PTR [rdi+rax]{k1},zmm11
vmovups ZMMWORD PTR [rbx]{k1},zmm13
vmovups ZMMWORD PTR [rbx+rax]{k1},zmm15
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,16
ja ProcessNextColumnLoop32x6
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN6:
lea rbx,[r10*2+r10]
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy16, 6
jmp Output16x6Block
;
; Process 3 rows of the matrices.
;
ProcessCountMLessThan6:
cmp r11,3
jb ProcessCountMLessThan3
mov r11d,3 ; return 3 rows handled
cmp rbp,16
jbe ProcessRemainingCountN3
ProcessNextColumnLoop32x3:
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, 3
add rdx,r12 ; advance matrix B by 16*CountK floats
IFDIFI <Mode>, <Add>
vmulps zmm4,zmm4,zmm31 ; multiply by alpha
vmulps zmm6,zmm6,zmm31
vmulps zmm8,zmm8,zmm31
ELSE
vfmadd213ps zmm4,zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm6,zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm8,zmm31,ZMMWORD PTR [r8+rax*2]
ENDIF
vmovups ZMMWORD PTR [r8],zmm4
vmovups ZMMWORD PTR [r8+rax],zmm6
vmovups ZMMWORD PTR [r8+rax*2],zmm8
add r8,16*4 ; advance matrix C by 16 columns
sub rbp,16
Output16x3Block:
sub rbp,16
jae Output16x3BlockWithMask
lea ecx,[ebp+16] ; correct for over-subtract above
mov edi,1
shl edi,cl
dec edi
kmovw k1,edi ; update mask for remaining columns
xor ebp,ebp ; no more columns remaining
Output16x3BlockWithMask:
IFDIFI <Mode>, <Add>
vmulps zmm5,zmm5,zmm31 ; multiply by alpha
vmulps zmm7,zmm7,zmm31
vmulps zmm9,zmm9,zmm31
ELSE
vfmadd213ps zmm5{k1},zmm31,ZMMWORD PTR [r8]
vfmadd213ps zmm7{k1},zmm31,ZMMWORD PTR [r8+rax]
vfmadd213ps zmm9{k1},zmm31,ZMMWORD PTR [r8+rax*2]
ENDIF
vmovups ZMMWORD PTR [r8]{k1},zmm5
vmovups ZMMWORD PTR [r8+rax]{k1},zmm7
vmovups ZMMWORD PTR [r8+rax*2]{k1},zmm9
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,16
ja ProcessNextColumnLoop32x3
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN3:
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy16, 3
jmp Output16x3Block
;
; Process 1 row of the matrices.
;
ProcessCountMLessThan3:
mov r11d,1 ; return 1 row handled
cmp rbp,16
jbe ProcessRemainingCountN1
ProcessNextColumnLoop32x1:
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, 1
add rdx,r12 ; advance matrix B by 16*CountK floats
IFDIFI <Mode>, <Add>
vmulps zmm4,zmm4,zmm31 ; multiply by alpha
ELSE
vfmadd213ps zmm4,zmm31,ZMMWORD PTR [r8]
ENDIF
vmovups ZMMWORD PTR [r8],zmm4
add r8,16*4 ; advance matrix C by 16 columns
sub rbp,16
Output16x1Block:
sub rbp,16
jae Output16x1BlockWithMask
lea ecx,[ebp+16] ; correct for over-subtract above
mov edi,1
shl edi,cl
dec edi
kmovw k1,edi ; update mask for remaining columns
xor ebp,ebp ; no more columns remaining
Output16x1BlockWithMask:
IFDIFI <Mode>, <Add>
vmulps zmm5,zmm5,zmm31 ; multiply by alpha
ELSE
vfmadd213ps zmm5{k1},zmm31,ZMMWORD PTR [r8]
ENDIF
vmovups ZMMWORD PTR [r8]{k1},zmm5
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,16
ja ProcessNextColumnLoop32x1
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN1:
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy16, 1
jmp Output16x1Block
ProcessCountM5:
ProcessCountMAvx512F Mode, 5
NESTED_END MlasSgemmKernel&Mode&Avx512F, _TEXT

59
onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc
View file

@ -74,29 +74,6 @@ SgemmKernelM1Frame STRUCT
SgemmKernelM1Frame ENDS
;
; save_xmm128_avx
;
; Macro Description:
;
; This macro uses AVX instructions to save a vector register as part of a
; function prologue as an alternative to save_xmm128.
;
; Arguments:
;
; Register - Supplies the vector register to be saved.
;
; Offset - Supplies the stack frame offset to store the contents of the
; vector register.
;
save_xmm128_avx MACRO Register, Offset
vmovaps Offset[rsp], Register
.savexmm128 Register, Offset
ENDM
;
; SgemmKernelAvxEntry
;
@ -212,7 +189,7 @@ ENDIF
ENDM
;
; ComputeBlockCommonLoop
; ComputeBlockLoop
;
; This macro generates code to execute the block compute macro multiple
; times and advancing the matrix A and matrix B data pointers.
@ -240,7 +217,7 @@ ENDIF
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockCommonLoop MACRO ComputeBlock, Count, AdvanceMatrixAPlusRows
ComputeBlockLoop MACRO ComputeBlock, Count, AdvanceMatrixAPlusRows
LOCAL ComputeBlockBy4Loop
LOCAL ProcessRemainingBlocks
@ -252,11 +229,11 @@ ComputeBlockCommonLoop MACRO ComputeBlock, Count, AdvanceMatrixAPlusRows
jb ProcessRemainingBlocks
ComputeBlockBy4Loop:
ComputeBlock Count, 0, 0
ComputeBlock Count, 16*4, 4
ComputeBlock Count, 0, 0, 64*4
ComputeBlock Count, 16*4, 4, 64*4
sub rdx,-32*4 ; advance matrix B by 32 columns
ComputeBlock Count, 0, 8
ComputeBlock Count, 16*4, 12
ComputeBlock Count, 0, 8, 64*4
ComputeBlock Count, 16*4, 12, 64*4
sub rdx,-32*4 ; advance matrix B by 32 columns
add rcx,4*4 ; advance matrix A by 4 columns
IF AdvanceMatrixAPlusRows
@ -290,27 +267,3 @@ ENDIF
OutputBlock:
ENDM
ComputeBlockAvxLoop MACRO ComputeBlock, Count
ComputeBlockCommonLoop ComputeBlock, Count, <Count GT 2>
ENDM
ComputeBlockFma3Loop MACRO ComputeBlock, Count
;
; Reload the alpha value which is lost after each vzeroall instruction.
;
vbroadcastss ymm2,DWORD PTR SgemmKernelFrame.Alpha[rsp]
ComputeBlockCommonLoop ComputeBlock, Count, <Count GT 3>
ENDM
ComputeBlockAvx512FLoop MACRO ComputeBlock, Count
ComputeBlockCommonLoop ComputeBlock, Count, <Count GT 3>
ENDM

694
onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm
View file

@ -24,48 +24,10 @@ INCLUDE SgemmKernelCommon.inc
EXTERN MlasMaskMoveAvx:NEAR
;
; ComputeBlockFma3By32
;
; This macro multiplies and accumulates for a 32xN block (where N is 1,3)
; of the output matrix.
;
; Arguments:
;
; Count - Supplies the number of rows to access from matrix A.
;
; VectorOffset - Supplies the byte offset from matrix B to fetch elements.
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; Implicit Arguments:
;
; rcx - Supplies the address into the matrix A data.
;
; rdx - Supplies the address into the matrix B data.
;
; r10 - Supplies the length in bytes of a row from matrix A.
;
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockFma3By32 MACRO Count, VectorOffset, BroadcastOffset
IF Count EQ 1
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps ymm4,ymm3,YMMWORD PTR [rdx+VectorOffset]
vfmadd231ps ymm5,ymm3,YMMWORD PTR [rdx+VectorOffset+32]
vfmadd231ps ymm6,ymm3,YMMWORD PTR [rdx+rbx+VectorOffset]
vfmadd231ps ymm7,ymm3,YMMWORD PTR [rdx+rbx+VectorOffset+32]
ENDIF
ENDM
;
; ComputeBlockFma3By16
;
; This macro multiplies and accumulates for a 16xN block (where N is 1,3,6)
; of the output matrix.
; This macro multiplies and accumulates for a 16xN block of the output matrix.
;
; Arguments:
;
@ -75,6 +37,9 @@ ENDIF
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 3 rows.
@ -88,8 +53,11 @@ ENDIF
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockFma3By16 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockFma3By16 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IFNB <PrefetchOffset>
prefetcht0 [rdx+VectorOffset+PrefetchOffset]
ENDIF
IF Count EQ 1
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps ymm4,ymm3,YMMWORD PTR [rdx+VectorOffset]
@ -97,28 +65,24 @@ IF Count EQ 1
ELSE
vmovaps ymm0,YMMWORD PTR [rdx+VectorOffset]
vmovaps ymm1,YMMWORD PTR [rdx+VectorOffset+32]
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps ymm4,ymm3,ymm0
vfmadd231ps ymm5,ymm3,ymm1
IF Count GE 3
vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]
vfmadd231ps ymm6,ymm3,ymm0
vfmadd231ps ymm7,ymm3,ymm1
vbroadcastss ymm3,DWORD PTR [rcx+r10*2+BroadcastOffset]
vfmadd231ps ymm8,ymm3,ymm0
vfmadd231ps ymm9,ymm3,ymm1
ENDIF
IF Count GE 6
vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]
vfmadd231ps ymm10,ymm3,ymm0
vfmadd231ps ymm11,ymm3,ymm1
vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]
vfmadd231ps ymm12,ymm3,ymm0
vfmadd231ps ymm13,ymm3,ymm1
vbroadcastss ymm3,DWORD PTR [rbx+r10*2+BroadcastOffset]
vfmadd231ps ymm14,ymm3,ymm0
vfmadd231ps ymm15,ymm3,ymm1
ENDIF
EmitIfCountGE Count, 1, <vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]>
EmitIfCountGE Count, 1, <vfmadd231ps ymm4,ymm3,ymm0>
EmitIfCountGE Count, 1, <vfmadd231ps ymm5,ymm3,ymm1>
EmitIfCountGE Count, 2, <vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 2, <vfmadd231ps ymm6,ymm3,ymm0>
EmitIfCountGE Count, 2, <vfmadd231ps ymm7,ymm3,ymm1>
EmitIfCountGE Count, 3, <vbroadcastss ymm3,DWORD PTR [rcx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 3, <vfmadd231ps ymm8,ymm3,ymm0>
EmitIfCountGE Count, 3, <vfmadd231ps ymm9,ymm3,ymm1>
EmitIfCountGE Count, 4, <vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]>
EmitIfCountGE Count, 4, <vfmadd231ps ymm10,ymm3,ymm0>
EmitIfCountGE Count, 4, <vfmadd231ps ymm11,ymm3,ymm1>
EmitIfCountGE Count, 5, <vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 5, <vfmadd231ps ymm12,ymm3,ymm0>
EmitIfCountGE Count, 5, <vfmadd231ps ymm13,ymm3,ymm1>
EmitIfCountGE Count, 6, <vbroadcastss ymm3,DWORD PTR [rbx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 6, <vfmadd231ps ymm14,ymm3,ymm0>
EmitIfCountGE Count, 6, <vfmadd231ps ymm15,ymm3,ymm1>
ENDIF
ENDM
@ -128,8 +92,7 @@ ENDIF
;
; Macro Description:
;
; This macro multiplies and accumulates for a 8xN block (where N is 1,3,6)
; of the output matrix.
; This macro multiplies and accumulates for a 8xN block of the output matrix.
;
; Arguments:
;
@ -139,6 +102,9 @@ ENDIF
;
; BroadcastOffset - Supplies the byte offset from matrix A to fetch elements.
;
; PrefetchOffset - Optionally supplies the byte offset from matrix B to
; prefetch elements.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus 3 rows.
@ -152,29 +118,249 @@ ENDIF
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockFma3By8 MACRO Count, VectorOffset, BroadcastOffset
ComputeBlockFma3By8 MACRO Count, VectorOffset, BroadcastOffset, PrefetchOffset
IFNB <PrefetchOffset>
prefetcht0 [rdx+VectorOffset+PrefetchOffset]
ENDIF
IF Count EQ 1
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps ymm5,ymm3,YMMWORD PTR [rdx+VectorOffset]
ELSE
vmovaps ymm0,YMMWORD PTR [rdx+VectorOffset]
vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]
vfmadd231ps ymm5,ymm3,ymm0
IF Count GE 3
vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]
vfmadd231ps ymm7,ymm3,ymm0
vbroadcastss ymm3,DWORD PTR [rcx+r10*2+BroadcastOffset]
vfmadd231ps ymm9,ymm3,ymm0
EmitIfCountGE Count, 1, <vbroadcastss ymm3,DWORD PTR [rcx+BroadcastOffset]>
EmitIfCountGE Count, 1, <vfmadd231ps ymm5,ymm3,ymm0>
EmitIfCountGE Count, 2, <vbroadcastss ymm3,DWORD PTR [rcx+r10+BroadcastOffset]>
EmitIfCountGE Count, 2, <vfmadd231ps ymm7,ymm3,ymm0>
EmitIfCountGE Count, 3, <vbroadcastss ymm3,DWORD PTR [rcx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 3, <vfmadd231ps ymm9,ymm3,ymm0>
EmitIfCountGE Count, 4, <vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]>
EmitIfCountGE Count, 4, <vfmadd231ps ymm11,ymm3,ymm0>
EmitIfCountGE Count, 5, <vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]>
EmitIfCountGE Count, 5, <vfmadd231ps ymm13,ymm3,ymm0>
EmitIfCountGE Count, 6, <vbroadcastss ymm3,DWORD PTR [rbx+r10*2+BroadcastOffset]>
EmitIfCountGE Count, 6, <vfmadd231ps ymm15,ymm3,ymm0>
ENDIF
IF Count GE 6
vbroadcastss ymm3,DWORD PTR [rbx+BroadcastOffset]
vfmadd231ps ymm11,ymm3,ymm0
vbroadcastss ymm3,DWORD PTR [rbx+r10+BroadcastOffset]
vfmadd231ps ymm13,ymm3,ymm0
vbroadcastss ymm3,DWORD PTR [rbx+r10*2+BroadcastOffset]
vfmadd231ps ymm15,ymm3,ymm0
ENDM
;
; ComputeBlockFma3Loop
;
; This macro generates code to execute the block compute macro multiple
; times and advancing the matrix A and matrix B data pointers.
;
; Arguments:
;
; ComputeBlock - Supplies the macro to compute a single block.
;
; Count - Supplies the number of rows to access from matrix A.
;
; Implicit Arguments:
;
; rbx - Supplies the address into the matrix A data plus N rows.
;
; rcx - Supplies the address into the matrix A data.
;
; rdx - Supplies the address into the matrix B data.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; ymm4-ymm15 - Supplies the block accumulators.
;
ComputeBlockFma3Loop MACRO ComputeBlock, Count
IF Count GT 3
lea rbx,[r10*2+r10]
add rbx,rcx ; compute matrix A plus 3 rows
ENDIF
ComputeBlockLoop ComputeBlock, Count, <Count GT 3>
vbroadcastss ymm2,DWORD PTR SgemmKernelFrame.Alpha[rsp]
IF Count GT 3
lea rbx,[rax*2+rax]
add rbx,r8 ; compute matrix C plus 3 rows
ENDIF
ENDM
;
; ProcessCountMFma3
;
; Macro Description:
;
; This macro generates code to compute matrix multiplication for a fixed set
; of rows.
;
; Arguments:
;
; Mode - Supplies the mode of operation for updating the contents of matrix C.
;
; Count - Supplies the number of rows to process.
;
; Implicit Arguments:
;
; rcx - Supplies the address of matrix A.
;
; rdx - Supplies the address of matrix B.
;
; rsi - Supplies the address of matrix A.
;
; rbp - Supplies the number of columns from matrix B and matrix C to iterate
; over.
;
; r8 - Supplies the address of matrix C.
;
; r9 - Supplies the number of columns from matrix A and the number of rows
; from matrix B to iterate over.
;
; r10 - Supplies the length in bytes of a row from matrix A.
;
ProcessCountMFma3 MACRO Mode, Count, Fallthrough
LOCAL ProcessNextColumnLoop16xN
LOCAL ProcessRemainingCountN
LOCAL OutputMasked8xNBlock
LOCAL OutputMasked16xNBlock
cmp rbp,8
jbe ProcessRemainingCountN
ProcessNextColumnLoop16xN:
ComputeBlockFma3Loop ComputeBlockFma3By16, Count
EmitIfCountGE Count, 1, <prefetcht0 [r8+64]>
EmitIfCountGE Count, 2, <prefetcht0 [r8+rax+64]>
EmitIfCountGE Count, 3, <prefetcht0 [r8+rax*2+64]>
EmitIfCountGE Count, 4, <prefetcht0 [rbx+64]>
EmitIfCountGE Count, 5, <prefetcht0 [rbx+rax+64]>
EmitIfCountGE Count, 6, <prefetcht0 [rbx+rax*2+64]>
IFDIFI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmulps ymm4,ymm4,ymm2>
EmitIfCountGE Count, 1, <vmulps ymm5,ymm5,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm6,ymm6,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm7,ymm7,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm8,ymm8,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm9,ymm9,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm10,ymm10,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm11,ymm11,ymm2>
EmitIfCountGE Count, 5, <vmulps ymm12,ymm12,ymm2>
EmitIfCountGE Count, 5, <vmulps ymm13,ymm13,ymm2>
EmitIfCountGE Count, 6, <vmulps ymm14,ymm14,ymm2>
EmitIfCountGE Count, 6, <vmulps ymm15,ymm15,ymm2>
ENDIF
sub rbp,16
jb OutputMasked16xNBlock
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]>
EmitIfCountGE Count, 1, <vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8+32]>
EmitIfCountGE Count, 2, <vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 2, <vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax+32]>
EmitIfCountGE Count, 3, <vfmadd213ps ymm8,ymm2,YMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 3, <vfmadd213ps ymm9,ymm2,YMMWORD PTR [r8+rax*2+32]>
EmitIfCountGE Count, 4, <vfmadd213ps ymm10,ymm2,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 4, <vfmadd213ps ymm11,ymm2,YMMWORD PTR [rbx+32]>
EmitIfCountGE Count, 5, <vfmadd213ps ymm12,ymm2,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 5, <vfmadd213ps ymm13,ymm2,YMMWORD PTR [rbx+rax+32]>
EmitIfCountGE Count, 6, <vfmadd213ps ymm14,ymm2,YMMWORD PTR [rbx+rax*2]>
EmitIfCountGE Count, 6, <vfmadd213ps ymm15,ymm2,YMMWORD PTR [rbx+rax*2+32]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm4>
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8+32],ymm5>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm6>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax+32],ymm7>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [r8+rax*2],ymm8>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [r8+rax*2+32],ymm9>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx],ymm10>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx+32],ymm11>
EmitIfCountGE Count, 5, <vmovups YMMWORD PTR [rbx+rax],ymm12>
EmitIfCountGE Count, 5, <vmovups YMMWORD PTR [rbx+rax+32],ymm13>
EmitIfCountGE Count, 6, <vmovups YMMWORD PTR [rbx+rax*2],ymm14>
EmitIfCountGE Count, 6, <vmovups YMMWORD PTR [rbx+rax*2+32],ymm15>
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,8
ja ProcessNextColumnLoop16xN
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN:
ComputeBlockFma3Loop ComputeBlockFma3By8, Count
IFDIFI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmulps ymm5,ymm5,ymm2>
EmitIfCountGE Count, 2, <vmulps ymm7,ymm7,ymm2>
EmitIfCountGE Count, 3, <vmulps ymm9,ymm9,ymm2>
EmitIfCountGE Count, 4, <vmulps ymm11,ymm11,ymm2>
EmitIfCountGE Count, 5, <vmulps ymm13,ymm13,ymm2>
EmitIfCountGE Count, 6, <vmulps ymm15,ymm15,ymm2>
ENDIF
cmp rbp,8
jb OutputMasked8xNBlock
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vfmadd213ps ymm9,ymm2,YMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 4, <vfmadd213ps ymm11,ymm2,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 5, <vfmadd213ps ymm13,ymm2,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 6, <vfmadd213ps ymm15,ymm2,YMMWORD PTR [rbx+rax*2]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm5>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm7>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [r8+rax*2],ymm9>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx],ymm11>
EmitIfCountGE Count, 5, <vmovups YMMWORD PTR [rbx+rax],ymm13>
EmitIfCountGE Count, 6, <vmovups YMMWORD PTR [rbx+rax*2],ymm15>
jmp ExitKernelAndZeroUpper
OutputMasked16xNBlock:
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vfmadd213ps ymm8,ymm2,YMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 4, <vfmadd213ps ymm10,ymm2,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 5, <vfmadd213ps ymm12,ymm2,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 6, <vfmadd213ps ymm14,ymm2,YMMWORD PTR [rbx+rax*2]>
ENDIF
EmitIfCountGE Count, 1, <vmovups YMMWORD PTR [r8],ymm4>
EmitIfCountGE Count, 2, <vmovups YMMWORD PTR [r8+rax],ymm6>
EmitIfCountGE Count, 3, <vmovups YMMWORD PTR [r8+rax*2],ymm8>
EmitIfCountGE Count, 4, <vmovups YMMWORD PTR [rbx],ymm10>
EmitIfCountGE Count, 5, <vmovups YMMWORD PTR [rbx+rax],ymm12>
EmitIfCountGE Count, 6, <vmovups YMMWORD PTR [rbx+rax*2],ymm14>
add r8,8*4 ; advance matrix C by 8 columns
IF Count GT 3
add rbx,8*4 ; advance matrix C plus 3 rows by 8 columns
ENDIF
add rbp,8 ; correct for over-subtract above
OutputMasked8xNBlock:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss ymm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd ymm0,ymm0,YMMWORD PTR [MlasMaskMoveAvx]
IFIDNI <Mode>, <Add>
EmitIfCountGE Count, 1, <vmaskmovps ymm4,ymm0,YMMWORD PTR [r8]>
EmitIfCountGE Count, 2, <vmaskmovps ymm6,ymm0,YMMWORD PTR [r8+rax]>
EmitIfCountGE Count, 3, <vmaskmovps ymm8,ymm0,YMMWORD PTR [r8+rax*2]>
EmitIfCountGE Count, 4, <vmaskmovps ymm10,ymm0,YMMWORD PTR [rbx]>
EmitIfCountGE Count, 5, <vmaskmovps ymm12,ymm0,YMMWORD PTR [rbx+rax]>
EmitIfCountGE Count, 6, <vmaskmovps ymm14,ymm0,YMMWORD PTR [rbx+rax*2]>
EmitIfCountGE Count, 1, <vfmadd213ps ymm5,ymm2,ymm4>
EmitIfCountGE Count, 2, <vfmadd213ps ymm7,ymm2,ymm6>
EmitIfCountGE Count, 3, <vfmadd213ps ymm9,ymm2,ymm8>
EmitIfCountGE Count, 4, <vfmadd213ps ymm11,ymm2,ymm10>
EmitIfCountGE Count, 5, <vfmadd213ps ymm13,ymm2,ymm12>
EmitIfCountGE Count, 6, <vfmadd213ps ymm15,ymm2,ymm14>
ENDIF
EmitIfCountGE Count, 1, <vmaskmovps YMMWORD PTR [r8],ymm0,ymm5>
EmitIfCountGE Count, 2, <vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm7>
EmitIfCountGE Count, 3, <vmaskmovps YMMWORD PTR [r8+rax*2],ymm0,ymm9>
EmitIfCountGE Count, 4, <vmaskmovps YMMWORD PTR [rbx],ymm0,ymm11>
EmitIfCountGE Count, 5, <vmaskmovps YMMWORD PTR [rbx+rax],ymm0,ymm13>
EmitIfCountGE Count, 6, <vmaskmovps YMMWORD PTR [rbx+rax*2],ymm0,ymm15>
IFB <Fallthrough>
jmp ExitKernelAndZeroUpper
ENDIF
ENDM
@ -211,6 +397,8 @@ ENDIF
;
; Alpha - Supplies the scaler multiplier (see SGEMM definition).
;
; Beta - Supplies the scaler multiplier (see SGEMM definition).
;
; Return Value:
;
; Returns the number of rows handled.
@ -224,147 +412,27 @@ SgemmKernelFma3Function MACRO Mode
SgemmKernelAvxEntry
;
; Process 6 rows of the matrices.
; Process N rows of the matrices.
;
cmp r11,6
jb ProcessCountMLessThan6
cmp r11,5
ja ProcessCountM6
je ProcessCountM5
cmp r11,3
ja ProcessCountM4
je ProcessCountM3
cmp r11,1
je ProcessCountM1
ProcessCountM2:
ProcessCountMFma3 Mode, 2
ProcessCountM4:
ProcessCountMFma3 Mode, 4
ProcessCountM6:
mov r11d,6 ; return 6 rows handled
cmp rbp,8
jbe ProcessRemainingCountN6
ProcessNextColumnLoop16x6:
lea rbx,[r10*2+r10]
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockFma3Loop ComputeBlockFma3By16, 6
lea rcx,[r8+rax*2] ; compute matrix C plus 2 rows
lea rbx,[r8+rax*4] ; compute matrix C plus 4 rows
IFDIFI <Mode>, <Add>
vmulps ymm4,ymm4,ymm2 ; multiply by alpha
vmulps ymm5,ymm5,ymm2
vmulps ymm6,ymm6,ymm2
vmulps ymm7,ymm7,ymm2
vmulps ymm8,ymm8,ymm2
vmulps ymm9,ymm9,ymm2
vmulps ymm10,ymm10,ymm2
vmulps ymm11,ymm11,ymm2
vmulps ymm12,ymm12,ymm2
vmulps ymm13,ymm13,ymm2
vmulps ymm14,ymm14,ymm2
vmulps ymm15,ymm15,ymm2
ENDIF
sub rbp,16
jb OutputMasked16x6Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8+32]
vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax+32]
vfmadd213ps ymm8,ymm2,YMMWORD PTR [rcx]
vfmadd213ps ymm9,ymm2,YMMWORD PTR [rcx+32]
vfmadd213ps ymm10,ymm2,YMMWORD PTR [rcx+rax]
vfmadd213ps ymm11,ymm2,YMMWORD PTR [rcx+rax+32]
vfmadd213ps ymm12,ymm2,YMMWORD PTR [rbx]
vfmadd213ps ymm13,ymm2,YMMWORD PTR [rbx+32]
vfmadd213ps ymm14,ymm2,YMMWORD PTR [rbx+rax]
vfmadd213ps ymm15,ymm2,YMMWORD PTR [rbx+rax+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+32],ymm5
vmovups YMMWORD PTR [r8+rax],ymm6
vmovups YMMWORD PTR [r8+rax+32],ymm7
vmovups YMMWORD PTR [rcx],ymm8
vmovups YMMWORD PTR [rcx+32],ymm9
vmovups YMMWORD PTR [rcx+rax],ymm10
vmovups YMMWORD PTR [rcx+rax+32],ymm11
vmovups YMMWORD PTR [rbx],ymm12
vmovups YMMWORD PTR [rbx+32],ymm13
vmovups YMMWORD PTR [rbx+rax],ymm14
vmovups YMMWORD PTR [rbx+rax+32],ymm15
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,8
ja ProcessNextColumnLoop16x6
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN6:
lea rbx,[r10*2+r10]
add rbx,rcx ; compute matrix A plus 3 rows
ComputeBlockFma3Loop ComputeBlockFma3By8, 6
lea rcx,[r8+rax*2] ; compute matrix C plus 2 rows
lea rbx,[r8+rax*4] ; compute matrix C plus 4 rows
IFDIFI <Mode>, <Add>
vmulps ymm5,ymm5,ymm2 ; multiply by alpha
vmulps ymm7,ymm7,ymm2
vmulps ymm9,ymm9,ymm2
vmulps ymm11,ymm11,ymm2
vmulps ymm13,ymm13,ymm2
vmulps ymm15,ymm15,ymm2
ENDIF
cmp rbp,8
jb OutputMasked8x6Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm9,ymm2,YMMWORD PTR [rcx]
vfmadd213ps ymm11,ymm2,YMMWORD PTR [rcx+rax]
vfmadd213ps ymm13,ymm2,YMMWORD PTR [rbx]
vfmadd213ps ymm15,ymm2,YMMWORD PTR [rbx+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm5
vmovups YMMWORD PTR [r8+rax],ymm7
vmovups YMMWORD PTR [rcx],ymm9
vmovups YMMWORD PTR [rcx+rax],ymm11
vmovups YMMWORD PTR [rbx],ymm13
vmovups YMMWORD PTR [rbx+rax],ymm15
jmp ExitKernelAndZeroUpper
OutputMasked16x6Block:
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm8,ymm2,YMMWORD PTR [rcx]
vfmadd213ps ymm10,ymm2,YMMWORD PTR [rcx+rax]
vfmadd213ps ymm12,ymm2,YMMWORD PTR [rbx]
vfmadd213ps ymm14,ymm2,YMMWORD PTR [rbx+rax]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+rax],ymm6
vmovups YMMWORD PTR [rcx],ymm8
vmovups YMMWORD PTR [rcx+rax],ymm10
vmovups YMMWORD PTR [rbx],ymm12
vmovups YMMWORD PTR [rbx+rax],ymm14
add r8,8*4 ; advance matrix C by 8 columns
add rcx,8*4 ; advance matrix C plus 2 rows by 8 columns
add rbx,8*4 ; advance matrix C plus 4 rows by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x6Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss ymm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd ymm0,ymm0,YMMWORD PTR [MlasMaskMoveAvx]
IFIDNI <Mode>, <Add>
vmaskmovps ymm4,ymm0,YMMWORD PTR [r8]
vmaskmovps ymm6,ymm0,YMMWORD PTR [r8+rax]
vmaskmovps ymm8,ymm0,YMMWORD PTR [rcx]
vmaskmovps ymm10,ymm0,YMMWORD PTR [rcx+rax]
vmaskmovps ymm12,ymm0,YMMWORD PTR [rbx]
vmaskmovps ymm14,ymm0,YMMWORD PTR [rbx+rax]
vfmadd213ps ymm5,ymm2,ymm4
vfmadd213ps ymm7,ymm2,ymm6
vfmadd213ps ymm9,ymm2,ymm8
vfmadd213ps ymm11,ymm2,ymm10
vfmadd213ps ymm13,ymm2,ymm12
vfmadd213ps ymm15,ymm2,ymm14
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm5
vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm7
vmaskmovps YMMWORD PTR [rcx],ymm0,ymm9
vmaskmovps YMMWORD PTR [rcx+rax],ymm0,ymm11
vmaskmovps YMMWORD PTR [rbx],ymm0,ymm13
vmaskmovps YMMWORD PTR [rbx+rax],ymm0,ymm15
ProcessCountMFma3 Mode, 6, Fallthrough
;
; Restore non-volatile registers and return.
@ -376,194 +444,14 @@ ExitKernelAndZeroUpper:
ExitKernel:
SgemmKernelAvxExit
;
; Process 3 rows of the matrices.
;
ProcessCountM1:
ProcessCountMFma3 Mode, 1
ProcessCountMLessThan6:
cmp r11,3
jb ProcessCountMLessThan3
mov r11d,3 ; return 3 rows handled
cmp rbp,8
jbe ProcessRemainingCountN3
ProcessCountM3:
ProcessCountMFma3 Mode, 3
ProcessNextColumnLoop16x3:
ComputeBlockFma3Loop ComputeBlockFma3By16, 3
IFDIFI <Mode>, <Add>
vmulps ymm4,ymm4,ymm2 ; multiply by alpha
vmulps ymm5,ymm5,ymm2
vmulps ymm6,ymm6,ymm2
vmulps ymm7,ymm7,ymm2
vmulps ymm8,ymm8,ymm2
vmulps ymm9,ymm9,ymm2
ENDIF
sub rbp,16
jb OutputMasked16x3Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8+32]
vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax+32]
vfmadd213ps ymm8,ymm2,YMMWORD PTR [r8+rax*2]
vfmadd213ps ymm9,ymm2,YMMWORD PTR [r8+rax*2+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+32],ymm5
vmovups YMMWORD PTR [r8+rax],ymm6
vmovups YMMWORD PTR [r8+rax+32],ymm7
vmovups YMMWORD PTR [r8+rax*2],ymm8
vmovups YMMWORD PTR [r8+rax*2+32],ymm9
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,8
ja ProcessNextColumnLoop16x3
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN3:
ComputeBlockFma3Loop ComputeBlockFma3By8, 3
IFDIFI <Mode>, <Add>
vmulps ymm5,ymm5,ymm2 ; multiply by alpha
vmulps ymm7,ymm7,ymm2
vmulps ymm9,ymm9,ymm2
ENDIF
cmp rbp,8
jb OutputMasked8x3Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm9,ymm2,YMMWORD PTR [r8+rax*2]
ENDIF
vmovups YMMWORD PTR [r8],ymm5
vmovups YMMWORD PTR [r8+rax],ymm7
vmovups YMMWORD PTR [r8+rax*2],ymm9
jmp ExitKernelAndZeroUpper
OutputMasked16x3Block:
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+rax]
vfmadd213ps ymm8,ymm2,YMMWORD PTR [r8+rax*2]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+rax],ymm6
vmovups YMMWORD PTR [r8+rax*2],ymm8
add r8,8*4 ; advance matrix C by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x3Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss ymm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd ymm0,ymm0,YMMWORD PTR [MlasMaskMoveAvx]
IFIDNI <Mode>, <Add>
vmaskmovps ymm4,ymm0,YMMWORD PTR [r8]
vmaskmovps ymm6,ymm0,YMMWORD PTR [r8+rax]
vmaskmovps ymm8,ymm0,YMMWORD PTR [r8+rax*2]
vfmadd213ps ymm5,ymm2,ymm4
vfmadd213ps ymm7,ymm2,ymm6
vfmadd213ps ymm9,ymm2,ymm8
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm5
vmaskmovps YMMWORD PTR [r8+rax],ymm0,ymm7
vmaskmovps YMMWORD PTR [r8+rax*2],ymm0,ymm9
jmp ExitKernelAndZeroUpper
;
; Process 1 row of the matrices.
;
ProcessCountMLessThan3:
mov r11d,1 ; return 1 row handled
cmp rbp,32
jb ProcessRemainingCountN1LessThan32
mov rbx,r9
shl rbx,6 ; compute 16*CountK*sizeof(float)
ProcessNextColumnLoop32x1:
ComputeBlockFma3Loop ComputeBlockFma3By32, 1
add rdx,rbx ; advance matrix B by 16*CountK floats
IFDIFI <Mode>, <Add>
vmulps ymm4,ymm4,ymm2 ; multiply by alpha
vmulps ymm5,ymm5,ymm2
vmulps ymm6,ymm6,ymm2
vmulps ymm7,ymm7,ymm2
ELSE
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8+32]
vfmadd213ps ymm6,ymm2,YMMWORD PTR [r8+64]
vfmadd213ps ymm7,ymm2,YMMWORD PTR [r8+96]
ENDIF
sub rbp,32
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+32],ymm5
vmovups YMMWORD PTR [r8+64],ymm6
vmovups YMMWORD PTR [r8+96],ymm7
add r8,32*4 ; advance matrix C by 32 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,32
jae ProcessNextColumnLoop32x1
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN1LessThan32:
cmp rbp,8
jbe ProcessRemainingCountN1
ProcessNextColumnLoop16x1:
ComputeBlockFma3Loop ComputeBlockFma3By16, 1
IFDIFI <Mode>, <Add>
vmulps ymm4,ymm4,ymm2 ; multiply by alpha
vmulps ymm5,ymm5,ymm2
ENDIF
sub rbp,16
jb OutputMasked16x1Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8+32]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
vmovups YMMWORD PTR [r8+32],ymm5
add r8,16*4 ; advance matrix C by 16 columns
mov rcx,rsi ; reload matrix A
vzeroall
cmp rbp,8
ja ProcessNextColumnLoop16x1
test rbp,rbp
jz ExitKernel
ProcessRemainingCountN1:
ComputeBlockFma3Loop ComputeBlockFma3By8, 1
IFDIFI <Mode>, <Add>
vmulps ymm5,ymm5,ymm2 ; multiply by alpha
ENDIF
cmp rbp,8
jb OutputMasked8x1Block
IFIDNI <Mode>, <Add>
vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8]
ENDIF
vmovups YMMWORD PTR [r8],ymm5
jmp ExitKernelAndZeroUpper
OutputMasked16x1Block:
IFIDNI <Mode>, <Add>
vfmadd213ps ymm4,ymm2,YMMWORD PTR [r8]
ENDIF
vmovups YMMWORD PTR [r8],ymm4
add r8,8*4 ; advance matrix C by 8 columns
add rbp,8 ; correct for over-subtract above
OutputMasked8x1Block:
mov DWORD PTR SgemmKernelFrame.CountN[rsp],ebp
vbroadcastss ymm0,DWORD PTR SgemmKernelFrame.CountN[rsp]
vpcmpgtd ymm0,ymm0,YMMWORD PTR [MlasMaskMoveAvx]
IFIDNI <Mode>, <Add>
vmaskmovps ymm4,ymm0,YMMWORD PTR [r8]
vfmadd213ps ymm5,ymm2,ymm4
ENDIF
vmaskmovps YMMWORD PTR [r8],ymm0,ymm5
jmp ExitKernelAndZeroUpper
ProcessCountM5:
ProcessCountMFma3 Mode, 5
NESTED_END MlasSgemmKernel&Mode&Fma3, _TEXT

1
onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm
View file

@ -19,7 +19,6 @@
.xlist
INCLUDE mlasi.inc
INCLUDE SgemmKernelCommon.inc
.list
EXTERN MlasMaskMoveAvx:NEAR

25
onnxruntime/core/mlas/lib/amd64/mlasi.inc
View file

@ -41,3 +41,28 @@ save_xmm128_avx MACRO Register, Offset
.savexmm128 Register, Offset
ENDM
;
; EmitIfCountGE
;
; Macro Description:
;
; This macro conditionally emits the statement if Count is greater than or
; equal to Value.
;
; Arguments:
;
; Count - Supplies the variable used in the comparison.
;
; Value - Supplies the static used in the comparison.
;
; Statement - Supplies the statement to conditionally emit.
;
EmitIfCountGE MACRO Count, Value, Statement
IF Count GE Value
Statement
ENDIF
ENDM