1
0
mirror of https://github.com/opencv/opencv.git synced 2026-07-30 15:53:03 +04:00
Files
opencv/3rdparty/mlas/lib/x86_64/FgemmKernelAvxCommon.h
Abhishek Gola bdf348c13a Merge pull request #28934 from abhishek-gola:mlas_gemm
Added MLAS third party module and integrated into GeMM path #28934

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
2026-05-22 20:22:15 +03:00

452 lines
16 KiB
C

/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
FgemmKernelAvxCommon.h
Abstract:
This module implements the kernels for the floating point matrix/matrix
multiply operation (SGEMM and DGEMM).
This implementation uses AVX instructions.
--*/
/*++
Macro Description:
This macro multiplies and accumulates for 2 YMMWORDs by N rows of the output
matrix.
Arguments:
RowCount - Supplies the number of rows to process.
VectorOffset - Supplies the byte offset from matrix B to fetch elements.
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:
rdi - Supplies the address into the matrix A data.
rbx - Supplies the address into the matrix A data plus 2 rows.
rsi - Supplies the address into the matrix B data.
r10 - Supplies the length in bytes of a row from matrix A.
ymm8-ymm15 - Supplies the block accumulators.
--*/
.macro ComputeBlockAvxBy16 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
.if \RowCount\() == 1
vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]
vmulpf ymm4,ymm3,YMMWORD PTR [rsi+\VectorOffset\()]
vaddpf ymm8,ymm8,ymm4
vmulpf ymm5,ymm3,YMMWORD PTR [rsi+\VectorOffset\()+32]
vaddpf ymm9,ymm9,ymm5
.else
vmovapf ymm0,YMMWORD PTR [rsi+\VectorOffset\()]
vmovapf ymm1,YMMWORD PTR [rsi+\VectorOffset\()+32]
EmitIfCountGE \RowCount\(), 1, "vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm4,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm8,ymm8,ymm4"
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm5,ymm3,ymm1"
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm9,ymm9,ymm5"
EmitIfCountGE \RowCount\(), 2, "vbroadcastsf ymm3,[rdi+r10+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm6,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm10,ymm10,ymm6"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm7,ymm3,ymm1"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm11,ymm11,ymm7"
EmitIfCountGE \RowCount\(), 3, "vbroadcastsf ymm3,[rbx+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm4,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm12,ymm12,ymm4"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm5,ymm3,ymm1"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm13,ymm13,ymm5"
EmitIfCountGE \RowCount\(), 4, "vbroadcastsf ymm3,[rbx+r10+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm6,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm14,ymm14,ymm6"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm7,ymm3,ymm1"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm15,ymm15,ymm7"
.endif
.endm
/*++
Macro Description:
This macro multiplies and accumulates for 1 YMMWORD by N rows of the output
matrix.
Arguments:
RowCount - Supplies the number of rows to process.
VectorOffset - Supplies the byte offset from matrix B to fetch elements.
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:
rdi - Supplies the address into the matrix A data.
rbx - Supplies the address into the matrix A data plus 2 rows.
rsi - Supplies the address into the matrix B data.
r10 - Supplies the length in bytes of a row from matrix A.
ymm8-ymm15 - Supplies the block accumulators.
--*/
.macro ComputeBlockAvxBy8 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
.if \RowCount\() == 1
vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]
vmulpf ymm5,ymm3,YMMWORD PTR [rsi+\VectorOffset\()]
vaddpf ymm9,ymm9,ymm5
.else
vmovapf ymm0,YMMWORD PTR [rsi+\VectorOffset\()]
EmitIfCountGE \RowCount\(), 1, "vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm5,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm9,ymm9,ymm5"
EmitIfCountGE \RowCount\(), 2, "vbroadcastsf ymm3,[rdi+r10+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm7,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm11,ymm11,ymm7"
EmitIfCountGE \RowCount\(), 3, "vbroadcastsf ymm3,[rbx+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm5,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm13,ymm13,ymm5"
EmitIfCountGE \RowCount\(), 4, "vbroadcastsf ymm3,[rbx+r10+\BroadcastOffset\()]"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm7,ymm3,ymm0"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm15,ymm15,ymm7"
.endif
.endm
/*++
Macro Description:
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.
RowCount - Supplies the number of rows to process.
Implicit Arguments:
rdi - Supplies the address into the matrix A data.
rsi - Supplies the address into the matrix B data.
rcx - 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.
ymm4-ymm15 - Supplies the block accumulators.
--*/
.macro ComputeBlockAvxLoop ComputeBlock, RowCount
.if \RowCount\() > 2
lea rbx,[rdi+r10*2] # compute matrix A plus 2 rows
.endif
ComputeBlockLoop \ComputeBlock\(), \RowCount\(), \RowCount\() > 2
.if \RowCount\() > 2
lea rbx,[rdx+rax*2] # compute matrix C plus 2 rows
.endif
.endm
/*++
Macro Description:
This macro generates code to compute matrix multiplication for a fixed set
of rows.
Arguments:
RowCount - Supplies the number of rows to process.
Fallthrough - Supplies a non-blank value if the macro may fall through to
the ExitKernel label.
Implicit Arguments:
rdi - Supplies the address of matrix A.
rsi - Supplies the address of matrix B.
r11 - Supplies the address of matrix A.
r9 - Supplies the number of columns from matrix B and matrix C to iterate
over.
rdx - Supplies the address of matrix C.
rcx - 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.
rax - Supplies the length in bytes of a row from matrix C.
r15 - Stores the ZeroMode argument from the stack frame.
--*/
.macro ProcessCountM RowCount, Fallthrough
cmp r9,.LFgemmYmmElementCount
jbe .LProcessRemainingCountN\@
.LProcessNextColumnLoop2xN\@:
EmitIfCountGE \RowCount\(), 1, "vxorpf xmm8,xmm8,xmm8"
EmitIfCountGE \RowCount\(), 1, "vxorpf xmm9,xmm9,xmm9"
EmitIfCountGE \RowCount\(), 2, "vxorpf xmm10,xmm10,xmm10"
EmitIfCountGE \RowCount\(), 2, "vxorpf xmm11,xmm11,xmm11"
EmitIfCountGE \RowCount\(), 3, "vxorpf xmm12,xmm12,xmm12"
EmitIfCountGE \RowCount\(), 3, "vxorpf xmm13,xmm13,xmm13"
EmitIfCountGE \RowCount\(), 4, "vxorpf xmm14,xmm14,xmm14"
EmitIfCountGE \RowCount\(), 4, "vxorpf xmm15,xmm15,xmm15"
ComputeBlockAvxLoop ComputeBlockAvxBy16, \RowCount\()
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm8,ymm8,ymm2"
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm9,ymm9,ymm2"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm10,ymm10,ymm2"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm11,ymm11,ymm2"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm12,ymm12,ymm2"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm13,ymm13,ymm2"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm14,ymm14,ymm2"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm15,ymm15,ymm2"
sub r9,2*.LFgemmYmmElementCount
jb .LOutputMasked2xNBlock\@
test r15b,r15b # ZeroMode?
jnz .LStore2xNBlock\@
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm8,ymm8,YMMWORD PTR [rdx]"
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm9,ymm9,YMMWORD PTR [rdx+32]"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm10,ymm10,YMMWORD PTR [rdx+rax]"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm11,ymm11,YMMWORD PTR [rdx+rax+32]"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm12,ymm12,YMMWORD PTR [rbx]"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm13,ymm13,YMMWORD PTR [rbx+32]"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm14,ymm14,YMMWORD PTR [rbx+rax]"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm15,ymm15,YMMWORD PTR [rbx+rax+32]"
.LStore2xNBlock\@:
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm8"
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx+32],ymm9"
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm10"
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax+32],ymm11"
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rbx],ymm12"
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rbx+32],ymm13"
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx+rax],ymm14"
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx+rax+32],ymm15"
add rdx,2*32 # advance matrix C by 2 YMMWORDs
mov rdi,r11 # reload matrix A
cmp r9,.LFgemmYmmElementCount
ja .LProcessNextColumnLoop2xN\@
test r9,r9
jz .LExitKernel
.LProcessRemainingCountN\@:
EmitIfCountGE \RowCount\(), 1, "vxorpf xmm9,xmm9,xmm9"
EmitIfCountGE \RowCount\(), 2, "vxorpf xmm11,xmm11,xmm11"
EmitIfCountGE \RowCount\(), 3, "vxorpf xmm13,xmm13,xmm13"
EmitIfCountGE \RowCount\(), 4, "vxorpf xmm15,xmm15,xmm15"
ComputeBlockAvxLoop ComputeBlockAvxBy8, \RowCount\()
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm9,ymm9,ymm2"
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm11,ymm11,ymm2"
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm13,ymm13,ymm2"
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm15,ymm15,ymm2"
cmp r9,.LFgemmYmmElementCount
jb .LOutputMasked1xNBlock\@
test r15b,r15b # ZeroMode?
jnz .LStore1xNBlock\@
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm9,ymm9,YMMWORD PTR [rdx]"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm11,ymm11,YMMWORD PTR [rdx+rax]"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm13,ymm13,YMMWORD PTR [rbx]"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm15,ymm15,YMMWORD PTR [rbx+rax]"
.LStore1xNBlock\@:
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm9"
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm11"
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rbx],ymm13"
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx+rax],ymm15"
jmp .LExitKernel
.LOutputMasked2xNBlock\@:
test r15b,r15b # ZeroMode?
jnz .LStoreMasked2xNBlock\@
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm8,ymm8,YMMWORD PTR [rdx]"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm10,ymm10,YMMWORD PTR [rdx+rax]"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm12,ymm12,YMMWORD PTR [rbx]"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm14,ymm14,YMMWORD PTR [rbx+rax]"
.LStoreMasked2xNBlock\@:
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm8"
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm10"
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rbx],ymm12"
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx+rax],ymm14"
add rdx,32 # advance matrix C by YMMWORD
.if \RowCount\() > 2
add rbx,32 # advance matrix C plus 2 rows by YMMWORD
.endif
add r9,.LFgemmYmmElementCount # correct for over-subtract above
.LOutputMasked1xNBlock\@:
neg r9
lea rdi,C_UNDERSCORE(MlasMaskMoveTableAvx)[rip+8*4]
vmovdqu ymm0,YMMWORD PTR [rdi+r9*.LFgemmElementSize]
test r15b,r15b # ZeroMode?
jnz .LStoreMasked1xNBlock\@
EmitIfCountGE \RowCount\(), 1, "vmaskmovpf ymm8,ymm0,YMMWORD PTR [rdx]"
EmitIfCountGE \RowCount\(), 2, "vmaskmovpf ymm10,ymm0,YMMWORD PTR [rdx+rax]"
EmitIfCountGE \RowCount\(), 3, "vmaskmovpf ymm12,ymm0,YMMWORD PTR [rbx]"
EmitIfCountGE \RowCount\(), 4, "vmaskmovpf ymm14,ymm0,YMMWORD PTR [rbx+rax]"
EmitIfCountGE \RowCount\(), 1, "vaddpf ymm9,ymm9,ymm8"
EmitIfCountGE \RowCount\(), 2, "vaddpf ymm11,ymm11,ymm10"
EmitIfCountGE \RowCount\(), 3, "vaddpf ymm13,ymm13,ymm12"
EmitIfCountGE \RowCount\(), 4, "vaddpf ymm15,ymm15,ymm14"
.LStoreMasked1xNBlock\@:
EmitIfCountGE \RowCount\(), 1, "vmaskmovpf YMMWORD PTR [rdx],ymm0,ymm9"
EmitIfCountGE \RowCount\(), 2, "vmaskmovpf YMMWORD PTR [rdx+rax],ymm0,ymm11"
EmitIfCountGE \RowCount\(), 3, "vmaskmovpf YMMWORD PTR [rbx],ymm0,ymm13"
EmitIfCountGE \RowCount\(), 4, "vmaskmovpf YMMWORD PTR [rbx+rax],ymm0,ymm15"
.ifb \Fallthrough\()
jmp .LExitKernel
.endif
.endm
/*++
Macro Description:
This macro generates the inner kernel to compute matrix multiplication.
Arguments:
FunctionName - Supplies the name for the generated function.
--*/
.macro FgemmKernelAvxFunction FunctionName
/*++
Routine Description:
This routine is an inner kernel to compute matrix multiplication for a
set of rows.
Arguments:
A (rdi) - Supplies the address of matrix A.
B (rsi) - Supplies the address of matrix B. The matrix data has been packed
using MlasSgemmCopyPackB or MlasSgemmTransposePackB.
C (rdx) - Supplies the address of matrix C.
CountK (rcx) - Supplies the number of columns from matrix A and the number
of rows from matrix B to iterate over.
CountM (r8) - Supplies the maximum number of rows that can be processed for
matrix A and matrix C. The actual number of rows handled for this
invocation depends on the kernel implementation.
CountN (r9) - Supplies the number of columns from matrix B and matrix C to
iterate over.
lda - Supplies the first dimension of matrix A.
ldc - Supplies the first dimension of matrix C.
Alpha (xmm0) - Supplies the scalar alpha multiplier (see GEMM definition).
ZeroMode - Supplies true if the output matrix must be zero initialized,
else false if the output matrix is accumulated into.
Return Value:
Returns the number of rows handled.
--*/
FUNCTION_ENTRY \FunctionName\()
push rbp
push rbx
push r15
mov r11,rdi
mov r10,.LFgemmKernelFrame_lda[rsp]
shl r10,.LFgemmElementShift # convert lda to bytes
mov rax,.LFgemmKernelFrame_ldc[rsp]
shl rax,.LFgemmElementShift # convert ldc to bytes
movzx r15,BYTE PTR .LFgemmKernelFrame_ZeroMode[rsp]
vmovsf .LFgemmKernelFrame_alpha[rsp],xmm0
vbroadcastsf ymm2,.LFgemmKernelFrame_alpha[rsp]
//
// Process 4 rows of the matrices.
//
cmp r8,4
jb .LProcessCountMLessThan4
mov r8d,4 # return 4 rows handled
ProcessCountM 4, Fallthrough
//
// Restore non-volatile registers and return.
//
.LExitKernel:
vzeroupper
mov eax,r8d
pop r15
pop rbx
pop rbp
ret
//
// Process 2 rows of the matrices.
//
.LProcessCountMLessThan4:
cmp r8,2
jb .LProcessCountMLessThan2
mov r8d,2 # return 2 rows handled
ProcessCountM 2
//
// Process 1 row of the matrices.
//
.LProcessCountMLessThan2:
ProcessCountM 1
.endm