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