From 20ef8b43a6df22e6effe71c61e11a0def9406cf0 Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Tue, 29 Jan 2019 17:20:32 -0800 Subject: [PATCH] Change GEMM kernels to natively handle broader range of row counts (#406) 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. --- .../mlas/lib/amd64/LogisticKernelFma3.asm | 1 - .../core/mlas/lib/amd64/SgemmKernelAvx.asm | 515 ++++++------- .../mlas/lib/amd64/SgemmKernelAvx512F.asm | 714 ++++++++---------- .../core/mlas/lib/amd64/SgemmKernelCommon.inc | 59 +- .../core/mlas/lib/amd64/SgemmKernelFma3.asm | 694 +++++++---------- .../core/mlas/lib/amd64/TanhKernelFma3.asm | 1 - onnxruntime/core/mlas/lib/amd64/mlasi.inc | 25 + 7 files changed, 866 insertions(+), 1143 deletions(-) diff --git a/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm index 28bc3507e3..5d81c23722 100644 --- a/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/LogisticKernelFma3.asm @@ -19,7 +19,6 @@ .xlist INCLUDE mlasi.inc -INCLUDE SgemmKernelCommon.inc .list EXTERN MlasMaskMoveAvx:NEAR diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm index 2cf23e0897..d5ab876440 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx.asm @@ -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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, 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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, 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, +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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + ComputeBlockAvxLoop ComputeBlockAvxBy16, Count + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + sub rbp,16 + jb OutputMasked16xNBlock +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + 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, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + ComputeBlockAvxLoop ComputeBlockAvxBy8, Count + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + cmp rbp,8 + jb OutputMasked8xNBlock +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + jmp ExitKernel + +OutputMasked16xNBlock: +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + 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 , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, +IFB + 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - vaddps ymm9,ymm9,YMMWORD PTR [r8] -ENDIF - vmovups YMMWORD PTR [r8],ymm9 - jmp ExitKernel - -OutputMasked16x1Block: -IFIDNI , - 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 , - 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 diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm index 1df94ce606..f7b85800b2 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelAvx512F.asm @@ -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 + 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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, 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 + 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, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + + 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, +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, + ; clear upper block accumulators + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + ComputeBlockAvx512FLoop ComputeBlockAvx512FBy32, Count + add rdx,r12 ; advance matrix B by 16*CountK floats +IFDIFI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, +ELSE + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + 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 , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, +ELSE + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + 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, + ; clear upper block accumulators + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + EmitIfCountGE Count, 12, + 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc b/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc index 86a7123b92..4f43688fd4 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelCommon.inc @@ -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, - - 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, - - ENDM - -ComputeBlockAvx512FLoop MACRO ComputeBlock, Count - - ComputeBlockCommonLoop ComputeBlock, Count, - - ENDM diff --git a/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm index 2343cc463a..d88556dd73 100644 --- a/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/SgemmKernelFma3.asm @@ -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 + 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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, 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 + 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, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, 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, + 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, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +IFDIFI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, +ENDIF + sub rbp,16 + jb OutputMasked16xNBlock +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 6, + 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 , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +ENDIF + cmp rbp,8 + jb OutputMasked8xNBlock +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + jmp ExitKernelAndZeroUpper + +OutputMasked16xNBlock: +IFIDNI , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + 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 , + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +ENDIF + EmitIfCountGE Count, 1, + EmitIfCountGE Count, 2, + EmitIfCountGE Count, 3, + EmitIfCountGE Count, 4, + EmitIfCountGE Count, 5, + EmitIfCountGE Count, 6, +IFB + 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - 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 , - vmulps ymm5,ymm5,ymm2 ; multiply by alpha - vmulps ymm7,ymm7,ymm2 - vmulps ymm9,ymm9,ymm2 -ENDIF - cmp rbp,8 - jb OutputMasked8x3Block -IFIDNI , - 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 , - 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 , - 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 , - 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 , - vmulps ymm4,ymm4,ymm2 ; multiply by alpha - vmulps ymm5,ymm5,ymm2 -ENDIF - sub rbp,16 - jb OutputMasked16x1Block -IFIDNI , - 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 , - vmulps ymm5,ymm5,ymm2 ; multiply by alpha -ENDIF - cmp rbp,8 - jb OutputMasked8x1Block -IFIDNI , - vfmadd213ps ymm5,ymm2,YMMWORD PTR [r8] -ENDIF - vmovups YMMWORD PTR [r8],ymm5 - jmp ExitKernelAndZeroUpper - -OutputMasked16x1Block: -IFIDNI , - 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 , - 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 diff --git a/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm index f8c423238c..ad5fc4e207 100644 --- a/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm +++ b/onnxruntime/core/mlas/lib/amd64/TanhKernelFma3.asm @@ -19,7 +19,6 @@ .xlist INCLUDE mlasi.inc -INCLUDE SgemmKernelCommon.inc .list EXTERN MlasMaskMoveAvx:NEAR diff --git a/onnxruntime/core/mlas/lib/amd64/mlasi.inc b/onnxruntime/core/mlas/lib/amd64/mlasi.inc index 3db76dc257..ba3dc6889e 100644 --- a/onnxruntime/core/mlas/lib/amd64/mlasi.inc +++ b/onnxruntime/core/mlas/lib/amd64/mlasi.inc @@ -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