vendor: OpenCV 5.0.0 snapshot at 40738fb16ceddb5fb3fea747585f7ce6abb0605b
This commit is contained in:
Gitea Mirror Bot
+482
@@ -0,0 +1,482 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelNeon.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.text
|
||||
|
||||
//
|
||||
// ClearRowAccumulators
|
||||
//
|
||||
// Generates the code to clear the accumulators for a single row of the output
|
||||
// block.
|
||||
//
|
||||
|
||||
.macro ClearRowAccumulators Columns, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
movi v\Vec1Reg\().16b,#0
|
||||
movi v\Vec2Reg\().16b,#0
|
||||
.if \Columns\() > 8
|
||||
movi v\Vec3Reg\().16b,#0
|
||||
movi v\Vec4Reg\().16b,#0
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// ClearBlockAccumulators
|
||||
//
|
||||
// Generates the code to clear the accumulators for a single row of the output
|
||||
// block.
|
||||
//
|
||||
|
||||
.macro ClearBlockAccumulators Columns, Rows
|
||||
|
||||
ClearRowAccumulators \Columns\(),16,17,18,19
|
||||
.if \Rows\() >= 2
|
||||
ClearRowAccumulators \Columns\(),20,21,22,23
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
ClearRowAccumulators \Columns\(),24,25,26,27
|
||||
ClearRowAccumulators \Columns\(),28,29,30,31
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// LoadMatrixAElementsBy4
|
||||
// LoadMatrixAElementsBy1
|
||||
//
|
||||
// Generates the code to load 1 or 4 elements from matrix A.
|
||||
//
|
||||
|
||||
.macro LoadMatrixAElementsBy4 Rows
|
||||
|
||||
ldr q8,[x0],#16
|
||||
.if \Rows\() >= 2
|
||||
ldr q9,[x10],#16
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
ldr q10,[x11],#16
|
||||
ldr q11,[x12],#16
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
.macro LoadMatrixAElementsBy1 Rows
|
||||
|
||||
ldr s8,[x0],#4
|
||||
.if \Rows\() >= 2
|
||||
ldr s9,[x10],#4
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
ldr s10,[x11],#4
|
||||
ldr s11,[x12],#4
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// MultiplyAccumulateRow
|
||||
//
|
||||
// Generates the code to multiply and accumulate a single row of the output
|
||||
// block.
|
||||
//
|
||||
|
||||
.macro MultiplyAccumulateRow Columns, MatrixAReg, Broadcast, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
fmla v\Vec1Reg\().4s,v4.4s,\MatrixAReg\().s[\Broadcast\()]
|
||||
fmla v\Vec2Reg\().4s,v5.4s,\MatrixAReg\().s[\Broadcast\()]
|
||||
.if \Columns\() > 8
|
||||
fmla v\Vec3Reg\().4s,v6.4s,\MatrixAReg\().s[\Broadcast\()]
|
||||
fmla v\Vec4Reg\().4s,v7.4s,\MatrixAReg\().s[\Broadcast\()]
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// MultiplyAccumulateBlock
|
||||
//
|
||||
// Generates the code to multiply and accumulate into the output block.
|
||||
//
|
||||
|
||||
.macro MultiplyAccumulateBlock Columns, Rows, Broadcast
|
||||
|
||||
MultiplyAccumulateRow \Columns\(),v8,\Broadcast\(),16,17,18,19
|
||||
.if \Rows\() >= 2
|
||||
MultiplyAccumulateRow \Columns\(),v9,\Broadcast\(),20,21,22,23
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
MultiplyAccumulateRow \Columns\(),v10,\Broadcast\(),24,25,26,27
|
||||
MultiplyAccumulateRow \Columns\(),v11,\Broadcast\(),28,29,30,31
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// ComputeBlockLoop
|
||||
//
|
||||
// Generates the code to loop over K entries of the input matrices to produce
|
||||
// the output block.
|
||||
//
|
||||
|
||||
.macro ComputeBlockLoop Mode, Columns, Rows
|
||||
|
||||
ClearBlockAccumulators \Columns\(),\Rows\()
|
||||
|
||||
.if \Rows\() >= 2
|
||||
add x10,x0,x6,lsl #2 // compute matrix A plus 1 row
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
add x11,x10,x6,lsl #2 // compute matrix A plus 2 rows
|
||||
add x12,x11,x6,lsl #2 // compute matrix A plus 3 rows
|
||||
.endif
|
||||
|
||||
sub x9,x3,#4 // decrement block count to process
|
||||
tbnz x9,#63,.L\Mode\().ProcessRemaining\Columns\().x\Rows\().Blocks
|
||||
|
||||
.L\Mode\().Compute\Columns\().x\Rows\().BlockBy4Loop:
|
||||
LoadMatrixAElementsBy4 \Rows\()
|
||||
ldp q4,q5,[x1],#64*4
|
||||
.if \Columns\() > 8
|
||||
ldp q6,q7,[x1,#-56*4]
|
||||
.endif
|
||||
MultiplyAccumulateBlock \Columns\(),\Rows\(),0
|
||||
ldp q4,q5,[x1,#-48*4]
|
||||
.if \Columns\() > 8
|
||||
ldp q6,q7,[x1,#-40*4]
|
||||
.endif
|
||||
MultiplyAccumulateBlock \Columns\(),\Rows\(),1
|
||||
ldp q4,q5,[x1,#-32*4]
|
||||
.if \Columns\() > 8
|
||||
ldp q6,q7,[x1,#-24*4]
|
||||
.endif
|
||||
MultiplyAccumulateBlock \Columns\(),\Rows\(),2
|
||||
ldp q4,q5,[x1,#-16*4]
|
||||
.if \Columns\() > 8
|
||||
ldp q6,q7,[x1,#-8*4]
|
||||
.endif
|
||||
MultiplyAccumulateBlock \Columns\(),\Rows\(),3
|
||||
sub x9,x9,#4
|
||||
tbz x9,#63,.L\Mode\().Compute\Columns\().x\Rows\().BlockBy4Loop
|
||||
|
||||
.L\Mode\().ProcessRemaining\Columns\().x\Rows\().Blocks:
|
||||
add x9,x9,#4 // correct for over-subtract above
|
||||
cbz x9,.L\Mode\().Output\Columns\().x\Rows\().Block
|
||||
|
||||
.L\Mode\().Compute\Columns\().x\Rows\().BlockBy1Loop:
|
||||
LoadMatrixAElementsBy1 \Rows\()
|
||||
ldp q4,q5,[x1],#16*4
|
||||
.if \Columns\() > 8
|
||||
ldp q6,q7,[x1,#-8*4]
|
||||
.endif
|
||||
MultiplyAccumulateBlock \Columns\(),\Rows\(),0
|
||||
sub x9,x9,#1
|
||||
cbnz x9,.L\Mode\().Compute\Columns\().x\Rows\().BlockBy1Loop
|
||||
|
||||
.L\Mode\().Output\Columns\().x\Rows\().Block:
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// MultiplyAlphaRow
|
||||
//
|
||||
// Generates the code to multiply a single row of the output block by the alpha
|
||||
// value.
|
||||
//
|
||||
|
||||
.macro MultiplyAlphaRow Columns, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.if \Columns\() <= 4
|
||||
fmul v\Vec1Reg\().4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
.elif \Columns\() <= 8
|
||||
fmul v\Vec1Reg\().4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
fmul v\Vec2Reg\().4s,v\Vec2Reg\().4s,v0.s[0]
|
||||
.elif \Columns\() <= 12
|
||||
fmul v\Vec1Reg\().4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
fmul v\Vec2Reg\().4s,v\Vec2Reg\().4s,v0.s[0]
|
||||
fmul v\Vec3Reg\().4s,v\Vec3Reg\().4s,v0.s[0]
|
||||
.else
|
||||
fmul v\Vec1Reg\().4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
fmul v\Vec2Reg\().4s,v\Vec2Reg\().4s,v0.s[0]
|
||||
fmul v\Vec3Reg\().4s,v\Vec3Reg\().4s,v0.s[0]
|
||||
fmul v\Vec4Reg\().4s,v\Vec4Reg\().4s,v0.s[0]
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// MultiplyAlphaBlock
|
||||
//
|
||||
// Generates the code to multiply the output block by the alpha value.
|
||||
//
|
||||
|
||||
.macro MultiplyAlphaBlock Columns, Rows
|
||||
|
||||
MultiplyAlphaRow \Columns\(),16,17,18,19
|
||||
.if \Rows\() >= 2
|
||||
MultiplyAlphaRow \Columns\(),20,21,22,23
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
MultiplyAlphaRow \Columns\(),24,25,26,27
|
||||
MultiplyAlphaRow \Columns\(),28,29,30,31
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// OutputRow1Element
|
||||
// OutputRow2Element
|
||||
// OutputRow4Element
|
||||
// OutputRow8Element
|
||||
// OutputRow16Element
|
||||
//
|
||||
// Generates the code to store elements to the output block.
|
||||
//
|
||||
|
||||
.macro OutputRow1Element Mode, AddrReg, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.ifeqs "\Mode\()","Add"
|
||||
ld1 {v4.s}[0],[\AddrReg\()]
|
||||
fmla v4.2s,v\Vec1Reg\().2s,v0.s[0]
|
||||
st1 {v4.s}[0],[\AddrReg\()] // post-increment not needed for last element
|
||||
.else
|
||||
st1 {v\Vec1Reg\().s}[0],[\AddrReg\()]// post-increment not needed for last element
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
.macro OutputRow2Element Mode, AddrReg, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.ifeqs "\Mode\()","Add"
|
||||
ld1 {v4.2s},[\AddrReg\()]
|
||||
fmla v4.2s,v\Vec1Reg\().2s,v0.s[0]
|
||||
st1 {v4.2s},[\AddrReg\()],#2*4
|
||||
.else
|
||||
st1 {v\Vec1Reg\().2s},[\AddrReg\()],#2*4
|
||||
.endif
|
||||
dup v\Vec1Reg\().4s,v\Vec1Reg\().s[2] // shift remaining elements down
|
||||
|
||||
.endm
|
||||
|
||||
.macro OutputRow4Element Mode, AddrReg, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.ifeqs "\Mode\()","Add"
|
||||
ld1 {v4.4s},[\AddrReg\()]
|
||||
fmla v4.4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
st1 {v4.4s},[\AddrReg\()],#4*4
|
||||
.else
|
||||
st1 {v\Vec1Reg\().4s},[\AddrReg\()],#4*4
|
||||
.endif
|
||||
mov v\Vec1Reg\().16b,v\Vec2Reg\().16b // shift remaining elements down
|
||||
|
||||
.endm
|
||||
|
||||
.macro OutputRow8Element Mode, AddrReg, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.ifeqs "\Mode\()","Add"
|
||||
ldp q4,q5,[\AddrReg\()]
|
||||
fmla v4.4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
fmla v5.4s,v\Vec2Reg\().4s,v0.s[0]
|
||||
stp q4,q5,[\AddrReg\()],#8*4
|
||||
.else
|
||||
stp q\Vec1Reg\(),q\Vec2Reg\(),[\AddrReg\()],#8*4
|
||||
.endif
|
||||
mov v\Vec1Reg\().16b,v\Vec3Reg\().16b // shift remaining elements down
|
||||
mov v\Vec2Reg\().16b,v\Vec4Reg\().16b
|
||||
|
||||
.endm
|
||||
|
||||
.macro OutputRow16Element Mode, AddrReg, Vec1Reg, Vec2Reg, Vec3Reg, Vec4Reg
|
||||
|
||||
.ifeqs "\Mode\()","Add"
|
||||
ldp q4,q5,[\AddrReg\()]
|
||||
ldp q6,q7,[\AddrReg\(),#8*4]
|
||||
fmla v4.4s,v\Vec1Reg\().4s,v0.s[0]
|
||||
fmla v5.4s,v\Vec2Reg\().4s,v0.s[0]
|
||||
fmla v6.4s,v\Vec3Reg\().4s,v0.s[0]
|
||||
fmla v7.4s,v\Vec4Reg\().4s,v0.s[0]
|
||||
stp q4,q5,[\AddrReg\()],#16*4
|
||||
stp q6,q7,[\AddrReg\(),#-8*4]
|
||||
.else
|
||||
stp q\Vec1Reg\(),q\Vec2Reg\(),[\AddrReg\()],#16*4
|
||||
stp q\Vec3Reg\(),q\Vec4Reg\(),[\AddrReg\(),#-8*4]
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// OutputBlock
|
||||
//
|
||||
// Generates the code to store the output block.
|
||||
//
|
||||
|
||||
.macro OutputBlock Mode, Columns, Rows
|
||||
|
||||
OutputRow\Columns\()Element \Mode\(),x2,16,17,18,19
|
||||
.if \Rows\() >= 2
|
||||
OutputRow\Columns\()Element \Mode\(),x13,20,21,22,23
|
||||
.endif
|
||||
.if \Rows\() >= 4
|
||||
OutputRow\Columns\()Element \Mode\(),x14,24,25,26,27
|
||||
OutputRow\Columns\()Element \Mode\(),x15,28,29,30,31
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// ProcessRows
|
||||
//
|
||||
// Generates the code to process a compute and store the output block for a
|
||||
// fixed number of rows.
|
||||
//
|
||||
|
||||
.macro ProcessRows Mode, Rows
|
||||
|
||||
mov x4,#\Rows\() // return number of rows handled
|
||||
cmp x5,#8
|
||||
ble .L\Mode\().ProcessRemainingCountN\Rows\()
|
||||
|
||||
.L\Mode\().ProcessNextColumnLoop16x\Rows\():
|
||||
ComputeBlockLoop \Mode\(),16,\Rows\()
|
||||
.ifeqs "\Mode\()","Zero"
|
||||
MultiplyAlphaBlock 16,\Rows\()
|
||||
.endif
|
||||
sub x5,x5,#16
|
||||
tbnz x5,#63,.L\Mode\().OutputMasked16x\Rows\().Block
|
||||
OutputBlock \Mode\(),16,\Rows\()
|
||||
mov x0,x8 // reload matrix A
|
||||
cmp x5,#8
|
||||
bgt .L\Mode\().ProcessNextColumnLoop16x\Rows\()
|
||||
cbz x5,.L\Mode\().ExitKernel
|
||||
|
||||
.L\Mode\().ProcessRemainingCountN\Rows\():
|
||||
ComputeBlockLoop \Mode\(),8,\Rows\()
|
||||
.ifeqs "\Mode\()","Zero"
|
||||
MultiplyAlphaBlock 8,\Rows\()
|
||||
.endif
|
||||
|
||||
.L\Mode\().OutputMasked16x\Rows\().Block:
|
||||
tbz x5,#3,.L\Mode\().OutputRemaining7x\Rows\().Block
|
||||
OutputBlock \Mode\(),8,\Rows\()
|
||||
|
||||
.L\Mode\().OutputRemaining7x\Rows\().Block:
|
||||
tbz x5,#2,.L\Mode\().OutputRemaining3x\Rows\().Block
|
||||
OutputBlock \Mode\(),4,\Rows\()
|
||||
|
||||
.L\Mode\().OutputRemaining3x\Rows\().Block:
|
||||
tbz x5,#1,.L\Mode\().OutputRemaining1x\Rows\().Block
|
||||
OutputBlock \Mode\(),2,\Rows\()
|
||||
|
||||
.L\Mode\().OutputRemaining1x\Rows\().Block:
|
||||
tbz x5,#0,.L\Mode\().ExitKernel
|
||||
OutputBlock \Mode\(),1,\Rows\()
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A (x0) - Supplies the address of matrix A.
|
||||
|
||||
B (x1) - Supplies the address of matrix B. The matrix data has been packed
|
||||
using MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C (x2) - Supplies the address of matrix C.
|
||||
|
||||
CountK (x3) - Supplies the number of columns from matrix A and the number
|
||||
of rows from matrix B to iterate over.
|
||||
|
||||
CountM (x4) - 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 (x5) - Supplies the number of columns from matrix B and matrix C to
|
||||
iterate over.
|
||||
|
||||
lda (x6) - Supplies the first dimension of matrix A.
|
||||
|
||||
ldc (x7) - Supplies the first dimension of matrix C.
|
||||
|
||||
Alpha (s0) - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
|
||||
.macro SgemmKernelNeonFunction Mode
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmKernel\Mode\()
|
||||
|
||||
stp d8,d9,[sp,#-32]!
|
||||
stp d10,d11,[sp,#16]
|
||||
|
||||
add x13,x2,x7,lsl #2 // compute matrix C plus 1 row
|
||||
add x14,x13,x7,lsl #2 // compute matrix C plus 2 rows
|
||||
add x15,x14,x7,lsl #2 // compute matrix C plus 3 rows
|
||||
mov x8,x0 // save matrix A
|
||||
|
||||
//
|
||||
// Process 4 rows of the matrices.
|
||||
//
|
||||
|
||||
cmp x4,#4
|
||||
blt .L\Mode\().ProcessCountMLessThan4
|
||||
ProcessRows \Mode\(),4
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.L\Mode\().ExitKernel:
|
||||
mov x0,x4
|
||||
ldp d10,d11,[sp,#16]
|
||||
ldp d8,d9,[sp],#32
|
||||
ret
|
||||
|
||||
//
|
||||
// Process 2 rows of the matrices.
|
||||
//
|
||||
|
||||
.L\Mode\().ProcessCountMLessThan4:
|
||||
cmp x4,#2
|
||||
blt .L\Mode\().ProcessCountMLessThan2
|
||||
ProcessRows \Mode\(),2
|
||||
b .L\Mode\().ExitKernel
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
.L\Mode\().ProcessCountMLessThan2:
|
||||
ProcessRows \Mode\(),1
|
||||
b .L\Mode\().ExitKernel
|
||||
|
||||
.endm
|
||||
|
||||
SgemmKernelNeonFunction Zero
|
||||
SgemmKernelNeonFunction Add
|
||||
|
||||
.end
|
||||
+303
@@ -0,0 +1,303 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemvKernelNeon.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/vector
|
||||
multiply operation (SGEMV).
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows. This handles the special case of M=1.
|
||||
|
||||
The elements in matrix B are not transposed.
|
||||
|
||||
Arguments:
|
||||
|
||||
A (x0) - Supplies the address of matrix A.
|
||||
|
||||
B (x1) - Supplies the address of matrix B.
|
||||
|
||||
C (x2) - Supplies the address of matrix C.
|
||||
|
||||
CountK (x3) - Supplies the number of columns from matrix A and the number
|
||||
of rows from matrix B to iterate over.
|
||||
|
||||
CountN (x4) - Supplies the number of columns from matrix B and matrix C to
|
||||
iterate over.
|
||||
|
||||
ldb (x5) - Supplies the first dimension of matrix B.
|
||||
|
||||
ZeroMode (x6) - Supplies true if the output matrix must be zero initialized,
|
||||
else false if the output matrix is accumulated into.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasGemvFloatKernel
|
||||
|
||||
cmp x4,#64
|
||||
blo .LSgemvN.ProcessRemainingCountN
|
||||
mov x14,x0 // preserve vector A
|
||||
|
||||
//
|
||||
// Process 64 columns at a time in a loop.
|
||||
//
|
||||
|
||||
.LSgemvN.ProcessColumnLoopBy64:
|
||||
ldr q4,[x1]
|
||||
add x15,x1,#256 // compute next matrix B
|
||||
ldr q5,[x1,#16]
|
||||
tst w6,0xFF // ZeroMode?
|
||||
mov x13,x3 // reload CountK
|
||||
ldr q6,[x1,#32]
|
||||
beq .LSgemvN.LoadOutputBy64
|
||||
movi v16.4s,#0
|
||||
movi v17.4s,#0
|
||||
movi v18.4s,#0
|
||||
movi v19.4s,#0
|
||||
movi v20.4s,#0
|
||||
movi v21.4s,#0
|
||||
movi v22.4s,#0
|
||||
movi v23.4s,#0
|
||||
movi v24.4s,#0
|
||||
movi v25.4s,#0
|
||||
movi v26.4s,#0
|
||||
movi v27.4s,#0
|
||||
movi v28.4s,#0
|
||||
movi v29.4s,#0
|
||||
movi v30.4s,#0
|
||||
movi v31.4s,#0
|
||||
b .LSgemvN.MultiplyAccumulateBy64
|
||||
|
||||
.LSgemvN.LoadOutputBy64:
|
||||
ldp q16,q17,[x2]
|
||||
ldp q18,q19,[x2,#32]
|
||||
ldp q20,q21,[x2,#64]
|
||||
ldp q22,q23,[x2,#96]
|
||||
ldp q24,q25,[x2,#128]
|
||||
ldp q26,q27,[x2,#160]
|
||||
ldp q28,q29,[x2,#192]
|
||||
ldp q30,q31,[x2,#224]
|
||||
|
||||
.LSgemvN.MultiplyAccumulateBy64:
|
||||
ld1r {v0.4s},[x0] // broadcast next vector A element
|
||||
add x0,x0,4 // advance vector A by 1 element
|
||||
sub x13,x13,#1 // decrement K remaining
|
||||
fmla v16.4s,v4.4s,v0.4s
|
||||
ldr q7,[x1,#48]
|
||||
fmla v17.4s,v5.4s,v0.4s
|
||||
ldr q4,[x1,#64]
|
||||
fmla v18.4s,v6.4s,v0.4s
|
||||
ldr q5,[x1,#80]
|
||||
fmla v19.4s,v7.4s,v0.4s
|
||||
ldr q6,[x1,#96]
|
||||
fmla v20.4s,v4.4s,v0.4s
|
||||
ldr q7,[x1,#112]
|
||||
fmla v21.4s,v5.4s,v0.4s
|
||||
ldr q4,[x1,#128]
|
||||
fmla v22.4s,v6.4s,v0.4s
|
||||
ldr q5,[x1,#144]
|
||||
fmla v23.4s,v7.4s,v0.4s
|
||||
ldr q6,[x1,#160]
|
||||
fmla v24.4s,v4.4s,v0.4s
|
||||
ldr q7,[x1,#176]
|
||||
fmla v25.4s,v5.4s,v0.4s
|
||||
ldr q4,[x1,#192]
|
||||
fmla v26.4s,v6.4s,v0.4s
|
||||
ldr q5,[x1,#208]
|
||||
fmla v27.4s,v7.4s,v0.4s
|
||||
ldr q6,[x1,#224]
|
||||
fmla v28.4s,v4.4s,v0.4s
|
||||
ldr q7,[x1,#240]
|
||||
add x1,x1,x5,lsl #2 // compute next matrix B row address
|
||||
cbz x13,.LSgemvN.StoreOutputBy64
|
||||
ldr q4,[x1] // load data for next iteration
|
||||
fmla v29.4s,v5.4s,v0.4s
|
||||
ldr q5,[x1,#16]
|
||||
fmla v30.4s,v6.4s,v0.4s
|
||||
ldr q6,[x1,#32]
|
||||
fmla v31.4s,v7.4s,v0.4s
|
||||
b .LSgemvN.MultiplyAccumulateBy64
|
||||
|
||||
.LSgemvN.StoreOutputBy64:
|
||||
stp q16,q17,[x2]
|
||||
fmla v29.4s,v5.4s,v0.4s // finish computing tail vectors
|
||||
stp q18,q19,[x2,#32]
|
||||
fmla v30.4s,v6.4s,v0.4s
|
||||
stp q20,q21,[x2,#64]
|
||||
fmla v31.4s,v7.4s,v0.4s
|
||||
stp q22,q23,[x2,#96]
|
||||
sub x4,x4,#64 // subtract 64 columns
|
||||
stp q24,q25,[x2,#128]
|
||||
mov x0,x14 // reload vector A
|
||||
stp q26,q27,[x2,#160]
|
||||
mov x1,x15 // load next matrix B
|
||||
stp q28,q29,[x2,#192]
|
||||
stp q30,q31,[x2,#224]
|
||||
add x2,x2,#256 // advance vector C by 64 columns
|
||||
cbz x4,.LSgemvN.ExitKernel
|
||||
cmp x4,#64
|
||||
bhs .LSgemvN.ProcessColumnLoopBy64
|
||||
|
||||
//
|
||||
// Process the remaining 1 to 63 columns.
|
||||
//
|
||||
|
||||
.LSgemvN.ProcessRemainingCountN:
|
||||
tst w6,0xFF // ZeroMode?
|
||||
beq .LSgemvN.LoadOutputPartial32
|
||||
movi v16.4s,#0
|
||||
movi v17.4s,#0
|
||||
movi v18.4s,#0
|
||||
movi v19.4s,#0
|
||||
movi v20.4s,#0
|
||||
movi v21.4s,#0
|
||||
movi v22.4s,#0
|
||||
movi v23.4s,#0
|
||||
movi v24.4s,#0
|
||||
movi v25.4s,#0
|
||||
movi v26.4s,#0
|
||||
movi v27.4s,#0
|
||||
movi v28.4s,#0
|
||||
movi v29.4s,#0
|
||||
movi v30.4s,#0
|
||||
movi v31.4s,#0 // trailing float[2]
|
||||
movi v1.4s,#0 // trailing float[1]
|
||||
b .LSgemvN.ProcessNextPartialRow
|
||||
|
||||
.LSgemvN.LoadOutputPartial32:
|
||||
mov x15,x2
|
||||
tbz x4,#5,.LSgemvN.LoadOutputPartial16
|
||||
ldp q16,q17,[x15],#128
|
||||
ldp q18,q19,[x15,#-96]
|
||||
ldp q20,q21,[x15,#-64]
|
||||
ldp q22,q23,[x15,#-32]
|
||||
|
||||
.LSgemvN.LoadOutputPartial16:
|
||||
tbz x4,#4,.LSgemvN.LoadOutputPartial8
|
||||
ldp q24,q25,[x15],#64
|
||||
ldp q26,q27,[x15,#-32]
|
||||
|
||||
.LSgemvN.LoadOutputPartial8:
|
||||
tbz x4,#3,.LSgemvN.LoadOutputPartial4
|
||||
ldp q28,q29,[x15],#32
|
||||
|
||||
.LSgemvN.LoadOutputPartial4:
|
||||
tbz x4,#2,.LSgemvN.LoadOutputPartial2
|
||||
ldr q30,[x15],#16
|
||||
|
||||
.LSgemvN.LoadOutputPartial2:
|
||||
tbz x4,#1,.LSgemvN.LoadOutputPartial1
|
||||
ldr d31,[x15],#8
|
||||
|
||||
.LSgemvN.LoadOutputPartial1:
|
||||
tbz x4,#0,.LSgemvN.ProcessNextPartialRow
|
||||
ldr s1,[x15]
|
||||
|
||||
.LSgemvN.ProcessNextPartialRow:
|
||||
ld1r {v0.4s},[x0]
|
||||
add x0,x0,4
|
||||
sub x3,x3,#1 // decrement K remaining
|
||||
mov x15,x1
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial32:
|
||||
tbz x4,#5,.LSgemvN.MultiplyAccumulatePartial16
|
||||
ldp q4,q5,[x15],#128
|
||||
fmla v16.4s,v4.4s,v0.4s
|
||||
ldp q6,q7,[x15,#-96]
|
||||
fmla v17.4s,v5.4s,v0.4s
|
||||
ldp q4,q5,[x15,#-64]
|
||||
fmla v18.4s,v6.4s,v0.4s
|
||||
fmla v19.4s,v7.4s,v0.4s
|
||||
ldp q6,q7,[x15,#-32]
|
||||
fmla v20.4s,v4.4s,v0.4s
|
||||
fmla v21.4s,v5.4s,v0.4s
|
||||
fmla v22.4s,v6.4s,v0.4s
|
||||
fmla v23.4s,v7.4s,v0.4s
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial16:
|
||||
tbz x4,#4,.LSgemvN.MultiplyAccumulatePartial8
|
||||
ldp q4,q5,[x15],#64
|
||||
fmla v24.4s,v4.4s,v0.4s
|
||||
ldp q6,q7,[x15,#-32]
|
||||
fmla v25.4s,v5.4s,v0.4s
|
||||
fmla v26.4s,v6.4s,v0.4s
|
||||
fmla v27.4s,v7.4s,v0.4s
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial8:
|
||||
tbz x4,#3,.LSgemvN.MultiplyAccumulatePartial4
|
||||
ldp q4,q5,[x15],#32
|
||||
fmla v28.4s,v4.4s,v0.4s
|
||||
fmla v29.4s,v5.4s,v0.4s
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial4:
|
||||
tbz x4,#2,.LSgemvN.MultiplyAccumulatePartial2
|
||||
ldr q4,[x15],#16
|
||||
fmla v30.4s,v4.4s,v0.4s
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial2:
|
||||
tbz x4,#1,.LSgemvN.MultiplyAccumulatePartial1
|
||||
ldr d4,[x15],#8
|
||||
fmla v31.4s,v4.4s,v0.4s
|
||||
|
||||
.LSgemvN.MultiplyAccumulatePartial1:
|
||||
tbz x4,#0,.LSgemvN.AdvancePartialRow
|
||||
ldr s4,[x15]
|
||||
fmla v1.4s,v4.4s,v0.4s
|
||||
|
||||
.LSgemvN.AdvancePartialRow:
|
||||
add x1,x1,x5,lsl #2 // compute next matrix B row address
|
||||
cbnz x3,.LSgemvN.ProcessNextPartialRow
|
||||
|
||||
.LSgemvN.StoreOutputPartial32:
|
||||
tbz x4,#5,.LSgemvN.StoreOutputPartial16
|
||||
stp q16,q17,[x2],#128
|
||||
stp q18,q19,[x2,#-96]
|
||||
stp q20,q21,[x2,#-64]
|
||||
stp q22,q23,[x2,#-32]
|
||||
|
||||
.LSgemvN.StoreOutputPartial16:
|
||||
tbz x4,#4,.LSgemvN.StoreOutputPartial8
|
||||
stp q24,q25,[x2],#64
|
||||
stp q26,q27,[x2,#-32]
|
||||
|
||||
.LSgemvN.StoreOutputPartial8:
|
||||
tbz x4,#3,.LSgemvN.StoreOutputPartial4
|
||||
stp q28,q29,[x2],#32
|
||||
|
||||
.LSgemvN.StoreOutputPartial4:
|
||||
tbz x4,#2,.LSgemvN.StoreOutputPartial2
|
||||
str q30,[x2],#16
|
||||
|
||||
.LSgemvN.StoreOutputPartial2:
|
||||
tbz x4,#1,.LSgemvN.StoreOutputPartial1
|
||||
str d31,[x2],#8
|
||||
|
||||
.LSgemvN.StoreOutputPartial1:
|
||||
tbz x4,#0,.LSgemvN.ExitKernel
|
||||
str s1,[x2]
|
||||
|
||||
.LSgemvN.ExitKernel:
|
||||
ret
|
||||
|
||||
.end
|
||||
Vendored
+95
@@ -0,0 +1,95 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
asmmacro.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements common macros for the assembly modules.
|
||||
|
||||
--*/
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the assembler directives to annotate a new function.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name of the function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FUNCTION_ENTRY FunctionName
|
||||
|
||||
.p2align 2
|
||||
#if defined(__APPLE__)
|
||||
.globl _\FunctionName\()
|
||||
_\FunctionName\():
|
||||
#else
|
||||
.globl \FunctionName\()
|
||||
.type \FunctionName\(),%function
|
||||
\FunctionName\():
|
||||
#endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
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.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCountGE Count1, Value1, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro conditionally emits the statement if Count1 is greater than or
|
||||
equal to Value1 and Count2 is greater than or equal to Value2.
|
||||
|
||||
Arguments:
|
||||
|
||||
Count1 - Supplies the variable used in the comparison.
|
||||
|
||||
Value1 - Supplies the static used in the comparison.
|
||||
|
||||
Count2 - Supplies the variable used in the comparison.
|
||||
|
||||
Value2 - Supplies the static used in the comparison.
|
||||
|
||||
Statement - Supplies the statement to conditionally emit.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCount2GE Count1, Value1, Count2, Value2, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\()) && (\Count2\() >= \Value2\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
Vendored
+531
@@ -0,0 +1,531 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
sgemmc.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
template<bool ZeroMode, bool ProcessTwoRows>
|
||||
size_t
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountN - 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 - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
float32x4_t Row0Block0;
|
||||
float32x4_t Row0Block1;
|
||||
float32x4_t Row0Block2;
|
||||
float32x4_t Row0Block3;
|
||||
|
||||
float32x4_t Row1Block0;
|
||||
float32x4_t Row1Block1;
|
||||
float32x4_t Row1Block2;
|
||||
float32x4_t Row1Block3;
|
||||
|
||||
#if defined(_WIN32)
|
||||
|
||||
if (!ProcessTwoRows) {
|
||||
UNREFERENCED_PARAMETER(lda);
|
||||
UNREFERENCED_PARAMETER(ldc);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
do {
|
||||
|
||||
float32x4_t BElements0;
|
||||
float32x4_t BElements1;
|
||||
float32x4_t BElements2;
|
||||
float32x4_t BElements3;
|
||||
|
||||
float32x2_t Row0AElements;
|
||||
float32x2_t Row1AElements;
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
|
||||
Row0Block0 = vdupq_n_f32(0.0f);
|
||||
Row0Block1 = vdupq_n_f32(0.0f);
|
||||
Row0Block2 = vdupq_n_f32(0.0f);
|
||||
Row0Block3 = vdupq_n_f32(0.0f);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0 = vdupq_n_f32(0.0f);
|
||||
Row1Block1 = vdupq_n_f32(0.0f);
|
||||
Row1Block2 = vdupq_n_f32(0.0f);
|
||||
Row1Block3 = vdupq_n_f32(0.0f);
|
||||
}
|
||||
|
||||
//
|
||||
// Compute the 16x1 or 16x2 output block.
|
||||
//
|
||||
|
||||
const float* a = A;
|
||||
size_t k = CountK;
|
||||
|
||||
while (k >= 2) {
|
||||
|
||||
Row0AElements = vld1_f32(a);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1AElements = vld1_f32(a + lda);
|
||||
}
|
||||
|
||||
BElements0 = vld1q_f32(B + 0);
|
||||
BElements1 = vld1q_f32(B + 4);
|
||||
BElements2 = vld1q_f32(B + 8);
|
||||
BElements3 = vld1q_f32(B + 12);
|
||||
|
||||
Row0Block0 = vmlaq_lane_f32(Row0Block0, BElements0, Row0AElements, 0);
|
||||
Row0Block1 = vmlaq_lane_f32(Row0Block1, BElements1, Row0AElements, 0);
|
||||
Row0Block2 = vmlaq_lane_f32(Row0Block2, BElements2, Row0AElements, 0);
|
||||
Row0Block3 = vmlaq_lane_f32(Row0Block3, BElements3, Row0AElements, 0);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0 = vmlaq_lane_f32(Row1Block0, BElements0, Row1AElements, 0);
|
||||
Row1Block1 = vmlaq_lane_f32(Row1Block1, BElements1, Row1AElements, 0);
|
||||
Row1Block2 = vmlaq_lane_f32(Row1Block2, BElements2, Row1AElements, 0);
|
||||
Row1Block3 = vmlaq_lane_f32(Row1Block3, BElements3, Row1AElements, 0);
|
||||
}
|
||||
|
||||
BElements0 = vld1q_f32(B + 16);
|
||||
BElements1 = vld1q_f32(B + 20);
|
||||
BElements2 = vld1q_f32(B + 24);
|
||||
BElements3 = vld1q_f32(B + 28);
|
||||
|
||||
Row0Block0 = vmlaq_lane_f32(Row0Block0, BElements0, Row0AElements, 1);
|
||||
Row0Block1 = vmlaq_lane_f32(Row0Block1, BElements1, Row0AElements, 1);
|
||||
Row0Block2 = vmlaq_lane_f32(Row0Block2, BElements2, Row0AElements, 1);
|
||||
Row0Block3 = vmlaq_lane_f32(Row0Block3, BElements3, Row0AElements, 1);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0 = vmlaq_lane_f32(Row1Block0, BElements0, Row1AElements, 1);
|
||||
Row1Block1 = vmlaq_lane_f32(Row1Block1, BElements1, Row1AElements, 1);
|
||||
Row1Block2 = vmlaq_lane_f32(Row1Block2, BElements2, Row1AElements, 1);
|
||||
Row1Block3 = vmlaq_lane_f32(Row1Block3, BElements3, Row1AElements, 1);
|
||||
}
|
||||
|
||||
a += 2;
|
||||
B += 32;
|
||||
k -= 2;
|
||||
}
|
||||
|
||||
if (k > 0) {
|
||||
|
||||
Row0AElements = vld1_dup_f32(a);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1AElements = vld1_dup_f32(a + lda);
|
||||
}
|
||||
|
||||
BElements0 = vld1q_f32(B + 0);
|
||||
BElements1 = vld1q_f32(B + 4);
|
||||
BElements2 = vld1q_f32(B + 8);
|
||||
BElements3 = vld1q_f32(B + 12);
|
||||
|
||||
Row0Block0 = vmlaq_lane_f32(Row0Block0, BElements0, Row0AElements, 0);
|
||||
Row0Block1 = vmlaq_lane_f32(Row0Block1, BElements1, Row0AElements, 0);
|
||||
Row0Block2 = vmlaq_lane_f32(Row0Block2, BElements2, Row0AElements, 0);
|
||||
Row0Block3 = vmlaq_lane_f32(Row0Block3, BElements3, Row0AElements, 0);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0 = vmlaq_lane_f32(Row1Block0, BElements0, Row1AElements, 0);
|
||||
Row1Block1 = vmlaq_lane_f32(Row1Block1, BElements1, Row1AElements, 0);
|
||||
Row1Block2 = vmlaq_lane_f32(Row1Block2, BElements2, Row1AElements, 0);
|
||||
Row1Block3 = vmlaq_lane_f32(Row1Block3, BElements3, Row1AElements, 0);
|
||||
}
|
||||
|
||||
B += 16;
|
||||
}
|
||||
|
||||
//
|
||||
// Multiply by the alpha value.
|
||||
//
|
||||
|
||||
Row0Block0 = vmulq_n_f32(Row0Block0, alpha);
|
||||
Row0Block1 = vmulq_n_f32(Row0Block1, alpha);
|
||||
Row0Block2 = vmulq_n_f32(Row0Block2, alpha);
|
||||
Row0Block3 = vmulq_n_f32(Row0Block3, alpha);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0 = vmulq_n_f32(Row1Block0, alpha);
|
||||
Row1Block1 = vmulq_n_f32(Row1Block1, alpha);
|
||||
Row1Block2 = vmulq_n_f32(Row1Block2, alpha);
|
||||
Row1Block3 = vmulq_n_f32(Row1Block3, alpha);
|
||||
}
|
||||
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block0 = vaddq_f32(Row0Block0, vld1q_f32(C));
|
||||
Row0Block1 = vaddq_f32(Row0Block1, vld1q_f32(C + 4));
|
||||
Row0Block2 = vaddq_f32(Row0Block2, vld1q_f32(C + 8));
|
||||
Row0Block3 = vaddq_f32(Row0Block3, vld1q_f32(C + 12));
|
||||
}
|
||||
|
||||
vst1q_f32(C, Row0Block0);
|
||||
vst1q_f32(C + 4, Row0Block1);
|
||||
vst1q_f32(C + 8, Row0Block2);
|
||||
vst1q_f32(C + 12, Row0Block3);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block0 = vaddq_f32(Row1Block0, vld1q_f32(C + ldc));
|
||||
Row1Block1 = vaddq_f32(Row1Block1, vld1q_f32(C + ldc + 4));
|
||||
Row1Block2 = vaddq_f32(Row1Block2, vld1q_f32(C + ldc + 8));
|
||||
Row1Block3 = vaddq_f32(Row1Block3, vld1q_f32(C + ldc + 12));
|
||||
}
|
||||
|
||||
vst1q_f32(C + ldc, Row1Block0);
|
||||
vst1q_f32(C + ldc + 4, Row1Block1);
|
||||
vst1q_f32(C + ldc + 8, Row1Block2);
|
||||
vst1q_f32(C + ldc + 12, Row1Block3);
|
||||
}
|
||||
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if ((CountN & 8) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block0 = vaddq_f32(Row0Block0, vld1q_f32(C));
|
||||
Row0Block1 = vaddq_f32(Row0Block1, vld1q_f32(C + 4));
|
||||
}
|
||||
|
||||
vst1q_f32(C, Row0Block0);
|
||||
vst1q_f32(C + 4, Row0Block1);
|
||||
Row0Block0 = Row0Block2;
|
||||
Row0Block1 = Row0Block3;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block0 = vaddq_f32(Row1Block0, vld1q_f32(C + ldc));
|
||||
Row1Block1 = vaddq_f32(Row1Block1, vld1q_f32(C + ldc + 4));
|
||||
}
|
||||
|
||||
vst1q_f32(C + ldc, Row1Block0);
|
||||
vst1q_f32(C + ldc + 4, Row1Block1);
|
||||
Row1Block0 = Row1Block2;
|
||||
Row1Block1 = Row1Block3;
|
||||
}
|
||||
|
||||
C += 8;
|
||||
}
|
||||
|
||||
if ((CountN & 4) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block0 = vaddq_f32(Row0Block0, vld1q_f32(C));
|
||||
}
|
||||
|
||||
vst1q_f32(C, Row0Block0);
|
||||
Row0Block0 = Row0Block1;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block0 = vaddq_f32(Row1Block0, vld1q_f32(C + ldc));
|
||||
}
|
||||
|
||||
vst1q_f32(C + ldc, Row1Block0);
|
||||
Row1Block0 = Row1Block1;
|
||||
}
|
||||
|
||||
C += 4;
|
||||
}
|
||||
|
||||
float32x2_t Row0Block0High;
|
||||
float32x2_t Row0Block0Low;
|
||||
|
||||
float32x2_t Row1Block0High;
|
||||
float32x2_t Row1Block0Low;
|
||||
|
||||
Row0Block0High = vget_high_f32(Row0Block0);
|
||||
Row0Block0Low = vget_low_f32(Row0Block0);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block0High = vget_high_f32(Row1Block0);
|
||||
Row1Block0Low = vget_low_f32(Row1Block0);
|
||||
}
|
||||
|
||||
if ((CountN & 2) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block0Low = vadd_f32(Row0Block0Low, vld1_f32(C));
|
||||
}
|
||||
|
||||
vst1_f32(C, Row0Block0Low);
|
||||
Row0Block0Low = Row0Block0High;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block0Low = vadd_f32(Row1Block0Low, vld1_f32(C + ldc));
|
||||
}
|
||||
|
||||
vst1_f32(C + ldc, Row1Block0Low);
|
||||
Row1Block0Low = Row1Block0High;
|
||||
}
|
||||
|
||||
C += 2;
|
||||
}
|
||||
|
||||
if ((CountN & 1) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block0Low = vadd_f32(Row0Block0Low, vld1_dup_f32(C));
|
||||
}
|
||||
|
||||
vst1_lane_f32(C, Row0Block0Low, 0);
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block0Low = vadd_f32(Row1Block0Low, vld1_dup_f32(C + ldc));
|
||||
}
|
||||
|
||||
vst1_lane_f32(C + ldc, Row1Block0Low, 0);
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return ProcessTwoRows ? 2 : 1;
|
||||
}
|
||||
|
||||
template<bool ZeroMode>
|
||||
size_t
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
|
||||
if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmKernel<ZeroMode, true>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmKernel<ZeroMode, false>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelZero(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
return MlasSgemmKernel<true>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelAdd(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
return MlasSgemmKernel<false>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
+502
@@ -0,0 +1,502 @@
|
||||
;++
|
||||
;
|
||||
; Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
;
|
||||
; Licensed under the MIT License.
|
||||
;
|
||||
; Module Name:
|
||||
;
|
||||
; SgemmKernelNeon.asm
|
||||
;
|
||||
; Abstract:
|
||||
;
|
||||
; This module implements the kernels for the single precision matrix/matrix
|
||||
; multiply operation (SGEMM).
|
||||
;
|
||||
;--
|
||||
|
||||
#include "kxarm64.h"
|
||||
|
||||
TEXTAREA
|
||||
|
||||
;
|
||||
; ClearRowAccumulators
|
||||
;
|
||||
; Generates the code to clear the accumulators for a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ClearRowAccumulators $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
movi $Vec1Reg..16b,#0
|
||||
movi $Vec2Reg..16b,#0
|
||||
IF $Columns > 8
|
||||
movi $Vec3Reg..16b,#0
|
||||
movi $Vec4Reg..16b,#0
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ClearBlockAccumulators
|
||||
;
|
||||
; Generates the code to clear the accumulators for a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ClearBlockAccumulators $Columns, $Rows
|
||||
|
||||
ClearRowAccumulators $Columns, v16, v17, v18, v19
|
||||
IF $Rows >= 2
|
||||
ClearRowAccumulators $Columns, v20, v21, v22, v23
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
ClearRowAccumulators $Columns, v24, v25, v26, v27
|
||||
ClearRowAccumulators $Columns, v28, v29, v30, v31
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; LoadMatrixAElementsBy4
|
||||
; LoadMatrixAElementsBy1
|
||||
;
|
||||
; Generates the code to load 1 or 4 elements from matrix A.
|
||||
;
|
||||
|
||||
MACRO
|
||||
LoadMatrixAElementsBy4 $Rows
|
||||
|
||||
ldr v8,[x0],#16
|
||||
IF $Rows >= 2
|
||||
ldr v9,[x10],#16
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
ldr v10,[x11],#16
|
||||
ldr v11,[x12],#16
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
LoadMatrixAElementsBy1 $Rows
|
||||
|
||||
ldr s8,[x0],#4
|
||||
IF $Rows >= 2
|
||||
ldr s9,[x10],#4
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
ldr s10,[x11],#4
|
||||
ldr s11,[x12],#4
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAccumulateRow
|
||||
;
|
||||
; Generates the code to multiply and accumulate a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAccumulateRow $Columns, $MatrixAReg, $Broadcast, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
fmla $Vec1Reg..4s,v4.4s,$MatrixAReg..s[$Broadcast]
|
||||
fmla $Vec2Reg..4s,v5.4s,$MatrixAReg..s[$Broadcast]
|
||||
IF $Columns > 8
|
||||
fmla $Vec3Reg..4s,v6.4s,$MatrixAReg..s[$Broadcast]
|
||||
fmla $Vec4Reg..4s,v7.4s,$MatrixAReg..s[$Broadcast]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAccumulateBlock
|
||||
;
|
||||
; Generates the code to multiply and accumulate into the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAccumulateBlock $Columns, $Rows, $Broadcast
|
||||
|
||||
MultiplyAccumulateRow $Columns, v8, $Broadcast, v16, v17, v18, v19
|
||||
IF $Rows >= 2
|
||||
MultiplyAccumulateRow $Columns, v9, $Broadcast, v20, v21, v22, v23
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
MultiplyAccumulateRow $Columns, v10, $Broadcast, v24, v25, v26, v27
|
||||
MultiplyAccumulateRow $Columns, v11, $Broadcast, v28, v29, v30, v31
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ComputeBlockLoop
|
||||
;
|
||||
; Generates the code to loop over K entries of the input matrices to produce
|
||||
; the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ComputeBlockLoop $Mode, $Columns, $Rows
|
||||
|
||||
ClearBlockAccumulators $Columns, $Rows
|
||||
|
||||
IF $Rows >= 2
|
||||
add x10,x0,x6 lsl #2 ; compute matrix A plus 1 row
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
add x11,x10,x6 lsl #2 ; compute matrix A plus 2 rows
|
||||
add x12,x11,x6 lsl #2 ; compute matrix A plus 3 rows
|
||||
ENDIF
|
||||
|
||||
sub x9,x3,#4 ; decrement block count to process
|
||||
tbnz x9,#63,$Mode.ProcessRemaining$Columns.x$Rows.Blocks
|
||||
|
||||
$Mode.Compute$Columns.x$Rows.BlockBy4Loop
|
||||
LoadMatrixAElementsBy4 $Rows
|
||||
ldp v4,v5,[x1],#64*4
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-56*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,0
|
||||
ldp v4,v5,[x1,#-48*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-40*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,1
|
||||
ldp v4,v5,[x1,#-32*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-24*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,2
|
||||
ldp v4,v5,[x1,#-16*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-8*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,3
|
||||
sub x9,x9,#4
|
||||
tbz x9,#63,$Mode.Compute$Columns.x$Rows.BlockBy4Loop
|
||||
|
||||
$Mode.ProcessRemaining$Columns.x$Rows.Blocks
|
||||
add x9,x9,#4 ; correct for over-subtract above
|
||||
cbz x9,$Mode.Output$Columns.x$Rows.Block
|
||||
|
||||
$Mode.Compute$Columns.x$Rows.BlockBy1Loop
|
||||
LoadMatrixAElementsBy1 $Rows
|
||||
ldp v4,v5,[x1],#16*4
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-8*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,0
|
||||
sub x9,x9,#1
|
||||
cbnz x9,$Mode.Compute$Columns.x$Rows.BlockBy1Loop
|
||||
|
||||
$Mode.Output$Columns.x$Rows.Block
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAlphaRow
|
||||
;
|
||||
; Generates the code to multiply a single row of the output block by the alpha
|
||||
; value.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAlphaRow $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF $Columns <= 4
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
ELIF $Columns <= 8
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
ELIF $Columns <= 12
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
|
||||
ELSE
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
|
||||
fmul $Vec4Reg..4s,$Vec4Reg..4s,v0.s[0]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAlphaBlock
|
||||
;
|
||||
; Generates the code to multiply the output block by the alpha value.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAlphaBlock $Columns, $Rows
|
||||
|
||||
MultiplyAlphaRow $Columns, v16, v17, v18, v19
|
||||
IF $Rows >= 2
|
||||
MultiplyAlphaRow $Columns, v20, v21, v22, v23
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
MultiplyAlphaRow $Columns, v24, v25, v26, v27
|
||||
MultiplyAlphaRow $Columns, v28, v29, v30, v31
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; OutputRow1Element
|
||||
; OutputRow2Element
|
||||
; OutputRow4Element
|
||||
; OutputRow8Element
|
||||
; OutputRow16Element
|
||||
;
|
||||
; Generates the code to store elements to the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
OutputRow1Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.s}[0],[$AddrReg]
|
||||
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
|
||||
st1 {v4.s}[0],[$AddrReg] ; post-increment not needed for last element
|
||||
ELSE
|
||||
st1 {$Vec1Reg..s}[0],[$AddrReg] ; post-increment not needed for last element
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow2Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.2s},[$AddrReg]
|
||||
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
|
||||
st1 {v4.2s},[$AddrReg],#2*4
|
||||
ELSE
|
||||
st1 {$Vec1Reg..2s},[$AddrReg],#2*4
|
||||
ENDIF
|
||||
dup $Vec1Reg..4s,$Vec1Reg..s[2] ; shift remaining elements down
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow4Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.4s},[$AddrReg]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
st1 {v4.4s},[$AddrReg],#4*4
|
||||
ELSE
|
||||
st1 {$Vec1Reg..4s},[$AddrReg],#4*4
|
||||
ENDIF
|
||||
mov $Vec1Reg..16b,$Vec2Reg..16b ; shift remaining elements down
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow8Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ldp v4,v5,[$AddrReg]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
|
||||
stp v4,v5,[$AddrReg],#8*4
|
||||
ELSE
|
||||
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#8*4
|
||||
ENDIF
|
||||
mov $Vec1Reg..16b,$Vec3Reg..16b ; shift remaining elements down
|
||||
mov $Vec2Reg..16b,$Vec4Reg..16b
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow16Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ldp v4,v5,[$AddrReg]
|
||||
ldp v6,v7,[$AddrReg,#8*4]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmla v6.4s,$Vec3Reg..4s,v0.s[0]
|
||||
fmla v7.4s,$Vec4Reg..4s,v0.s[0]
|
||||
stp v4,v5,[$AddrReg],#16*4
|
||||
stp v6,v7,[$AddrReg,#-8*4]
|
||||
ELSE
|
||||
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#16*4
|
||||
stp $Vec3Reg.,$Vec4Reg.,[$AddrReg,#-8*4]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; OutputBlock
|
||||
;
|
||||
; Generates the code to store the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
OutputBlock $Mode, $Columns, $Rows
|
||||
|
||||
OutputRow$Columns.Element $Mode, x2, v16, v17, v18, v19
|
||||
IF $Rows >= 2
|
||||
OutputRow$Columns.Element $Mode, x13, v20, v21, v22, v23
|
||||
ENDIF
|
||||
IF $Rows >= 4
|
||||
OutputRow$Columns.Element $Mode, x14, v24, v25, v26, v27
|
||||
OutputRow$Columns.Element $Mode, x15, v28, v29, v30, v31
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ProcessRows
|
||||
;
|
||||
; Generates the code to process a compute and store the output block for a
|
||||
; fixed number of rows.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ProcessRows $Mode, $Rows
|
||||
|
||||
mov x4,#$Rows ; return number of rows handled
|
||||
cmp x5,#8
|
||||
ble $Mode.ProcessRemainingCountN$Rows
|
||||
|
||||
$Mode.ProcessNextColumnLoop16x$Rows
|
||||
ComputeBlockLoop $Mode,16,$Rows
|
||||
IF "$Mode"=="Zero"
|
||||
MultiplyAlphaBlock 16,$Rows
|
||||
ENDIF
|
||||
sub x5,x5,#16
|
||||
tbnz x5,#63,$Mode.OutputMasked16x$Rows.Block
|
||||
OutputBlock $Mode,16,$Rows
|
||||
mov x0,x8 ; reload matrix A
|
||||
cmp x5,#8
|
||||
bgt $Mode.ProcessNextColumnLoop16x$Rows
|
||||
cbz x5,$Mode.ExitKernel
|
||||
|
||||
$Mode.ProcessRemainingCountN$Rows
|
||||
ComputeBlockLoop $Mode,8,$Rows
|
||||
IF "$Mode"=="Zero"
|
||||
MultiplyAlphaBlock 8,$Rows
|
||||
ENDIF
|
||||
|
||||
$Mode.OutputMasked16x$Rows.Block
|
||||
tbz x5,#3,$Mode.OutputRemaining7x$Rows.Block
|
||||
OutputBlock $Mode,8,$Rows
|
||||
|
||||
$Mode.OutputRemaining7x$Rows.Block
|
||||
tbz x5,#2,$Mode.OutputRemaining3x$Rows.Block
|
||||
OutputBlock $Mode,4,$Rows
|
||||
|
||||
$Mode.OutputRemaining3x$Rows.Block
|
||||
tbz x5,#1,$Mode.OutputRemaining1x$Rows.Block
|
||||
OutputBlock $Mode,2,$Rows
|
||||
|
||||
$Mode.OutputRemaining1x$Rows.Block
|
||||
tbz x5,#0,$Mode.ExitKernel
|
||||
OutputBlock $Mode,1,$Rows
|
||||
|
||||
MEND
|
||||
|
||||
SUBT "SGEMM kernel"
|
||||
;++
|
||||
;
|
||||
; Routine Description:
|
||||
;
|
||||
; This routine is an inner kernel to compute matrix multiplication for a
|
||||
; set of rows.
|
||||
;
|
||||
; Arguments:
|
||||
;
|
||||
; A (x0) - Supplies the address of matrix A.
|
||||
;
|
||||
; B (x1) - Supplies the address of matrix B. The matrix data has been packed
|
||||
; using MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
;
|
||||
; C (x2) - Supplies the address of matrix C.
|
||||
;
|
||||
; CountK (x3) - Supplies the number of columns from matrix A and the number
|
||||
; of rows from matrix B to iterate over.
|
||||
;
|
||||
; CountM (x4) - 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 (x5) - Supplies the number of columns from matrix B and matrix C to
|
||||
; iterate over.
|
||||
;
|
||||
; lda (x6) - Supplies the first dimension of matrix A.
|
||||
;
|
||||
; ldc (x7) - Supplies the first dimension of matrix C.
|
||||
;
|
||||
; Alpha (s0) - Supplies the scalar multiplier (see SGEMM definition).
|
||||
;
|
||||
; Return Value:
|
||||
;
|
||||
; Returns the number of rows handled.
|
||||
;
|
||||
;--
|
||||
|
||||
MACRO
|
||||
SgemmKernelNeonFunction $Mode
|
||||
|
||||
NESTED_ENTRY MlasSgemmKernel$Mode
|
||||
|
||||
PROLOG_SAVE_REG_PAIR d8,d9,#-32!
|
||||
PROLOG_SAVE_REG_PAIR d10,d11,#16
|
||||
|
||||
add x13,x2,x7 lsl #2 ; compute matrix C plus 1 row
|
||||
add x14,x13,x7 lsl #2 ; compute matrix C plus 2 rows
|
||||
add x15,x14,x7 lsl #2 ; compute matrix C plus 3 rows
|
||||
mov x8,x0 ; save matrix A
|
||||
|
||||
;
|
||||
; Process 4 rows of the matrices.
|
||||
;
|
||||
|
||||
cmp x4,#4
|
||||
blt $Mode.ProcessCountMLessThan4
|
||||
ProcessRows $Mode,4
|
||||
|
||||
;
|
||||
; Restore non-volatile registers and return.
|
||||
;
|
||||
|
||||
$Mode.ExitKernel
|
||||
mov x0,x4
|
||||
EPILOG_RESTORE_REG_PAIR d10,d11,#16
|
||||
EPILOG_RESTORE_REG_PAIR d8,d9,#32!
|
||||
EPILOG_RETURN
|
||||
|
||||
;
|
||||
; Process 2 rows of the matrices.
|
||||
;
|
||||
|
||||
$Mode.ProcessCountMLessThan4
|
||||
cmp x4,#2
|
||||
blt $Mode.ProcessCountMLessThan2
|
||||
ProcessRows $Mode,2
|
||||
b $Mode.ExitKernel
|
||||
|
||||
;
|
||||
; Process 1 row of the matrices.
|
||||
;
|
||||
|
||||
$Mode.ProcessCountMLessThan2
|
||||
ProcessRows $Mode,1
|
||||
b $Mode.ExitKernel
|
||||
|
||||
NESTED_END
|
||||
|
||||
MEND
|
||||
|
||||
SgemmKernelNeonFunction Zero
|
||||
SgemmKernelNeonFunction Add
|
||||
|
||||
END
|
||||
+466
@@ -0,0 +1,466 @@
|
||||
;++
|
||||
;
|
||||
; Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
;
|
||||
; Licensed under the MIT License.
|
||||
;
|
||||
; Module Name:
|
||||
;
|
||||
; SgemmKernelNeon.asm
|
||||
;
|
||||
; Abstract:
|
||||
;
|
||||
; This module implements the kernels for the single precision matrix/matrix
|
||||
; multiply operation (SGEMM).
|
||||
;
|
||||
;--
|
||||
|
||||
#include "kxarm64.h"
|
||||
|
||||
TEXTAREA
|
||||
|
||||
;
|
||||
; ClearRowAccumulators
|
||||
;
|
||||
; Generates the code to clear the accumulators for a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ClearRowAccumulators $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
movi $Vec1Reg..16b,#0
|
||||
movi $Vec2Reg..16b,#0
|
||||
IF $Columns > 8
|
||||
movi $Vec3Reg..16b,#0
|
||||
movi $Vec4Reg..16b,#0
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ClearBlockAccumulators
|
||||
;
|
||||
; Generates the code to clear the accumulators for a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ClearBlockAccumulators $Columns, $Rows
|
||||
|
||||
ClearRowAccumulators $Columns, v8, v9, v10, v11
|
||||
IF $Rows >= 2
|
||||
ClearRowAccumulators $Columns, v12, v13, v14, v15
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; LoadMatrixAElementsBy4
|
||||
; LoadMatrixAElementsBy1
|
||||
;
|
||||
; Generates the code to load 1 or 4 elements from matrix A.
|
||||
;
|
||||
|
||||
MACRO
|
||||
LoadMatrixAElementsBy4 $Rows
|
||||
|
||||
ldr v2,[x0],#16
|
||||
IF $Rows >= 2
|
||||
ldr v3,[x10],#16
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
LoadMatrixAElementsBy1 $Rows
|
||||
|
||||
ldr s2,[x0],#4
|
||||
IF $Rows >= 2
|
||||
ldr s3,[x10],#4
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAccumulateRow
|
||||
;
|
||||
; Generates the code to multiply and accumulate a single row of the output
|
||||
; block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAccumulateRow $Columns, $MatrixAReg, $Broadcast, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
fmla $Vec1Reg..4s,v4.4s,$MatrixAReg..s[$Broadcast]
|
||||
fmla $Vec2Reg..4s,v5.4s,$MatrixAReg..s[$Broadcast]
|
||||
IF $Columns > 8
|
||||
fmla $Vec3Reg..4s,v6.4s,$MatrixAReg..s[$Broadcast]
|
||||
fmla $Vec4Reg..4s,v7.4s,$MatrixAReg..s[$Broadcast]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAccumulateBlock
|
||||
;
|
||||
; Generates the code to multiply and accumulate into the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAccumulateBlock $Columns, $Rows, $Broadcast
|
||||
|
||||
MultiplyAccumulateRow $Columns, v2, $Broadcast, v8, v9, v10, v11
|
||||
IF $Rows >= 2
|
||||
MultiplyAccumulateRow $Columns, v3, $Broadcast, v12, v13, v14, v15
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ComputeBlockLoop
|
||||
;
|
||||
; Generates the code to loop over K entries of the input matrices to produce
|
||||
; the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ComputeBlockLoop $Mode, $Columns, $Rows
|
||||
|
||||
ClearBlockAccumulators $Columns, $Rows
|
||||
|
||||
IF $Rows >= 2
|
||||
add x10,x0,x6 lsl #2 ; compute matrix A plus 1 row
|
||||
ENDIF
|
||||
|
||||
sub x9,x3,#4 ; decrement block count to process
|
||||
tbnz x9,#63,$Mode.ProcessRemaining$Columns.x$Rows.Blocks
|
||||
|
||||
$Mode.Compute$Columns.x$Rows.BlockBy4Loop
|
||||
LoadMatrixAElementsBy4 $Rows
|
||||
ldp v4,v5,[x1],#64*4
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-56*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,0
|
||||
ldp v4,v5,[x1,#-48*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-40*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,1
|
||||
ldp v4,v5,[x1,#-32*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-24*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,2
|
||||
ldp v4,v5,[x1,#-16*4]
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-8*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,3
|
||||
sub x9,x9,#4
|
||||
tbz x9,#63,$Mode.Compute$Columns.x$Rows.BlockBy4Loop
|
||||
|
||||
$Mode.ProcessRemaining$Columns.x$Rows.Blocks
|
||||
add x9,x9,#4 ; correct for over-subtract above
|
||||
cbz x9,$Mode.Output$Columns.x$Rows.Block
|
||||
|
||||
$Mode.Compute$Columns.x$Rows.BlockBy1Loop
|
||||
LoadMatrixAElementsBy1 $Rows
|
||||
ldp v4,v5,[x1],#16*4
|
||||
IF $Columns > 8
|
||||
ldp v6,v7,[x1,#-8*4]
|
||||
ENDIF
|
||||
MultiplyAccumulateBlock $Columns,$Rows,0
|
||||
sub x9,x9,#1
|
||||
cbnz x9,$Mode.Compute$Columns.x$Rows.BlockBy1Loop
|
||||
|
||||
$Mode.Output$Columns.x$Rows.Block
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAlphaRow
|
||||
;
|
||||
; Generates the code to multiply a single row of the output block by the alpha
|
||||
; value.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAlphaRow $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF $Columns <= 4
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
ELIF $Columns <= 8
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
ELIF $Columns <= 12
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
|
||||
ELSE
|
||||
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
|
||||
fmul $Vec4Reg..4s,$Vec4Reg..4s,v0.s[0]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; MultiplyAlphaBlock
|
||||
;
|
||||
; Generates the code to multiply the output block by the alpha value.
|
||||
;
|
||||
|
||||
MACRO
|
||||
MultiplyAlphaBlock $Columns, $Rows
|
||||
|
||||
MultiplyAlphaRow $Columns, v8, v9, v10, v11
|
||||
IF $Rows >= 2
|
||||
MultiplyAlphaRow $Columns, v12, v13, v14, v15
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; OutputRow1Element
|
||||
; OutputRow2Element
|
||||
; OutputRow4Element
|
||||
; OutputRow8Element
|
||||
; OutputRow16Element
|
||||
;
|
||||
; Generates the code to store elements to the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
OutputRow1Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.s}[0],[$AddrReg]
|
||||
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
|
||||
st1 {v4.s}[0],[$AddrReg] ; post-increment not needed for last element
|
||||
ELSE
|
||||
st1 {$Vec1Reg..s}[0],[$AddrReg] ; post-increment not needed for last element
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow2Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.2s},[$AddrReg]
|
||||
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
|
||||
st1 {v4.2s},[$AddrReg],#2*4
|
||||
ELSE
|
||||
st1 {$Vec1Reg..2s},[$AddrReg],#2*4
|
||||
ENDIF
|
||||
dup $Vec1Reg..4s,$Vec1Reg..s[2] ; shift remaining elements down
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow4Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ld1 {v4.4s},[$AddrReg]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
st1 {v4.4s},[$AddrReg],#4*4
|
||||
ELSE
|
||||
st1 {$Vec1Reg..4s},[$AddrReg],#4*4
|
||||
ENDIF
|
||||
mov $Vec1Reg..16b,$Vec2Reg..16b ; shift remaining elements down
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow8Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ldp v4,v5,[$AddrReg]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
|
||||
stp v4,v5,[$AddrReg],#8*4
|
||||
ELSE
|
||||
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#8*4
|
||||
ENDIF
|
||||
mov $Vec1Reg..16b,$Vec3Reg..16b ; shift remaining elements down
|
||||
mov $Vec2Reg..16b,$Vec4Reg..16b
|
||||
|
||||
MEND
|
||||
|
||||
MACRO
|
||||
OutputRow16Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
|
||||
|
||||
IF "$Mode"=="Add"
|
||||
ldp v4,v5,[$AddrReg]
|
||||
ldp v6,v7,[$AddrReg,#8*4]
|
||||
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
|
||||
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
|
||||
fmla v6.4s,$Vec3Reg..4s,v0.s[0]
|
||||
fmla v7.4s,$Vec4Reg..4s,v0.s[0]
|
||||
stp v4,v5,[$AddrReg],#16*4
|
||||
stp v6,v7,[$AddrReg,#-8*4]
|
||||
ELSE
|
||||
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#16*4
|
||||
stp $Vec3Reg.,$Vec4Reg.,[$AddrReg,#-8*4]
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; OutputBlock
|
||||
;
|
||||
; Generates the code to store the output block.
|
||||
;
|
||||
|
||||
MACRO
|
||||
OutputBlock $Mode, $Columns, $Rows
|
||||
|
||||
OutputRow$Columns.Element $Mode, x2, v8, v9, v10, v11
|
||||
IF $Rows >= 2
|
||||
OutputRow$Columns.Element $Mode, x11, v12, v13, v14, v15
|
||||
ENDIF
|
||||
|
||||
MEND
|
||||
|
||||
;
|
||||
; ProcessRows
|
||||
;
|
||||
; Generates the code to process a compute and store the output block for a
|
||||
; fixed number of rows.
|
||||
;
|
||||
|
||||
MACRO
|
||||
ProcessRows $Mode, $Rows
|
||||
|
||||
mov x4,#$Rows ; return number of rows handled
|
||||
cmp x5,#8
|
||||
ble $Mode.ProcessRemainingCountN$Rows
|
||||
|
||||
$Mode.ProcessNextColumnLoop16x$Rows
|
||||
ComputeBlockLoop $Mode,16,$Rows
|
||||
IF "$Mode"=="Zero"
|
||||
MultiplyAlphaBlock 16,$Rows
|
||||
ENDIF
|
||||
sub x5,x5,#16
|
||||
tbnz x5,#63,$Mode.OutputMasked16x$Rows.Block
|
||||
OutputBlock $Mode,16,$Rows
|
||||
mov x0,x8 ; reload matrix A
|
||||
cmp x5,#8
|
||||
bgt $Mode.ProcessNextColumnLoop16x$Rows
|
||||
cbz x5,$Mode.ExitKernel
|
||||
|
||||
$Mode.ProcessRemainingCountN$Rows
|
||||
ComputeBlockLoop $Mode,8,$Rows
|
||||
IF "$Mode"=="Zero"
|
||||
MultiplyAlphaBlock 8,$Rows
|
||||
ENDIF
|
||||
|
||||
$Mode.OutputMasked16x$Rows.Block
|
||||
tbz x5,#3,$Mode.OutputRemaining7x$Rows.Block
|
||||
OutputBlock $Mode,8,$Rows
|
||||
|
||||
$Mode.OutputRemaining7x$Rows.Block
|
||||
tbz x5,#2,$Mode.OutputRemaining3x$Rows.Block
|
||||
OutputBlock $Mode,4,$Rows
|
||||
|
||||
$Mode.OutputRemaining3x$Rows.Block
|
||||
tbz x5,#1,$Mode.OutputRemaining1x$Rows.Block
|
||||
OutputBlock $Mode,2,$Rows
|
||||
|
||||
$Mode.OutputRemaining1x$Rows.Block
|
||||
tbz x5,#0,$Mode.ExitKernel
|
||||
OutputBlock $Mode,1,$Rows
|
||||
|
||||
MEND
|
||||
|
||||
SUBT "SGEMM kernel"
|
||||
;++
|
||||
;
|
||||
; Routine Description:
|
||||
;
|
||||
; This routine is an inner kernel to compute matrix multiplication for a
|
||||
; set of rows.
|
||||
;
|
||||
; Arguments:
|
||||
;
|
||||
; A (x0) - Supplies the address of matrix A.
|
||||
;
|
||||
; B (x1) - Supplies the address of matrix B. The matrix data has been packed
|
||||
; using MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
;
|
||||
; C (x2) - Supplies the address of matrix C.
|
||||
;
|
||||
; CountK (x3) - Supplies the number of columns from matrix A and the number
|
||||
; of rows from matrix B to iterate over.
|
||||
;
|
||||
; CountM (x4) - 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 (x5) - Supplies the number of columns from matrix B and matrix C to
|
||||
; iterate over.
|
||||
;
|
||||
; lda (x6) - Supplies the first dimension of matrix A.
|
||||
;
|
||||
; ldc (x7) - Supplies the first dimension of matrix C.
|
||||
;
|
||||
; Alpha (s0) - Supplies the scalar multiplier (see SGEMM definition).
|
||||
;
|
||||
; Return Value:
|
||||
;
|
||||
; Returns the number of rows handled.
|
||||
;
|
||||
;--
|
||||
|
||||
MACRO
|
||||
SgemmKernelNeonFunction $Mode
|
||||
|
||||
NESTED_ENTRY_COMDAT A64NAME(MlasSgemmKernel$Mode)
|
||||
|
||||
PROLOG_SAVE_REG_PAIR d8,d9,#-64!
|
||||
PROLOG_SAVE_REG_PAIR d10,d11,#16
|
||||
PROLOG_SAVE_REG_PAIR d12,d13,#32
|
||||
PROLOG_SAVE_REG_PAIR d14,d15,#48
|
||||
|
||||
add x11,x2,x7 lsl #2 ; compute matrix C plus 1 row
|
||||
mov x8,x0 ; save matrix A
|
||||
|
||||
;
|
||||
; Process 2 rows of the matrices.
|
||||
;
|
||||
|
||||
cmp x4,#2
|
||||
blt $Mode.ProcessCountMLessThan2
|
||||
ProcessRows $Mode,2
|
||||
|
||||
;
|
||||
; Restore non-volatile registers and return.
|
||||
;
|
||||
|
||||
$Mode.ExitKernel
|
||||
mov x0,x4
|
||||
EPILOG_RESTORE_REG_PAIR d14,d15,#48
|
||||
EPILOG_RESTORE_REG_PAIR d12,d13,#32
|
||||
EPILOG_RESTORE_REG_PAIR d10,d11,#16
|
||||
EPILOG_RESTORE_REG_PAIR d8,d9,#64!
|
||||
EPILOG_RETURN
|
||||
|
||||
;
|
||||
; Process 1 row of the matrices.
|
||||
;
|
||||
|
||||
$Mode.ProcessCountMLessThan2
|
||||
ProcessRows $Mode,1
|
||||
b $Mode.ExitKernel
|
||||
|
||||
NESTED_END
|
||||
|
||||
MEND
|
||||
|
||||
SgemmKernelNeonFunction Zero
|
||||
SgemmKernelNeonFunction Add
|
||||
|
||||
END
|
||||
Vendored
+1160
File diff suppressed because it is too large
Load Diff
+43
@@ -0,0 +1,43 @@
|
||||
// Shim for ORT's core/common/common.h. Provides the small subset of
|
||||
// macros that MLAS's q4_dq.cpp / q4common.h use: ORT_ENFORCE and ORT_THROW.
|
||||
// Upstream's common.h pulls in logging, status, exceptions, and lots more
|
||||
// — none of which MLAS itself needs. We map both macros to throwing
|
||||
// std::runtime_error since MLAS is built without exception-disable in
|
||||
// our CMake (see mlasi.h's MLAS_NO_EXCEPTION guard).
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <sstream>
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
// Concatenate stream-like arguments into a single string. Supports the
|
||||
// same `operator<<` chain that ORT_ENFORCE uses for its diagnostic.
|
||||
template <typename... Args>
|
||||
inline std::string MlasShimMakeMessage(const Args&... args) {
|
||||
std::ostringstream oss;
|
||||
using expand = int[];
|
||||
(void)expand{0, ((void)(oss << args), 0)...};
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
} // namespace onnxruntime
|
||||
|
||||
#define ORT_THROW(...) \
|
||||
do { \
|
||||
throw std::runtime_error( \
|
||||
::onnxruntime::MlasShimMakeMessage(__VA_ARGS__)); \
|
||||
} while (0)
|
||||
|
||||
#define ORT_ENFORCE(cond, ...) \
|
||||
do { \
|
||||
if (!(cond)) { \
|
||||
throw std::runtime_error( \
|
||||
::onnxruntime::MlasShimMakeMessage( \
|
||||
"ORT_ENFORCE(" #cond ") failed: ", ##__VA_ARGS__)); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
#define ORT_NOT_IMPLEMENTED(...) ORT_THROW("not implemented: ", ##__VA_ARGS__)
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
// Shim for ORT's core/common/narrow.h — used by the vendored MLAS (cast.cpp).
|
||||
// Upstream provides a checked narrowing cast a la gsl::narrow. The MLAS
|
||||
// translation units here only #include the header; they do not actually
|
||||
// invoke narrow<T>(...). We provide a minimal definition anyway so the file
|
||||
// compiles cleanly and any future MLAS update that does call narrow keeps
|
||||
// working.
|
||||
//
|
||||
// This file is intentionally tiny so OpenCV can keep a stable shim while
|
||||
// upstream MLAS evolves.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <stdexcept>
|
||||
#include <type_traits>
|
||||
|
||||
namespace onnxruntime {
|
||||
|
||||
template <typename T, typename U>
|
||||
constexpr T narrow(U u) {
|
||||
static_assert(std::is_arithmetic<T>::value && std::is_arithmetic<U>::value,
|
||||
"narrow<T>(U): T and U must be arithmetic types");
|
||||
const T t = static_cast<T>(u);
|
||||
if (static_cast<U>(t) != u ||
|
||||
((t < T{}) != (u < U{}))) {
|
||||
throw std::runtime_error("onnxruntime::narrow: narrowing failed");
|
||||
}
|
||||
return t;
|
||||
}
|
||||
|
||||
} // namespace onnxruntime
|
||||
Vendored
+27
@@ -0,0 +1,27 @@
|
||||
/*++
|
||||
|
||||
Copyright 2025 FUJITSU LIMITED
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
erf_neon_fp16.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module contains the procedure prototypes for the ERF NEON FP16 intrinsics.
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <arm_neon.h>
|
||||
|
||||
#include "mlasi.h"
|
||||
#include "fp16_common.h"
|
||||
#include "softmax_kernel_neon.h"
|
||||
#include <cstring>
|
||||
|
||||
void MlasNeonErfFP16Kernel(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N);
|
||||
Vendored
+167
@@ -0,0 +1,167 @@
|
||||
#include <numeric>
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
void
|
||||
MlasFlashAttentionThreaded(
|
||||
void* argptr,
|
||||
std::ptrdiff_t thread_id
|
||||
)
|
||||
{
|
||||
const MlasFlashAttentionThreadedArgs* args = reinterpret_cast<MlasFlashAttentionThreadedArgs*>(argptr);
|
||||
ptrdiff_t q_block_size = static_cast<ptrdiff_t>(args->q_block_size);
|
||||
ptrdiff_t kv_block_size = static_cast<ptrdiff_t>(args->kv_block_size);
|
||||
ptrdiff_t batch_size = static_cast<ptrdiff_t>(args->batch_size);
|
||||
ptrdiff_t num_heads = static_cast<ptrdiff_t>(args->num_heads);
|
||||
ptrdiff_t q_sequence_length = static_cast<ptrdiff_t>(args->q_sequence_length);
|
||||
ptrdiff_t kv_sequence_length = static_cast<ptrdiff_t>(args->kv_sequence_length);
|
||||
ptrdiff_t qk_head_size = static_cast<ptrdiff_t>(args->qk_head_size);
|
||||
ptrdiff_t v_head_size = static_cast<ptrdiff_t>(args->v_head_size);
|
||||
float* buffer = args->buffer;
|
||||
ptrdiff_t buffer_size_per_thread = static_cast<ptrdiff_t>(args->buffer_size_per_thread);
|
||||
ptrdiff_t thread_count = static_cast<ptrdiff_t>(args->thread_count);
|
||||
const float* query = args->query;
|
||||
const float* key = args->key;
|
||||
const float* value = args->value;
|
||||
float* output = args->output;
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64)
|
||||
auto&& mlas_platform = GetMlasPlatform();
|
||||
#endif
|
||||
|
||||
ptrdiff_t q_chunk_count = (q_sequence_length + (q_block_size - 1)) / q_block_size;
|
||||
|
||||
ptrdiff_t task_start = 0;
|
||||
ptrdiff_t task_end = 0;
|
||||
ptrdiff_t total_task_count = batch_size * num_heads * q_chunk_count;
|
||||
ptrdiff_t quotient = total_task_count / thread_count;
|
||||
ptrdiff_t remainder = total_task_count % thread_count;
|
||||
if (thread_id < remainder) {
|
||||
task_start = (quotient + 1) * thread_id;
|
||||
task_end = task_start + quotient + 1;
|
||||
} else {
|
||||
task_start = quotient * thread_id + remainder;
|
||||
task_end = task_start + quotient;
|
||||
}
|
||||
|
||||
for (ptrdiff_t task_index = task_start; task_index < task_end; ++task_index) {
|
||||
ptrdiff_t batch_idx = task_index;
|
||||
ptrdiff_t q_idx = (batch_idx % q_chunk_count) * q_block_size;
|
||||
batch_idx /= q_chunk_count;
|
||||
ptrdiff_t head_idx = batch_idx % num_heads;
|
||||
batch_idx /= num_heads;
|
||||
|
||||
char* buffer_current_thread = reinterpret_cast<char*>(buffer) + thread_id * buffer_size_per_thread;
|
||||
float* l = reinterpret_cast<float*>(buffer_current_thread);
|
||||
float* m = l + q_block_size;
|
||||
for (ptrdiff_t t = 0; t < q_block_size; ++t) {
|
||||
m[t] = std::numeric_limits<float>::lowest();
|
||||
}
|
||||
float* intermediate = m + q_block_size;
|
||||
float* temp_output = intermediate + q_block_size * kv_block_size;
|
||||
float negmax = 0;
|
||||
|
||||
for (ptrdiff_t ir = 0; ir < kv_sequence_length; ir += kv_block_size) {
|
||||
/*
|
||||
S = Q[batch_idx, head_idx, q_idx:q_idx+q_block_size, :] * (K[batch_idx, head_idx, ir:ir+kv_block_size, :]).T
|
||||
old_m = m
|
||||
m = max(m, rowmax(S))
|
||||
diff = old_m - m
|
||||
S = exp(S - m)
|
||||
l = exp(diff) * l + rowsum(S)
|
||||
O = diag(exp(diff)) * O + S * V[batch_idx, head_idx, ir:ir+kv_block_size, :]
|
||||
*/
|
||||
ptrdiff_t h = batch_idx * num_heads + head_idx;
|
||||
const float* inputQ = query + (h * q_sequence_length + q_idx) * qk_head_size;
|
||||
const float* inputK = key + (h * kv_sequence_length + ir) * qk_head_size;
|
||||
const float* inputV = value + (h * kv_sequence_length + ir) * v_head_size;
|
||||
|
||||
size_t row_size_q_capped = static_cast<size_t>(std::min(q_block_size, q_sequence_length - q_idx));
|
||||
size_t row_size_kv_capped = static_cast<size_t>(std::min(kv_block_size, kv_sequence_length - ir));
|
||||
|
||||
MlasSgemmOperation(CBLAS_TRANSPOSE::CblasNoTrans,
|
||||
CBLAS_TRANSPOSE::CblasTrans,
|
||||
row_size_q_capped,
|
||||
row_size_kv_capped,
|
||||
static_cast<size_t>(qk_head_size),
|
||||
args->scale,
|
||||
inputQ,
|
||||
static_cast<size_t>(qk_head_size),
|
||||
inputK,
|
||||
static_cast<size_t>(qk_head_size),
|
||||
0.0f,
|
||||
intermediate,
|
||||
row_size_kv_capped);
|
||||
|
||||
for (ptrdiff_t irow = 0; irow < static_cast<ptrdiff_t>(row_size_q_capped); ++irow) {
|
||||
float* p = intermediate + irow * row_size_kv_capped;
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64)
|
||||
float rowmax = mlas_platform.ReduceMaximumF32Kernel(p, row_size_kv_capped);
|
||||
#else
|
||||
float rowmax = MlasReduceMaximumF32Kernel(p, row_size_kv_capped);
|
||||
#endif
|
||||
float m_diff = m[irow];
|
||||
m[irow] = std::max(m[irow], rowmax); // new m
|
||||
negmax = -m[irow];
|
||||
m_diff -= m[irow]; // old - new (less than 0)
|
||||
|
||||
#if defined(MLAS_TARGET_AMD64)
|
||||
float rowsum = mlas_platform.ComputeSumExpF32Kernel(p, p, row_size_kv_capped, &negmax);
|
||||
#else
|
||||
float rowsum = MlasComputeSumExpF32Kernel(p, p, row_size_kv_capped, &negmax);
|
||||
#endif
|
||||
|
||||
// Note: for ir == 0, there is actually no need to calculate exp_diff
|
||||
if (ir != 0) {
|
||||
float exp_diff = std::exp(m_diff);
|
||||
l[irow] = exp_diff * l[irow] + rowsum;
|
||||
|
||||
for (ptrdiff_t icol = 0; icol < v_head_size; ++icol) {
|
||||
temp_output[irow * v_head_size + icol] = exp_diff * temp_output[irow * v_head_size + icol];
|
||||
}
|
||||
} else {
|
||||
l[irow] = rowsum;
|
||||
// When ir == 0, there is no need to scale the old result because it is zero.
|
||||
}
|
||||
}
|
||||
MlasSgemmOperation(CBLAS_TRANSPOSE::CblasNoTrans,
|
||||
CBLAS_TRANSPOSE::CblasNoTrans,
|
||||
row_size_q_capped,
|
||||
static_cast<size_t>(v_head_size),
|
||||
row_size_kv_capped,
|
||||
1.0f,
|
||||
intermediate,
|
||||
row_size_kv_capped,
|
||||
inputV,
|
||||
static_cast<size_t>(v_head_size),
|
||||
ir == 0 ? 0.0f : 1.0f,
|
||||
temp_output,
|
||||
static_cast<size_t>(v_head_size));
|
||||
}
|
||||
|
||||
float* output_row = output + ((batch_idx * q_sequence_length + q_idx) * num_heads + head_idx) * v_head_size;
|
||||
ptrdiff_t row_size_q_valid = std::min(q_block_size, q_sequence_length - q_idx);
|
||||
// TODO: leverage advanced instruction sets
|
||||
for (ptrdiff_t irow = 0; irow < row_size_q_valid; ++irow) {
|
||||
for (ptrdiff_t icol = 0; icol < v_head_size; ++icol) {
|
||||
output_row[icol] = temp_output[irow * v_head_size + icol] / l[irow];
|
||||
}
|
||||
output_row += num_heads * v_head_size;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasFlashAttention(
|
||||
MlasFlashAttentionThreadedArgs* args,
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
)
|
||||
{
|
||||
MlasExecuteThreaded(
|
||||
MlasFlashAttentionThreaded,
|
||||
static_cast<void *>(args),
|
||||
static_cast<std::ptrdiff_t>(args->thread_count),
|
||||
ThreadPool);
|
||||
}
|
||||
Vendored
+31
@@ -0,0 +1,31 @@
|
||||
/*++
|
||||
|
||||
Copyright 2025 FUJITSU LIMITED
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
gelu_neon_fp16.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module contains Gelu helper functions .
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "fp16_common.h"
|
||||
#include "erf_neon_fp16.h"
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasNeonGeluFP16Kernel(
|
||||
const MLAS_FP16* input,
|
||||
MLAS_FP16* output,
|
||||
MLAS_FP16* temp,
|
||||
size_t count,
|
||||
MLAS_GELU_ALGORITHM algo
|
||||
);
|
||||
+236
@@ -0,0 +1,236 @@
|
||||
//
|
||||
// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates <open-source-office@arm.com>
|
||||
//
|
||||
// SPDX-License-Identifier: MIT
|
||||
//
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../mlasi.h"
|
||||
#include <iostream>
|
||||
|
||||
// Fix to ensure compatibility with MSVC build
|
||||
#if defined(_MSC_VER)
|
||||
#define RESTRICT __restrict
|
||||
#else
|
||||
#define RESTRICT __restrict__
|
||||
#endif
|
||||
|
||||
// Logging macros.
|
||||
#ifndef KLEIDIAI_DEBUG_LOGGING
|
||||
#define KLEIDIAI_DEBUG_LOGGING 0
|
||||
#endif
|
||||
#ifndef KLEIDIAI_KERNEL_LOGGING
|
||||
#define KLEIDIAI_KERNEL_LOGGING 0
|
||||
#endif
|
||||
|
||||
#if KLEIDIAI_DEBUG_LOGGING ||KLEIDIAI_KERNEL_LOGGING
|
||||
#define KLEIDIAI_LOG(tag, msg) \
|
||||
do { \
|
||||
std::cout << "[KLEIDIAI " << tag << "]: " << __FILE__ << " : " << __LINE__ << " : " << msg << std::endl; \
|
||||
} while(false)
|
||||
#endif
|
||||
|
||||
// General logging. "tag" is expected to qualify the type of message.
|
||||
#if KLEIDIAI_DEBUG_LOGGING
|
||||
// General debug messages.
|
||||
#define KLEIDIAI_DEBUG_LOG(msg) KLEIDIAI_LOG("DEBUG", msg)
|
||||
#else
|
||||
#define KLEIDIAI_DEBUG_LOG(msg)
|
||||
#endif
|
||||
|
||||
#if KLEIDIAI_KERNEL_LOGGING
|
||||
// Messages specifically written before a call to kai_run.
|
||||
// Note: In cases where a kernel is called in multiple threads, for example MlasTrySimpleParallel,
|
||||
// the output order can be inconsistient. The solution is to set the intra-node thread size to 1.
|
||||
// If using onnxruntime_perf_test this is done with "--x 1".
|
||||
#define KLEIDIAI_KERNEL_LOG(kernel_name) KLEIDIAI_LOG("KERNEL", kernel_name)
|
||||
#else
|
||||
#define KLEIDIAI_KERNEL_LOG(msg)
|
||||
#endif
|
||||
|
||||
namespace ArmKleidiAI {
|
||||
|
||||
// By default we should try for SME2 first before falling back to SME.
|
||||
inline const bool UseSME2 = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME2();
|
||||
inline const bool UseSME = MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME();
|
||||
inline const std::string_view vendor_name = MLAS_CPUIDINFO::GetCPUIDInfo().GetCPUVendor();
|
||||
|
||||
// Buffer packing routines.
|
||||
//
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasGemmPackBSize(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t N,
|
||||
size_t K
|
||||
);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasGemmPackB(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t N,
|
||||
size_t K,
|
||||
const float* B,
|
||||
size_t ldb,
|
||||
void* PackedB
|
||||
);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasGemvBatch(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t M,
|
||||
size_t N,
|
||||
size_t K,
|
||||
const MLAS_SGEMM_DATA_PARAMS* Data,
|
||||
size_t BatchSize
|
||||
);
|
||||
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasGemmBatch(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t M,
|
||||
size_t N,
|
||||
size_t K,
|
||||
const MLAS_SGEMM_DATA_PARAMS* Data,
|
||||
size_t BatchSize,
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
);
|
||||
|
||||
#if defined(__aarch64__) && defined(__linux__)
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSBGemmPackBSize(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t N,
|
||||
size_t K
|
||||
);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasSBGemmPackB(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t N,
|
||||
size_t K,
|
||||
const float* B,
|
||||
size_t ldb,
|
||||
void* PackedB
|
||||
);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasSBGemmBatch(
|
||||
CBLAS_TRANSPOSE TransA,
|
||||
CBLAS_TRANSPOSE TransB,
|
||||
size_t M,
|
||||
size_t N,
|
||||
size_t K,
|
||||
const MLAS_SBGEMM_DATA_PARAMS* Data,
|
||||
size_t BatchSize,
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
);
|
||||
#endif
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasDynamicQGemmPackBSize(
|
||||
size_t N,
|
||||
size_t K
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasDynamicQGemmPackB(
|
||||
size_t N,
|
||||
size_t K,
|
||||
const int8_t* B,
|
||||
const float* Scales,
|
||||
const float* Bias,
|
||||
void* PackedB
|
||||
);
|
||||
|
||||
//pack symmetric quantized B and dynamic quantized A
|
||||
void
|
||||
MLASCALL
|
||||
MlasDynamicQGemmBatch(
|
||||
const MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS& Shape,
|
||||
const MLAS_GEMM_DYN_QUANT_DATA_PARAMS* DataParams,
|
||||
const size_t BatchN,
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasConvPrepare(MLAS_CONV_PARAMETERS* Parameters,
|
||||
size_t Dimensions,
|
||||
size_t BatchCount,
|
||||
size_t GroupCount,
|
||||
size_t InputChannels,
|
||||
const int64_t* InputShape,
|
||||
const int64_t* KernelShape,
|
||||
const int64_t* DilationShape,
|
||||
const int64_t* Padding,
|
||||
const int64_t* StrideShape,
|
||||
const int64_t* OutputShape,
|
||||
size_t FilterCount,
|
||||
const MLAS_ACTIVATION* Activation,
|
||||
size_t* WorkingBufferSize,
|
||||
float Beta,
|
||||
MLAS_THREADPOOL* ThreadPool);
|
||||
|
||||
bool
|
||||
MLASCALL
|
||||
MlasConv(
|
||||
const MLAS_CONV_PARAMETERS* Parameters,
|
||||
const float* Input,
|
||||
const float* Filter,
|
||||
const float* Bias,
|
||||
float* WorkingBuffer,
|
||||
float* Output,
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
);
|
||||
}
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine determines if a wraparound will occur when multiplying two size_t variables
|
||||
Uses __builtin_mul_overflow if available on the current system and if not falls back
|
||||
to a default implementation to check this wraparound.
|
||||
|
||||
Arguments:
|
||||
|
||||
a - Supplies the first number to be muliplied.
|
||||
|
||||
b - Supplies the second number to be muliplied.
|
||||
|
||||
out - pointer to a size_t which acts as the return value in success cases.
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns false if the operation was successful
|
||||
Returns true if wraparound of size_t was detected
|
||||
|
||||
--*/
|
||||
inline bool mul_overflow_size_t_builtin(size_t a, size_t b, size_t* out) {
|
||||
#if defined(__has_builtin)
|
||||
# if __has_builtin(__builtin_mul_overflow)
|
||||
return __builtin_mul_overflow(a, b, out);
|
||||
# endif
|
||||
#endif
|
||||
// Fallback to manual check if builtin not available
|
||||
if (b != 0 && a > SIZE_MAX / b) return true;
|
||||
if (out) *out = a * b;
|
||||
return false;
|
||||
}
|
||||
+33
@@ -0,0 +1,33 @@
|
||||
/*++
|
||||
|
||||
Copyright (C) 2023 Loongson Technology Corporation Limited. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelLasx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses LASX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "SgemmKernelCommon.h"
|
||||
#include "FgemmKernelLasxCommon.h"
|
||||
|
||||
|
||||
.text
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelLasxFunction MlasGemmFloatKernelLasx
|
||||
|
||||
.end
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
/*++
|
||||
|
||||
Copyright (C) 2023 Loongson Technology Corporation Limited. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelLsx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses Lsx instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "FgemmKernelLsxCommon.h"
|
||||
|
||||
FGEMM_TYPED_INSTRUCTION(vfadd, vfadd.s)
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for a 16xN block of the output matrix.
|
||||
|
||||
Arguments:
|
||||
|
||||
RowCount - Supplies the number of rows to process.
|
||||
|
||||
VectorOffset - Supplies the byte offset from matrix B to fetch elements.
|
||||
|
||||
Shuffle - Supplies the shuffle mask to extract the element from matrix A.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
a1 - Supplies the address into the matrix B data.
|
||||
|
||||
vr0-vr1 - Supplies up to four elements loaded from matrix A and matrix A
|
||||
plus one row.
|
||||
|
||||
vr8-vr15 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockSseBy16 RowCount, VectorOffset, Shuffle
|
||||
vld $vr4, $a1, \VectorOffset
|
||||
vld $vr5, $a1, \VectorOffset + 16
|
||||
vreplvei.w $vr2, $vr0, \Shuffle
|
||||
.if \RowCount\() == 2
|
||||
vreplvei.w $vr3, $vr1, \Shuffle
|
||||
vmove $vr6, $vr4
|
||||
vmove $vr7, $vr5
|
||||
.endif
|
||||
vfmadd.s $vr8, $vr4, $vr2, $vr8
|
||||
vfmadd.s $vr9, $vr5, $vr2, $vr9
|
||||
.if \RowCount\() == 2
|
||||
vfmadd.s $vr12, $vr6, $vr3, $vr12
|
||||
vfmadd.s $vr13, $vr7, $vr3, $vr13
|
||||
.endif
|
||||
vld $vr4, $a1, \VectorOffset + 32
|
||||
vld $vr5, $a1, \VectorOffset + 48
|
||||
.if \RowCount\() == 2
|
||||
vmove $vr6, $vr4
|
||||
vmove $vr7, $vr5
|
||||
.endif
|
||||
vfmadd.s $vr10, $vr4, $vr2, $vr10
|
||||
vfmadd.s $vr11, $vr5, $vr2, $vr11
|
||||
.if \RowCount\() == 2
|
||||
vfmadd.s $vr14, $vr6, $vr3, $vr14
|
||||
vfmadd.s $vr15, $vr7, $vr3, $vr15
|
||||
.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:
|
||||
|
||||
a0 - Supplies the address of matrix A.
|
||||
|
||||
a1 - Supplies the address of matrix B.
|
||||
|
||||
t8 - Supplies the address of matrix A.
|
||||
|
||||
a5 - Supplies the number of columns from matrix B and matrix C to iterate
|
||||
over.
|
||||
|
||||
a2 - Supplies the address of matrix C.
|
||||
|
||||
a3 - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
t7 - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
t5 - Supplies the length in bytes of a row from matrix C.
|
||||
|
||||
s3 - Stores the ZeroMode argument from the stack frame.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ProcessCountM RowCount, Fallthrough
|
||||
.LProcessNextColumnLoop16xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vxor.v $vr8, $vr8,$vr8"
|
||||
EmitIfCountGE \RowCount\(), 1, "vxor.v $vr9, $vr9,$vr9"
|
||||
EmitIfCountGE \RowCount\(), 1, "vxor.v $vr10, $vr10,$vr10"
|
||||
EmitIfCountGE \RowCount\(), 1, "vxor.v $vr11, $vr11,$vr11"
|
||||
EmitIfCountGE \RowCount\(), 2, "vxor.v $vr12, $vr12,$vr12"
|
||||
EmitIfCountGE \RowCount\(), 2, "vxor.v $vr13, $vr13,$vr13"
|
||||
EmitIfCountGE \RowCount\(), 2, "vxor.v $vr14, $vr14,$vr14"
|
||||
EmitIfCountGE \RowCount\(), 2, "vxor.v $vr15, $vr15,$vr15"
|
||||
move $t8, $a3
|
||||
li.d $s0, 4
|
||||
blt $t8, $s0, .LProcessRemaining16xNBlocks\@
|
||||
.LCompute16xNBlockBy4Loop\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vld $vr0, $a0, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vldx $vr1, $a0, $t0" #second line of A
|
||||
ComputeBlockSseBy16 2, 0, 0x0
|
||||
ComputeBlockSseBy16 2, 16*4, 0x1
|
||||
addi.d $a1, $a1, 32*4 # advance matrix B by 32 columns
|
||||
ComputeBlockSseBy16 2, 0, 0x2
|
||||
ComputeBlockSseBy16 2, 16*4, 0x3
|
||||
addi.d $a1, $a1, 32*4 # advance matrix B by 32 columns
|
||||
addi.d $a0, $a0, 4*4 # advance matrix A by 4 columns
|
||||
addi.d $t8, $t8, -4
|
||||
li.d $s0, 4 #check matrix A remaining less than 4
|
||||
bge $t8, $s0, .LCompute16xNBlockBy4Loop\@
|
||||
|
||||
.LProcessRemaining16xNBlocks\@:
|
||||
beqz $t8, .LOutput16xNBlock\@
|
||||
|
||||
.LCompute16xNBlockBy1Loop\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "ld.w $s0, $a0, 0"
|
||||
EmitIfCountGE \RowCount\(), 1, "vinsgr2vr.w $vr0, $s0, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "ldx.w $s0,$a0, $t0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vinsgr2vr.w $vr1,$s0, 0"
|
||||
ComputeBlockSseBy16 2, 0, 0x00
|
||||
addi.d $a1, $a1, 16*4 #advance matrix B by 16 columns
|
||||
addi.d $a0, $a0, 1*4 #advance matrix A by 1 column
|
||||
addi.d $t8, $t8, -1
|
||||
bnez $t8, .LCompute16xNBlockBy1Loop\@
|
||||
|
||||
.LOutput16xNBlock\@:
|
||||
movfr2gr.s $s0, $f24
|
||||
vreplgr2vr.w $vr2, $s0
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmul.s $vr8,$vr8,$vr2"
|
||||
# multiply by alpha
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmul.s $vr9,$vr9,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmul.s $vr10,$vr10,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmul.s $vr11,$vr11,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmul.s $vr12,$vr12,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmul.s $vr13,$vr13,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmul.s $vr14,$vr14,$vr2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmul.s $vr15,$vr15,$vr2"
|
||||
li.d $s0, 16
|
||||
blt $a5, $s0, .LOutputPartial16xNBlock\@
|
||||
sub.d $a5, $a5, $s0
|
||||
AccumulateAndStoreBlock \RowCount\(), 4
|
||||
addi.d $a2, $a2, 16*4 # advance matrix C by 16 columns
|
||||
move $a0, $t1 # reload matrix A
|
||||
bnez $a5, .LProcessNextColumnLoop16xN\@
|
||||
b .LExitKernel
|
||||
|
||||
//
|
||||
// Output a partial 16xN block to the matrix.
|
||||
//
|
||||
|
||||
.LOutputPartial16xNBlock\@:
|
||||
li.d $s0, 4
|
||||
blt $a5, $s0, .LOutputPartialLessThan4xNBlock\@
|
||||
li.d $s0, 8
|
||||
blt $a5, $s0, .LOutputPartialLessThan8xNBlock\@
|
||||
li.d $s0, 12
|
||||
blt $a5, $s0, .LOutputPartialLessThan12xNBlock\@
|
||||
AccumulateAndStoreBlock \RowCount\(), 3
|
||||
andi $a5, $a5, 3
|
||||
beqz $a5, .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "vmove $vr8, $vr11"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "vmove $vr12, $vr15"
|
||||
addi.d $a2, $a2,12*4 # advance matrix C by 12 columns
|
||||
b .LOutputPartialLessThan4xNBlock\@
|
||||
|
||||
.LOutputPartialLessThan12xNBlock\@:
|
||||
AccumulateAndStoreBlock \RowCount\(), 2
|
||||
andi $a5, $a5, 3
|
||||
beqz $a5, .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "vmove $vr8, $vr10"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "vmove $vr12, $vr14"
|
||||
addi.d $a2, $a2,8*4 # advance matrix C by 8 columns
|
||||
b .LOutputPartialLessThan4xNBlock\@
|
||||
|
||||
.LOutputPartialLessThan8xNBlock\@:
|
||||
AccumulateAndStoreBlock \RowCount\(), 1
|
||||
andi $a5, $a5, 3
|
||||
beqz $a5, .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "vmove $vr8, $vr9"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "vmove $vr12, $vr13"
|
||||
addi.d $a2, $a2, 4*4 # advance matrix C by 4 columns
|
||||
|
||||
.LOutputPartialLessThan4xNBlock\@:
|
||||
andi $s0, $a5, 2
|
||||
beqz $s0, .LOutputPartial1xNBlock\@
|
||||
and $s0, $t5, $t5 # ZeroMode?
|
||||
bnez $s0, .LSkipAccumulateOutput2xN\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vxor.v $vr0, $vr0, $vr0"
|
||||
EmitIfCountGE \RowCount\(), 1, "ld.d $s0, $a2, 0"
|
||||
EmitIfCountGE \RowCount\(), 1, "vinsgr2vr.d $vr0, $s0, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vxor.v $vr1, $vr1, $vr1"
|
||||
EmitIfCountGE \RowCount\(), 2, "ldx.d $s0, $a2, $t6"
|
||||
EmitIfCountGE \RowCount\(), 2, "vinsgr2vr.d $vr1, $s0, 0"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfadd.s $vr8, $vr8, $vr0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfadd.s $vr12, $vr12, $vr1"
|
||||
|
||||
.LSkipAccumulateOutput2xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vstelm.d $vr8, $a2, 0, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vpickve2gr.d $s0, $vr12, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "stx.d $s0, $a2, $t6"
|
||||
andi $s0, $a5, 1
|
||||
beqz $s0, .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "vpermi.w $vr8, $vr8, 0xee"
|
||||
# shift third element down
|
||||
EmitIfCountGE \RowCount\(), 2, "vpermi.w $vr12, $vr12, 0xee"
|
||||
addi.d $a2, $a2, 2*4 # advance matrix C by 2 columns
|
||||
|
||||
.LOutputPartial1xNBlock\@:
|
||||
and $s0, $t5, $t5 # ZeroMode?
|
||||
bnez $s0, .LSkipAccumulateOutput1xN\@
|
||||
|
||||
EmitIfCountGE \RowCount\(), 1, "fld.s $f16, $a2, 0"
|
||||
EmitIfCountGE \RowCount\(), 1, "fadd.s $f8, $f16, $f8"
|
||||
EmitIfCountGE \RowCount\(), 2, "fldx.s $f17, $a2, $t6"
|
||||
EmitIfCountGE \RowCount\(), 2, "fadd.s $f12, $f12, $f17"
|
||||
|
||||
.LSkipAccumulateOutput1xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "fst.s $f8, $a2, 0"
|
||||
EmitIfCountGE \RowCount\(), 2, "fstx.s $f12, $a2, $t6"
|
||||
.ifb \Fallthrough\()
|
||||
b .LExitKernel
|
||||
.endif
|
||||
.endm
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelLsxFunction MlasGemmFloatKernelLSX
|
||||
|
||||
.end
|
||||
@@ -0,0 +1,89 @@
|
||||
/*++
|
||||
|
||||
Copyright (C) 2023 Loongson Technology Corporation Limited. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmTransposePackB16x4LSX.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements routines for packing buffers for the single precision
|
||||
matrix/matrix multiply operation (SGEMM).
|
||||
|
||||
This implementation uses Lsx instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine transposes elements from the source matrix to the destination
|
||||
packed buffer.
|
||||
|
||||
4 columns of 16 rows from the source matrix are transposed to 16 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
D (a0) - Supplies the address of the destination packed buffer.
|
||||
|
||||
B (a1) - Supplies the address of the source matrix.
|
||||
|
||||
ldb (a2) - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmTransposePackB16x4LSX
|
||||
addi.d $sp, $sp, -64
|
||||
st.d $s0, $sp, 0*8
|
||||
st.d $s1, $sp, 1*8
|
||||
slli.d $a2, $a2, 2 # convert ldb to bytes
|
||||
ori $a3, $zero, 4 # transpose four 4x4 blocks
|
||||
vxor.v $vr7, $vr7, $vr7
|
||||
.LTransposeBlockLoop:
|
||||
slli.d $s0, $a2, 1
|
||||
add.d $s1, $a1, $s0
|
||||
vld $vr0, $a1, 0
|
||||
vldx $vr1, $a1, $a2
|
||||
vld $vr2, $s1, 0
|
||||
vldx $vr3, $s1, $a2
|
||||
|
||||
vor.v $vr4, $vr0, $vr7
|
||||
vilvl.w $vr4, $vr1, $vr4
|
||||
vilvh.w $vr0, $vr1, $vr0
|
||||
vor.v $vr5, $vr2, $vr7
|
||||
vilvl.w $vr5, $vr3, $vr5
|
||||
vilvh.w $vr2, $vr3, $vr2
|
||||
vor.v $vr1, $vr4, $vr7
|
||||
vilvl.d $vr1, $vr5, $vr1
|
||||
vilvh.d $vr4, $vr5, $vr4
|
||||
vor.v $vr3, $vr0, $vr7
|
||||
vilvl.d $vr3, $vr2, $vr3
|
||||
vilvh.d $vr0, $vr2, $vr0
|
||||
vst $vr1, $a0, 0
|
||||
vst $vr4, $a0, 0x40
|
||||
vst $vr3, $a0, 0x80
|
||||
vst $vr0, $a0, 0xc0
|
||||
addi.d $a0, $a0, 0x10
|
||||
slli.d $s0, $a2, 1
|
||||
add.d $a1, $s0, $s1
|
||||
addi.d $a3, $a3, -1
|
||||
bnez $a3, .LTransposeBlockLoop
|
||||
ld.d $s0, $sp, 0*8
|
||||
ld.d $s1, $sp, 1*8
|
||||
addi.d $sp, $sp, 64
|
||||
jr $ra
|
||||
|
||||
.end
|
||||
@@ -0,0 +1,126 @@
|
||||
/*++
|
||||
|
||||
Copyright (C) 2023 Loongson Technology Corporation Limited. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmTransposePackB16x4Lasx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements routines for packing buffers for the single precision
|
||||
matrix/matrix multiply operation (SGEMM).
|
||||
|
||||
This implementation uses Lasx instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
4 columns of 8 rows from the source matrix are transposed to 8 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
StoreOffset - Supplies the relative byte offset into the destination packed
|
||||
buffer.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
a0 - Supplies the address of the destination packed buffer.
|
||||
|
||||
a1 - Supplies the address of the source matrix.
|
||||
|
||||
a2 - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
--*/
|
||||
|
||||
.macro TransposePackB8x4BlockLasx StoreOffset
|
||||
|
||||
//
|
||||
// Load 4 columns from 8 rows of the source matrix into the lower and upper
|
||||
// halves of 4 XR registers.
|
||||
//
|
||||
|
||||
add.d $t0, $a2, $a2
|
||||
add.d $t6, $a1, $t0
|
||||
vld $vr0, $a1, 0
|
||||
vldx $vr1, $a1, $a2
|
||||
add.d $t0, $a2, $a2
|
||||
add.d $a1, $t6, $t0
|
||||
vld $vr2, $t6, 0
|
||||
vldx $vr3, $t6, $a2
|
||||
add.d $t0, $a2, $a2
|
||||
add.d $t6, $a1, $t0
|
||||
|
||||
vld $vr4, $a1, 0
|
||||
xvpermi.q $xr0, $xr4, 0x2
|
||||
vldx $vr5, $a1, $a2
|
||||
xvpermi.q $xr1, $xr5, 0x2
|
||||
vld $vr4, $t6, 0
|
||||
xvpermi.q $xr2, $xr4, 0x2
|
||||
vldx $vr5, $t6, $a2
|
||||
xvpermi.q $xr3, $xr5, 0x2
|
||||
|
||||
//
|
||||
// Transpose the lower and upper halves of the 4 XR registers as two 4x4
|
||||
// matrices and store the output to the destination packed buffer.
|
||||
//
|
||||
|
||||
xvilvl.w $xr4, $xr1, $xr0
|
||||
xvilvh.w $xr5, $xr1, $xr0
|
||||
xvilvl.w $xr0, $xr3, $xr2
|
||||
xvilvh.w $xr1, $xr3, $xr2
|
||||
xvilvl.d $xr2, $xr0, $xr4
|
||||
xvilvh.d $xr3, $xr0, $xr4
|
||||
xvst $xr2, $a0, \StoreOffset\()
|
||||
xvst $xr3, $a0, 0x40+\StoreOffset\()
|
||||
xvilvl.d $xr0, $xr1, $xr5
|
||||
xvilvh.d $xr4, $xr1, $xr5
|
||||
xvst $xr0, $a0, 0x80+\StoreOffset\()
|
||||
xvst $xr4, $a0, 0xc0+\StoreOffset\()
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine transposes elements from the source matrix to the destination
|
||||
packed buffer.
|
||||
|
||||
4 columns of 16 rows from the source matrix are transposed to 16 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
D (a0) - Supplies the address of the destination packed buffer.
|
||||
|
||||
B (a1) - Supplies the address of the source matrix.
|
||||
|
||||
ldb (a2) - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmTransposePackB16x4Lasx
|
||||
|
||||
slli.d $a2, $a2, 2 # convert ldb to bytes
|
||||
TransposePackB8x4BlockLasx 0*4
|
||||
add.d $t0, $a2, $a2
|
||||
add.d $a1, $t0, $t6
|
||||
TransposePackB8x4BlockLasx 8*4
|
||||
jr $ra
|
||||
|
||||
.end
|
||||
+144
@@ -0,0 +1,144 @@
|
||||
/*++
|
||||
|
||||
Copyright (C) 2023 Loongson Technology Corporation Limited. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
asmmacro.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements common macros for the assembly modules.
|
||||
|
||||
--*/
|
||||
|
||||
#define C_UNDERSCORE(symbol) symbol
|
||||
|
||||
.macro vmove dst src
|
||||
vand.v \dst, \src, \src
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the assembler directives to annotate a new function.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name of the function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FUNCTION_ENTRY FunctionName
|
||||
.align 2
|
||||
.globl \FunctionName\()
|
||||
.type \FunctionName\(),@function
|
||||
\FunctionName\():
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro generates an optimization for "add reg,128" which can instead
|
||||
be encoded as "sub reg,-128" to reduce code size by using a signed 8-bit
|
||||
value.
|
||||
|
||||
Arguments:
|
||||
|
||||
Register - Supplies the register to be added to.
|
||||
|
||||
Immediate - Supplies the immediate to add to the register.
|
||||
|
||||
--*/
|
||||
|
||||
.macro add_immed Register, Immediate
|
||||
|
||||
.if (\Immediate\() != 128)
|
||||
addi.d \Register\(),\Register\(),\Immediate\()
|
||||
.else
|
||||
addi.d \Register\(),\Register\(),\Immediate\() # smaller encoding
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
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.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCountGE Count1, Value1, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro conditionally emits the statement if Count1 is greater than or
|
||||
equal to Value1 and Count2 is greater than or equal to Value2.
|
||||
|
||||
Arguments:
|
||||
|
||||
Count1 - Supplies the variable used in the comparison.
|
||||
|
||||
Value1 - Supplies the static used in the comparison.
|
||||
|
||||
Count2 - Supplies the variable used in the comparison.
|
||||
|
||||
Value2 - Supplies the static used in the comparison.
|
||||
|
||||
Statement - Supplies the statement to conditionally emit.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCount2GE Count1, Value1, Count2, Value2, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\()) && (\Count2\() >= \Value2\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the statement for each register listed in the register
|
||||
list. The statement can use RegItem to access the current register.
|
||||
|
||||
Arguments:
|
||||
|
||||
RegList - Supplies the list of registers.
|
||||
|
||||
Statement - Supplies the statement to emit.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitForEachRegister RegList, Statement
|
||||
|
||||
.irp RegItem, \RegList\()
|
||||
\Statement\()
|
||||
.endr
|
||||
|
||||
.endm
|
||||
Vendored
+3102
File diff suppressed because it is too large
Load Diff
Vendored
+1095
File diff suppressed because it is too large
Load Diff
+697
@@ -0,0 +1,697 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelPower.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#define PREFETCH_ADDR(addr) \
|
||||
asm volatile("dcbt 0, %0" ::"r"(addr) : "memory");
|
||||
|
||||
#include "SgemmKernelpower.h"
|
||||
extern "C" void
|
||||
PackAKernelPOWER10(__vector float* D, const float* A, size_t lda, size_t k, size_t RowCount);
|
||||
struct MlasSgemmBroadcastAElementsMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
ABroadcast[0] = vec_insert(A[Row * lda], ABroadcast[0], Row);
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeAElements(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount]
|
||||
)
|
||||
{
|
||||
__vector float a1,a2;
|
||||
a1 = vec_mergee (AElements[0], AElements[1]);
|
||||
a2 = vec_mergee (AElements[2], AElements[3]);
|
||||
ABroadcast[0] =vec_xxpermdi(a1,a2,0);
|
||||
ABroadcast[2] =vec_xxpermdi(a1,a2,3);
|
||||
a1 = vec_mergeo (AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo (AElements[2], AElements[3]);
|
||||
ABroadcast[1] =vec_xxpermdi(a1,a2,0);
|
||||
ABroadcast[3] =vec_xxpermdi(a1,a2,3);
|
||||
}
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeBlockMMA(
|
||||
__vector_quad acc[8],
|
||||
MLAS_FLOAT32X4 ABroadcast,
|
||||
MLAS_FLOAT32X4 A2Broadcast,
|
||||
const float* B,
|
||||
size_t CountM
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 BElements[4];
|
||||
typedef __vector unsigned char vec_t;
|
||||
|
||||
BElements[0] = MlasLoadFloat32x4(B);
|
||||
BElements[1] = MlasLoadFloat32x4(B + 4);
|
||||
BElements[2] = MlasLoadFloat32x4(B + 8);
|
||||
BElements[3] = MlasLoadFloat32x4(B + 12);
|
||||
__builtin_mma_xvf32gerpp (&acc[0], reinterpret_cast<vec_t>(ABroadcast), reinterpret_cast<vec_t>(BElements[0]));
|
||||
__builtin_mma_xvf32gerpp (&acc[1], reinterpret_cast<vec_t>(ABroadcast), reinterpret_cast<vec_t>(BElements[1]));
|
||||
__builtin_mma_xvf32gerpp (&acc[2], reinterpret_cast<vec_t>(ABroadcast), reinterpret_cast<vec_t>(BElements[2]));
|
||||
__builtin_mma_xvf32gerpp (&acc[3], reinterpret_cast<vec_t>(ABroadcast), reinterpret_cast<vec_t>(BElements[3]));
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_xvf32gerpp (&acc[4], reinterpret_cast<vec_t>(A2Broadcast), reinterpret_cast<vec_t>(BElements[0]));
|
||||
__builtin_mma_xvf32gerpp (&acc[5], reinterpret_cast<vec_t>(A2Broadcast), reinterpret_cast<vec_t>(BElements[1]));
|
||||
__builtin_mma_xvf32gerpp (&acc[6], reinterpret_cast<vec_t>(A2Broadcast), reinterpret_cast<vec_t>(BElements[2]));
|
||||
__builtin_mma_xvf32gerpp (&acc[7], reinterpret_cast<vec_t>(A2Broadcast), reinterpret_cast<vec_t>(BElements[3]));
|
||||
}
|
||||
}
|
||||
template<size_t VectorCount>
|
||||
struct MlasSgemmStoreVectorMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Result[4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 *rowC;
|
||||
if (ZeroMode) {
|
||||
rowC = reinterpret_cast<MLAS_FLOAT32X4 *>(&C[Row * ldc + VectorCount]);
|
||||
rowC[0] = Result[Row] * AlphaBroadcast;
|
||||
} else {
|
||||
rowC = reinterpret_cast<MLAS_FLOAT32X4 *>(&C[Row * ldc + VectorCount]);
|
||||
rowC[0] += Result[Row] * AlphaBroadcast;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaTrailingMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmStoreScalarMMA
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc + Lane;
|
||||
float Value = Accumulators[Row][Lane];
|
||||
if (!ZeroMode) {
|
||||
Value += *c;
|
||||
}
|
||||
|
||||
*c = Value;
|
||||
}
|
||||
};
|
||||
|
||||
template <size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmMMAProcessCount(
|
||||
__vector float* Pa,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountM,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
do {
|
||||
__vector float* pa1 = Pa;
|
||||
size_t k = CountK;
|
||||
|
||||
MLAS_FLOAT32X4 Accumulators[2][RowCount] = {{0}};
|
||||
MLAS_FLOAT32X4 Result[RowCount];
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount] = {0};
|
||||
__vector_quad acc[8];
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
__builtin_mma_xxsetaccz(&acc[0]);
|
||||
__builtin_mma_xxsetaccz(&acc[1]);
|
||||
__builtin_mma_xxsetaccz(&acc[2]);
|
||||
__builtin_mma_xxsetaccz(&acc[3]);
|
||||
__builtin_mma_xxsetaccz(&acc[4]);
|
||||
__builtin_mma_xxsetaccz(&acc[5]);
|
||||
__builtin_mma_xxsetaccz(&acc[6]);
|
||||
__builtin_mma_xxsetaccz(&acc[7]);
|
||||
|
||||
//
|
||||
// Compute the output block.
|
||||
//
|
||||
while (k >= 8) {
|
||||
if (CountM == 8) {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], pa1[4], B, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[1], pa1[5], B + 16, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[2], pa1[6], B + 32, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[3], pa1[7], B + 48, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[8], pa1[12], B + 64, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[9], pa1[13], B + 80, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[10], pa1[14], B + 96, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[11], pa1[15], B + 112, CountM);
|
||||
B += 128;
|
||||
pa1 += 16;
|
||||
k -= 8;
|
||||
} else {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], ABroadcast[0], B, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[1], ABroadcast[1], B + 16, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[2], ABroadcast[2], B + 32, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[3], ABroadcast[3], B + 48, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[4], ABroadcast[0], B + 64, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[5], ABroadcast[1], B + 80, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[6], ABroadcast[2], B + 96, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[7], ABroadcast[3], B + 112, CountM);
|
||||
B += 128;
|
||||
pa1 += 8;
|
||||
k -= 8;
|
||||
}
|
||||
}
|
||||
|
||||
while (k >= 4) {
|
||||
if (CountM == 8) {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], pa1[4], B, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[1], pa1[5], B + 16, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[2], pa1[6], B + 32, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[3], pa1[7], B + 48, CountM);
|
||||
B += 16 * 4;
|
||||
pa1 += 8;
|
||||
k -= 4;
|
||||
} else {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], ABroadcast[0], B, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[1], ABroadcast[1], B + 16, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[2], ABroadcast[2], B + 32, CountM);
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[3], ABroadcast[3], B + 48, CountM);
|
||||
B += 16 * 4;
|
||||
pa1 += 4;
|
||||
k -= 4;
|
||||
}
|
||||
}
|
||||
while (k > 0) {
|
||||
if (CountM == 8) {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], pa1[1], B, CountM);
|
||||
pa1 += 2;
|
||||
} else {
|
||||
MlasSgemmComputeBlockMMA<RowCount>(&acc[0], pa1[0], ABroadcast[0], B, CountM);
|
||||
pa1 += 1;
|
||||
}
|
||||
B += 16;
|
||||
k -= 1;
|
||||
}
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
__builtin_mma_disassemble_acc (Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[2]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[3]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<12>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[6]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[7]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<12>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if (CountN >= 12) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[2]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[6]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<8>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 12 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[1], &acc[7]);
|
||||
}
|
||||
}
|
||||
if (CountN - 12 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[0], &acc[3]);
|
||||
}
|
||||
} else if (CountN >= 8) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[1]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
__builtin_mma_disassemble_acc (Result, &acc[5]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<4>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 8 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[1], &acc[6]);
|
||||
}
|
||||
}
|
||||
if (CountN - 8 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[0], &acc[2]);
|
||||
}
|
||||
} else if (CountN >= 4) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[0]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc (Result, &acc[4]);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorMMA<0>>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 4 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[1], &acc[5]);
|
||||
}
|
||||
}
|
||||
if (CountN - 4 > 0) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[0], &acc[1]);
|
||||
}
|
||||
} else {
|
||||
__builtin_mma_disassemble_acc (Accumulators[0], &acc[0]);
|
||||
if (CountM == 8) {
|
||||
__builtin_mma_disassemble_acc (Accumulators[1], &acc[4]);
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Store the remaining unaligned columns.
|
||||
//
|
||||
|
||||
C += (CountN & ~3);
|
||||
CountN &= 3;
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingMMA>()(Accumulators[0], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<0>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingMMA>()(Accumulators[1], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<0>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
if (CountN >= 2) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<1>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<1>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
if (CountN >= 3) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<2>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarMMA<2>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return CountM;
|
||||
}
|
||||
|
||||
template <size_t RowCount>
|
||||
MLAS_FORCEINLINE void
|
||||
MlasSgemmPackA(
|
||||
__vector float* D,
|
||||
const float* A,
|
||||
size_t lda,
|
||||
size_t k
|
||||
)
|
||||
{
|
||||
__vector float a1, a2;
|
||||
const float* a = A;
|
||||
MLAS_FLOAT32X4 AElements[RowCount] = {};
|
||||
MLAS_FLOAT32X4 A2Elements[RowCount] = {};
|
||||
while (k >= 16)
|
||||
|
||||
{
|
||||
PREFETCH_ADDR(a);
|
||||
PREFETCH_ADDR(a + lda);
|
||||
PREFETCH_ADDR(a + 2 * lda);
|
||||
PREFETCH_ADDR(a + 3 * lda);
|
||||
PREFETCH_ADDR(a + 4 * lda);
|
||||
PREFETCH_ADDR(a + 5 * lda);
|
||||
PREFETCH_ADDR(a + 6 * lda);
|
||||
PREFETCH_ADDR(a + 7 * lda);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[0] = vec_xxpermdi(a1, a2, 0);
|
||||
D[2] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[1] = vec_xxpermdi(a1, a2, 0);
|
||||
D[3] = vec_xxpermdi(a1, a2, 3);
|
||||
if (RowCount == 8) {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 4, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[8] = vec_xxpermdi(a1, a2, 0);
|
||||
D[10] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[9] = vec_xxpermdi(a1, a2, 0);
|
||||
D[11] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 8, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[16] = vec_xxpermdi(a1, a2, 0);
|
||||
D[18] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[17] = vec_xxpermdi(a1, a2, 0);
|
||||
D[19] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 12, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[24] = vec_xxpermdi(a1, a2, 0);
|
||||
D[26] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[25] = vec_xxpermdi(a1, a2, 0);
|
||||
D[27] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, a + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[4] = vec_xxpermdi(a1, a2, 0);
|
||||
D[6] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[5] = vec_xxpermdi(a1, a2, 0);
|
||||
D[7] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, (a + 4) + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[12] = vec_xxpermdi(a1, a2, 0);
|
||||
D[14] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[13] = vec_xxpermdi(a1, a2, 0);
|
||||
D[15] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, (a + 8) + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[20] = vec_xxpermdi(a1, a2, 0);
|
||||
D[22] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[21] = vec_xxpermdi(a1, a2, 0);
|
||||
D[23] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, (a + 12) + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[28] = vec_xxpermdi(a1, a2, 0);
|
||||
D[30] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[29] = vec_xxpermdi(a1, a2, 0);
|
||||
D[31] = vec_xxpermdi(a1, a2, 3);
|
||||
D += 32;
|
||||
} else {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 4, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[4] = vec_xxpermdi(a1, a2, 0);
|
||||
D[6] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[5] = vec_xxpermdi(a1, a2, 0);
|
||||
D[7] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 8, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[8] = vec_xxpermdi(a1, a2, 0);
|
||||
D[10] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[9] = vec_xxpermdi(a1, a2, 0);
|
||||
D[11] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 12, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[12] = vec_xxpermdi(a1, a2, 0);
|
||||
D[14] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[13] = vec_xxpermdi(a1, a2, 0);
|
||||
D[15] = vec_xxpermdi(a1, a2, 3);
|
||||
D += 16;
|
||||
}
|
||||
k -= 16;
|
||||
a += 16;
|
||||
}
|
||||
|
||||
while (k >= 8) {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[0] = vec_xxpermdi(a1, a2, 0);
|
||||
D[2] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[1] = vec_xxpermdi(a1, a2, 0);
|
||||
D[3] = vec_xxpermdi(a1, a2, 3);
|
||||
if (RowCount == 8) {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 4, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[8] = vec_xxpermdi(a1, a2, 0);
|
||||
D[10] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[9] = vec_xxpermdi(a1, a2, 0);
|
||||
D[11] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, a + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[4] = vec_xxpermdi(a1, a2, 0);
|
||||
D[6] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[5] = vec_xxpermdi(a1, a2, 0);
|
||||
D[7] = vec_xxpermdi(a1, a2, 3);
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, (a + 4) + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[12] = vec_xxpermdi(a1, a2, 0);
|
||||
D[14] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[13] = vec_xxpermdi(a1, a2, 0);
|
||||
D[15] = vec_xxpermdi(a1, a2, 3);
|
||||
D += 16;
|
||||
} else {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a + 4, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[4] = vec_xxpermdi(a1, a2, 0);
|
||||
D[6] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[5] = vec_xxpermdi(a1, a2, 0);
|
||||
D[7] = vec_xxpermdi(a1, a2, 3);
|
||||
D += 8;
|
||||
}
|
||||
a += 8;
|
||||
k -= 8;
|
||||
}
|
||||
|
||||
while (k >= 4) {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(AElements, a, lda);
|
||||
a1 = vec_mergee(AElements[0], AElements[1]);
|
||||
a2 = vec_mergee(AElements[2], AElements[3]);
|
||||
D[0] = vec_xxpermdi(a1, a2, 0);
|
||||
D[2] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(AElements[0], AElements[1]);
|
||||
a2 = vec_mergeo(AElements[2], AElements[3]);
|
||||
D[1] = vec_xxpermdi(a1, a2, 0);
|
||||
D[3] = vec_xxpermdi(a1, a2, 3);
|
||||
if (RowCount == 8) {
|
||||
MlasLoopUnroll<4, MlasFgemmLoadAElements>()(A2Elements, a + (lda * 4), lda);
|
||||
a1 = vec_mergee(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergee(A2Elements[2], A2Elements[3]);
|
||||
D[4] = vec_xxpermdi(a1, a2, 0);
|
||||
D[6] = vec_xxpermdi(a1, a2, 3);
|
||||
a1 = vec_mergeo(A2Elements[0], A2Elements[1]);
|
||||
a2 = vec_mergeo(A2Elements[2], A2Elements[3]);
|
||||
D[5] = vec_xxpermdi(a1, a2, 0);
|
||||
D[7] = vec_xxpermdi(a1, a2, 3);
|
||||
D += 8;
|
||||
} else
|
||||
D += 4;
|
||||
a += 4;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
/* When k is less than 4, copy a single element from each row. */
|
||||
while (k > 0) {
|
||||
MlasLoopUnroll<4, MlasSgemmBroadcastAElementsMMA>()(AElements, a, lda);
|
||||
D[0] = AElements[0];
|
||||
if (RowCount == 8) {
|
||||
MlasLoopUnroll<4, MlasSgemmBroadcastAElementsMMA>()(A2Elements, a + (lda * 4), lda);
|
||||
D[1] = A2Elements[0];
|
||||
D += 2;
|
||||
} else {
|
||||
D += 1;
|
||||
}
|
||||
a += 1;
|
||||
k -= 1;
|
||||
}
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelPOWER10(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
size_t index = CountK * 2;
|
||||
|
||||
MLAS_FLOAT32X4* PackA =
|
||||
reinterpret_cast<MLAS_FLOAT32X4*>(alloca(sizeof(MLAS_FLOAT32X4) * index));
|
||||
MLAS_FLOAT32X4 AlphaBroadcast = MlasBroadcastFloat32x4(alpha);
|
||||
|
||||
if (CountM >= 8) {
|
||||
#ifdef _AIX
|
||||
MlasSgemmPackA<8>(PackA, A, lda, CountK);
|
||||
#else
|
||||
if (CountK >= 16 && !(CountK % 16)) {
|
||||
PackAKernelPOWER10(PackA, A, lda, CountK, 8);
|
||||
} else {
|
||||
MlasSgemmPackA<8>(PackA, A, lda, CountK);
|
||||
}
|
||||
#endif
|
||||
|
||||
RowsHandled = MlasSgemmMMAProcessCount<4>(PackA, B, C, 8, CountK, CountN, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 4) {
|
||||
memset(PackA + CountK, 0, sizeof(MLAS_FLOAT32X4) * CountK);
|
||||
#ifdef _AIX
|
||||
MlasSgemmPackA<4>(PackA, A, lda, CountK);
|
||||
#else
|
||||
if (CountK >= 16 && !(CountK % 16)) {
|
||||
PackAKernelPOWER10(PackA, A, lda, CountK, 4);
|
||||
} else {
|
||||
MlasSgemmPackA<4>(PackA, A, lda, CountK);
|
||||
}
|
||||
#endif
|
||||
RowsHandled = MlasSgemmMMAProcessCount<4>(PackA, B, C, 4, CountK, CountN, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
return RowsHandled;
|
||||
}
|
||||
+247
@@ -0,0 +1,247 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelPackA.S
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the POWER10 kernel for packing matrix A for single precision SGEMM.
|
||||
|
||||
This implementation targets power10 using VSX instructions.
|
||||
|
||||
--*/
|
||||
/*++
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to pack matrix A for rows 4 or 8.
|
||||
|
||||
Arguments:
|
||||
|
||||
D (r3) - Supplies the address of Packed A.
|
||||
|
||||
A (r4) - Supplies the address of matrix A.
|
||||
|
||||
lda (r5) - LDA.
|
||||
|
||||
k (r6) - Supplies the number of columns from matrix A.
|
||||
|
||||
RowCount (r7) - Supplies the number of rows to process.
|
||||
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
--*/
|
||||
#include "asmmacro.h"
|
||||
.text
|
||||
FUNCTION_ENTRY PackAKernelPOWER10
|
||||
slwi 9,5,2
|
||||
cmpldi 7,8
|
||||
add 8,4,9
|
||||
add 10,8,9
|
||||
add 11,10,9
|
||||
dcbt 0,4
|
||||
dcbt 0,8
|
||||
dcbt 0,10
|
||||
dcbt 0,11
|
||||
blt L_loop
|
||||
L_Rows8:
|
||||
lxvp 32,0(4)
|
||||
lxvp 42,32(4)
|
||||
addi 4,4,64
|
||||
dcbt 0,4
|
||||
lxvp 34,0(8) //a+lda
|
||||
lxvp 44,32(8) //a+32+lda
|
||||
lxvp 36,0(10) //a+2*lda
|
||||
lxvp 46,32(10) //a+32+2*lda
|
||||
lxvp 38,0(11) //a+3*lda
|
||||
lxvp 48,32(11) //a+32+3*lda
|
||||
add 7,11,9
|
||||
dcbt 0,7
|
||||
add 8,7,9
|
||||
dcbt 0,8
|
||||
add 10,8,9
|
||||
dcbt 0,10
|
||||
add 11,10,9
|
||||
dcbt 0,11
|
||||
vmrgow 8,3,1
|
||||
vmrgew 18,3,1
|
||||
vmrgow 9,7,5
|
||||
vmrgew 19,7,5
|
||||
xxpermdi 1,41,40,3
|
||||
xxpermdi 0,51,50,3
|
||||
xxpermdi 3,41,40,0
|
||||
xxpermdi 2,51,50,0
|
||||
vmrgow 8,2,0
|
||||
vmrgow 9,6,4
|
||||
vmrgew 18,2,0
|
||||
vmrgew 19,6,4
|
||||
stxvp 0,0(3)
|
||||
xxpermdi 9,41,40,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,41,40,0
|
||||
xxpermdi 10,51,50,0
|
||||
stxvp 2,32(3)
|
||||
stxvp 8,128(3)
|
||||
stxvp 10,160(3)
|
||||
|
||||
vmrgow 0,13,11
|
||||
vmrgow 1,17,15
|
||||
vmrgew 18,13,11
|
||||
vmrgew 19,17,15
|
||||
xxpermdi 5,33,32,3
|
||||
xxpermdi 7,33,32,0
|
||||
xxpermdi 4,51,50,3
|
||||
xxpermdi 6,51,50,0
|
||||
vmrgow 0,12,10
|
||||
vmrgow 1,16,14
|
||||
stxvp 4,256(3)
|
||||
stxvp 6,288(3)
|
||||
vmrgew 18,12,10
|
||||
vmrgew 19,16,14
|
||||
xxpermdi 9,33,32,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,33,32,0
|
||||
xxpermdi 10,51,50,0
|
||||
lxvp 32,0(7) //a+4*lda
|
||||
lxvp 34,0(8) //a+5*lda
|
||||
lxvp 36,0(10) //a+6*lda
|
||||
lxvp 38, 0(11) //a+7*lda
|
||||
|
||||
stxvp 8,384(3)
|
||||
stxvp 10,416(3)
|
||||
lxvp 42,32(7) //a+32+4*lda
|
||||
|
||||
vmrgow 8,3,1
|
||||
vmrgew 18,3,1
|
||||
vmrgow 9,7,5
|
||||
vmrgew 19,7,5
|
||||
lxvp 44,32(8) //a+32+5*lda
|
||||
lxvp 46,32(10) //a+32+6*lda
|
||||
|
||||
xxpermdi 1,41,40,3
|
||||
xxpermdi 0,51,50,3
|
||||
xxpermdi 3,41,40,0
|
||||
xxpermdi 2,51,50,0
|
||||
lxvp 48,32(11) //a+32+7*lda
|
||||
add 8,4,9
|
||||
dcbt 0,8
|
||||
add 10,8,9
|
||||
dcbt 0,10
|
||||
add 11,10,9
|
||||
dcbt 0,11
|
||||
stxvp 0,64(3)
|
||||
stxvp 2,96(3)
|
||||
vmrgow 8,2,0
|
||||
vmrgow 9,6,4
|
||||
vmrgew 18,2,0
|
||||
vmrgew 19,6,4
|
||||
vmrgow 0,13,11
|
||||
vmrgow 1,17,15
|
||||
|
||||
xxpermdi 9,41,40,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,41,40,0
|
||||
xxpermdi 10,51,50,0
|
||||
vmrgew 18,13,11
|
||||
vmrgew 19,17,15
|
||||
stxvp 8,192(3)
|
||||
stxvp 10,224(3)
|
||||
xxpermdi 5,33,32,3
|
||||
xxpermdi 4,51,50,3
|
||||
xxpermdi 7,33,32,0
|
||||
xxpermdi 6,51,50,0
|
||||
vmrgow 0,12,10
|
||||
vmrgow 1,16,14
|
||||
stxvp 4,320(3)
|
||||
stxvp 6,352(3)
|
||||
vmrgew 18,12,10
|
||||
vmrgew 19,16,14
|
||||
xxpermdi 9,33,32,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,33,32,0
|
||||
xxpermdi 10,51,50,0
|
||||
|
||||
stxvp 8,448(3)
|
||||
stxvp 10,480(3)
|
||||
addi 6,6,-16
|
||||
cmpldi 6,16
|
||||
|
||||
addi 3,3,512
|
||||
bge L_Rows8
|
||||
b L_exit
|
||||
|
||||
L_loop:
|
||||
lxvp 32,0(4)
|
||||
lxvp 42,32(4)
|
||||
addi 4,4,64
|
||||
dcbt 0,4
|
||||
lxvp 34,0(8) //a+lda
|
||||
lxvp 44,32(8) //a+32+lda
|
||||
lxvp 36,0(10) //a+2*lda
|
||||
lxvp 46,32(10) //a+32+2*lda
|
||||
lxvp 38,0(11) //a+3*lda
|
||||
lxvp 48,32(11) //a+32+3*lda
|
||||
vmrgow 8,3,1
|
||||
vmrgew 18,3,1
|
||||
vmrgow 9,7,5
|
||||
vmrgew 19,7,5
|
||||
|
||||
add 8,4,9
|
||||
dcbt 0,8
|
||||
add 10,8,9
|
||||
dcbt 0,10
|
||||
add 11,10,9
|
||||
dcbt 0,11
|
||||
|
||||
xxpermdi 1,41,40,3
|
||||
xxpermdi 0,51,50,3
|
||||
xxpermdi 3,41,40,0
|
||||
xxpermdi 2,51,50,0
|
||||
vmrgow 8,2,0
|
||||
vmrgow 9,6,4
|
||||
vmrgew 18,2,0
|
||||
vmrgew 19,6,4
|
||||
stxvp 0,0(3)
|
||||
|
||||
xxpermdi 9,41,40,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,41,40,0
|
||||
xxpermdi 10,51,50,0
|
||||
stxvp 2,32(3)
|
||||
stxvp 8,64(3)
|
||||
stxvp 10,96(3)
|
||||
|
||||
vmrgow 0,13,11
|
||||
vmrgow 1,17,15
|
||||
vmrgew 18,13,11
|
||||
vmrgew 19,17,15
|
||||
xxpermdi 5,33,32,3
|
||||
xxpermdi 7,33,32,0
|
||||
xxpermdi 4,51,50,3
|
||||
xxpermdi 6,51,50,0
|
||||
vmrgow 0,12,10
|
||||
vmrgow 1,16,14
|
||||
stxvp 4,128(3)
|
||||
stxvp 6,160(3)
|
||||
vmrgew 18,12,10
|
||||
vmrgew 19,16,14
|
||||
xxpermdi 9,33,32,3
|
||||
xxpermdi 8,51,50,3
|
||||
xxpermdi 11,33,32,0
|
||||
xxpermdi 10,51,50,0
|
||||
stxvp 8,192(3)
|
||||
stxvp 10,224(3)
|
||||
|
||||
addi 3,3,256
|
||||
addi 6,6,-16
|
||||
cmpldi 6,16
|
||||
bge L_loop
|
||||
|
||||
L_exit:
|
||||
blr
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelPower.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
#include "SgemmKernelpower.h"
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
|
||||
MLAS_FLOAT32X4 AlphaBroadcast = MlasBroadcastFloat32x4(alpha);
|
||||
|
||||
if (CountM >= 4) {
|
||||
RowsHandled = MlasSgemmProcessCount<4>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
Vendored
+930
@@ -0,0 +1,930 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
qgemm.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module defines the set of template functions to implement a kernel of
|
||||
quantized integer matrix/matrix multiply operation (QGEMM).
|
||||
|
||||
To implement a new kernel, template functions below need to be specialized:
|
||||
MlasGemmQuantFixupZeroPointA
|
||||
MlasGemmQuantFixupZeroPointB
|
||||
MlasGemmQuantCopyPackA
|
||||
MlasGemmQuantCopyPackB
|
||||
MlasGemmQuantKernel
|
||||
Specialization of MlasGemmQuantTryGemvKernel is optional.
|
||||
|
||||
MlasGemmQuantOperation and MlasGemmQuantPackedOperation are shared kernel drivers.
|
||||
MlasGemmQuantScaleSumBuffer is a helper function.
|
||||
|
||||
It also includes the dispatcher logics.
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <cstdlib>
|
||||
|
||||
//
|
||||
// Define the default striding parameters used for the quantized integer
|
||||
// matrix/matrix multiply operation.
|
||||
//
|
||||
|
||||
struct MLAS_GEMM_QUANT_STRIDES {
|
||||
size_t M;
|
||||
size_t N;
|
||||
size_t K;
|
||||
};
|
||||
|
||||
template<typename KernelType>
|
||||
MLAS_FORCEINLINE
|
||||
bool
|
||||
MlasGemmQuantTryGemvKernel(
|
||||
const uint8_t* A,
|
||||
const uint8_t* B,
|
||||
size_t ldb,
|
||||
int32_t* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
bool AIsSigned,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
MLAS_UNREFERENCED_PARAMETER(A);
|
||||
MLAS_UNREFERENCED_PARAMETER(B);
|
||||
MLAS_UNREFERENCED_PARAMETER(ldb);
|
||||
MLAS_UNREFERENCED_PARAMETER(C);
|
||||
MLAS_UNREFERENCED_PARAMETER(CountK);
|
||||
MLAS_UNREFERENCED_PARAMETER(CountN);
|
||||
MLAS_UNREFERENCED_PARAMETER(AIsSigned);
|
||||
MLAS_UNREFERENCED_PARAMETER(BIsSigned);
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename KernelType>
|
||||
MLAS_FORCEINLINE constexpr
|
||||
int32_t
|
||||
MlasGemmQuantFixupZeroPointA(
|
||||
int32_t ZeroPointA,
|
||||
bool AIsSigned)
|
||||
{
|
||||
MLAS_UNREFERENCED_PARAMETER(AIsSigned);
|
||||
return ZeroPointA;
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
int32_t constexpr
|
||||
MlasGemmQuantFixupZeroPointB(
|
||||
int32_t ZeroPointB,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
MLAS_UNREFERENCED_PARAMETER(BIsSigned);
|
||||
|
||||
return ZeroPointB;
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasGemmQuantFixupZeroPointB(
|
||||
const uint8_t* PackedZeroPointB,
|
||||
int32_t* ZeroPointBBuffer,
|
||||
size_t N,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
int32_t ZeroPointB;
|
||||
|
||||
for (size_t n = 0; n < N; n++) {
|
||||
|
||||
ZeroPointB = typename KernelType::OffsetBType(PackedZeroPointB[n]);
|
||||
ZeroPointB = MlasGemmQuantFixupZeroPointB<KernelType>(ZeroPointB, BIsSigned);
|
||||
|
||||
ZeroPointBBuffer[n] = -ZeroPointB;
|
||||
}
|
||||
|
||||
//
|
||||
// Fill the misaligned slots of the zero point buffer with zeros to guard
|
||||
// against tools that check for uninitialized data usage.
|
||||
//
|
||||
|
||||
size_t AlignedN = (N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1);
|
||||
|
||||
for (size_t n = N; n < AlignedN; n++) {
|
||||
ZeroPointBBuffer[n] = 0;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
void
|
||||
MlasGemmQuantCopyPackA(
|
||||
typename KernelType::PackedAType* D,
|
||||
const uint8_t* A,
|
||||
size_t lda,
|
||||
size_t CountM,
|
||||
size_t CountK,
|
||||
int32_t* RowSumBuffer,
|
||||
bool AIsSigned
|
||||
);
|
||||
|
||||
template<typename KernelType>
|
||||
void
|
||||
MlasGemmQuantCopyPackB(
|
||||
typename KernelType::PackedBType* D,
|
||||
const uint8_t* B,
|
||||
size_t ldb,
|
||||
size_t CountN,
|
||||
size_t CountK,
|
||||
int32_t* ColumnSumBuffer,
|
||||
bool BIsSigned
|
||||
);
|
||||
|
||||
template<typename KernelType>
|
||||
size_t
|
||||
MlasGemmQuantKernel(
|
||||
const typename KernelType::PackedAType* A,
|
||||
const typename KernelType::PackedBType* B,
|
||||
int32_t* C,
|
||||
size_t PackedCountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t ldc,
|
||||
const int32_t* RowSumBuffer,
|
||||
const int32_t* ColumnSumBuffer,
|
||||
const int32_t* ZeroPointB,
|
||||
bool ZeroMode
|
||||
);
|
||||
|
||||
/**
|
||||
* @brief Usually a wrapper of assembly/intrinsic kernel
|
||||
* of symmetric quant gemm
|
||||
* @tparam KernelType
|
||||
* @param A Left hand side matrix
|
||||
* @param B Prepacked right hand side matrix
|
||||
* @param C Result matrix
|
||||
* @param PackedCountK Number of packed rows from B
|
||||
* @param CountM Number of rows to process
|
||||
* @param CountN Number of columns to process
|
||||
* @param ldc Row stride of C
|
||||
* @param lda Row stride of A
|
||||
* @param ColumnSumVector Column sum of B scaled by zero point A
|
||||
* @return Number of rows processed
|
||||
*/
|
||||
template<typename KernelType>
|
||||
size_t
|
||||
MlasSymmQGemmKernel(
|
||||
const int8_t* A,
|
||||
const int8_t* B,
|
||||
int32_t* C,
|
||||
size_t PackedCountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t ldc,
|
||||
size_t lda,
|
||||
const int32_t* ColumnSumVector
|
||||
);
|
||||
|
||||
inline
|
||||
void
|
||||
MlasGemmQuantScaleSumBuffer(
|
||||
int32_t* Output,
|
||||
const int32_t* Input,
|
||||
size_t N,
|
||||
int32_t Scale
|
||||
)
|
||||
{
|
||||
for (size_t n = 0; n < N; n++) {
|
||||
Output[n] = Input[n] * Scale;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasGemmQuantScaleSumBuffer(
|
||||
int32_t* SumBuffer,
|
||||
size_t N,
|
||||
int32_t Scale
|
||||
)
|
||||
{
|
||||
return MlasGemmQuantScaleSumBuffer(SumBuffer, SumBuffer, N, Scale);
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasGemmQuantThreadInit()
|
||||
{
|
||||
constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::Strides;
|
||||
constexpr size_t packASize =
|
||||
UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType));
|
||||
constexpr size_t packBSize =
|
||||
UpAlignSize(Strides.N * Strides.K * sizeof(typename KernelType::PackedBType));
|
||||
constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t));
|
||||
constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t));
|
||||
constexpr size_t zpbSize = UpAlignSize(Strides.N * sizeof(int32_t));
|
||||
|
||||
constexpr MLAS_GEMM_QUANT_STRIDES PackedStrides = KernelType::PackedStrides;
|
||||
constexpr size_t packedASize =
|
||||
UpAlignSize(PackedStrides.M * PackedStrides.K * sizeof(typename KernelType::PackedAType));
|
||||
|
||||
constexpr size_t bufsize = std::max(packASize + packBSize, packedASize) + rowSumSize + colSumSize + zpbSize;
|
||||
|
||||
MlasThreadedBufAlloc(bufsize);
|
||||
}
|
||||
|
||||
template<typename KernelType>
|
||||
void
|
||||
MlasGemmQuantOperation(
|
||||
const MLAS_GEMM_QUANT_SHAPE_PARAMS* Shape,
|
||||
const MLAS_GEMM_QUANT_DATA_PARAMS* Data,
|
||||
const size_t RangeStartM,
|
||||
const size_t RangeCountM,
|
||||
const size_t RangeStartN,
|
||||
const size_t RangeCountN
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine implements the quantized integer matrix/matrix multiply
|
||||
operation (QGEMM).
|
||||
|
||||
Arguments:
|
||||
|
||||
Shape - Supplies the structure containing the GEMM input and output shapes.
|
||||
|
||||
Data - Supplies the structure containing the GEMM input and output data layout
|
||||
|
||||
RangeStartM - Supplies the starting row index to output.
|
||||
|
||||
RangeCountM - Supplies the number of rows to output.
|
||||
|
||||
RangeStartN - Supplies the starting column index to output.
|
||||
|
||||
RangeCountN - Supplies the number of columns to output.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
{
|
||||
constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::Strides;
|
||||
constexpr size_t packASize =
|
||||
UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType));
|
||||
constexpr size_t packBSize =
|
||||
UpAlignSize(Strides.N * Strides.K * sizeof(typename KernelType::PackedBType));
|
||||
constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t));
|
||||
constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t));
|
||||
|
||||
MlasGemmQuantThreadInit<KernelType>();
|
||||
|
||||
uint8_t* p = ThreadedBufHolder.get();
|
||||
typename KernelType::PackedAType* PanelA =
|
||||
reinterpret_cast<typename KernelType::PackedAType*>(p);
|
||||
p += packASize;
|
||||
typename KernelType::PackedBType* PanelB =
|
||||
reinterpret_cast<typename KernelType::PackedBType*>(p);
|
||||
p += packBSize;
|
||||
int32_t* RowSumBuffer = reinterpret_cast<int32_t*>(p);
|
||||
p += rowSumSize;
|
||||
int32_t* ColumnSumBuffer = reinterpret_cast<int32_t*>(p);
|
||||
p += colSumSize;
|
||||
int32_t* ZeroPointBBuffer = reinterpret_cast<int32_t*>(p);
|
||||
|
||||
|
||||
const size_t K = Shape->K;
|
||||
|
||||
const size_t lda = Data->lda;
|
||||
const size_t ldb = Data->ldb;
|
||||
const size_t ldc = Data->ldc;
|
||||
|
||||
const uint8_t* A = Data->A + RangeStartM * lda;
|
||||
const uint8_t* B = (const uint8_t*)Data->B + RangeStartN;
|
||||
int32_t* C = Data->C + RangeStartM * ldc + RangeStartN;
|
||||
const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ?
|
||||
Data->ZeroPointB + RangeStartN : nullptr;
|
||||
bool IsAccumulateMode = Shape->IsAccumulateMode;
|
||||
|
||||
int32_t ZeroPointA = typename KernelType::OffsetAType(Data->ZeroPointA);
|
||||
int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB);
|
||||
|
||||
//
|
||||
// Try to use a GEMV kernel if supported by this kernel type.
|
||||
//
|
||||
|
||||
if ((RangeCountM == 1) &&
|
||||
(ZeroPointA == 0) && (PackedZeroPointB == nullptr) && (ZeroPointB == 0) &&
|
||||
(Data->OutputProcessor == nullptr)) {
|
||||
if (MlasGemmQuantTryGemvKernel<KernelType>(A, B, ldb, C, K, RangeCountN, Shape->AIsSigned, Shape->BIsSigned)) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-matrix zero point offset of matrix A if the
|
||||
// kernel requires opposite-signed data.
|
||||
//
|
||||
|
||||
ZeroPointA = MlasGemmQuantFixupZeroPointA<KernelType>(ZeroPointA, Shape->AIsSigned);
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-matrix zero point offset of matrix B if the
|
||||
// data is the opposite format of the kernel implementation. This value is
|
||||
// ignored if per-column zero point offsets are used instead.
|
||||
//
|
||||
|
||||
ZeroPointB = MlasGemmQuantFixupZeroPointB<KernelType>(ZeroPointB, Shape->BIsSigned);
|
||||
|
||||
//
|
||||
// Step through each slice of matrix B along the K dimension.
|
||||
//
|
||||
|
||||
size_t CountK;
|
||||
|
||||
for (size_t k = 0; k < K; k += CountK) {
|
||||
|
||||
CountK = std::min(K - k, Strides.K);
|
||||
|
||||
const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / KernelType::PackedK;
|
||||
|
||||
//
|
||||
// Step through each slice of matrix B along the N dimension.
|
||||
//
|
||||
|
||||
size_t CountN;
|
||||
|
||||
for (size_t n = 0; n < RangeCountN; n += CountN) {
|
||||
|
||||
CountN = std::min(RangeCountN - n, Strides.N);
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-column zero point offsets of matrix B
|
||||
// if the data is the opposite format of the kernel implementation.
|
||||
//
|
||||
|
||||
if (PackedZeroPointB != nullptr) {
|
||||
MlasGemmQuantFixupZeroPointB<KernelType>(
|
||||
PackedZeroPointB + n,
|
||||
ZeroPointBBuffer,
|
||||
CountN,
|
||||
Shape->BIsSigned);
|
||||
}
|
||||
|
||||
//
|
||||
// Copy a panel of matrix B to a local packed buffer.
|
||||
//
|
||||
|
||||
MlasGemmQuantCopyPackB<KernelType>(
|
||||
PanelB,
|
||||
B + n,
|
||||
ldb,
|
||||
CountN,
|
||||
CountK,
|
||||
ColumnSumBuffer,
|
||||
Shape->BIsSigned);
|
||||
|
||||
MlasGemmQuantScaleSumBuffer(ColumnSumBuffer, CountN, -ZeroPointA);
|
||||
|
||||
//
|
||||
// Step through each slice of matrix A along the M dimension.
|
||||
//
|
||||
|
||||
int32_t* c = C + n;
|
||||
size_t CountM;
|
||||
|
||||
for (size_t m = 0; m < RangeCountM; m += CountM) {
|
||||
|
||||
CountM = std::min(RangeCountM - m, Strides.M);
|
||||
|
||||
//
|
||||
// Copy a panel of matrix A to a local packed buffer.
|
||||
//
|
||||
|
||||
MlasGemmQuantCopyPackA<KernelType>(
|
||||
PanelA,
|
||||
A + m * lda,
|
||||
lda,
|
||||
CountM,
|
||||
CountK,
|
||||
RowSumBuffer,
|
||||
Shape->AIsSigned);
|
||||
|
||||
//
|
||||
// Apply the global depth value constant without the ZeroPointB scaling from:
|
||||
//
|
||||
// (A[i] - ZeroPointA) * (B[i] - ZeroPointB)
|
||||
// ==>
|
||||
// A[i] * B[i] - A[i] * ZeroPointB - B[i] * ZeroPointA + ZeroPointA * ZeroPointB
|
||||
//
|
||||
// The ZeroPointB term is factored out and either applied below for per-matrix
|
||||
// quantization or inside the kernel for per-column quantization.
|
||||
//
|
||||
|
||||
for (size_t mm = 0; mm < CountM; mm++) {
|
||||
RowSumBuffer[mm] -= int32_t(CountK) * ZeroPointA;
|
||||
}
|
||||
|
||||
//
|
||||
// Scale the row sums by the per-matrix zero point offset of matrix B.
|
||||
//
|
||||
|
||||
if (PackedZeroPointB == nullptr) {
|
||||
MlasGemmQuantScaleSumBuffer(RowSumBuffer, CountM, -ZeroPointB);
|
||||
}
|
||||
|
||||
//
|
||||
// Step through the rows of the local packed buffer.
|
||||
//
|
||||
|
||||
typename KernelType::PackedAType* pa = PanelA;
|
||||
int32_t* RowSums = RowSumBuffer;
|
||||
size_t RowsRemaining = CountM;
|
||||
|
||||
bool ZeroMode = (k == 0) && !IsAccumulateMode;
|
||||
bool PostProcess = (k + CountK == K);
|
||||
|
||||
while (RowsRemaining > 0) {
|
||||
|
||||
size_t RowsHandled = MlasGemmQuantKernel<KernelType>(
|
||||
pa,
|
||||
PanelB,
|
||||
c,
|
||||
PackedCountK,
|
||||
RowsRemaining,
|
||||
CountN,
|
||||
ldc,
|
||||
RowSums,
|
||||
ColumnSumBuffer,
|
||||
(PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr,
|
||||
ZeroMode);
|
||||
|
||||
if (PostProcess && Data->OutputProcessor != nullptr) {
|
||||
Data->OutputProcessor->Process(
|
||||
Data->C,
|
||||
RangeStartM + m + CountM - RowsRemaining,
|
||||
RangeStartN + n,
|
||||
RowsHandled,
|
||||
CountN,
|
||||
Data->ldc);
|
||||
}
|
||||
|
||||
c += ldc * RowsHandled;
|
||||
pa += KernelType::PackedK * PackedCountK * RowsHandled;
|
||||
RowSums += RowsHandled;
|
||||
RowsRemaining -= RowsHandled;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
A += CountK;
|
||||
B += CountK * ldb;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template<typename KernelType>
|
||||
void
|
||||
MlasGemmQuantPackedOperation(
|
||||
const MLAS_GEMM_QUANT_SHAPE_PARAMS* Shape,
|
||||
const MLAS_GEMM_QUANT_DATA_PARAMS* Data,
|
||||
const size_t RangeStartM,
|
||||
const size_t RangeCountM,
|
||||
const size_t RangeStartN,
|
||||
const size_t RangeCountN
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine implements the quantized integer matrix/matrix multiply
|
||||
operation (QGEMM).
|
||||
|
||||
Arguments:
|
||||
|
||||
Shape - Supplies the structure containing the GEMM input and output shapes.
|
||||
|
||||
Data - Supplies the structure containing the GEMM input and output data layout
|
||||
|
||||
RangeStartM - Supplies the starting row index to output.
|
||||
|
||||
RangeCountM - Supplies the number of rows to output.
|
||||
|
||||
RangeStartN - Supplies the starting column index to output.
|
||||
|
||||
RangeCountN - Supplies the number of columns to output.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
{
|
||||
constexpr MLAS_GEMM_QUANT_STRIDES Strides = KernelType::PackedStrides;
|
||||
constexpr size_t packASize =
|
||||
UpAlignSize(Strides.M * Strides.K * sizeof(typename KernelType::PackedAType));
|
||||
constexpr size_t rowSumSize = UpAlignSize(Strides.M * sizeof(int32_t));
|
||||
constexpr size_t colSumSize = UpAlignSize(Strides.N * sizeof(int32_t));
|
||||
|
||||
MlasGemmQuantThreadInit<KernelType>();
|
||||
|
||||
uint8_t* p = ThreadedBufHolder.get();
|
||||
typename KernelType::PackedAType* PanelA =
|
||||
reinterpret_cast<typename KernelType::PackedAType*>(p);
|
||||
p += packASize;
|
||||
int32_t* RowSumBuffer = reinterpret_cast<int32_t*>(p);
|
||||
p += rowSumSize;
|
||||
int32_t* ColumnSumBuffer = reinterpret_cast<int32_t*>(p);
|
||||
p += colSumSize;
|
||||
int32_t* ZeroPointBBuffer = reinterpret_cast<int32_t*>(p);
|
||||
|
||||
const size_t K = Shape->K;
|
||||
|
||||
const size_t lda = Data->lda;
|
||||
const size_t ldc = Data->ldc;
|
||||
|
||||
const uint8_t* A = Data->A + RangeStartM * lda;
|
||||
const uint8_t* PackedB = (const uint8_t*)Data->B;
|
||||
int32_t* C = Data->C + RangeStartM * ldc + RangeStartN;
|
||||
const uint8_t* PackedZeroPointB = Data->PerColumnZeroPoints ?
|
||||
Data->ZeroPointB + RangeStartN : nullptr;
|
||||
bool IsAccumulateMode = Shape->IsAccumulateMode;
|
||||
|
||||
int32_t ZeroPointA = typename KernelType::OffsetAType(Data->ZeroPointA);
|
||||
int32_t ZeroPointB = typename KernelType::OffsetBType(*Data->ZeroPointB);
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-matrix zero point offset of matrix A if the
|
||||
// kernel requires signed data.
|
||||
//
|
||||
|
||||
ZeroPointA = MlasGemmQuantFixupZeroPointA<KernelType>(ZeroPointA, Shape->AIsSigned);
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-matrix zero point offset of matrix B if the
|
||||
// data is the opposite format of the kernel implementation. This value is
|
||||
// ignored if per-column zero point offsets are used instead.
|
||||
//
|
||||
|
||||
ZeroPointB = MlasGemmQuantFixupZeroPointB<KernelType>(ZeroPointB, Shape->BIsSigned);
|
||||
|
||||
//
|
||||
// Extract the pointer to the column sum buffer from the packed matrix.
|
||||
//
|
||||
|
||||
const size_t AlignedN =
|
||||
(Shape->N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1);
|
||||
const int32_t* PackedColumnSumBuffer = (const int32_t*)PackedB;
|
||||
PackedB = (const uint8_t*)(PackedColumnSumBuffer + AlignedN);
|
||||
PackedColumnSumBuffer += RangeStartN;
|
||||
|
||||
//
|
||||
// Step through each slice of matrix B along the K dimension.
|
||||
//
|
||||
|
||||
size_t CountK;
|
||||
|
||||
for (size_t k = 0; k < K; k += CountK) {
|
||||
|
||||
CountK = std::min(K - k, Strides.K);
|
||||
|
||||
const size_t PackedCountK = (CountK + KernelType::PackedK - 1) / KernelType::PackedK;
|
||||
|
||||
if (k > 0) {
|
||||
std::fill_n(ColumnSumBuffer, Strides.N, 0);
|
||||
}
|
||||
|
||||
//
|
||||
// Step through each slice of matrix B along the N dimension.
|
||||
//
|
||||
|
||||
size_t CountN;
|
||||
|
||||
for (size_t n = 0; n < RangeCountN; n += CountN) {
|
||||
|
||||
CountN = std::min(RangeCountN - n, Strides.N);
|
||||
|
||||
if (k == 0) {
|
||||
MlasGemmQuantScaleSumBuffer(ColumnSumBuffer, PackedColumnSumBuffer + n,
|
||||
CountN, -ZeroPointA);
|
||||
}
|
||||
|
||||
//
|
||||
// Fixup the sign bit of the per-column zero point offsets of matrix B
|
||||
// if the data is the opposite format of the kernel implementation.
|
||||
//
|
||||
|
||||
if (PackedZeroPointB != nullptr) {
|
||||
MlasGemmQuantFixupZeroPointB<KernelType>(
|
||||
PackedZeroPointB + n,
|
||||
ZeroPointBBuffer,
|
||||
CountN,
|
||||
Shape->BIsSigned);
|
||||
}
|
||||
|
||||
//
|
||||
// Step through each slice of matrix A along the M dimension.
|
||||
//
|
||||
|
||||
const uint8_t* b = PackedB + (RangeStartN + n) *
|
||||
KernelType::PackedK * PackedCountK;
|
||||
int32_t* c = C + n;
|
||||
size_t CountM;
|
||||
|
||||
for (size_t m = 0; m < RangeCountM; m += CountM) {
|
||||
|
||||
CountM = std::min(RangeCountM - m, Strides.M);
|
||||
|
||||
//
|
||||
// Copy a panel of matrix A to a local packed buffer.
|
||||
//
|
||||
|
||||
MlasGemmQuantCopyPackA<KernelType>(
|
||||
PanelA,
|
||||
A + m * lda,
|
||||
lda,
|
||||
CountM,
|
||||
CountK,
|
||||
RowSumBuffer,
|
||||
Shape->AIsSigned);
|
||||
|
||||
//
|
||||
// Apply the global depth value constant without the ZeroPointB scaling from:
|
||||
//
|
||||
// (A[i] - ZeroPointA) * (B[i] - ZeroPointB)
|
||||
// ==>
|
||||
// A[i] * B[i] - A[i] * ZeroPointB - B[i] * ZeroPointA + ZeroPointA * ZeroPointB
|
||||
//
|
||||
// The ZeroPointB term is factored out and either applied below for per-matrix
|
||||
// quantization or inside the kernel for per-column quantization.
|
||||
//
|
||||
|
||||
for (size_t mm = 0; mm < CountM; mm++) {
|
||||
RowSumBuffer[mm] -= int32_t(CountK) * ZeroPointA;
|
||||
}
|
||||
|
||||
//
|
||||
// Scale the row sums by the per-matrix zero point offset of matrix B.
|
||||
//
|
||||
|
||||
if (PackedZeroPointB == nullptr) {
|
||||
MlasGemmQuantScaleSumBuffer(RowSumBuffer, CountM, -ZeroPointB);
|
||||
}
|
||||
|
||||
//
|
||||
// Step through the rows of the local packed buffer.
|
||||
//
|
||||
|
||||
typename KernelType::PackedAType* pa = PanelA;
|
||||
int32_t* RowSums = RowSumBuffer;
|
||||
size_t RowsRemaining = CountM;
|
||||
|
||||
bool ZeroMode = (k == 0) && !IsAccumulateMode;
|
||||
bool PostProcess = (k + CountK == K);
|
||||
|
||||
while (RowsRemaining > 0) {
|
||||
|
||||
size_t RowsHandled = MlasGemmQuantKernel<KernelType>(
|
||||
pa,
|
||||
b,
|
||||
c,
|
||||
PackedCountK,
|
||||
RowsRemaining,
|
||||
CountN,
|
||||
ldc,
|
||||
RowSums,
|
||||
ColumnSumBuffer,
|
||||
(PackedZeroPointB != nullptr) ? ZeroPointBBuffer : nullptr,
|
||||
ZeroMode);
|
||||
|
||||
if (PostProcess && Data->OutputProcessor != nullptr) {
|
||||
Data->OutputProcessor->Process(
|
||||
Data->C,
|
||||
RangeStartM + m + CountM - RowsRemaining,
|
||||
RangeStartN + n,
|
||||
RowsHandled,
|
||||
CountN,
|
||||
Data->ldc);
|
||||
}
|
||||
|
||||
c += ldc * RowsHandled;
|
||||
pa += KernelType::PackedK * PackedCountK * RowsHandled;
|
||||
RowSums += RowsHandled;
|
||||
RowsRemaining -= RowsHandled;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
A += CountK;
|
||||
PackedB = (const uint8_t*)PackedB + AlignedN * CountK;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Operation for Quantized GEMM where B is symmetrically
|
||||
* quantized and packed matrix
|
||||
* @param Shape
|
||||
* @param Data
|
||||
* @param RangeStartM
|
||||
* @param RangeCountM
|
||||
* @param RangeStartN
|
||||
* @param RangeCountN
|
||||
*/
|
||||
template<typename KernelType>
|
||||
void
|
||||
MlasSymmQGemmPackedOperation(
|
||||
const MLAS_GEMM_QUANT_SHAPE_PARAMS* Shape,
|
||||
const MLAS_SYMM_QGEMM_DATA_PARAMS* Data,
|
||||
const size_t RangeStartM,
|
||||
const size_t RangeCountM,
|
||||
const size_t RangeStartN,
|
||||
const size_t RangeCountN
|
||||
)
|
||||
{
|
||||
|
||||
const size_t K = Shape->K;
|
||||
|
||||
const size_t lda = Data->lda;
|
||||
const size_t ldc = Data->ldc;
|
||||
|
||||
const int8_t* PanelA = (const int8_t*)(Data->A) + RangeStartM * lda;
|
||||
const int8_t* PackedB = (const int8_t*)Data->B;
|
||||
int32_t* C = (int32_t*)(Data->C) + RangeStartM * ldc + RangeStartN;
|
||||
|
||||
//
|
||||
// Extract the pointer to the column sum buffer from the packed matrix.
|
||||
//
|
||||
const size_t AlignedN =
|
||||
(Shape->N + MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_QGEMM_STRIDEN_THREAD_ALIGN - 1);
|
||||
const int32_t* PackedColumnSumBuffer = (const int32_t*)PackedB;
|
||||
PackedB = (const int8_t*)(PackedColumnSumBuffer + AlignedN);
|
||||
PackedColumnSumBuffer += RangeStartN;
|
||||
|
||||
const size_t PackedCountK = (K + KernelType::PackedK - 1) / KernelType::PackedK;
|
||||
|
||||
//
|
||||
// Apply the global depth value constant without the ZeroPointB scaling from:
|
||||
//
|
||||
// (A[i] - ZeroPointA) * (B[i] - ZeroPointB)
|
||||
// ==>
|
||||
// A[i] * B[i] - A[i] * ZeroPointB - B[i] * ZeroPointA + ZeroPointA * ZeroPointB
|
||||
//
|
||||
// ZeroPointB is zero, which makes this much simpler
|
||||
//
|
||||
|
||||
const int8_t* b = PackedB + RangeStartN * KernelType::PackedK * PackedCountK;
|
||||
int32_t* c = C;
|
||||
|
||||
auto pa = PanelA;
|
||||
size_t RowsRemaining = RangeCountM;
|
||||
|
||||
while (RowsRemaining > 0) {
|
||||
size_t RowsHandled = MlasSymmQGemmKernel<KernelType>(
|
||||
pa, b, c, PackedCountK, RowsRemaining, RangeCountN, ldc, lda, PackedColumnSumBuffer);
|
||||
|
||||
c += ldc * RowsHandled;
|
||||
pa += lda * RowsHandled;
|
||||
RowsRemaining -= RowsHandled;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
//
|
||||
// Quantized integer matrix/matrix dispatch structure.
|
||||
//
|
||||
|
||||
typedef
|
||||
void
|
||||
(MLAS_GEMM_QUANT_OPERATION)(
|
||||
const MLAS_GEMM_QUANT_SHAPE_PARAMS* Shape,
|
||||
const MLAS_GEMM_QUANT_DATA_PARAMS* Data,
|
||||
const size_t RangeStartM,
|
||||
const size_t RangeCountM,
|
||||
const size_t RangeStartN,
|
||||
const size_t RangeCountN
|
||||
);
|
||||
|
||||
typedef
|
||||
void
|
||||
(MLAS_SYMM_QGEMM_OPERATION)(
|
||||
const MLAS_GEMM_QUANT_SHAPE_PARAMS* Shape,
|
||||
const MLAS_SYMM_QGEMM_DATA_PARAMS* Data,
|
||||
const size_t RangeStartM,
|
||||
const size_t RangeCountM,
|
||||
const size_t RangeStartN,
|
||||
const size_t RangeCountN
|
||||
);
|
||||
|
||||
typedef
|
||||
void
|
||||
(MLAS_GEMM_QUANT_COPY_PACKB_ROUTINE)(
|
||||
uint8_t* D,
|
||||
const uint8_t* B,
|
||||
size_t ldb,
|
||||
size_t CountN,
|
||||
size_t CountK,
|
||||
int32_t* ColumnSumBuffer,
|
||||
bool BIsSigned
|
||||
);
|
||||
|
||||
struct MLAS_GEMM_QUANT_DISPATCH {
|
||||
MLAS_GEMM_QUANT_OPERATION* Operation;
|
||||
MLAS_GEMM_QUANT_OPERATION* PackedOperation;
|
||||
MLAS_GEMM_QUANT_COPY_PACKB_ROUTINE* CopyPackBRoutine;
|
||||
size_t PackedK;
|
||||
size_t PackedStrideK;
|
||||
size_t StrideM;
|
||||
};
|
||||
|
||||
struct MLAS_SYMM_QGEMM_DISPATCH {
|
||||
MLAS_SYMM_QGEMM_OPERATION* LitOperation; /// running on little cores with narrow memory load
|
||||
MLAS_SYMM_QGEMM_OPERATION* BigOperation; /// running on big cores with wider memory load
|
||||
MLAS_GEMM_QUANT_COPY_PACKB_ROUTINE* CopyPackBRoutine;
|
||||
size_t StrideM; /**< num of rows processed by kernel at a time */
|
||||
size_t PackedK;
|
||||
};
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
const MLAS_GEMM_QUANT_DISPATCH*
|
||||
MlasGemmQuantGetDispatch(
|
||||
bool AIsSigned,
|
||||
bool BIsSigned
|
||||
)
|
||||
{
|
||||
const MLAS_GEMM_QUANT_DISPATCH* GemmQuantDispatch = &MlasGemmQuantDispatchDefault;
|
||||
|
||||
#if !defined(FORCE_GENERIC_ALGORITHMS)
|
||||
#if defined(MLAS_TARGET_AMD64_IX86)
|
||||
if (AIsSigned) {
|
||||
GemmQuantDispatch =
|
||||
BIsSigned ? GetMlasPlatform().GemmS8S8Dispatch : GetMlasPlatform().GemmS8U8Dispatch;
|
||||
} else {
|
||||
GemmQuantDispatch =
|
||||
BIsSigned ? GetMlasPlatform().GemmU8S8Dispatch : GetMlasPlatform().GemmU8U8Dispatch;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_ARM64)
|
||||
if(BIsSigned) {
|
||||
GemmQuantDispatch = AIsSigned ? GetMlasPlatform().GemmS8S8Dispatch : GetMlasPlatform().GemmU8S8Dispatch;
|
||||
} else if(!AIsSigned) {
|
||||
GemmQuantDispatch = GetMlasPlatform().GemmU8U8Dispatch;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_ARM64EC) || (defined(MLAS_TARGET_ARM) && !defined(_MSC_VER))
|
||||
if(BIsSigned || !AIsSigned) {
|
||||
GemmQuantDispatch = &MlasGemmU8X8DispatchNeon;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_WASM_RELAXED_SIMD)
|
||||
if (!AIsSigned) {
|
||||
if (HasUSDot()) {
|
||||
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmRelaxedSimd;
|
||||
} else {
|
||||
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmSimd;
|
||||
}
|
||||
}
|
||||
#elif defined(MLAS_TARGET_WASM_SIMD)
|
||||
if (!AIsSigned) {
|
||||
GemmQuantDispatch = &MlasGemmU8X8DispatchWasmSimd;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_POWER) && (defined(__linux__) || defined(_AIX)) && defined(POWER10) && \
|
||||
((defined(__GNUC__) && ((__GNUC__ > 10) || (__GNUC__== 10 && __GNUC_MINOR__ >= 2))) || \
|
||||
(defined(__clang__) && (__clang_major__ >= 12)))
|
||||
if (GetMlasPlatform().GemmU8X8Dispatch == &MlasGemm8X8DispatchPOWER10) {
|
||||
GemmQuantDispatch = GetMlasPlatform().GemmU8X8Dispatch;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_LARCH64)
|
||||
if (AIsSigned) {
|
||||
GemmQuantDispatch =
|
||||
BIsSigned ? GetMlasPlatform().GemmS8S8Dispatch : GetMlasPlatform().GemmS8U8Dispatch;
|
||||
} else { // !AIsSigned
|
||||
GemmQuantDispatch =
|
||||
BIsSigned ? GetMlasPlatform().GemmU8S8Dispatch : GetMlasPlatform().GemmU8U8Dispatch;
|
||||
}
|
||||
#elif defined(MLAS_TARGET_S390X)
|
||||
if (GetMlasPlatform().GemmU8X8Dispatch == &MlasGemm8X8DispatchZVECTOR) {
|
||||
GemmQuantDispatch = GetMlasPlatform().GemmU8X8Dispatch;
|
||||
}
|
||||
#endif
|
||||
#endif // !defined(FORCE_GENERIC_ALGORITHMS)
|
||||
|
||||
if (nullptr == GemmQuantDispatch) {
|
||||
std::stringstream ss;
|
||||
ss << "Quant GEMM format: AIsSigned(" << AIsSigned << "), BIsSigned(" << BIsSigned
|
||||
<< ") is not supported on this device";
|
||||
MLAS_THROW_EX(std::invalid_argument, ss.str());
|
||||
}
|
||||
|
||||
return GemmQuantDispatch;
|
||||
}
|
||||
+275
@@ -0,0 +1,275 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
sgemm_kernel_rvv.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements an RVV kernel for the single precision matrix/matrix
|
||||
multiply operation (SGEMM) on riscv64.
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
#if defined(MLAS_USE_RVV)
|
||||
|
||||
#include <riscv_vector.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// The packed B layout stays 16 columns wide to match MLAS, but each tile is
|
||||
// consumed in runtime-sized RVV chunks so the kernel is not tied to a fixed
|
||||
// VLEN such as 128 or 256 bits.
|
||||
constexpr size_t kPackedCountN = 16;
|
||||
|
||||
template<bool ZeroMode, bool AlphaIsOne>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasStoreAccumulatorRvv(
|
||||
float* C,
|
||||
vfloat32m4_t Accumulator,
|
||||
size_t vl,
|
||||
float alpha
|
||||
)
|
||||
{
|
||||
#if defined(_WIN32)
|
||||
|
||||
if constexpr (AlphaIsOne) {
|
||||
UNREFERENCED_PARAMETER(alpha);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
if constexpr (!AlphaIsOne) {
|
||||
Accumulator = __riscv_vfmul_vf_f32m4(Accumulator, alpha, vl);
|
||||
}
|
||||
|
||||
if constexpr (!ZeroMode) {
|
||||
Accumulator = __riscv_vfadd_vv_f32m4(Accumulator, __riscv_vle32_v_f32m4(C, vl), vl);
|
||||
}
|
||||
|
||||
__riscv_vse32_v_f32m4(C, Accumulator, vl);
|
||||
}
|
||||
|
||||
template<bool ZeroMode, bool AlphaIsOne, size_t Rows>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmKernelRvv(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
{
|
||||
static_assert(Rows >= 1 && Rows <= 4, "unsupported RVV SGEMM tile height");
|
||||
|
||||
#if defined(_WIN32)
|
||||
|
||||
if constexpr (Rows == 1) {
|
||||
UNREFERENCED_PARAMETER(lda);
|
||||
UNREFERENCED_PARAMETER(ldc);
|
||||
}
|
||||
|
||||
if constexpr (AlphaIsOne) {
|
||||
UNREFERENCED_PARAMETER(alpha);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
const float* packed_b_block = B;
|
||||
float* c_block = C;
|
||||
size_t remaining_n_total = CountN;
|
||||
|
||||
do {
|
||||
const size_t count_n_block = remaining_n_total >= kPackedCountN ? kPackedCountN : remaining_n_total;
|
||||
size_t remaining_n_block = count_n_block;
|
||||
size_t column_offset = 0;
|
||||
float* c = c_block;
|
||||
|
||||
while (remaining_n_block > 0) {
|
||||
// Split a packed 16-column tile into however many lanes the current
|
||||
// machine exposes for e32,m4. This keeps the kernel VLEN-agnostic.
|
||||
const size_t vl = __riscv_vsetvl_e32m4(remaining_n_block);
|
||||
vfloat32m4_t row0_block = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
||||
vfloat32m4_t row1_block;
|
||||
vfloat32m4_t row2_block;
|
||||
vfloat32m4_t row3_block;
|
||||
|
||||
if constexpr (Rows >= 2) {
|
||||
row1_block = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
||||
}
|
||||
if constexpr (Rows >= 3) {
|
||||
row2_block = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
||||
}
|
||||
if constexpr (Rows >= 4) {
|
||||
row3_block = __riscv_vfmv_v_f_f32m4(0.0f, vl);
|
||||
}
|
||||
|
||||
const float* a = A;
|
||||
const float* b = packed_b_block + column_offset;
|
||||
size_t k = CountK;
|
||||
|
||||
while (k >= 2) {
|
||||
const float row0_a0 = a[0];
|
||||
const float row0_a1 = a[1];
|
||||
vfloat32m4_t b_elements = __riscv_vle32_v_f32m4(b, vl);
|
||||
row0_block = __riscv_vfmacc_vf_f32m4(row0_block, row0_a0, b_elements, vl);
|
||||
|
||||
if constexpr (Rows >= 2) {
|
||||
row1_block = __riscv_vfmacc_vf_f32m4(row1_block, a[lda], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 3) {
|
||||
row2_block = __riscv_vfmacc_vf_f32m4(row2_block, a[lda * 2], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 4) {
|
||||
row3_block = __riscv_vfmacc_vf_f32m4(row3_block, a[lda * 3], b_elements, vl);
|
||||
}
|
||||
|
||||
b_elements = __riscv_vle32_v_f32m4(b + kPackedCountN, vl);
|
||||
row0_block = __riscv_vfmacc_vf_f32m4(row0_block, row0_a1, b_elements, vl);
|
||||
|
||||
if constexpr (Rows >= 2) {
|
||||
row1_block = __riscv_vfmacc_vf_f32m4(row1_block, a[lda + 1], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 3) {
|
||||
row2_block = __riscv_vfmacc_vf_f32m4(row2_block, a[lda * 2 + 1], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 4) {
|
||||
row3_block = __riscv_vfmacc_vf_f32m4(row3_block, a[lda * 3 + 1], b_elements, vl);
|
||||
}
|
||||
|
||||
a += 2;
|
||||
b += kPackedCountN * 2;
|
||||
k -= 2;
|
||||
}
|
||||
|
||||
if (k > 0) {
|
||||
vfloat32m4_t b_elements = __riscv_vle32_v_f32m4(b, vl);
|
||||
row0_block = __riscv_vfmacc_vf_f32m4(row0_block, a[0], b_elements, vl);
|
||||
|
||||
if constexpr (Rows >= 2) {
|
||||
row1_block = __riscv_vfmacc_vf_f32m4(row1_block, a[lda], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 3) {
|
||||
row2_block = __riscv_vfmacc_vf_f32m4(row2_block, a[lda * 2], b_elements, vl);
|
||||
}
|
||||
if constexpr (Rows >= 4) {
|
||||
row3_block = __riscv_vfmacc_vf_f32m4(row3_block, a[lda * 3], b_elements, vl);
|
||||
}
|
||||
}
|
||||
|
||||
MlasStoreAccumulatorRvv<ZeroMode, AlphaIsOne>(c, row0_block, vl, alpha);
|
||||
|
||||
if constexpr (Rows >= 2) {
|
||||
MlasStoreAccumulatorRvv<ZeroMode, AlphaIsOne>(c + ldc, row1_block, vl, alpha);
|
||||
}
|
||||
if constexpr (Rows >= 3) {
|
||||
MlasStoreAccumulatorRvv<ZeroMode, AlphaIsOne>(c + ldc * 2, row2_block, vl, alpha);
|
||||
}
|
||||
if constexpr (Rows >= 4) {
|
||||
MlasStoreAccumulatorRvv<ZeroMode, AlphaIsOne>(c + ldc * 3, row3_block, vl, alpha);
|
||||
}
|
||||
|
||||
c += vl;
|
||||
column_offset += vl;
|
||||
remaining_n_block -= vl;
|
||||
}
|
||||
|
||||
c_block += count_n_block;
|
||||
packed_b_block += CountK * kPackedCountN;
|
||||
remaining_n_total -= count_n_block;
|
||||
|
||||
} while (remaining_n_total > 0);
|
||||
|
||||
return Rows;
|
||||
}
|
||||
|
||||
template<bool ZeroMode, bool AlphaIsOne>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasGemmFloatKernelRvvDispatchRows(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
{
|
||||
if (CountM >= 4) {
|
||||
return MlasSgemmKernelRvv<ZeroMode, AlphaIsOne, 4>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
if (CountM == 3) {
|
||||
return MlasSgemmKernelRvv<ZeroMode, AlphaIsOne, 3>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
if (CountM >= 2) {
|
||||
return MlasSgemmKernelRvv<ZeroMode, AlphaIsOne, 2>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
return MlasSgemmKernelRvv<ZeroMode, AlphaIsOne, 1>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
template<bool ZeroMode>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasGemmFloatKernelRvvDispatch(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
{
|
||||
if (alpha == 1.0f) {
|
||||
return MlasGemmFloatKernelRvvDispatchRows<ZeroMode, true>(
|
||||
A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
return MlasGemmFloatKernelRvvDispatchRows<ZeroMode, false>(
|
||||
A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasGemmFloatKernelRvv(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
if (ZeroMode) {
|
||||
return MlasGemmFloatKernelRvvDispatch<true>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
return MlasGemmFloatKernelRvvDispatch<false>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
#endif // defined(MLAS_USE_RVV)
|
||||
+115
@@ -0,0 +1,115 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
sgemm_pack_b_rvv.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements an RVV packing helper for the single precision
|
||||
matrix/matrix multiply operation (SGEMM) on riscv64.
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
#if defined(MLAS_USE_RVV)
|
||||
|
||||
#include <riscv_vector.h>
|
||||
|
||||
namespace {
|
||||
|
||||
// Keep MLAS packing in 16-column tiles, but let RVV decide the actual chunk
|
||||
// size at runtime via vsetvl so the same code works across different VLENs.
|
||||
constexpr size_t kPackedCountN = 16;
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasStoreZeroPaddedBlock(
|
||||
float* D,
|
||||
const float* B,
|
||||
size_t CountX
|
||||
)
|
||||
{
|
||||
size_t remaining = kPackedCountN;
|
||||
size_t offset = 0;
|
||||
|
||||
while (remaining > 0) {
|
||||
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
||||
__riscv_vse32_v_f32m4(D + offset, __riscv_vfmv_v_f_f32m4(0.0f, vl), vl);
|
||||
offset += vl;
|
||||
remaining -= vl;
|
||||
}
|
||||
|
||||
remaining = CountX;
|
||||
offset = 0;
|
||||
|
||||
while (remaining > 0) {
|
||||
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
||||
__riscv_vse32_v_f32m4(D + offset, __riscv_vle32_v_f32m4(B + offset, vl), vl);
|
||||
offset += vl;
|
||||
remaining -= vl;
|
||||
}
|
||||
}
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasStoreFullBlock(
|
||||
float* D,
|
||||
const float* B
|
||||
)
|
||||
{
|
||||
size_t remaining = kPackedCountN;
|
||||
size_t offset = 0;
|
||||
|
||||
while (remaining > 0) {
|
||||
const size_t vl = __riscv_vsetvl_e32m4(remaining);
|
||||
__riscv_vse32_v_f32m4(D + offset, __riscv_vle32_v_f32m4(B + offset, vl), vl);
|
||||
offset += vl;
|
||||
remaining -= vl;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
void
|
||||
MlasSgemmCopyPackBRvv(
|
||||
float* D,
|
||||
const float* B,
|
||||
size_t ldb,
|
||||
size_t CountX,
|
||||
size_t CountY
|
||||
)
|
||||
{
|
||||
while (CountX >= kPackedCountN) {
|
||||
const float* b = B;
|
||||
size_t y = CountY;
|
||||
|
||||
do {
|
||||
MlasStoreFullBlock(D, b);
|
||||
D += kPackedCountN;
|
||||
b += ldb;
|
||||
y--;
|
||||
} while (y > 0);
|
||||
|
||||
B += kPackedCountN;
|
||||
CountX -= kPackedCountN;
|
||||
}
|
||||
|
||||
if (CountX > 0) {
|
||||
size_t y = CountY;
|
||||
|
||||
do {
|
||||
MlasStoreZeroPaddedBlock(D, B, CountX);
|
||||
D += kPackedCountN;
|
||||
B += ldb;
|
||||
y--;
|
||||
} while (y > 0);
|
||||
}
|
||||
}
|
||||
|
||||
#endif // defined(MLAS_USE_RVV)
|
||||
+87
@@ -0,0 +1,87 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernel.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
#include "SgemmKernelZVECTOR.h"
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
|
||||
MLAS_FLOAT32X4 AlphaBroadcast = MlasBroadcastFloat32x4(alpha);
|
||||
|
||||
if (CountM >= 4) {
|
||||
RowsHandled = MlasSgemmProcessCount<4>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
+451
@@ -0,0 +1,451 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelZVECTOR.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "SgemmKernelZVECTOR.h"
|
||||
|
||||
#include <vecintrin.h>
|
||||
|
||||
struct MlasSgemmBroadcastAElementsZVECTOR
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount],
|
||||
const float* A,
|
||||
size_t lda
|
||||
)
|
||||
{
|
||||
ABroadcast[0][Row] = A [Row * lda];
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeAElements(
|
||||
MLAS_FLOAT32X4 AElements[RowCount],
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount]
|
||||
)
|
||||
{
|
||||
const __vector unsigned char mask0 = { 0, 1, 2, 3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 22, 23 };
|
||||
const __vector unsigned char mask3 = { 8, 9, 10, 11, 12, 13, 14, 15, 24, 25, 26, 27, 28, 29, 30, 31 };
|
||||
const __vector unsigned char mask_even = { 0, 1, 2, 3, 16, 17, 18, 19, 8, 9, 10, 11, 24, 25, 26, 27 };
|
||||
const __vector unsigned char mask_odd = { 4, 5, 6, 7, 20, 21, 22, 23, 12, 13, 14, 15, 28, 29, 30, 31 };
|
||||
|
||||
__vector float a1,a2;
|
||||
|
||||
a1 = vec_perm(AElements[0], AElements[1], mask_even);
|
||||
a2 = vec_perm(AElements[2], AElements[3], mask_even);
|
||||
ABroadcast[0] = vec_perm(a1, a2, mask0);
|
||||
ABroadcast[2] = vec_perm(a1, a2, mask3);
|
||||
a1 = vec_perm(AElements[0], AElements[1], mask_odd);
|
||||
a2 = vec_perm(AElements[2], AElements[3], mask_odd);
|
||||
ABroadcast[1] = vec_perm(a1, a2, mask0);
|
||||
ABroadcast[3] = vec_perm(a1, a2, mask3);
|
||||
}
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSgemmComputeBlockZVECTOR(
|
||||
MLAS_FLOAT32X4 acc[32],
|
||||
MLAS_FLOAT32X4 ABroadcast,
|
||||
MLAS_FLOAT32X4 A2Broadcast,
|
||||
const float* B,
|
||||
size_t CountM
|
||||
)
|
||||
{
|
||||
|
||||
MLAS_FLOAT32X4 AElements[8];
|
||||
|
||||
AElements[0] = vec_splats(ABroadcast[0]);
|
||||
AElements[1] = vec_splats(ABroadcast[1]);
|
||||
AElements[2] = vec_splats(ABroadcast[2]);
|
||||
AElements[3] = vec_splats(ABroadcast[3]);
|
||||
|
||||
if (CountM == 8) {
|
||||
AElements[4] = vec_splats(A2Broadcast[0]);
|
||||
AElements[5] = vec_splats(A2Broadcast[1]);
|
||||
AElements[6] = vec_splats(A2Broadcast[2]);
|
||||
AElements[7] = vec_splats(A2Broadcast[3]);
|
||||
}
|
||||
|
||||
MLAS_FLOAT32X4 BElements[4];
|
||||
|
||||
BElements[0] = MlasLoadFloat32x4(B);
|
||||
BElements[1] = MlasLoadFloat32x4(B + 4);
|
||||
BElements[2] = MlasLoadFloat32x4(B + 8);
|
||||
BElements[3] = MlasLoadFloat32x4(B + 12);
|
||||
|
||||
acc[0] = __builtin_s390_vfmasb(AElements[0], BElements[0], acc[0]);
|
||||
acc[1] = __builtin_s390_vfmasb(AElements[1], BElements[0], acc[1]);
|
||||
acc[2] = __builtin_s390_vfmasb(AElements[2], BElements[0], acc[2]);
|
||||
acc[3] = __builtin_s390_vfmasb(AElements[3], BElements[0], acc[3]);
|
||||
|
||||
acc[4] = __builtin_s390_vfmasb(AElements[0], BElements[1], acc[4]);
|
||||
acc[5] = __builtin_s390_vfmasb(AElements[1], BElements[1], acc[5]);
|
||||
acc[6] = __builtin_s390_vfmasb(AElements[2], BElements[1], acc[6]);
|
||||
acc[7] = __builtin_s390_vfmasb(AElements[3], BElements[1], acc[7]);
|
||||
|
||||
acc[8] = __builtin_s390_vfmasb(AElements[0], BElements[2], acc[8]);
|
||||
acc[9] = __builtin_s390_vfmasb(AElements[1], BElements[2], acc[9]);
|
||||
acc[10] = __builtin_s390_vfmasb(AElements[2], BElements[2], acc[10]);
|
||||
acc[11] = __builtin_s390_vfmasb(AElements[3], BElements[2], acc[11]);
|
||||
|
||||
acc[12] = __builtin_s390_vfmasb(AElements[0], BElements[3], acc[12]);
|
||||
acc[13] = __builtin_s390_vfmasb(AElements[1], BElements[3], acc[13]);
|
||||
acc[14] = __builtin_s390_vfmasb(AElements[2], BElements[3], acc[14]);
|
||||
acc[15] = __builtin_s390_vfmasb(AElements[3], BElements[3], acc[15]);
|
||||
|
||||
if (CountM == 8) {
|
||||
acc[16] = __builtin_s390_vfmasb(AElements[4], BElements[0], acc[16]);
|
||||
acc[17] = __builtin_s390_vfmasb(AElements[5], BElements[0], acc[17]);
|
||||
acc[18] = __builtin_s390_vfmasb(AElements[6], BElements[0], acc[18]);
|
||||
acc[19] = __builtin_s390_vfmasb(AElements[7], BElements[0], acc[19]);
|
||||
|
||||
acc[20] = __builtin_s390_vfmasb(AElements[4], BElements[1], acc[20]);
|
||||
acc[21] = __builtin_s390_vfmasb(AElements[5], BElements[1], acc[21]);
|
||||
acc[22] = __builtin_s390_vfmasb(AElements[6], BElements[1], acc[22]);
|
||||
acc[23] = __builtin_s390_vfmasb(AElements[7], BElements[1], acc[23]);
|
||||
|
||||
acc[24] = __builtin_s390_vfmasb(AElements[4], BElements[2], acc[24]);
|
||||
acc[25] = __builtin_s390_vfmasb(AElements[5], BElements[2], acc[25]);
|
||||
acc[26] = __builtin_s390_vfmasb(AElements[6], BElements[2], acc[26]);
|
||||
acc[27] = __builtin_s390_vfmasb(AElements[7], BElements[2], acc[27]);
|
||||
|
||||
acc[28] = __builtin_s390_vfmasb(AElements[4], BElements[3], acc[28]);
|
||||
acc[29] = __builtin_s390_vfmasb(AElements[5], BElements[3], acc[29]);
|
||||
acc[30] = __builtin_s390_vfmasb(AElements[6], BElements[3], acc[30]);
|
||||
acc[31] = __builtin_s390_vfmasb(AElements[7], BElements[3], acc[31]);
|
||||
}
|
||||
}
|
||||
template<size_t VectorCount>
|
||||
struct MlasSgemmStoreVectorZVECTOR
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Result[4],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
MLAS_FLOAT32X4 *rowC;
|
||||
if (ZeroMode) {
|
||||
rowC = reinterpret_cast<MLAS_FLOAT32X4 *>(&C[Row * ldc + VectorCount]);
|
||||
rowC[0] = Result[Row] * AlphaBroadcast;
|
||||
} else {
|
||||
rowC = reinterpret_cast<MLAS_FLOAT32X4 *>(&C[Row * ldc + VectorCount]);
|
||||
rowC[0] += Result[Row] * AlphaBroadcast;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
struct MlasSgemmMultiplyAlphaTrailingZVECTOR
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
MLAS_FLOAT32X4 AlphaBroadcast
|
||||
)
|
||||
{
|
||||
Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast);
|
||||
}
|
||||
};
|
||||
template<unsigned Lane>
|
||||
struct MlasSgemmStoreScalarZVECTOR
|
||||
{
|
||||
template<size_t RowCount, size_t Row>
|
||||
MLAS_FORCEINLINE
|
||||
static
|
||||
void
|
||||
Iteration(
|
||||
MLAS_FLOAT32X4 Accumulators[RowCount],
|
||||
float* C,
|
||||
size_t ldc,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
float* c = C + Row * ldc + Lane;
|
||||
float Value = Accumulators[Row][Lane];
|
||||
if (!ZeroMode) {
|
||||
Value += *c;
|
||||
}
|
||||
|
||||
*c = Value;
|
||||
}
|
||||
};
|
||||
|
||||
template<size_t RowCount>
|
||||
MLAS_FORCEINLINE
|
||||
size_t
|
||||
MlasSgemmZVECTORProcessCount(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountM,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
MLAS_FLOAT32X4 AlphaBroadcast,
|
||||
bool ZeroMode
|
||||
)
|
||||
{
|
||||
do {
|
||||
|
||||
const float* a = A;
|
||||
size_t k = CountK;
|
||||
|
||||
MLAS_FLOAT32X4 AElements[RowCount];
|
||||
MLAS_FLOAT32X4 ABroadcast[RowCount] = { 0 };
|
||||
MLAS_FLOAT32X4 A2Broadcast[RowCount] = { 0 };
|
||||
MLAS_FLOAT32X4 acc[32] = { 0 };
|
||||
MLAS_FLOAT32X4 Accumulators[2][RowCount] = {{0}};
|
||||
|
||||
//
|
||||
// Compute the output block.
|
||||
//
|
||||
while (k >= 4) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasFgemmLoadAElements>()(AElements, a, lda);
|
||||
MlasSgemmComputeAElements<RowCount>(AElements, ABroadcast);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasFgemmLoadAElements>()(AElements, a + ( lda * 4), lda);
|
||||
MlasSgemmComputeAElements<RowCount>(AElements, A2Broadcast);
|
||||
}
|
||||
MlasSgemmComputeBlockZVECTOR<RowCount>(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM);
|
||||
MlasSgemmComputeBlockZVECTOR<RowCount>(&acc[0], ABroadcast[1], A2Broadcast[1], B+16, CountM);
|
||||
MlasSgemmComputeBlockZVECTOR<RowCount>(&acc[0], ABroadcast[2], A2Broadcast[2], B+32, CountM);
|
||||
MlasSgemmComputeBlockZVECTOR<RowCount>(&acc[0], ABroadcast[3], A2Broadcast[3], B+48, CountM);
|
||||
B += 16 * 4;
|
||||
a += 4;
|
||||
k -= 4;
|
||||
}
|
||||
|
||||
while (k > 0) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElementsZVECTOR>()(ABroadcast, a, lda);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmBroadcastAElementsZVECTOR>()(A2Broadcast, a + (lda * 4), lda);
|
||||
}
|
||||
MlasSgemmComputeBlockZVECTOR<RowCount>(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM);
|
||||
a += 1;
|
||||
B += 16;
|
||||
k -= 1;
|
||||
}
|
||||
if (CountN >= 16) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<8>>()(acc + 8, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<12>>()(acc + 12, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 20, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<8>>()(acc + 24, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<12>>()(acc + 28, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
|
||||
if (CountN >= 12) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<8>>()(acc + 8, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 20, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<8>>()(acc + 24, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 12 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[1][i] = acc[i + 28];
|
||||
}
|
||||
}
|
||||
}
|
||||
if (CountN - 12 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[0][i] = acc[i + 12];
|
||||
}
|
||||
}
|
||||
} else if (CountN >= 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<4>>()(acc + 20, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 8 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[1][i] = acc[i + 24];
|
||||
}
|
||||
}
|
||||
}
|
||||
if (CountN - 8 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[0][i] = acc[i + 8];
|
||||
}
|
||||
}
|
||||
} else if (CountN >= 4) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc, C, ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreVectorZVECTOR<0>>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode);
|
||||
if (CountN - 4 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[1][i] = acc[i + 20];
|
||||
}
|
||||
}
|
||||
}
|
||||
if (CountN - 4 > 0) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[0][i] = acc[i + 4];
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[0][i] = acc[i];
|
||||
}
|
||||
|
||||
if (CountM == 8) {
|
||||
for (size_t i = 0; i < 4; ++i) {
|
||||
Accumulators[1][i] = acc[i + 16];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Store the remaining unaligned columns.
|
||||
//
|
||||
|
||||
C += (CountN & ~3);
|
||||
CountN &= 3;
|
||||
|
||||
if (CountN > 0) {
|
||||
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingZVECTOR>()(Accumulators[0], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<0>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmMultiplyAlphaTrailingZVECTOR>()(Accumulators[1], AlphaBroadcast);
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<0>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
if (CountN >= 2) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<1>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<1>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
if (CountN >= 3) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<2>>()(Accumulators[0], C, ldc, ZeroMode);
|
||||
if (CountM == 8) {
|
||||
MlasLoopUnroll<RowCount, MlasSgemmStoreScalarZVECTOR<2>>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
C += 16;
|
||||
CountN -= 16;
|
||||
|
||||
} while (CountN > 0);
|
||||
|
||||
return CountM;
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelZVECTOR(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
MLAS_FLOAT32X4 AlphaBroadcast = MlasBroadcastFloat32x4(alpha);
|
||||
|
||||
if (CountM >= 8) {
|
||||
RowsHandled = MlasSgemmZVECTORProcessCount<4>(A, B, C, 8 ,CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 4) {
|
||||
RowsHandled = MlasSgemmZVECTORProcessCount<4>(A, B, C, 4, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmProcessCount<2>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmProcessCount<1>(A, B, C, CountK, CountN, lda, ldc, AlphaBroadcast, ZeroMode);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SconvDepthwiseKernelScalar.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision direct
|
||||
convolution kernels.
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
static
|
||||
void
|
||||
MlasConv2dSingleChannel_CHW_Kernel3x3_Pad01_Dilation1(
|
||||
const MLAS_CONV_PARAMETERS* Parameters,
|
||||
const float* Input,
|
||||
const float* Filter,
|
||||
float* Output,
|
||||
const float* Zeros
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute convolution on one channel input with one filter channel.
|
||||
|
||||
Arguments:
|
||||
|
||||
Parameters - conv parameters calculated based on conv parameters like padding, strides, dilations, etc.
|
||||
|
||||
Input - input channel data start. Input is NCHW, so this pointer point to single H x W image data.
|
||||
|
||||
Filter - Whole filters are of F x CpG x FH x FW, this filter point to single FH x FW filter data.
|
||||
|
||||
Output - whole output are of N x F x OH x OW. This pointer point to single OH x OW output image data.
|
||||
|
||||
Zeroes - Point to working buffer where all 0.0f are filled.
|
||||
|
||||
--*/
|
||||
{
|
||||
const size_t W = Parameters->InputShape[1];
|
||||
const float beta = Parameters->Beta;
|
||||
|
||||
if (W > 1) {
|
||||
|
||||
const float w00 = Filter[0];
|
||||
const float w01 = Filter[1];
|
||||
const float w02 = Filter[2];
|
||||
const float w10 = Filter[3];
|
||||
const float w11 = Filter[4];
|
||||
const float w12 = Filter[5];
|
||||
const float w20 = Filter[6];
|
||||
const float w21 = Filter[7];
|
||||
const float w22 = Filter[8];
|
||||
|
||||
const size_t H = Parameters->InputShape[0];
|
||||
const size_t pad_top = Parameters->Padding[0];
|
||||
const size_t pad_left = Parameters->Padding[1];
|
||||
const size_t stride_h = Parameters->StrideShape[0];
|
||||
const size_t stride_w = Parameters->StrideShape[1];
|
||||
|
||||
// We treat pad_left, pad_top are hard require.
|
||||
// While pad_right and pad_bottom could be adjusted if they do not 100% match other parameters.
|
||||
const size_t pad_right = (((Parameters->OutputShape[1] - 1) * stride_w + 3) > (pad_left + W)) ? 1 : 0;
|
||||
|
||||
const float* row0 = (pad_top > 0) ? Zeros : (Input - pad_left);
|
||||
// Need to handle effective pad_bottom is 2 when H == 1
|
||||
const float* row1 = (H + pad_top <= 1) ? Zeros : (Input + (1 - pad_top) * W) - pad_left;
|
||||
const float* row2 = (H + pad_top <= 2) ? Zeros : (row1 + W);
|
||||
|
||||
for (size_t h = 0, out_row = Parameters->OutputShape[0]; out_row > 0; --out_row) {
|
||||
auto out_col = Parameters->OutputShape[1];
|
||||
|
||||
if (pad_left == 1) {
|
||||
float dotsum = w01 * row0[1] + w02 * row0[2] + w11 * row1[1] + w12 * row1[2] +
|
||||
w21 * row2[1] + w22 * row2[2] + (beta == 0.f ? 0.f : *Output * beta);
|
||||
*Output++ = dotsum;
|
||||
out_col--;
|
||||
row0 += stride_w;
|
||||
row1 += stride_w;
|
||||
row2 += stride_w;
|
||||
}
|
||||
|
||||
for (; out_col > pad_right; out_col--) {
|
||||
float dotsum = w00 * row0[0] + w01 * row0[1] + w02 * row0[2] + w10 * row1[0] +
|
||||
w11 * row1[1] + w12 * row1[2] + w20 * row2[0] + w21 * row2[1] +
|
||||
w22 * row2[2] + (beta == 0.f ? 0.f : *Output * beta);
|
||||
*Output++ = dotsum;
|
||||
row0 += stride_w;
|
||||
row1 += stride_w;
|
||||
row2 += stride_w;
|
||||
}
|
||||
|
||||
if (out_col == 1) { // pad_right == 1
|
||||
float dotsum = w00 * row0[0] + w01 * row0[1] + w10 * row1[0] + w11 * row1[1] +
|
||||
w20 * row2[0] + w21 * row2[1] + (beta == 0.f ? 0.f : *Output * beta);
|
||||
*Output++ = dotsum;
|
||||
}
|
||||
|
||||
h += stride_h;
|
||||
row0 = (Input + (h - pad_top) * W) - pad_left;
|
||||
row1 = row0 + W;
|
||||
row2 = (h + 2 >= H + pad_top) ? Zeros : (row1 + W);
|
||||
}
|
||||
|
||||
} else { // W == 1
|
||||
|
||||
const size_t H = Parameters->InputShape[0];
|
||||
const size_t pad_left = Parameters->Padding[1];
|
||||
const size_t pad_top = Parameters->Padding[0];
|
||||
const size_t stride_h = Parameters->StrideShape[0];
|
||||
size_t out_row = Parameters->OutputShape[0];
|
||||
|
||||
// Make sure pad_bottom is consistent with other parameters.
|
||||
size_t pad_bottom = ((out_row - 1) * stride_h + 3) > (pad_top + H) ?
|
||||
((out_row - 1) * stride_h + 3) - (pad_top + H) : 0;
|
||||
|
||||
const float w0 = Filter[pad_left ? 1 : 0];
|
||||
const float w1 = Filter[pad_left ? 4 : 3];
|
||||
const float w2 = Filter[pad_left ? 7 : 6];
|
||||
auto init_v = (beta == 0.f ? 0.f : *Output * beta);
|
||||
|
||||
if (pad_top == 1) {
|
||||
*Output++ = w1 * Input[0] + w2 * ((H + pad_top <= 2) ? 0.0f : Input[1]) + init_v;
|
||||
out_row--;
|
||||
}
|
||||
|
||||
for (const float* row = Input + pad_top * stride_h - pad_top; out_row > pad_bottom; --out_row) {
|
||||
// All pixels are in the input col
|
||||
auto init = (beta == 0.f ? 0.f : *Output * beta);
|
||||
*Output++ = w0 * row[0] + w1 * row[1] + w2 * row[2] + init;
|
||||
row += stride_h;
|
||||
}
|
||||
|
||||
if (out_row > 0) {
|
||||
// last 1 or 2 rows are from the padding zero row.
|
||||
// out_row == 1 when arrive here
|
||||
if (pad_bottom == 1) {
|
||||
const float* row = Input + H - 2;
|
||||
*Output++ = w0 * row[0] + w1 * row[1] + init_v;
|
||||
} else { // pad_bottom == 2 and H == 1 and padding_top == 0
|
||||
*Output++ = w0 * Input[0] + init_v;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
void
|
||||
MlasConvDepthwiseFloat_CHW(
|
||||
const MLAS_CONV_PARAMETERS* Parameters,
|
||||
const float* Input,
|
||||
const float* Filter,
|
||||
float* Output,
|
||||
const float* Zeros
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute depthwise convolution for one filter channel on one input channel.
|
||||
|
||||
Arguments:
|
||||
|
||||
Parameters - conv parameters calculated based on conv parameters like padding, strides, dilations, etc.
|
||||
|
||||
Input - input channel data start. Input is NCHW, so this pointer point to single H x W image data.
|
||||
|
||||
Filter - Whole filters are of F x CpG x FH x FW, this filter point to single FH x FW filter data.
|
||||
|
||||
Output - whole output are of N x F x OH x OW. This pointer point to single OH x OW output image data.
|
||||
|
||||
Zeroes - Point to working buffer where all 0.0f are filled.
|
||||
|
||||
Note:
|
||||
No checking here as it is inner loop. Logic in generating Parameters controls the check.
|
||||
|
||||
Currently only support 2d kernel 3x3.
|
||||
Will add general case and more special case if needed later.
|
||||
|
||||
--*/
|
||||
{
|
||||
MlasConv2dSingleChannel_CHW_Kernel3x3_Pad01_Dilation1(Parameters, Input, Filter, Output, Zeros);
|
||||
}
|
||||
+480
@@ -0,0 +1,480 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelScalar.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
template<bool ZeroMode, bool ProcessTwoRows>
|
||||
size_t
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB with a packing width
|
||||
of 16.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountN - 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 - Supplies the scaler multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
float Row0Block00;
|
||||
float Row0Block01;
|
||||
float Row0Block02;
|
||||
float Row0Block03;
|
||||
|
||||
float Row1Block00;
|
||||
float Row1Block01;
|
||||
float Row1Block02;
|
||||
float Row1Block03;
|
||||
|
||||
#if defined(_WIN32)
|
||||
|
||||
if (!ProcessTwoRows) {
|
||||
UNREFERENCED_PARAMETER(lda);
|
||||
UNREFERENCED_PARAMETER(ldc);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
int countb = 0;
|
||||
|
||||
do {
|
||||
|
||||
float BElements00;
|
||||
float BElements01;
|
||||
float BElements02;
|
||||
float BElements03;
|
||||
|
||||
float Row0AElements0;
|
||||
float Row0AElements1;
|
||||
float Row1AElements0;
|
||||
float Row1AElements1;
|
||||
|
||||
//
|
||||
// Clear the block accumulators.
|
||||
//
|
||||
|
||||
Row0Block00 = 0.0f;
|
||||
Row0Block01 = 0.0f;
|
||||
Row0Block02 = 0.0f;
|
||||
Row0Block03 = 0.0f;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block00 = 0.0f;
|
||||
Row1Block01 = 0.0f;
|
||||
Row1Block02 = 0.0f;
|
||||
Row1Block03 = 0.0f;
|
||||
}
|
||||
|
||||
//
|
||||
// Compute the 4x1 or 4x2 output block.
|
||||
//
|
||||
|
||||
const float* a = A;
|
||||
const float* b = B;
|
||||
size_t k = CountK;
|
||||
|
||||
while (k >= 2) {
|
||||
|
||||
Row0AElements0 = a[0];
|
||||
Row0AElements1 = a[1];
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1AElements0 = a[lda];
|
||||
Row1AElements1 = a[lda + 1];
|
||||
}
|
||||
|
||||
BElements00 = b[0];
|
||||
BElements01 = b[1];
|
||||
BElements02 = b[2];
|
||||
BElements03 = b[3];
|
||||
Row0Block00 = Row0Block00 + BElements00 * Row0AElements0;
|
||||
Row0Block01 = Row0Block01 + BElements01 * Row0AElements0;
|
||||
Row0Block02 = Row0Block02 + BElements02 * Row0AElements0;
|
||||
Row0Block03 = Row0Block03 + BElements03 * Row0AElements0;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block00 = Row1Block00 + BElements00 * Row1AElements0;
|
||||
Row1Block01 = Row1Block01 + BElements01 * Row1AElements0;
|
||||
Row1Block02 = Row1Block02 + BElements02 * Row1AElements0;
|
||||
Row1Block03 = Row1Block03 + BElements03 * Row1AElements0;
|
||||
}
|
||||
|
||||
BElements00 = b[16];
|
||||
BElements01 = b[17];
|
||||
BElements02 = b[18];
|
||||
BElements03 = b[19];
|
||||
Row0Block00 = Row0Block00 + BElements00 * Row0AElements1;
|
||||
Row0Block01 = Row0Block01 + BElements01 * Row0AElements1;
|
||||
Row0Block02 = Row0Block02 + BElements02 * Row0AElements1;
|
||||
Row0Block03 = Row0Block03 + BElements03 * Row0AElements1;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block00 = Row1Block00 + BElements00 * Row1AElements1;
|
||||
Row1Block01 = Row1Block01 + BElements01 * Row1AElements1;
|
||||
Row1Block02 = Row1Block02 + BElements02 * Row1AElements1;
|
||||
Row1Block03 = Row1Block03 + BElements03 * Row1AElements1;
|
||||
}
|
||||
|
||||
a += 2;
|
||||
b += 32;
|
||||
k -= 2;
|
||||
}
|
||||
|
||||
if (k > 0) {
|
||||
|
||||
Row0AElements0 = a[0];
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1AElements0 = a[lda];
|
||||
}
|
||||
|
||||
BElements00 = b[0];
|
||||
BElements01 = b[1];
|
||||
BElements02 = b[2];
|
||||
BElements03 = b[3];
|
||||
Row0Block00 = Row0Block00 + BElements00 * Row0AElements0;
|
||||
Row0Block01 = Row0Block01 + BElements01 * Row0AElements0;
|
||||
Row0Block02 = Row0Block02 + BElements02 * Row0AElements0;
|
||||
Row0Block03 = Row0Block03 + BElements03 * Row0AElements0;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block00 = Row1Block00 + BElements00 * Row1AElements0;
|
||||
Row1Block01 = Row1Block01 + BElements01 * Row1AElements0;
|
||||
Row1Block02 = Row1Block02 + BElements02 * Row1AElements0;
|
||||
Row1Block03 = Row1Block03 + BElements03 * Row1AElements0;
|
||||
}
|
||||
}
|
||||
|
||||
//
|
||||
// Multiply by the alpha value.
|
||||
//
|
||||
|
||||
Row0Block00 = Row0Block00 * alpha;
|
||||
Row0Block01 = Row0Block01 * alpha;
|
||||
Row0Block02 = Row0Block02 * alpha;
|
||||
Row0Block03 = Row0Block03 * alpha;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
Row1Block00 = Row1Block00 * alpha;
|
||||
Row1Block01 = Row1Block01 * alpha;
|
||||
Row1Block02 = Row1Block02 * alpha;
|
||||
Row1Block03 = Row1Block03 * alpha;
|
||||
}
|
||||
|
||||
if (CountN >= 4) {
|
||||
|
||||
//
|
||||
// Store the entire output block.
|
||||
//
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block00 = Row0Block00 + C[0];
|
||||
Row0Block01 = Row0Block01 + C[1];
|
||||
Row0Block02 = Row0Block02 + C[2];
|
||||
Row0Block03 = Row0Block03 + C[3];
|
||||
}
|
||||
|
||||
C[0] = Row0Block00;
|
||||
C[1] = Row0Block01;
|
||||
C[2] = Row0Block02;
|
||||
C[3] = Row0Block03;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block00 = Row1Block00 + C[ldc];
|
||||
Row1Block01 = Row1Block01 + C[ldc + 1];
|
||||
Row1Block02 = Row1Block02 + C[ldc + 2];
|
||||
Row1Block03 = Row1Block03 + C[ldc + 3];
|
||||
}
|
||||
|
||||
C[ldc] = Row1Block00;
|
||||
C[ldc + 1] = Row1Block01;
|
||||
C[ldc + 2] = Row1Block02;
|
||||
C[ldc + 3] = Row1Block03;
|
||||
}
|
||||
|
||||
} else {
|
||||
|
||||
//
|
||||
// Store the partial output block.
|
||||
//
|
||||
if ((CountN & 2) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block00 = Row0Block00 + C[0];
|
||||
Row0Block01 = Row0Block01 + C[1];
|
||||
}
|
||||
|
||||
C[0] = Row0Block00;
|
||||
C[1] = Row0Block01;
|
||||
Row0Block00 = Row0Block02;
|
||||
Row0Block01 = Row0Block03;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block00 = Row1Block00 + C[ldc];
|
||||
Row1Block01 = Row1Block01 + C[ldc + 1];
|
||||
}
|
||||
|
||||
C[ldc] = Row1Block00;
|
||||
C[ldc + 1] = Row1Block01;
|
||||
Row1Block00 = Row1Block02;
|
||||
Row1Block01 = Row1Block03;
|
||||
}
|
||||
|
||||
C += 2;
|
||||
}
|
||||
|
||||
if ((CountN & 1) != 0) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row0Block00 = Row0Block00 + C[0];
|
||||
}
|
||||
|
||||
C[0] = Row0Block00;
|
||||
|
||||
if (ProcessTwoRows) {
|
||||
|
||||
if (!ZeroMode) {
|
||||
Row1Block00 = Row1Block00 + C[ldc];
|
||||
}
|
||||
|
||||
C[ldc] = Row1Block00;
|
||||
}
|
||||
}
|
||||
|
||||
break;
|
||||
}
|
||||
|
||||
B += 4;
|
||||
C += 4;
|
||||
CountN -= 4;
|
||||
|
||||
countb = (countb + 1) % 4;
|
||||
if (countb == 0) {
|
||||
B += CountK * 16 - 16;
|
||||
}
|
||||
} while (CountN > 0);
|
||||
|
||||
return ProcessTwoRows ? 2 : 1;
|
||||
}
|
||||
|
||||
template<bool ZeroMode>
|
||||
size_t
|
||||
MlasSgemmKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scaler multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
size_t RowsHandled;
|
||||
|
||||
if (CountM >= 2) {
|
||||
RowsHandled = MlasSgemmKernel<ZeroMode, true>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
} else {
|
||||
RowsHandled = MlasSgemmKernel<ZeroMode, false>(A, B, C, CountK, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
return RowsHandled;
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelZero(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scaler multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
return MlasSgemmKernel<true>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
MlasSgemmKernelAdd(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountM,
|
||||
size_t CountN,
|
||||
size_t lda,
|
||||
size_t ldc,
|
||||
float alpha
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scaler multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
Returns the number of rows handled.
|
||||
|
||||
--*/
|
||||
{
|
||||
return MlasSgemmKernel<false>(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
|
||||
}
|
||||
+169
@@ -0,0 +1,169 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemvKernelScalar.cpp
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/vector
|
||||
multiply operation (SGEMV).
|
||||
|
||||
--*/
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasGemvFloatKernel(
|
||||
const float* A,
|
||||
const float* B,
|
||||
float* C,
|
||||
size_t CountK,
|
||||
size_t CountN,
|
||||
size_t ldb,
|
||||
bool ZeroMode
|
||||
)
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows. This handles the special case of M=1.
|
||||
|
||||
The elements in matrix B are not transposed.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number
|
||||
of rows from matrix B to iterate over.
|
||||
|
||||
CountN - Supplies the number of columns from matrix B and matrix C to
|
||||
iterate over.
|
||||
|
||||
ldb - Supplies the first dimension of matrix B.
|
||||
|
||||
ZeroMode - Supplies true if the output matrix must be zero initialized,
|
||||
else false if the output matrix is accumulated into.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
{
|
||||
if (ZeroMode && CountK > 0) {
|
||||
float* c = C;
|
||||
const float* b = B;
|
||||
const float A0 = A[0];
|
||||
auto N = CountN;
|
||||
constexpr size_t kWidth = 4;
|
||||
for (; N >= kWidth; N -= kWidth) {
|
||||
c[0] = A0 * b[0];
|
||||
c[1] = A0 * b[1];
|
||||
c[2] = A0 * b[2];
|
||||
c[3] = A0 * b[3];
|
||||
c += kWidth;
|
||||
b += kWidth;
|
||||
}
|
||||
|
||||
for (; N > 0; N--) {
|
||||
c[0] = A0 * b[0];
|
||||
c++;
|
||||
b++;
|
||||
}
|
||||
A++;
|
||||
B += ldb;
|
||||
|
||||
CountK--;
|
||||
}
|
||||
|
||||
for (; CountK >= 4; CountK -= 4) {
|
||||
float* c = C;
|
||||
const float* b = B;
|
||||
const float* b2 = B + ldb * 2;
|
||||
|
||||
const float A0 = A[0];
|
||||
const float A1 = A[1];
|
||||
const float A2 = A[2];
|
||||
const float A3 = A[3];
|
||||
|
||||
constexpr size_t kWidth = 4;
|
||||
auto N = CountN;
|
||||
for (; N >= kWidth; N -= kWidth) {
|
||||
float c0 = c[0] + A0 * b[0];
|
||||
float c1 = c[1] + A0 * b[1];
|
||||
float c2 = c[2] + A0 * b[2];
|
||||
float c3 = c[3] + A0 * b[3];
|
||||
|
||||
c0 += A1 * b[ldb + 0];
|
||||
c1 += A1 * b[ldb + 1];
|
||||
c2 += A1 * b[ldb + 2];
|
||||
c3 += A1 * b[ldb + 3];
|
||||
|
||||
c0 += A2 * b2[0];
|
||||
c1 += A2 * b2[1];
|
||||
c2 += A2 * b2[2];
|
||||
c3 += A2 * b2[3];
|
||||
|
||||
c0 += A3 * b2[ldb + 0];
|
||||
c1 += A3 * b2[ldb + 1];
|
||||
c2 += A3 * b2[ldb + 2];
|
||||
c3 += A3 * b2[ldb + 3];
|
||||
|
||||
c[0] = c0;
|
||||
c[1] = c1;
|
||||
c[2] = c2;
|
||||
c[3] = c3;
|
||||
|
||||
c += kWidth;
|
||||
b += kWidth;
|
||||
b2 += kWidth;
|
||||
}
|
||||
|
||||
for (; N > 0; N--) {
|
||||
c[0] += A0 * b[0] + A1 * b[ldb] + A2 * b2[0] + A3 * b2[ldb];
|
||||
c++;
|
||||
b++;
|
||||
b2++;
|
||||
}
|
||||
|
||||
B += 4 * ldb;
|
||||
A += 4;
|
||||
}
|
||||
|
||||
for (; CountK > 0; CountK--) {
|
||||
float* c = C;
|
||||
const float* b = B;
|
||||
const float A0 = A[0];
|
||||
constexpr size_t kWidth = 4;
|
||||
auto N = CountN;
|
||||
for (; N >= kWidth; N -= kWidth) {
|
||||
c[0] += A0 * b[0];
|
||||
c[1] += A0 * b[1];
|
||||
c[2] += A0 * b[2];
|
||||
c[3] += A0 * b[3];
|
||||
|
||||
c += kWidth;
|
||||
b += kWidth;
|
||||
}
|
||||
|
||||
for (; N > 0; N--) {
|
||||
c[0] += A0 * b[0];
|
||||
c++;
|
||||
b++;
|
||||
}
|
||||
B += ldb;
|
||||
A++;
|
||||
}
|
||||
}
|
||||
Vendored
+1740
File diff suppressed because it is too large
Load Diff
Vendored
+129
@@ -0,0 +1,129 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
softmax.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module includes kernel function prototypes and helper functions for
|
||||
softmax.
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "mlasi.h"
|
||||
|
||||
struct MLAS_SOFTMAX_DISPATCH {
|
||||
/**
|
||||
* @brief Compute the hyperbolic tangent function for each element of the input array
|
||||
* @param Input Address of the input array. Valid in [-3.51562, 3.51562].
|
||||
* @param Output Address of the output array. Could be the same as the input array.
|
||||
* @param N Number of elements in the input array
|
||||
*/
|
||||
typedef void(Tanh_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
Tanh_Fp16_Fn* Tanh_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Compute the softcap function for each element of the input array. Use tanh activation.
|
||||
* @param Input Address of the input array. Valid if input / softcap in [-3.51562, 3.51562].
|
||||
* @param Output Address of the output array. Could be the same as the input array.
|
||||
* @param N Number of elements in the input array
|
||||
* @param Softcap The softcap value
|
||||
*/
|
||||
typedef void(Softcap_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N,
|
||||
const MLAS_FP16 Softcap
|
||||
);
|
||||
|
||||
Softcap_Fp16_Fn* Softcap_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Compute the exponential function for each element of the input array.
|
||||
* @param Input Address of the input array. Valid in [-17.3287, 11.0904].
|
||||
* @param Output Address of the output array. Could be the same as the input array.
|
||||
* @param N Number of elements in the input array
|
||||
*/
|
||||
typedef void(Exp_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
Exp_Fp16_Fn* Exp_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Find the max value among the input array
|
||||
* @param Input Address of the input array
|
||||
* @param N Number of elements in the input array
|
||||
*/
|
||||
typedef MLAS_FP16(ReduceMax_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
size_t N
|
||||
);
|
||||
|
||||
ReduceMax_Fp16_Fn* ReduceMax_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Compute the expotential function for each element of the input array and returnt he sum. It has smaller
|
||||
* dynamic range for the input than Exp_Fp16_Fn thus is faster.
|
||||
* @param Input Address of the input array. Valid in [-10.7438, 10.7438]
|
||||
* @param Output Address of the output array. Could be the same as the input array or nullptr.
|
||||
* @param N Number of elements in the input array
|
||||
* @param NegativeMaximum The negative of the maximum value in the input array
|
||||
*/
|
||||
typedef MLAS_FP16(SumExp_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N,
|
||||
const MLAS_FP16 NegativeMaximum
|
||||
);
|
||||
|
||||
SumExp_Fp16_Fn* SumExp_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Compute the softmax output for each element of the input array. input / sum.
|
||||
* @param Input Address of the input array. Values of exp(x)
|
||||
* @param Output Address of the output array. Could be the same as the input array.
|
||||
* @param N Number of elements in the input array
|
||||
* @param Sum Sum of exp(input)
|
||||
*/
|
||||
typedef void(Softmax_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N,
|
||||
const MLAS_FP16 Sum
|
||||
);
|
||||
|
||||
Softmax_Fp16_Fn* Softmax_Fp16 = nullptr;
|
||||
|
||||
/**
|
||||
* @brief Compute the log softmax output for each element of the input array. input - max - logSum
|
||||
* @param Input Address of the input array
|
||||
* @param Output Address of the output array. Could be the same as the input array.
|
||||
* @param N Number of elements in the input array
|
||||
* @param NagativeMaximum The negative of the maximum value in the input array
|
||||
* @param LogSum The logarithm of the sum of the exponential function of the input array
|
||||
*/
|
||||
typedef void(LogSoftmax_Fp16_Fn)(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N,
|
||||
const MLAS_FP16 NagativeMaximum,
|
||||
const MLAS_FP16 LogSum
|
||||
);
|
||||
|
||||
LogSoftmax_Fp16_Fn* LogSoftmax_Fp16 = nullptr;
|
||||
};
|
||||
Vendored
+680
@@ -0,0 +1,680 @@
|
||||
/*++
|
||||
|
||||
Copyright 2025 FUJITSU LIMITED
|
||||
|
||||
Module Name:
|
||||
|
||||
mlasi_sve.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module contains the procedure prototypes for the SVE intrinsics.
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "../mlasi.h"
|
||||
#include <arm_sve.h> // SVE intrinsic header
|
||||
|
||||
#ifndef __clang__
|
||||
#pragma GCC push_options
|
||||
#pragma GCC target("arch=armv8.2-a+sve")
|
||||
|
||||
// Use Clang-specific per-function attribute
|
||||
#ifdef __clang__
|
||||
#define MLAS_SVE_TARGET __attribute__((target("arch=armv8.2-a+sve")))
|
||||
#else
|
||||
#define MLAS_SVE_TARGET
|
||||
#endif
|
||||
|
||||
typedef svfloat32_t MLAS_SVFLOAT32;
|
||||
typedef svint32_t MLAS_SVINT32;
|
||||
typedef svuint32_t MLAS_SVUINT32;
|
||||
typedef svbool_t MLAS_SVBOOL;
|
||||
typedef svfloat16_t MLAS_SVFLOAT16;
|
||||
typedef svuint16_t MLAS_SVUINT16;
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveErfFP16Kernel(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveTanhFP16Kernel(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveGeluFP16Kernel(
|
||||
const MLAS_FP16* Input,
|
||||
MLAS_FP16* Output,
|
||||
MLAS_FP16* Temp,
|
||||
size_t N,
|
||||
MLAS_GELU_ALGORITHM Algo
|
||||
);
|
||||
// function declarations
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveComputeExpVector(
|
||||
MLAS_SVBOOL Pred,
|
||||
MLAS_SVFLOAT32 Vector
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveComputeExpF32Kernel(
|
||||
const float* Input,
|
||||
float* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveComputeSumExpVector(
|
||||
MLAS_SVBOOL Pred,
|
||||
MLAS_SVFLOAT32 Vector,
|
||||
MLAS_SVFLOAT32 NegativeMaximumVector
|
||||
);
|
||||
|
||||
float
|
||||
MLASCALL
|
||||
MlasSveComputeSumExpF32Kernel(
|
||||
const float* Input,
|
||||
float* Output,
|
||||
size_t N,
|
||||
const float* NegativeMaximum
|
||||
);
|
||||
|
||||
float MLASCALL
|
||||
MlasSveReduceMaximumF32Kernel(
|
||||
const float* Input,
|
||||
size_t N
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveReduceMinimumMaximumF32Kernel(
|
||||
const float* Input,
|
||||
float* Min,
|
||||
float* Max,
|
||||
size_t N
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveComputeSoftmaxOutputF32Kernel(
|
||||
float* Output,
|
||||
size_t N,
|
||||
const float* Parameters
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveComputeLogSoftmaxOutputF32Kernel(
|
||||
const float* Input,
|
||||
float* Output,
|
||||
size_t N,
|
||||
const float* Parameters
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveErfKernel(
|
||||
const float* Input,
|
||||
float* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
void
|
||||
MLASCALL
|
||||
MlasSveLogisticKernel(
|
||||
const float* Input,
|
||||
float* Output,
|
||||
size_t N
|
||||
);
|
||||
|
||||
//MLAS API for SVE intrinsics
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveReinterpretAsInt32(MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svreinterpret_s32_f32(Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVUINT32
|
||||
MlasSveReinterpretAsUInt32(MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svreinterpret_u32_f32(Vector);
|
||||
}
|
||||
|
||||
// Reinterprets an unsigned 32-bit vector as a 32-bit floating-point vector.
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveReinterpretAsFLOAT32(MLAS_SVUINT32 Vector)
|
||||
{
|
||||
return svreinterpret_f32_u32(Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveCastToInt32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svcvt_s32_f32_z(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveCastToFloat32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector)
|
||||
{
|
||||
return svcvt_f32_s32_z(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveBroadcastInt32(int32_t Value)
|
||||
{
|
||||
return svdup_n_s32(Value);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveLoadInt32(MLAS_SVBOOL Pred, const int32_t* Buffer)
|
||||
{
|
||||
return svld1_s32(Pred, Buffer);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSveStoreInt32(MLAS_SVBOOL Pred, int32_t* Buffer, MLAS_SVINT32 Vector)
|
||||
{
|
||||
svst1_s32(Pred, Buffer, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveAddInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svadd_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveSubtractInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svsub_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveAndInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svand_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVUINT32
|
||||
MlasSveAndUInt32(MLAS_SVBOOL Pred, MLAS_SVUINT32 Vector1, MLAS_SVUINT32 Vector2)
|
||||
{
|
||||
return svand_u32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveOrInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svorr_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveAndNotInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 VectorNot, MLAS_SVINT32 Vector)
|
||||
{
|
||||
return svand_s32_m(Pred, svnot_s32_z(Pred, VectorNot), Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveXorInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return sveor_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveBlendInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2, MLAS_SVINT32 Selection)
|
||||
{
|
||||
return MlasSveOrInt32(
|
||||
Pred,
|
||||
MlasSveAndInt32(Pred, Vector2, Selection),
|
||||
MlasSveAndNotInt32(Pred, Selection, Vector1)
|
||||
);
|
||||
}
|
||||
|
||||
template<unsigned ShiftCount>
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVUINT32
|
||||
MlasSveShiftLeftUInt32(MLAS_SVBOOL Pred, MLAS_SVUINT32 Vector)
|
||||
{
|
||||
return svlsl_n_u32_z(Pred, Vector, ShiftCount);
|
||||
}
|
||||
|
||||
template<unsigned ShiftCount>
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveShiftLeftInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector)
|
||||
{
|
||||
return svlsl_n_s32_z(Pred, Vector, ShiftCount);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVUINT32
|
||||
MlasSveShiftRightInt32(MLAS_SVBOOL Pred, MLAS_SVUINT32 Vector, uint ShiftCount)
|
||||
{
|
||||
return svlsr_n_u32_m(Pred, Vector, ShiftCount);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveMaximumInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svmax_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVINT32
|
||||
MlasSveMinimumInt32(MLAS_SVBOOL Pred, MLAS_SVINT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svmin_s32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveReinterpretAsFloat32(MLAS_SVINT32 Vector)
|
||||
{
|
||||
return svreinterpret_f32_s32(Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveBroadcastFloat32(float Value)
|
||||
{
|
||||
return svdup_n_f32(Value);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVUINT32
|
||||
MlasSveBroadcastUINT32(uint Value)
|
||||
{
|
||||
return svdup_n_u32(Value);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveBroadcastFloat32(const float* Value)
|
||||
{
|
||||
return svld1_f32(svptrue_b32(), Value);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveZeroFloat32(void)
|
||||
{
|
||||
return svdup_n_f32(0.0f);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveLoadFloat32(MLAS_SVBOOL Pred, const float* Buffer)
|
||||
{
|
||||
return svld1_f32(Pred, Buffer);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSveStoreFloat32(MLAS_SVBOOL Pred, float* Buffer, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
svst1_f32(Pred, Buffer, Vector);
|
||||
}
|
||||
|
||||
template<unsigned Lane>
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSveStoreLaneFloat32(float* Buffer, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
svbool_t Pred = svwhilelt_b32(Lane, Lane + 1);
|
||||
svst1_f32(Pred, Buffer, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
void
|
||||
MlasSveStoreLowHalfFloat32(float* Buffer, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
svbool_t Pred = svwhilelt_b32(0, (int32_t)svcntw() / 2);
|
||||
svst1_f32(Pred, Buffer, Vector);
|
||||
}
|
||||
|
||||
template<unsigned Lane>
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
float
|
||||
MlasSveExtractLaneFloat32(MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
float TmpBuffer[1];
|
||||
svbool_t Pred = svwhilelt_b32(Lane, Lane + 1);
|
||||
svst1_f32(Pred, TmpBuffer, Vector);
|
||||
return TmpBuffer[0];
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveInterleaveLowFloat32(MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svzip1_f32(Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveInterleaveHighFloat32(MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svzip2_f32(Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveAddFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svadd_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveSubtractFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svsub_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMultiplyFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svmul_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveExpFloat32(MLAS_SVUINT32 Vector)
|
||||
{
|
||||
return svexpa_f32(Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveScaleFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVINT32 Vector2)
|
||||
{
|
||||
return svscale_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveRoundINTFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svrintm_f32_z(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMultiplyAddFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2, MLAS_SVFLOAT32 Vector3)
|
||||
{
|
||||
return svmla_f32_m(Pred, Vector3, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMultiplyAddFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, float Scalar2, MLAS_SVFLOAT32 Vector3)
|
||||
{
|
||||
return MlasSveMultiplyAddFloat32(Pred, Vector1, MlasSveBroadcastFloat32(Scalar2), Vector3);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMultiplyAddFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2, float Scalar3)
|
||||
{
|
||||
return MlasSveMultiplyAddFloat32(Pred, Vector1, Vector2, MlasSveBroadcastFloat32(Scalar3));
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveDivideFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svdiv_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveGreaterThanFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
// Compare Vector1 and Vector2, return a predicate vector
|
||||
svbool_t cmp_mask = svcmpgt_f32(Pred, Vector1, Vector2);
|
||||
|
||||
//Convert predicate to uint32_t mask
|
||||
svuint32_t mask_bits = svdup_u32_z(cmp_mask, 0xFFFFFFFF);
|
||||
|
||||
//Reinterpret to float32
|
||||
return svreinterpret_f32_u32(mask_bits);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveAndFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return MlasSveReinterpretAsFloat32(
|
||||
MlasSveAndInt32(
|
||||
Pred,
|
||||
MlasSveReinterpretAsInt32(Vector1),
|
||||
MlasSveReinterpretAsInt32(Vector2)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveOrFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return MlasSveReinterpretAsFloat32(
|
||||
MlasSveOrInt32(
|
||||
Pred,
|
||||
MlasSveReinterpretAsInt32(Vector1),
|
||||
MlasSveReinterpretAsInt32(Vector2)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveAndNotFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return MlasSveReinterpretAsFloat32(
|
||||
MlasSveAndNotInt32(
|
||||
Pred,
|
||||
MlasSveReinterpretAsInt32(Vector1),
|
||||
MlasSveReinterpretAsInt32(Vector2)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveXorFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return MlasSveReinterpretAsFloat32(
|
||||
MlasSveXorInt32(
|
||||
Pred,
|
||||
MlasSveReinterpretAsInt32(Vector1),
|
||||
MlasSveReinterpretAsInt32(Vector2)
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveBlendFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2, MLAS_SVFLOAT32 Selection)
|
||||
{
|
||||
return MlasSveOrFloat32(
|
||||
Pred,
|
||||
MlasSveAndFloat32(Pred, Vector2, Selection),
|
||||
MlasSveAndFloat32(Pred, Vector1, Selection)
|
||||
);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMaximumFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svmax_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveMinimumFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector1, MLAS_SVFLOAT32 Vector2)
|
||||
{
|
||||
return svmin_f32_m(Pred, Vector1, Vector2);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveClampFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Value, float LowerRange, float UpperRange)
|
||||
{
|
||||
Value = MlasSveMaximumFloat32(Pred, MlasSveBroadcastFloat32(LowerRange), Value);
|
||||
Value = MlasSveMinimumFloat32(Pred, MlasSveBroadcastFloat32(UpperRange), Value);
|
||||
return Value;
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
float
|
||||
MlasSveReduceAddFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svaddv_f32(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
float
|
||||
MlasSveReduceMaximumFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svmaxv_f32(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
float
|
||||
MlasSveReduceMinimumFloat32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
return svminv_f32(Pred, Vector);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSvePowerOf2Float32(MLAS_SVBOOL Pred, MLAS_SVFLOAT32 Vector)
|
||||
{
|
||||
MLAS_SVINT32 emm0 = MlasSveAddInt32(
|
||||
Pred,
|
||||
MlasSveCastToInt32(Pred, Vector),
|
||||
MlasSveBroadcastInt32(127)
|
||||
);
|
||||
return MlasSveReinterpretAsFloat32(MlasSveShiftLeftInt32<23>(Pred, emm0));
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVFLOAT32
|
||||
MlasSveSelect(svbool_t Pred, MLAS_SVFLOAT32 TrueValue, MLAS_SVFLOAT32 FalseValue)
|
||||
{
|
||||
return svsel_f32(Pred, TrueValue, FalseValue);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVBOOL
|
||||
MlasSveCompareLessThan(svbool_t Pred, MLAS_SVFLOAT32 A, MLAS_SVFLOAT32 B)
|
||||
{
|
||||
return svcmplt_f32(Pred, A, B);
|
||||
}
|
||||
|
||||
MLAS_SVE_TARGET
|
||||
MLAS_FORCEINLINE
|
||||
MLAS_SVBOOL
|
||||
MlasSveCompareGreaterThan(svbool_t Pred, MLAS_SVFLOAT32 A, MLAS_SVFLOAT32 B)
|
||||
{
|
||||
return svcmpgt_f32(Pred, A, B);
|
||||
}
|
||||
|
||||
// GCC: Pop options after SVE-specific functions
|
||||
#ifndef __clang__
|
||||
#pragma GCC pop_options
|
||||
#endif
|
||||
|
||||
#endif
|
||||
|
||||
+435
@@ -0,0 +1,435 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelAvx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses AVX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
//
|
||||
// Stack frame layout for the SGEMM kernel.
|
||||
//
|
||||
|
||||
.equ .LSgemmKernelFrame_SavedEdi, 0
|
||||
.equ .LSgemmKernelFrame_SavedEsi, 4
|
||||
.equ .LSgemmKernelFrame_SavedEbx, 8
|
||||
.equ .LSgemmKernelFrame_SavedEbp, 12
|
||||
.equ .LSgemmKernelFrame_ReturnAddress, 16
|
||||
.equ .LSgemmKernelFrame_MatrixA, 20
|
||||
.equ .LSgemmKernelFrame_MatrixB, 24
|
||||
.equ .LSgemmKernelFrame_MatrixC, 28
|
||||
.equ .LSgemmKernelFrame_CountK, 32
|
||||
.equ .LSgemmKernelFrame_CountM, 36
|
||||
.equ .LSgemmKernelFrame_CountN, 40
|
||||
.equ .LSgemmKernelFrame_lda, 44
|
||||
.equ .LSgemmKernelFrame_ldc, 48
|
||||
.equ .LSgemmKernelFrame_alpha, 52
|
||||
.equ .LSgemmKernelFrame_ZeroMode, 56
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for a 16xN block 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.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
ebx - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
ecx - Supplies the address into the matrix A data.
|
||||
|
||||
edx - Supplies the address into the matrix B data.
|
||||
|
||||
ymm4-ymm7 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvxBy16 RowCount, VectorOffset, BroadcastOffset
|
||||
|
||||
.if \RowCount\() == 1
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+\BroadcastOffset\()]
|
||||
vmulps ymm1,ymm3,YMMWORD PTR [edx+\VectorOffset\()]
|
||||
vaddps ymm4,ymm1,ymm4
|
||||
vmulps ymm3,ymm3,YMMWORD PTR [edx+\VectorOffset\()+32]
|
||||
vaddps ymm5,ymm3,ymm5
|
||||
.else
|
||||
vmovaps ymm0,YMMWORD PTR [edx+\VectorOffset\()]
|
||||
vmovaps ymm1,YMMWORD PTR [edx+\VectorOffset\()+32]
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+\BroadcastOffset\()]
|
||||
vmulps ymm2,ymm3,ymm0
|
||||
vaddps ymm4,ymm2,ymm4
|
||||
vmulps ymm2,ymm3,ymm1
|
||||
vaddps ymm5,ymm2,ymm5
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+ebx+\BroadcastOffset\()]
|
||||
vmulps ymm2,ymm3,ymm0
|
||||
vaddps ymm6,ymm2,ymm6
|
||||
vmulps ymm2,ymm3,ymm1
|
||||
vaddps ymm7,ymm2,ymm7
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for a 8xN block 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.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
ebx - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
ecx - Supplies the address into the matrix A data.
|
||||
|
||||
edx - Supplies the address into the matrix B data.
|
||||
|
||||
ymm4-ymm7 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvxBy8 RowCount, VectorOffset, BroadcastOffset
|
||||
|
||||
.if \RowCount\() == 1
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+\BroadcastOffset\()]
|
||||
vmulps ymm3,ymm3,YMMWORD PTR [edx+\VectorOffset\()]
|
||||
vaddps ymm5,ymm3,ymm5
|
||||
.else
|
||||
vmovaps ymm0,YMMWORD PTR [edx+\VectorOffset\()]
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+\BroadcastOffset\()]
|
||||
vmulps ymm3,ymm3,ymm0
|
||||
vaddps ymm5,ymm3,ymm5
|
||||
vbroadcastss ymm3,DWORD PTR [ecx+ebx+\BroadcastOffset\()]
|
||||
vmulps ymm3,ymm3,ymm0
|
||||
vaddps ymm7,ymm3,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:
|
||||
|
||||
ebx - Supplies the number of bytes to the next row of matrix A.
|
||||
|
||||
ecx - Supplies the address into the matrix A data.
|
||||
|
||||
edx - Supplies the address into the matrix B data.
|
||||
|
||||
edi - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
ymm4-ymm7 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvxLoop ComputeBlock, RowCount
|
||||
|
||||
sub edi,4
|
||||
jb .LProcessRemainingBlocks\@
|
||||
|
||||
.LComputeBlockBy4Loop\@:
|
||||
\ComputeBlock\() \RowCount\(), 0, 0
|
||||
\ComputeBlock\() \RowCount\(), 16*4, 4
|
||||
sub edx,-32*4 # advance matrix B by 32 columns
|
||||
\ComputeBlock\() \RowCount\(), 0, 8
|
||||
\ComputeBlock\() \RowCount\(), 16*4, 12
|
||||
sub edx,-32*4 # advance matrix B by 32 columns
|
||||
add ecx,4*4 # advance matrix A by 4 columns
|
||||
sub edi,4
|
||||
jae .LComputeBlockBy4Loop\@
|
||||
|
||||
.LProcessRemainingBlocks\@:
|
||||
add edi,4 # correct for over-subtract above
|
||||
jz .LOutputBlock\@
|
||||
|
||||
.LComputeBlockBy1Loop\@:
|
||||
\ComputeBlock\() \RowCount\(), 0, 0
|
||||
add edx,16*4 # advance matrix B by 16 columns
|
||||
add ecx,4 # advance matrix A by 1 column
|
||||
dec edi
|
||||
jne .LComputeBlockBy1Loop\@
|
||||
|
||||
.LOutputBlock\@:
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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 MlasGemmFloatKernelAvx
|
||||
|
||||
push ebp
|
||||
push ebx
|
||||
push esi
|
||||
push edi
|
||||
mov edx,.LSgemmKernelFrame_MatrixB[esp]
|
||||
mov esi,.LSgemmKernelFrame_MatrixC[esp]
|
||||
mov ebp,.LSgemmKernelFrame_CountN[esp]
|
||||
|
||||
//
|
||||
// Process 2 rows of the matrices.
|
||||
//
|
||||
|
||||
cmp DWORD PTR .LSgemmKernelFrame_CountM[esp],2
|
||||
jb .LProcessCountMLessThan2
|
||||
mov BYTE PTR .LSgemmKernelFrame_CountM[esp],2
|
||||
mov eax,.LSgemmKernelFrame_ldc[esp]
|
||||
mov ebx,.LSgemmKernelFrame_lda[esp]
|
||||
shl eax,2 # convert ldc to bytes
|
||||
shl ebx,2 # convert lda to bytes
|
||||
cmp ebp,8
|
||||
jbe .LProcessRemainingCountN2
|
||||
|
||||
.LProcessNextColumnLoop16x2:
|
||||
mov edi,.LSgemmKernelFrame_CountK[esp]
|
||||
mov ecx,.LSgemmKernelFrame_MatrixA[esp]
|
||||
vxorps xmm4,xmm4,xmm4 # clear block accumulators
|
||||
vxorps xmm5,xmm5,xmm5
|
||||
vxorps xmm6,xmm6,xmm6
|
||||
vxorps xmm7,xmm7,xmm7
|
||||
ComputeBlockAvxLoop ComputeBlockAvxBy16, 2
|
||||
vbroadcastss ymm2,DWORD PTR .LSgemmKernelFrame_alpha[esp]
|
||||
vmulps ymm4,ymm4,ymm2 # multiply by alpha
|
||||
vmulps ymm5,ymm5,ymm2
|
||||
vmulps ymm6,ymm6,ymm2
|
||||
vmulps ymm7,ymm7,ymm2
|
||||
sub ebp,16
|
||||
jb .LOutputMasked16x2Block
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateOutput16x2
|
||||
vaddps ymm4,ymm4,YMMWORD PTR [esi]
|
||||
vaddps ymm5,ymm5,YMMWORD PTR [esi+32]
|
||||
vaddps ymm6,ymm6,YMMWORD PTR [esi+eax]
|
||||
vaddps ymm7,ymm7,YMMWORD PTR [esi+eax+32]
|
||||
|
||||
.LSkipAccumulateOutput16x2:
|
||||
vmovups YMMWORD PTR [esi],ymm4
|
||||
vmovups YMMWORD PTR [esi+32],ymm5
|
||||
vmovups YMMWORD PTR [esi+eax],ymm6
|
||||
vmovups YMMWORD PTR [esi+eax+32],ymm7
|
||||
add esi,16*4 # advance matrix C by 16 columns
|
||||
cmp ebp,8
|
||||
ja .LProcessNextColumnLoop16x2
|
||||
test ebp,ebp
|
||||
jz .LExitKernel
|
||||
|
||||
.LProcessRemainingCountN2:
|
||||
mov edi,.LSgemmKernelFrame_CountK[esp]
|
||||
mov ecx,.LSgemmKernelFrame_MatrixA[esp]
|
||||
vxorps xmm5,xmm5,xmm5 # clear block accumulators
|
||||
vxorps xmm7,xmm7,xmm7
|
||||
ComputeBlockAvxLoop ComputeBlockAvxBy8, 2
|
||||
vbroadcastss ymm2,DWORD PTR .LSgemmKernelFrame_alpha[esp]
|
||||
vmulps ymm5,ymm5,ymm2 # multiply by alpha
|
||||
vmulps ymm7,ymm7,ymm2
|
||||
cmp ebp,8
|
||||
jb .LOutputMasked8x2Block
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateOutput8x2
|
||||
vaddps ymm5,ymm5,YMMWORD PTR [esi]
|
||||
vaddps ymm7,ymm7,YMMWORD PTR [esi+eax]
|
||||
|
||||
.LSkipAccumulateOutput8x2:
|
||||
vmovups YMMWORD PTR [esi],ymm5
|
||||
vmovups YMMWORD PTR [esi+eax],ymm7
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.LExitKernel:
|
||||
movzx eax,BYTE PTR .LSgemmKernelFrame_CountM[esp]
|
||||
vzeroupper
|
||||
pop edi
|
||||
pop esi
|
||||
pop ebx
|
||||
pop ebp
|
||||
ret
|
||||
|
||||
.LOutputMasked16x2Block:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateMasked16x2Block
|
||||
vaddps ymm4,ymm4,YMMWORD PTR [esi]
|
||||
vaddps ymm6,ymm6,YMMWORD PTR [esi+eax]
|
||||
|
||||
.LSkipAccumulateMasked16x2Block:
|
||||
vmovups YMMWORD PTR [esi],ymm4
|
||||
vmovups YMMWORD PTR [esi+eax],ymm6
|
||||
add esi,8*4 # advance matrix C by 8 columns
|
||||
add ebp,8 # correct for over-subtract above
|
||||
|
||||
.LOutputMasked8x2Block:
|
||||
neg ebp
|
||||
LoadGlobalOffsetTable bx
|
||||
mov ebx,DWORD PTR C_UNDERSCORE(MlasMaskMoveTableAvx)@GOT[ebx]
|
||||
vmovdqu ymm0,YMMWORD PTR [ebx+ebp*4+8*4]
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateMasked8x2Block
|
||||
vmaskmovps ymm4,ymm0,YMMWORD PTR [esi]
|
||||
vmaskmovps ymm6,ymm0,YMMWORD PTR [esi+eax]
|
||||
vaddps ymm5,ymm5,ymm4
|
||||
vaddps ymm7,ymm7,ymm6
|
||||
|
||||
.LSkipAccumulateMasked8x2Block:
|
||||
vmaskmovps YMMWORD PTR [esi],ymm0,ymm5
|
||||
vmaskmovps YMMWORD PTR [esi+eax],ymm0,ymm7
|
||||
jmp .LExitKernel
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
.LProcessCountMLessThan2:
|
||||
mov BYTE PTR .LSgemmKernelFrame_CountM[esp],1
|
||||
mov ebx,.LSgemmKernelFrame_MatrixA[esp]
|
||||
vbroadcastss ymm2,DWORD PTR .LSgemmKernelFrame_alpha[esp]
|
||||
cmp ebp,8
|
||||
jbe .LProcessRemainingCountN1
|
||||
|
||||
.LProcessNextColumnLoop16x1:
|
||||
mov edi,.LSgemmKernelFrame_CountK[esp]
|
||||
mov ecx,ebx # reload matrix A
|
||||
vxorps xmm4,xmm4,xmm4 # clear block accumulators
|
||||
vxorps xmm5,xmm5,xmm5
|
||||
ComputeBlockAvxLoop ComputeBlockAvxBy16, 1
|
||||
vmulps ymm4,ymm4,ymm2 # multiply by alpha
|
||||
vmulps ymm5,ymm5,ymm2
|
||||
sub ebp,16
|
||||
jb .LOutputMasked16x1Block
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulate16x1Block
|
||||
vaddps ymm4,ymm4,YMMWORD PTR [esi]
|
||||
vaddps ymm5,ymm5,YMMWORD PTR [esi+32]
|
||||
|
||||
.LSkipAccumulate16x1Block:
|
||||
vmovups YMMWORD PTR [esi],ymm4
|
||||
vmovups YMMWORD PTR [esi+32],ymm5
|
||||
add esi,16*4 # advance matrix C by 16 columns
|
||||
cmp ebp,8
|
||||
ja .LProcessNextColumnLoop16x1
|
||||
test ebp,ebp
|
||||
jz .LExitKernel
|
||||
|
||||
.LProcessRemainingCountN1:
|
||||
mov edi,.LSgemmKernelFrame_CountK[esp]
|
||||
mov ecx,ebx # reload matrix A
|
||||
vxorps xmm5,xmm5,xmm5 # clear block accumulators
|
||||
ComputeBlockAvxLoop ComputeBlockAvxBy8, 1
|
||||
vmulps ymm5,ymm5,ymm2 # multiply by alpha
|
||||
cmp ebp,8
|
||||
jb .LOutputMasked8x1Block
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulate8x1Block
|
||||
vaddps ymm5,ymm5,YMMWORD PTR [esi]
|
||||
|
||||
.LSkipAccumulate8x1Block:
|
||||
vmovups YMMWORD PTR [esi],ymm5
|
||||
jmp .LExitKernel
|
||||
|
||||
.LOutputMasked16x1Block:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateMasked16x1Block
|
||||
vaddps ymm4,ymm4,YMMWORD PTR [esi]
|
||||
|
||||
.LSkipAccumulateMasked16x1Block:
|
||||
vmovups YMMWORD PTR [esi],ymm4
|
||||
add esi,8*4 # advance matrix C by 8 columns
|
||||
add ebp,8 # correct for over-subtract above
|
||||
|
||||
.LOutputMasked8x1Block:
|
||||
neg ebp
|
||||
LoadGlobalOffsetTable bx
|
||||
mov ebx,DWORD PTR C_UNDERSCORE(MlasMaskMoveTableAvx)@GOT[ebx]
|
||||
vmovdqu ymm0,YMMWORD PTR [ebx+ebp*4+8*4]
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateMasked8x1Block
|
||||
vmaskmovps ymm4,ymm0,YMMWORD PTR [esi]
|
||||
vaddps ymm5,ymm5,ymm4
|
||||
|
||||
.LSkipAccumulateMasked8x1Block:
|
||||
vmaskmovps YMMWORD PTR [esi],ymm0,ymm5
|
||||
jmp .LExitKernel
|
||||
|
||||
.end
|
||||
+406
@@ -0,0 +1,406 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelSse2.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses SSE2 instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
//
|
||||
// Stack frame layout for the SGEMM kernel.
|
||||
//
|
||||
|
||||
.equ .LSgemmKernelFrame_SavedEdi, 0
|
||||
.equ .LSgemmKernelFrame_SavedEsi, 4
|
||||
.equ .LSgemmKernelFrame_SavedEbx, 8
|
||||
.equ .LSgemmKernelFrame_SavedEbp, 12
|
||||
.equ .LSgemmKernelFrame_ReturnAddress, 16
|
||||
.equ .LSgemmKernelFrame_MatrixA, 20
|
||||
.equ .LSgemmKernelFrame_MatrixB, 24
|
||||
.equ .LSgemmKernelFrame_MatrixC, 28
|
||||
.equ .LSgemmKernelFrame_CountK, 32
|
||||
.equ .LSgemmKernelFrame_CountM, 36
|
||||
.equ .LSgemmKernelFrame_CountN, 40
|
||||
.equ .LSgemmKernelFrame_lda, 44
|
||||
.equ .LSgemmKernelFrame_ldc, 48
|
||||
.equ .LSgemmKernelFrame_alpha, 52
|
||||
.equ .LSgemmKernelFrame_ZeroMode, 56
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for a Nx1 block of the output matrix.
|
||||
|
||||
Arguments:
|
||||
|
||||
VectorOffset - Supplies the byte offset from matrix B to fetch elements.
|
||||
|
||||
Shuffle - Supplies the shuffle mask to extract the element from matrix A.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
ebx - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
ecx - Supplies the address into the matrix A data.
|
||||
|
||||
edx - Supplies the address into the matrix B data.
|
||||
|
||||
xmm2 - Supplies up to four elements loaded from matrix A.
|
||||
|
||||
xmm4-xmm7 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockSseBy4 VectorOffset, Shuffle
|
||||
|
||||
pshufd xmm3,xmm1,\Shuffle\()
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm4,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+16]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm5,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+32]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm6,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+48]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm7,xmm0
|
||||
|
||||
.endm
|
||||
|
||||
.macro ComputeBlockSseBy3 VectorOffset, Shuffle
|
||||
|
||||
pshufd xmm3,xmm1,\Shuffle\()
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm5,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+16]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm6,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+32]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm7,xmm0
|
||||
|
||||
.endm
|
||||
|
||||
.macro ComputeBlockSseBy2 VectorOffset, Shuffle
|
||||
|
||||
pshufd xmm3,xmm1,\Shuffle\()
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm6,xmm0
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()+16]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm7,xmm0
|
||||
|
||||
.endm
|
||||
|
||||
.macro ComputeBlockSseBy1 VectorOffset, Shuffle
|
||||
|
||||
pshufd xmm3,xmm1,\Shuffle\()
|
||||
movaps xmm0,XMMWORD PTR [edx+\VectorOffset\()]
|
||||
mulps xmm0,xmm3
|
||||
addps xmm7,xmm0
|
||||
|
||||
.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:
|
||||
|
||||
ebx - Supplies the number of bytes to the next row of matrix A.
|
||||
|
||||
ecx - Supplies the address into the matrix A data.
|
||||
|
||||
edx - Supplies the address into the matrix B data.
|
||||
|
||||
edi - Supplies the number of columns from matrix A and the number of rows
|
||||
from matrix B to iterate over.
|
||||
|
||||
xmm4-xmm7 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockSseLoop RowCount
|
||||
|
||||
sub edi,4
|
||||
jb .LProcessRemainingBlocks\@
|
||||
|
||||
.LComputeBlockBy4Loop\@:
|
||||
movups xmm1,XMMWORD PTR [ecx]
|
||||
ComputeBlockSseBy\RowCount\() 0, 0x00
|
||||
ComputeBlockSseBy\RowCount\() 16*4, 0x55
|
||||
sub edx,-32*4 # advance matrix B by 32 columns
|
||||
ComputeBlockSseBy\RowCount\() 0, 0xAA
|
||||
ComputeBlockSseBy\RowCount\() 16*4, 0xFF
|
||||
sub edx,-32*4 # advance matrix B by 32 columns
|
||||
add ecx,4*4 # advance matrix A by 4 columns
|
||||
sub edi,4
|
||||
jae .LComputeBlockBy4Loop\@
|
||||
|
||||
.LProcessRemainingBlocks\@:
|
||||
add edi,4 # correct for over-subtract above
|
||||
jz .LOutputBlock\@
|
||||
|
||||
.LComputeBlockBy1Loop\@:
|
||||
movss xmm1,DWORD PTR [ecx]
|
||||
ComputeBlockSseBy\RowCount\() 0, 0x00
|
||||
add edx,16*4 # advance matrix B by 16 columns
|
||||
add ecx,4 # advance matrix A by 1 column
|
||||
dec edi
|
||||
jne .LComputeBlockBy1Loop\@
|
||||
|
||||
.LOutputBlock\@:
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows.
|
||||
|
||||
Arguments:
|
||||
|
||||
A - Supplies the address of matrix A.
|
||||
|
||||
B - Supplies the address of matrix B. The matrix data has been packed using
|
||||
MlasSgemmCopyPackB or MlasSgemmTransposePackB.
|
||||
|
||||
C - Supplies the address of matrix C.
|
||||
|
||||
CountK - Supplies the number of columns from matrix A and the number of
|
||||
rows from matrix B to iterate over.
|
||||
|
||||
CountM - 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 - 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 - Supplies the scalar multiplier (see SGEMM 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 MlasGemmFloatKernelSse
|
||||
|
||||
push ebp
|
||||
push ebx
|
||||
push esi
|
||||
push edi
|
||||
mov edx,.LSgemmKernelFrame_MatrixB[esp]
|
||||
mov esi,.LSgemmKernelFrame_MatrixC[esp]
|
||||
mov ebp,.LSgemmKernelFrame_CountN[esp]
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
mov eax,.LSgemmKernelFrame_CountK[esp]
|
||||
mov ebx,.LSgemmKernelFrame_MatrixA[esp]
|
||||
cmp ebp,12
|
||||
jbe .LProcessRemainingCountN
|
||||
|
||||
.LProcessNextColumnLoop16x1:
|
||||
mov edi,eax # reload CountK
|
||||
mov ecx,ebx # reload matrix A
|
||||
xorps xmm4,xmm4 # clear block accumulators
|
||||
xorps xmm5,xmm5
|
||||
xorps xmm6,xmm6
|
||||
xorps xmm7,xmm7
|
||||
ComputeBlockSseLoop 4
|
||||
movss xmm2,DWORD PTR .LSgemmKernelFrame_alpha[esp]
|
||||
shufps xmm2,xmm2,0
|
||||
mulps xmm4,xmm2 # multiply by alpha
|
||||
mulps xmm5,xmm2
|
||||
mulps xmm6,xmm2
|
||||
mulps xmm7,xmm2
|
||||
sub ebp,16
|
||||
jb .LOutputMasked16x1Block
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateOutput16x1
|
||||
movups xmm0,XMMWORD PTR [esi]
|
||||
movups xmm1,XMMWORD PTR [esi+16]
|
||||
movups xmm2,XMMWORD PTR [esi+32]
|
||||
movups xmm3,XMMWORD PTR [esi+48]
|
||||
addps xmm4,xmm0
|
||||
addps xmm5,xmm1
|
||||
addps xmm6,xmm2
|
||||
addps xmm7,xmm3
|
||||
|
||||
.LSkipAccumulateOutput16x1:
|
||||
movups XMMWORD PTR [esi],xmm4
|
||||
movups XMMWORD PTR [esi+16],xmm5
|
||||
movups XMMWORD PTR [esi+32],xmm6
|
||||
movups XMMWORD PTR [esi+48],xmm7
|
||||
add esi,16*4 # advance matrix C by 16 columns
|
||||
cmp ebp,12
|
||||
ja .LProcessNextColumnLoop16x1
|
||||
test ebp,ebp
|
||||
jnz .LProcessRemainingCountN
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.LExitKernel:
|
||||
mov eax,1 # return 1 row handled
|
||||
pop edi
|
||||
pop esi
|
||||
pop ebx
|
||||
pop ebp
|
||||
ret
|
||||
|
||||
//
|
||||
// Process the remaining 1 to 12 columns of the matrices.
|
||||
//
|
||||
|
||||
.LProcessRemainingCountN:
|
||||
mov edi,eax # reload CountK
|
||||
mov ecx,ebx # reload matrix A
|
||||
movss xmm4,DWORD PTR .LSgemmKernelFrame_alpha[esp]
|
||||
shufps xmm4,xmm4,0
|
||||
xorps xmm5,xmm5 # clear block accumulators
|
||||
xorps xmm6,xmm6
|
||||
xorps xmm7,xmm7
|
||||
cmp ebp,4
|
||||
jbe .LProcessRemainingCountN4OrLess
|
||||
cmp ebp,8
|
||||
jbe .LProcessRemainingCountN8OrLess
|
||||
|
||||
.LProcessRemainingCountN12OrLess:
|
||||
ComputeBlockSseLoop 3
|
||||
mulps xmm5,xmm4 # multiply by alpha
|
||||
mulps xmm6,xmm4
|
||||
mulps xmm7,xmm4
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateLeadingN12OrLess
|
||||
movups xmm0,XMMWORD PTR [esi]
|
||||
movups xmm1,XMMWORD PTR [esi+16]
|
||||
addps xmm5,xmm0
|
||||
addps xmm6,xmm1
|
||||
|
||||
.LSkipAccumulateLeadingN12OrLess:
|
||||
movups XMMWORD PTR [esi],xmm5
|
||||
movups XMMWORD PTR [esi+16],xmm6
|
||||
add esi,8*4 # advance matrix C by 8 columns
|
||||
jmp .LOutputTrailingBlock
|
||||
|
||||
.LProcessRemainingCountN8OrLess:
|
||||
ComputeBlockSseLoop 2
|
||||
mulps xmm6,xmm4 # multiply by alpha
|
||||
mulps xmm7,xmm4
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateLeadingN8OrLess
|
||||
movups xmm0,XMMWORD PTR [esi]
|
||||
addps xmm6,xmm0
|
||||
|
||||
.LSkipAccumulateLeadingN8OrLess:
|
||||
movups XMMWORD PTR [esi],xmm6
|
||||
add esi,4*4 # advance matrix C by 4 columns
|
||||
jmp .LOutputTrailingBlock
|
||||
|
||||
.LProcessRemainingCountN4OrLess:
|
||||
ComputeBlockSseLoop 1
|
||||
mulps xmm7,xmm4 # multiply by alpha
|
||||
jmp .LOutputTrailingBlock
|
||||
|
||||
.LOutputMasked16x1Block:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateLeading16x1Block
|
||||
movups xmm0,XMMWORD PTR [esi]
|
||||
movups xmm1,XMMWORD PTR [esi+16]
|
||||
movups xmm2,XMMWORD PTR [esi+32]
|
||||
addps xmm4,xmm0
|
||||
addps xmm5,xmm1
|
||||
addps xmm6,xmm2
|
||||
|
||||
.LSkipAccumulateLeading16x1Block:
|
||||
movups XMMWORD PTR [esi],xmm4
|
||||
movups XMMWORD PTR [esi+16],xmm5
|
||||
movups XMMWORD PTR [esi+32],xmm6
|
||||
add esi,12*4 # advance matrix C by 12 columns
|
||||
|
||||
.LOutputTrailingBlock:
|
||||
test ebp,3
|
||||
jz .LOutputTrailingBlock4Elements
|
||||
test ebp,2
|
||||
jz .LOutputTrailingBlock1Element
|
||||
|
||||
.LOutputTrailingBlock2Elements:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateTrailingBlock2Elements
|
||||
movsd xmm0,QWORD PTR [esi]
|
||||
addps xmm7,xmm0
|
||||
|
||||
.LSkipAccumulateTrailingBlock2Elements:
|
||||
movsd QWORD PTR [esi],xmm7
|
||||
test ebp,1
|
||||
jz .LExitKernel
|
||||
shufps xmm7,xmm7,0xAA # shuffle third float down
|
||||
add esi,2*4 # advance matrix C by 2 columns
|
||||
|
||||
.LOutputTrailingBlock1Element:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateTrailingBlock1Element
|
||||
movss xmm0,DWORD PTR [esi]
|
||||
addss xmm7,xmm0
|
||||
|
||||
.LSkipAccumulateTrailingBlock1Element:
|
||||
movss DWORD PTR [esi],xmm7
|
||||
jmp .LExitKernel
|
||||
|
||||
.LOutputTrailingBlock4Elements:
|
||||
cmp BYTE PTR .LSgemmKernelFrame_ZeroMode[esp],0
|
||||
jnz .LSkipAccumulateTrailingBlock4Elements
|
||||
movups xmm0,XMMWORD PTR [esi]
|
||||
addps xmm7,xmm0
|
||||
|
||||
.LSkipAccumulateTrailingBlock4Elements:
|
||||
movups XMMWORD PTR [esi],xmm7
|
||||
jmp .LExitKernel
|
||||
|
||||
.end
|
||||
Vendored
+79
@@ -0,0 +1,79 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
asmmacro.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements common macros for the assembly modules.
|
||||
|
||||
--*/
|
||||
|
||||
#if defined(__APPLE__)
|
||||
#define C_UNDERSCORE(symbol) _##symbol
|
||||
#else
|
||||
#define C_UNDERSCORE(symbol) symbol
|
||||
#endif
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the assembler directives to annotate a new function.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name of the function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FUNCTION_ENTRY FunctionName
|
||||
|
||||
.p2align 4
|
||||
#if defined(__APPLE__)
|
||||
.globl _\FunctionName\()
|
||||
_\FunctionName\():
|
||||
#else
|
||||
.globl \FunctionName\()
|
||||
.type \FunctionName\(),@function
|
||||
\FunctionName\():
|
||||
#endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the code to load the global offset table address into the
|
||||
supplied register.
|
||||
|
||||
Arguments:
|
||||
|
||||
TargetReg - Specifies the target register.
|
||||
|
||||
--*/
|
||||
|
||||
.macro LoadGlobalOffsetTable, TargetReg
|
||||
|
||||
//
|
||||
// The LLVM integrated assembler doesn't support the Intel syntax for OFFSET:
|
||||
//
|
||||
// add ebx,OFFSET _GLOBAL_OFFSET_TABLE_
|
||||
//
|
||||
// Workaround this by temporarily switching to AT&T syntax.
|
||||
//
|
||||
|
||||
.att_syntax
|
||||
|
||||
calll __x86.get_pc_thunk.\TargetReg\()
|
||||
addl $_GLOBAL_OFFSET_TABLE_,%e\TargetReg\()
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.endm
|
||||
+36
@@ -0,0 +1,36 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
x86.get_pc_thunk.S
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements __x86.get_pc_thunk.* to avoid external dependency.
|
||||
|
||||
--*/
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
The routine loads its return address -- which is the address of the
|
||||
instruction that immediately follows -- into the ebx register.
|
||||
|
||||
--*/
|
||||
|
||||
.p2align 4
|
||||
.weak __x86.get_pc_thunk.bx
|
||||
// Hidden = non-preemptible, so the PC-relative call in the SGEMM
|
||||
// kernels binds locally in the .so (no R_386_PC32 link error).
|
||||
.hidden __x86.get_pc_thunk.bx
|
||||
.type __x86.get_pc_thunk.bx,@function
|
||||
__x86.get_pc_thunk.bx:
|
||||
mov ebx, [esp]
|
||||
ret
|
||||
+529
@@ -0,0 +1,529 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
FgemmKernelAvx512FCommon.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the floating point matrix/matrix
|
||||
multiply operation (SGEMM and DGEMM).
|
||||
|
||||
This implementation uses AVX512F instructions.
|
||||
|
||||
--*/
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for 2 ZMMWORDs 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 3 rows.
|
||||
|
||||
r13 - Supplies the address into the matrix A data plus 6 rows.
|
||||
|
||||
r14 - Supplies the address into the matrix A data plus 9 rows.
|
||||
|
||||
rsi - Supplies the address into the matrix B data.
|
||||
|
||||
r10 - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
zmm4-zmm27 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvx512FBy2 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
|
||||
|
||||
.ifnb \PrefetchOffset\()
|
||||
prefetcht0 [rsi+\VectorOffset\()+\PrefetchOffset\()]
|
||||
prefetcht0 [rsi+r12+\VectorOffset\()+\PrefetchOffset\()]
|
||||
.endif
|
||||
.if \RowCount\() == 1
|
||||
vbroadcastsf zmm3,[rdi+\BroadcastOffset\()]
|
||||
vfmadd231pf zmm4,zmm3,ZMMWORD PTR [rsi+\VectorOffset\()]
|
||||
vfmadd231pf zmm5,zmm3,ZMMWORD PTR [rsi+r12+\VectorOffset\()]
|
||||
.else
|
||||
vmovapf zmm0,ZMMWORD PTR [rsi+\VectorOffset\()]
|
||||
vmovapf zmm1,ZMMWORD PTR [rsi+r12+\VectorOffset\()]
|
||||
EmitIfCountGE \RowCount\(), 1, "vbroadcastsf zmm3,[rdi+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf zmm4,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf zmm5,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 2, "vbroadcastsf zmm3,[rdi+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf zmm6,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf zmm7,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 3, "vbroadcastsf zmm3,[rdi+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf zmm8,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf zmm9,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 4, "vbroadcastsf zmm3,[rbx+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf zmm10,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf zmm11,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 5, "vbroadcastsf zmm3,[rbx+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf zmm12,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf zmm13,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 6, "vbroadcastsf zmm3,[rbx+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf zmm14,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf zmm15,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r13+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm16,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm17,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r13+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm18,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm19,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r13+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm20,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm21,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r14+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm22,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm23,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r14+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm24,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm25,zmm3,zmm1"
|
||||
EmitIfCountGE \RowCount\(), 12, "vbroadcastsf zmm3,[r14+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm26,zmm3,zmm0"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf zmm27,zmm3,zmm1"
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for 1 ZMMWORD 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 3 rows.
|
||||
|
||||
r13 - Supplies the address into the matrix A data plus 6 rows.
|
||||
|
||||
r14 - Supplies the address into the matrix A data plus 9 rows.
|
||||
|
||||
rsi - Supplies the address into the matrix B data.
|
||||
|
||||
r10 - Supplies the length in bytes of a row from matrix A.
|
||||
|
||||
zmm4-zmm27 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvx512FBy1 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
|
||||
|
||||
.ifnb \PrefetchOffset\()
|
||||
prefetcht0 [rsi+\VectorOffset\()+\PrefetchOffset\()]
|
||||
.endif
|
||||
vmovapf zmm0,ZMMWORD PTR [rsi+\VectorOffset\()]
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf_bcst zmm5,zmm0,[rdi+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf_bcst zmm7,zmm0,[rdi+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf_bcst zmm9,zmm0,[rdi+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf_bcst zmm11,zmm0,[rbx+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf_bcst zmm13,zmm0,[rbx+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf_bcst zmm15,zmm0,[rbx+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm17,zmm0,[r13+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm19,zmm0,[r13+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm21,zmm0,[r13+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm23,zmm0,[r14+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm25,zmm0,[r14+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd231pf_bcst zmm27,zmm0,[r14+r10*2+\BroadcastOffset\()]"
|
||||
|
||||
.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.
|
||||
|
||||
zmm4-zmm27 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockAvx512FLoop ComputeBlock, RowCount
|
||||
|
||||
.if \RowCount\() > 3
|
||||
lea rbx,[r10*2+r10]
|
||||
.if \RowCount\() == 12
|
||||
lea r13,[rdi+rbx*2] # compute matrix A plus 6 rows
|
||||
lea r14,[r13+rbx] # compute matrix A plus 9 rows
|
||||
.endif
|
||||
add rbx,rdi # compute matrix A plus 3 rows
|
||||
.endif
|
||||
ComputeBlockLoop \ComputeBlock\(), \RowCount\(), \RowCount\() > 3
|
||||
.if \RowCount\() > 3
|
||||
lea rbx,[rax*2+rax]
|
||||
.if \RowCount\() == 12
|
||||
lea r13,[rdx+rbx*2] # compute matrix C plus 6 rows
|
||||
lea r14,[r13+rbx] # compute matrix C plus 9 rows
|
||||
.endif
|
||||
add rbx,rdx # compute matrix C plus 3 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.
|
||||
|
||||
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
|
||||
|
||||
cmp r9,.LFgemmZmmElementCount
|
||||
jbe .LProcessRemainingCountN\@
|
||||
|
||||
.LProcessNextColumnLoop2xN\@:
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm16,zmm4"
|
||||
# clear upper block accumulators
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm17,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm18,zmm4"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm19,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm20,zmm4"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm21,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm22,zmm4"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm23,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm24,zmm4"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm25,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm26,zmm4"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm27,zmm5"
|
||||
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy2, \RowCount\()
|
||||
add rsi,r12 # advance matrix B by 64*CountK bytes
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LMultiplyAlpha2xNBlock\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf zmm4,zmm31,ZMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf zmm6,zmm31,ZMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf zmm8,zmm31,ZMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf zmm10,zmm31,ZMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf zmm12,zmm31,ZMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf zmm14,zmm31,ZMMWORD PTR [rbx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm16,zmm31,ZMMWORD PTR [r13]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm18,zmm31,ZMMWORD PTR [r13+rax]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm20,zmm31,ZMMWORD PTR [r13+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm22,zmm31,ZMMWORD PTR [r14]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm24,zmm31,ZMMWORD PTR [r14+rax]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm26,zmm31,ZMMWORD PTR [r14+rax*2]"
|
||||
jmp .LStore2xNBlock\@
|
||||
|
||||
.LMultiplyAlpha2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf zmm4,zmm4,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf zmm6,zmm6,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf zmm8,zmm8,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf zmm10,zmm10,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf zmm12,zmm12,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf zmm14,zmm14,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm16,zmm16,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm18,zmm18,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm20,zmm20,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm22,zmm22,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm24,zmm24,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm26,zmm26,zmm31"
|
||||
|
||||
.LStore2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf ZMMWORD PTR [rdx],zmm4"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf ZMMWORD PTR [rdx+rax],zmm6"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf ZMMWORD PTR [rdx+rax*2],zmm8"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf ZMMWORD PTR [rbx],zmm10"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf ZMMWORD PTR [rbx+rax],zmm12"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf ZMMWORD PTR [rbx+rax*2],zmm14"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13],zmm16"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13+rax],zmm18"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13+rax*2],zmm20"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14],zmm22"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14+rax],zmm24"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14+rax*2],zmm26"
|
||||
add rdx,64 # advance matrix C by ZMMWORD
|
||||
.if \RowCount\() > 3
|
||||
add rbx,64 # advance matrix C plus 3 rows by ZMMWORD
|
||||
.if \RowCount\() == 12
|
||||
add r13,64 # advance matrix C plus 6 rows by ZMMWORD
|
||||
add r14,64 # advance matrix C plus 9 rows by ZMMWORD
|
||||
.endif
|
||||
.endif
|
||||
sub r9,.LFgemmZmmElementCount
|
||||
|
||||
.LOutput1xNBlock\@:
|
||||
sub r9,.LFgemmZmmElementCount
|
||||
jae .LOutput1xNBlockWithMask\@
|
||||
lea rcx,[r9+.LFgemmZmmElementCount]
|
||||
# correct for over-subtract above
|
||||
mov ebp,1
|
||||
shl ebp,cl
|
||||
dec ebp
|
||||
kmovw k1,ebp # update mask for remaining columns
|
||||
xor r9,r9 # no more columns remaining
|
||||
|
||||
.LOutput1xNBlockWithMask\@:
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LMultiplyAlpha1xNBlockWithMask\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf zmm5{k1},zmm31,ZMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf zmm7{k1},zmm31,ZMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf zmm9{k1},zmm31,ZMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf zmm11{k1},zmm31,ZMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf zmm13{k1},zmm31,ZMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf zmm15{k1},zmm31,ZMMWORD PTR [rbx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm17{k1},zmm31,ZMMWORD PTR [r13]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm19{k1},zmm31,ZMMWORD PTR [r13+rax]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm21{k1},zmm31,ZMMWORD PTR [r13+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm23{k1},zmm31,ZMMWORD PTR [r14]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm25{k1},zmm31,ZMMWORD PTR [r14+rax]"
|
||||
EmitIfCountGE \RowCount\(), 12, "vfmadd213pf zmm27{k1},zmm31,ZMMWORD PTR [r14+rax*2]"
|
||||
jmp .LStore1xNBlockWithMask\@
|
||||
|
||||
.LMultiplyAlpha1xNBlockWithMask\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf zmm5,zmm5,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf zmm7,zmm7,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf zmm9,zmm9,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf zmm11,zmm11,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf zmm13,zmm13,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf zmm15,zmm15,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm17,zmm17,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm19,zmm19,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm21,zmm21,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm23,zmm23,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm25,zmm25,zmm31"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmulpf zmm27,zmm27,zmm31"
|
||||
|
||||
.LStore1xNBlockWithMask\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf ZMMWORD PTR [rdx]{k1},zmm5"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf ZMMWORD PTR [rdx+rax]{k1},zmm7"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf ZMMWORD PTR [rdx+rax*2]{k1},zmm9"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf ZMMWORD PTR [rbx]{k1},zmm11"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf ZMMWORD PTR [rbx+rax]{k1},zmm13"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf ZMMWORD PTR [rbx+rax*2]{k1},zmm15"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13]{k1},zmm17"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13+rax]{k1},zmm19"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r13+rax*2]{k1},zmm21"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14]{k1},zmm23"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14+rax]{k1},zmm25"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovupf ZMMWORD PTR [r14+rax*2]{k1},zmm27"
|
||||
add rdx,64 # advance matrix C by ZMMWORD
|
||||
mov rdi,r11 # reload matrix A
|
||||
vzeroall
|
||||
cmp r9,.LFgemmZmmElementCount
|
||||
ja .LProcessNextColumnLoop2xN\@
|
||||
test r9,r9
|
||||
jz .LExitKernel
|
||||
|
||||
.LProcessRemainingCountN\@:
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm17,zmm5"
|
||||
# clear upper block accumulators
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm19,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm21,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm23,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm25,zmm5"
|
||||
EmitIfCountGE \RowCount\(), 12, "vmovapf zmm27,zmm5"
|
||||
ComputeBlockAvx512FLoop ComputeBlockAvx512FBy1, \RowCount\()
|
||||
jmp .LOutput1xNBlock\@
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro generates the inner kernel to compute matrix multiplication.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name for the generated function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FgemmKernelAvx512FFunction 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 .LFgemmKernelFrame_SavedR12[rsp],r12
|
||||
mov .LFgemmKernelFrame_SavedR13[rsp],r13
|
||||
mov .LFgemmKernelFrame_SavedR14[rsp],r14
|
||||
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
|
||||
mov r12,rcx
|
||||
shl r12,6 # compute 64*CountK bytes
|
||||
mov ebp,-1
|
||||
kmovw k1,ebp # update mask to write all columns
|
||||
movzx r15,BYTE PTR .LFgemmKernelFrame_ZeroMode[rsp]
|
||||
vbroadcastsf zmm31,xmm0
|
||||
vzeroall
|
||||
|
||||
//
|
||||
// Process CountM rows of the matrices.
|
||||
//
|
||||
|
||||
cmp r8,12
|
||||
jb .LProcessCountMLessThan12
|
||||
mov r8d,12 # return 12 rows handled
|
||||
ProcessCountM 12
|
||||
|
||||
.LProcessCountMLessThan12:
|
||||
cmp r8,5
|
||||
ja .LProcessCountM6
|
||||
je .LProcessCountM5
|
||||
cmp r8,3
|
||||
ja .LProcessCountM4
|
||||
je .LProcessCountM3
|
||||
cmp r8,1
|
||||
je .LProcessCountM1
|
||||
|
||||
.LProcessCountM2:
|
||||
ProcessCountM 2
|
||||
|
||||
.LProcessCountM4:
|
||||
ProcessCountM 4
|
||||
|
||||
.LProcessCountM6:
|
||||
mov r8d,6 # return 6 rows handled
|
||||
ProcessCountM 6
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.LExitKernel:
|
||||
mov eax,r8d
|
||||
mov r12,.LFgemmKernelFrame_SavedR12[rsp]
|
||||
mov r13,.LFgemmKernelFrame_SavedR13[rsp]
|
||||
mov r14,.LFgemmKernelFrame_SavedR14[rsp]
|
||||
pop r15
|
||||
pop rbx
|
||||
pop rbp
|
||||
ret
|
||||
|
||||
.LProcessCountM1:
|
||||
ProcessCountM 1
|
||||
|
||||
.LProcessCountM3:
|
||||
ProcessCountM 3
|
||||
|
||||
.LProcessCountM5:
|
||||
ProcessCountM 5
|
||||
|
||||
.endm
|
||||
+451
@@ -0,0 +1,451 @@
|
||||
/*++
|
||||
|
||||
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
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
FgemmKernelCommon.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module contains common kernel macros and structures for the floating
|
||||
point matrix/matrix multiply operation (SGEMM and DGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
//
|
||||
// Stack frame layout for the floating point kernels.
|
||||
//
|
||||
|
||||
.equ .LFgemmKernelFrame_SavedR12, -32
|
||||
.equ .LFgemmKernelFrame_SavedR13, -24
|
||||
.equ .LFgemmKernelFrame_SavedR14, -16
|
||||
.equ .LFgemmKernelFrame_alpha, -8
|
||||
.equ .LFgemmKernelFrame_SavedR15, 0
|
||||
.equ .LFgemmKernelFrame_SavedRbx, 8
|
||||
.equ .LFgemmKernelFrame_SavedRbp, 16
|
||||
.equ .LFgemmKernelFrame_ReturnAddress, 24
|
||||
.equ .LFgemmKernelFrame_lda, 32
|
||||
.equ .LFgemmKernelFrame_ldc, 40
|
||||
.equ .LFgemmKernelFrame_ZeroMode, 48
|
||||
|
||||
//
|
||||
// Define the number of elements per vector register.
|
||||
//
|
||||
|
||||
.equ .LFgemmXmmElementCount, 16 / .LFgemmElementSize
|
||||
.equ .LFgemmYmmElementCount, 32 / .LFgemmElementSize
|
||||
.equ .LFgemmZmmElementCount, 64 / .LFgemmElementSize
|
||||
|
||||
//
|
||||
// Define the typed instruction template.
|
||||
//
|
||||
|
||||
#define FGEMM_TYPED_INSTRUCTION(Untyped, Typed) \
|
||||
.macro Untyped Operand:vararg; Typed \Operand\(); .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.
|
||||
|
||||
AdvanceMatrixAPlusRows - Supplies a non-zero value if the data pointer
|
||||
in rbx should also be advanced as part of the loop.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
rdi - Supplies the address into the matrix A data.
|
||||
|
||||
rbx - Supplies the address into the matrix A data plus 3 rows.
|
||||
|
||||
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.
|
||||
|
||||
ymm4-ymm15 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockLoop ComputeBlock, RowCount, AdvanceMatrixAPlusRows
|
||||
|
||||
mov rbp,rcx # reload CountK
|
||||
sub rbp,4
|
||||
jb .LProcessRemainingBlocks\@
|
||||
|
||||
.LComputeBlockBy4Loop\@:
|
||||
\ComputeBlock\() \RowCount\(), 0, .LFgemmElementSize*0, 64*4
|
||||
\ComputeBlock\() \RowCount\(), 2*32, .LFgemmElementSize*1, 64*4
|
||||
add_immed rsi,2*2*32 # advance matrix B by 128 bytes
|
||||
\ComputeBlock\() \RowCount\(), 0, .LFgemmElementSize*2, 64*4
|
||||
\ComputeBlock\() \RowCount\(), 2*32, .LFgemmElementSize*3, 64*4
|
||||
add_immed rsi,2*2*32 # advance matrix B by 128 bytes
|
||||
add rdi,4*.LFgemmElementSize # advance matrix A by 4 elements
|
||||
.if \RowCount\() > 3
|
||||
add rbx,4*.LFgemmElementSize # advance matrix A plus rows by 4 elements
|
||||
.if \RowCount\() == 12
|
||||
add r13,4*.LFgemmElementSize
|
||||
add r14,4*.LFgemmElementSize
|
||||
.endif
|
||||
.endif
|
||||
sub rbp,4
|
||||
jae .LComputeBlockBy4Loop\@
|
||||
|
||||
.LProcessRemainingBlocks\@:
|
||||
add rbp,4 # correct for over-subtract above
|
||||
jz .LOutputBlock\@
|
||||
|
||||
.LComputeBlockBy1Loop\@:
|
||||
\ComputeBlock\() \RowCount\(), 0, 0
|
||||
add rsi,2*32 # advance matrix B by 64 bytes
|
||||
add rdi,.LFgemmElementSize # advance matrix A by 1 element
|
||||
.if \RowCount\() > 3
|
||||
add rbx,.LFgemmElementSize # advance matrix A plus rows by 1 element
|
||||
.if \RowCount\() == 12
|
||||
add r13,.LFgemmElementSize
|
||||
add r14,.LFgemmElementSize
|
||||
.endif
|
||||
.endif
|
||||
dec rbp
|
||||
jne .LComputeBlockBy1Loop\@
|
||||
|
||||
.LOutputBlock\@:
|
||||
|
||||
.endm
|
||||
+512
@@ -0,0 +1,512 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
FgemmKernelFma3Common.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the floating point matrix/matrix
|
||||
multiply operation (SGEMM and DGEMM).
|
||||
|
||||
This implementation uses AVX fused multiply/add 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 3 rows.
|
||||
|
||||
rsi - 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.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockFma3By2 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
|
||||
|
||||
.ifnb \PrefetchOffset\()
|
||||
prefetcht0 [rsi+\VectorOffset\()+\PrefetchOffset\()]
|
||||
.endif
|
||||
.if \RowCount\() == 1
|
||||
vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]
|
||||
vfmadd231pf ymm4,ymm3,YMMWORD PTR [rsi+\VectorOffset\()]
|
||||
vfmadd231pf ymm5,ymm3,YMMWORD PTR [rsi+\VectorOffset\()+32]
|
||||
.else
|
||||
vmovapf ymm0,YMMWORD PTR [rsi+\VectorOffset\()]
|
||||
vmovapf ymm1,YMMWORD PTR [rsi+\VectorOffset\()+32]
|
||||
EmitIfCountGE \RowCount\(), 1, "vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf ymm4,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf ymm5,ymm3,ymm1"
|
||||
EmitIfCountGE \RowCount\(), 2, "vbroadcastsf ymm3,[rdi+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf ymm6,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf ymm7,ymm3,ymm1"
|
||||
EmitIfCountGE \RowCount\(), 3, "vbroadcastsf ymm3,[rdi+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf ymm8,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf ymm9,ymm3,ymm1"
|
||||
EmitIfCountGE \RowCount\(), 4, "vbroadcastsf ymm3,[rbx+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf ymm10,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf ymm11,ymm3,ymm1"
|
||||
EmitIfCountGE \RowCount\(), 5, "vbroadcastsf ymm3,[rbx+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf ymm12,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf ymm13,ymm3,ymm1"
|
||||
EmitIfCountGE \RowCount\(), 6, "vbroadcastsf ymm3,[rbx+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf ymm14,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf ymm15,ymm3,ymm1"
|
||||
.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 3 rows.
|
||||
|
||||
rsi - 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.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockFma3By1 RowCount, VectorOffset, BroadcastOffset, PrefetchOffset
|
||||
|
||||
.ifnb \PrefetchOffset\()
|
||||
prefetcht0 [rsi+\VectorOffset\()+\PrefetchOffset\()]
|
||||
.endif
|
||||
.if \RowCount\() == 1
|
||||
vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]
|
||||
vfmadd231pf ymm5,ymm3,YMMWORD PTR [rsi+\VectorOffset\()]
|
||||
.else
|
||||
vmovapf ymm0,YMMWORD PTR [rsi+\VectorOffset\()]
|
||||
EmitIfCountGE \RowCount\(), 1, "vbroadcastsf ymm3,[rdi+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd231pf ymm5,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 2, "vbroadcastsf ymm3,[rdi+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd231pf ymm7,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 3, "vbroadcastsf ymm3,[rdi+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd231pf ymm9,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 4, "vbroadcastsf ymm3,[rbx+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd231pf ymm11,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 5, "vbroadcastsf ymm3,[rbx+r10+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd231pf ymm13,ymm3,ymm0"
|
||||
EmitIfCountGE \RowCount\(), 6, "vbroadcastsf ymm3,[rbx+r10*2+\BroadcastOffset\()]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd231pf ymm15,ymm3,ymm0"
|
||||
.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 ComputeBlockFma3Loop ComputeBlock, RowCount
|
||||
|
||||
.if \RowCount\() > 3
|
||||
lea rbx,[r10*2+r10]
|
||||
add rbx,rdi # compute matrix A plus 3 rows
|
||||
.endif
|
||||
ComputeBlockLoop \ComputeBlock\(), \RowCount\(), \RowCount\() > 3
|
||||
vbroadcastsf ymm2,[rsp+.LFgemmKernelFrame_alpha]
|
||||
.if \RowCount\() > 3
|
||||
lea rbx,[rax*2+rax]
|
||||
add rbx,rdx # compute matrix C plus 3 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 ExitKernelAndZeroUpper 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\@:
|
||||
ComputeBlockFma3Loop ComputeBlockFma3By2, \RowCount\()
|
||||
EmitIfCountGE \RowCount\(), 1, "prefetcht0 [rdx+64]"
|
||||
EmitIfCountGE \RowCount\(), 2, "prefetcht0 [rdx+rax+64]"
|
||||
EmitIfCountGE \RowCount\(), 3, "prefetcht0 [rdx+rax*2+64]"
|
||||
EmitIfCountGE \RowCount\(), 4, "prefetcht0 [rbx+64]"
|
||||
EmitIfCountGE \RowCount\(), 5, "prefetcht0 [rbx+rax+64]"
|
||||
EmitIfCountGE \RowCount\(), 6, "prefetcht0 [rbx+rax*2+64]"
|
||||
sub r9,2*.LFgemmYmmElementCount
|
||||
jb .LOutputMasked2xNBlock\@
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LMultiplyAlpha2xNBlock\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf ymm4,ymm2,YMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf ymm5,ymm2,YMMWORD PTR [rdx+32]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf ymm6,ymm2,YMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf ymm7,ymm2,YMMWORD PTR [rdx+rax+32]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf ymm8,ymm2,YMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf ymm9,ymm2,YMMWORD PTR [rdx+rax*2+32]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf ymm10,ymm2,YMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf ymm11,ymm2,YMMWORD PTR [rbx+32]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf ymm12,ymm2,YMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf ymm13,ymm2,YMMWORD PTR [rbx+rax+32]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf ymm14,ymm2,YMMWORD PTR [rbx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf ymm15,ymm2,YMMWORD PTR [rbx+rax*2+32]"
|
||||
jmp .LStore2xNBlock\@
|
||||
|
||||
.LMultiplyAlpha2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm4,ymm4,ymm2"
|
||||
# multiply by alpha
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm5,ymm5,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm6,ymm6,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm7,ymm7,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm8,ymm8,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm9,ymm9,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm10,ymm10,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm11,ymm11,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf ymm12,ymm12,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf ymm13,ymm13,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf ymm14,ymm14,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf ymm15,ymm15,ymm2"
|
||||
|
||||
.LStore2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm4"
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx+32],ymm5"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm6"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax+32],ymm7"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rdx+rax*2],ymm8"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rdx+rax*2+32],ymm9"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx],ymm10"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx+32],ymm11"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf YMMWORD PTR [rbx+rax],ymm12"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf YMMWORD PTR [rbx+rax+32],ymm13"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf YMMWORD PTR [rbx+rax*2],ymm14"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf YMMWORD PTR [rbx+rax*2+32],ymm15"
|
||||
add rdx,2*32 # advance matrix C by 2 YMMWORDs
|
||||
mov rdi,r11 # reload matrix A
|
||||
vzeroall
|
||||
cmp r9,.LFgemmYmmElementCount
|
||||
ja .LProcessNextColumnLoop2xN\@
|
||||
test r9,r9
|
||||
jz .LExitKernel
|
||||
|
||||
.LProcessRemainingCountN\@:
|
||||
ComputeBlockFma3Loop ComputeBlockFma3By1, \RowCount\()
|
||||
cmp r9,.LFgemmYmmElementCount
|
||||
jb .LOutputMasked1xNBlock\@
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LMultiplyAlpha1xNBlock\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf ymm5,ymm2,YMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf ymm7,ymm2,YMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf ymm9,ymm2,YMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf ymm11,ymm2,YMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf ymm13,ymm2,YMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf ymm15,ymm2,YMMWORD PTR [rbx+rax*2]"
|
||||
jmp .LStore1xNBlock\@
|
||||
|
||||
.LMultiplyAlpha1xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm5,ymm5,ymm2"
|
||||
# multiply by alpha
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm7,ymm7,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm9,ymm9,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm11,ymm11,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf ymm13,ymm13,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf ymm15,ymm15,ymm2"
|
||||
|
||||
.LStore1xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm5"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm7"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rdx+rax*2],ymm9"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx],ymm11"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf YMMWORD PTR [rbx+rax],ymm13"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf YMMWORD PTR [rbx+rax*2],ymm15"
|
||||
jmp .LExitKernelAndZeroUpper
|
||||
|
||||
.LOutputMasked2xNBlock\@:
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LMultiplyAlphaMasked2xNBlock\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf ymm4,ymm2,YMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf ymm6,ymm2,YMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf ymm8,ymm2,YMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf ymm10,ymm2,YMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf ymm12,ymm2,YMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf ymm14,ymm2,YMMWORD PTR [rbx+rax*2]"
|
||||
jmp .LStoreMasked2xNBlock\@
|
||||
|
||||
.LMultiplyAlphaMasked2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm4,ymm4,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm6,ymm6,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm8,ymm8,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm10,ymm10,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf ymm12,ymm12,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf ymm14,ymm14,ymm2"
|
||||
|
||||
.LStoreMasked2xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmovupf YMMWORD PTR [rdx],ymm4"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmovupf YMMWORD PTR [rdx+rax],ymm6"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmovupf YMMWORD PTR [rdx+rax*2],ymm8"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmovupf YMMWORD PTR [rbx],ymm10"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmovupf YMMWORD PTR [rbx+rax],ymm12"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmovupf YMMWORD PTR [rbx+rax*2],ymm14"
|
||||
add rdx,32 # advance matrix C by YMMWORD
|
||||
.if \RowCount\() > 3
|
||||
add rbx,32 # advance matrix C plus 3 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 .LMultiplyAlphaMasked1xNBlock\@
|
||||
EmitIfCountGE \RowCount\(), 1, "vmaskmovpf ymm4,ymm0,YMMWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmaskmovpf ymm6,ymm0,YMMWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmaskmovpf ymm8,ymm0,YMMWORD PTR [rdx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmaskmovpf ymm10,ymm0,YMMWORD PTR [rbx]"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmaskmovpf ymm12,ymm0,YMMWORD PTR [rbx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmaskmovpf ymm14,ymm0,YMMWORD PTR [rbx+rax*2]"
|
||||
EmitIfCountGE \RowCount\(), 1, "vfmadd213pf ymm5,ymm2,ymm4"
|
||||
EmitIfCountGE \RowCount\(), 2, "vfmadd213pf ymm7,ymm2,ymm6"
|
||||
EmitIfCountGE \RowCount\(), 3, "vfmadd213pf ymm9,ymm2,ymm8"
|
||||
EmitIfCountGE \RowCount\(), 4, "vfmadd213pf ymm11,ymm2,ymm10"
|
||||
EmitIfCountGE \RowCount\(), 5, "vfmadd213pf ymm13,ymm2,ymm12"
|
||||
EmitIfCountGE \RowCount\(), 6, "vfmadd213pf ymm15,ymm2,ymm14"
|
||||
jmp .LStoreMasked1xNBlock\@
|
||||
|
||||
.LMultiplyAlphaMasked1xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmulpf ymm5,ymm5,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmulpf ymm7,ymm7,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmulpf ymm9,ymm9,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmulpf ymm11,ymm11,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmulpf ymm13,ymm13,ymm2"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmulpf ymm15,ymm15,ymm2"
|
||||
|
||||
.LStoreMasked1xNBlock\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "vmaskmovpf YMMWORD PTR [rdx],ymm0,ymm5"
|
||||
EmitIfCountGE \RowCount\(), 2, "vmaskmovpf YMMWORD PTR [rdx+rax],ymm0,ymm7"
|
||||
EmitIfCountGE \RowCount\(), 3, "vmaskmovpf YMMWORD PTR [rdx+rax*2],ymm0,ymm9"
|
||||
EmitIfCountGE \RowCount\(), 4, "vmaskmovpf YMMWORD PTR [rbx],ymm0,ymm11"
|
||||
EmitIfCountGE \RowCount\(), 5, "vmaskmovpf YMMWORD PTR [rbx+rax],ymm0,ymm13"
|
||||
EmitIfCountGE \RowCount\(), 6, "vmaskmovpf YMMWORD PTR [rbx+rax*2],ymm0,ymm15"
|
||||
.ifb \Fallthrough\()
|
||||
jmp .LExitKernelAndZeroUpper
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro generates the inner kernel to compute matrix multiplication.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name for the generated function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FgemmKernelFma3Function 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
|
||||
vzeroall
|
||||
|
||||
//
|
||||
// Process CountM rows of the matrices.
|
||||
//
|
||||
|
||||
cmp r8,5
|
||||
ja .LProcessCountM6
|
||||
je .LProcessCountM5
|
||||
cmp r8,3
|
||||
ja .LProcessCountM4
|
||||
je .LProcessCountM3
|
||||
cmp r8,1
|
||||
je .LProcessCountM1
|
||||
|
||||
.LProcessCountM2:
|
||||
ProcessCountM 2
|
||||
|
||||
.LProcessCountM4:
|
||||
ProcessCountM 4
|
||||
|
||||
.LProcessCountM6:
|
||||
mov r8d,6 # return 6 rows handled
|
||||
ProcessCountM 6, Fallthrough
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.LExitKernelAndZeroUpper:
|
||||
vzeroupper
|
||||
|
||||
.LExitKernel:
|
||||
mov eax,r8d
|
||||
pop r15
|
||||
pop rbx
|
||||
pop rbp
|
||||
ret
|
||||
|
||||
.LProcessCountM1:
|
||||
ProcessCountM 1
|
||||
|
||||
.LProcessCountM3:
|
||||
ProcessCountM 3
|
||||
|
||||
.LProcessCountM5:
|
||||
ProcessCountM 5
|
||||
|
||||
.endm
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
FgemmKernelSse2Common.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the floating point matrix/matrix
|
||||
multiply operation (SGEMM and DGEMM).
|
||||
|
||||
This implementation uses SSE2 instructions.
|
||||
|
||||
--*/
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This stores the block accumulators to the output matrix with an optional
|
||||
accumulation of the existing contents of the output matrix.
|
||||
|
||||
Arguments:
|
||||
|
||||
RowCount - Supplies the number of rows to process.
|
||||
|
||||
VectorCount - Supplies the number of vector columns to process.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
rax - Supplies the length in bytes of a row from matrix C.
|
||||
|
||||
rdx - Supplies the address of matrix C.
|
||||
|
||||
r15 - Stores the ZeroMode argument from the stack frame.
|
||||
|
||||
xmm8-xmm15 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro AccumulateAndStoreBlock RowCount, VectorCount
|
||||
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LSkipAccumulateOutput\@
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 1, "movupf xmm0,XMMWORD PTR [rdx]"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 2, "movupf xmm1,XMMWORD PTR [rdx+16]"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 3, "movupf xmm2,XMMWORD PTR [rdx+32]"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 4, "movupf xmm3,XMMWORD PTR [rdx+48]"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 1, "movupf xmm4,XMMWORD PTR [rdx+rax]"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 2, "movupf xmm5,XMMWORD PTR [rdx+rax+16]"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 3, "movupf xmm6,XMMWORD PTR [rdx+rax+32]"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 4, "movupf xmm7,XMMWORD PTR [rdx+rax+48]"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 1, "addpf xmm8,xmm0"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 2, "addpf xmm9,xmm1"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 3, "addpf xmm10,xmm2"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 4, "addpf xmm11,xmm3"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 1, "addpf xmm12,xmm4"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 2, "addpf xmm13,xmm5"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 3, "addpf xmm14,xmm6"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 4, "addpf xmm15,xmm7"
|
||||
|
||||
.LSkipAccumulateOutput\@:
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 1, "movupf XMMWORD PTR [rdx],xmm8"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 2, "movupf XMMWORD PTR [rdx+16],xmm9"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 3, "movupf XMMWORD PTR [rdx+32],xmm10"
|
||||
EmitIfCount2GE \RowCount\(), 1, \VectorCount\(), 4, "movupf XMMWORD PTR [rdx+48],xmm11"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 1, "movupf XMMWORD PTR [rdx+rax],xmm12"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 2, "movupf XMMWORD PTR [rdx+rax+16],xmm13"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 3, "movupf XMMWORD PTR [rdx+rax+32],xmm14"
|
||||
EmitIfCount2GE \RowCount\(), 2, \VectorCount\(), 4, "movupf XMMWORD PTR [rdx+rax+48],xmm15"
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro generates the inner kernel to compute matrix multiplication.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name for the generated function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FgemmKernelSse2Function 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]
|
||||
movsf .LFgemmKernelFrame_alpha[rsp],xmm0
|
||||
|
||||
//
|
||||
// Process CountM rows of the matrices.
|
||||
//
|
||||
|
||||
cmp r8,2
|
||||
jb .LProcessCountM1
|
||||
mov r8d,2 # return 2 rows handled
|
||||
ProcessCountM 2, Fallthrough
|
||||
|
||||
//
|
||||
// Restore non-volatile registers and return.
|
||||
//
|
||||
|
||||
.LExitKernel:
|
||||
mov eax,r8d
|
||||
pop r15
|
||||
pop rbx
|
||||
pop rbp
|
||||
ret
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
.LProcessCountM1:
|
||||
ProcessCountM 1
|
||||
|
||||
.endm
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelAvx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses AVX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "SgemmKernelCommon.h"
|
||||
#include "FgemmKernelAvxCommon.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelAvxFunction MlasGemmFloatKernelAvx
|
||||
|
||||
.end
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelAvx512F.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses AVX512F instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "SgemmKernelCommon.h"
|
||||
#include "FgemmKernelAvx512FCommon.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelAvx512FFunction MlasGemmFloatKernelAvx512F
|
||||
|
||||
.end
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelCommon.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module contains common kernel macros and structures for the single
|
||||
precision matrix/matrix multiply operation (SGEMM).
|
||||
|
||||
--*/
|
||||
|
||||
//
|
||||
// Define the single precision parameters.
|
||||
//
|
||||
|
||||
.equ .LFgemmElementShift, 2
|
||||
.equ .LFgemmElementSize, 1 << .LFgemmElementShift
|
||||
|
||||
#include "FgemmKernelCommon.h"
|
||||
|
||||
//
|
||||
// Define the typed instructions for single precision.
|
||||
//
|
||||
|
||||
FGEMM_TYPED_INSTRUCTION(addpf, addps)
|
||||
FGEMM_TYPED_INSTRUCTION(movsf, movss)
|
||||
FGEMM_TYPED_INSTRUCTION(movupf, movups)
|
||||
|
||||
FGEMM_TYPED_INSTRUCTION(vaddpf, vaddps)
|
||||
FGEMM_TYPED_INSTRUCTION(vbroadcastsf, vbroadcastss)
|
||||
FGEMM_TYPED_INSTRUCTION(vfmadd213pf, vfmadd213ps)
|
||||
FGEMM_TYPED_INSTRUCTION(vfmadd231pf, vfmadd231ps)
|
||||
FGEMM_TYPED_INSTRUCTION(vmaskmovpf, vmaskmovps)
|
||||
FGEMM_TYPED_INSTRUCTION(vmovapf, vmovaps)
|
||||
FGEMM_TYPED_INSTRUCTION(vmovsf, vmovss)
|
||||
FGEMM_TYPED_INSTRUCTION(vmovupf, vmovups)
|
||||
FGEMM_TYPED_INSTRUCTION(vmulpf, vmulps)
|
||||
FGEMM_TYPED_INSTRUCTION(vxorpf, vxorps)
|
||||
|
||||
.macro vfmadd231pf_bcst DestReg, SrcReg, Address
|
||||
|
||||
vfmadd231ps \DestReg\(), \SrcReg\(), \Address\(){1to16}
|
||||
|
||||
.endm
|
||||
+34
@@ -0,0 +1,34 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelFma3.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses AVX fused multiply/add instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "SgemmKernelCommon.h"
|
||||
#include "FgemmKernelFma3Common.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelFma3Function MlasGemmFloatKernelFma3
|
||||
|
||||
.end
|
||||
+267
@@ -0,0 +1,267 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelM1Avx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM). This handles the special case of M=1.
|
||||
|
||||
This implementation uses AVX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows. This handles the special case of M=1.
|
||||
|
||||
The elements in matrix B are not transposed.
|
||||
|
||||
Arguments:
|
||||
|
||||
A (rdi) - Supplies the address of matrix A.
|
||||
|
||||
B (rsi) - Supplies the address of matrix B.
|
||||
|
||||
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.
|
||||
|
||||
CountN (r8) - Supplies the number of columns from matrix B and matrix C to
|
||||
iterate over.
|
||||
|
||||
ldb (r9) - Supplies the first dimension of matrix B.
|
||||
|
||||
Beta (xmm0) - Supplies the scalar beta multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmKernelM1Avx
|
||||
|
||||
push rbx
|
||||
shl r9,2 # convert ldb to bytes
|
||||
mov r10,rdx
|
||||
mov r11,rsi
|
||||
|
||||
//
|
||||
// Compute the initial results mask for zeroing or accumulate mode.
|
||||
//
|
||||
|
||||
vxorps xmm1,xmm1,xmm1
|
||||
vcmpeqss xmm0,xmm1,xmm0
|
||||
vshufps xmm0,xmm0,xmm0,0
|
||||
vinsertf128 ymm0,ymm0,xmm0,1
|
||||
|
||||
//
|
||||
// Compute the conditional load/store mask for an unaligned CountN.
|
||||
//
|
||||
|
||||
mov eax,r8d
|
||||
and eax,7
|
||||
vmovd xmm7,eax
|
||||
vshufps xmm7,xmm7,xmm7,0
|
||||
vpcmpgtd xmm6,xmm7,XMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip+16]
|
||||
vpcmpgtd xmm7,xmm7,XMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip]
|
||||
vinsertf128 ymm7,ymm7,xmm6,1
|
||||
|
||||
//
|
||||
// Process 4 rows of the matrices in a loop.
|
||||
//
|
||||
|
||||
sub rcx,4
|
||||
jb .LProcessRemainingCountK
|
||||
|
||||
.LProcessRowLoop4:
|
||||
vbroadcastss ymm2,DWORD PTR [rdi]
|
||||
mov rax,r8 # reload CountN
|
||||
vbroadcastss ymm3,DWORD PTR [rdi+4]
|
||||
mov rsi,r11 # reload matrix B
|
||||
vbroadcastss ymm4,DWORD PTR [rdi+8]
|
||||
mov rdx,r10 # reload matrix C
|
||||
vbroadcastss ymm5,DWORD PTR [rdi+12]
|
||||
add rdi,4*4 # advance matrix A by 4 columns
|
||||
lea r11,[rsi+r9*4] # advance matrix B by 4 rows
|
||||
sub rax,16
|
||||
jb .LProcessRemainingCountN4
|
||||
|
||||
.LProcessColumnLoop4:
|
||||
lea rbx,[rsi+r9*2] # compute matrix B plus 2 rows
|
||||
vmulps ymm1,ymm2,YMMWORD PTR [rsi]
|
||||
vmulps ymm6,ymm2,YMMWORD PTR [rsi+32]
|
||||
vmulps ymm8,ymm3,YMMWORD PTR [rsi+r9]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmulps ymm8,ymm3,YMMWORD PTR [rsi+r9+32]
|
||||
vaddps ymm6,ymm6,ymm8
|
||||
vmulps ymm8,ymm4,YMMWORD PTR [rbx]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmulps ymm8,ymm4,YMMWORD PTR [rbx+32]
|
||||
vaddps ymm6,ymm6,ymm8
|
||||
vmulps ymm8,ymm5,YMMWORD PTR [rbx+r9]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmulps ymm8,ymm5,YMMWORD PTR [rbx+r9+32]
|
||||
vaddps ymm6,ymm6,ymm8
|
||||
vandnps ymm8,ymm0,YMMWORD PTR [rdx]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vandnps ymm8,ymm0,YMMWORD PTR [rdx+32]
|
||||
vaddps ymm6,ymm6,ymm8
|
||||
vmovups YMMWORD PTR [rdx],ymm1
|
||||
vmovups YMMWORD PTR [rdx+32],ymm6
|
||||
add rsi,16*4 # advance matrix B by 16 columns
|
||||
add rdx,16*4 # advance matrix C by 16 columns
|
||||
sub rax,16
|
||||
jae .LProcessColumnLoop4
|
||||
|
||||
.LProcessRemainingCountN4:
|
||||
test al,15 # test for unaligned columns
|
||||
jz .LProcessedRemainingCountN4
|
||||
test al,8 # CountN >= 8?
|
||||
jz .LProcessRemainingCountNSmall4
|
||||
lea rbx,[rsi+r9*2] # compute matrix B plus 2 rows
|
||||
vmulps ymm1,ymm2,YMMWORD PTR [rsi]
|
||||
vmulps ymm8,ymm3,YMMWORD PTR [rsi+r9]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmulps ymm8,ymm4,YMMWORD PTR [rbx]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmulps ymm8,ymm5,YMMWORD PTR [rbx+r9]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vandnps ymm8,ymm0,YMMWORD PTR [rdx]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmovups YMMWORD PTR [rdx],ymm1
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
add rdx,8*4 # advance matrix C by 8 columns
|
||||
test al,7
|
||||
jz .LProcessedRemainingCountN4
|
||||
|
||||
.LProcessRemainingCountNSmall4:
|
||||
lea rbx,[rsi+r9*2] # compute matrix B plus 2 rows
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm1,ymm2,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi+r9]
|
||||
vmulps ymm8,ymm3,ymm6
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rbx]
|
||||
vmulps ymm8,ymm4,ymm6
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rbx+r9]
|
||||
vmulps ymm8,ymm5,ymm6
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rdx]
|
||||
vandnps ymm6,ymm0,ymm6
|
||||
vaddps ymm1,ymm1,ymm6
|
||||
vmaskmovps YMMWORD PTR [rdx],ymm7,ymm1
|
||||
|
||||
.LProcessedRemainingCountN4:
|
||||
vxorps xmm0,xmm0,xmm0 # switch to accumulate mode
|
||||
sub rcx,4
|
||||
jae .LProcessRowLoop4
|
||||
|
||||
.LProcessRemainingCountK:
|
||||
test cl,2
|
||||
jnz .LProcessRowLoop2
|
||||
test cl,1
|
||||
jnz .LProcessRowLoop1
|
||||
|
||||
.LExitKernel:
|
||||
vzeroupper
|
||||
pop rbx
|
||||
ret
|
||||
|
||||
//
|
||||
// Process 2 rows of the matrices.
|
||||
//
|
||||
|
||||
.LProcessRowLoop2:
|
||||
vbroadcastss ymm2,DWORD PTR [rdi]
|
||||
mov rax,r8 # reload CountN
|
||||
vbroadcastss ymm3,DWORD PTR [rdi+4]
|
||||
mov rsi,r11 # reload matrix B
|
||||
mov rdx,r10 # reload matrix C
|
||||
add rdi,2*4 # advance matrix A by 2 columns
|
||||
lea r11,[rsi+r9*2] # advance matrix B by 2 rows
|
||||
sub rax,8
|
||||
jb .LProcessRemainingCountN2
|
||||
|
||||
.LProcessColumnLoop2:
|
||||
vmulps ymm1,ymm2,YMMWORD PTR [rsi]
|
||||
vmulps ymm8,ymm3,YMMWORD PTR [rsi+r9]
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vandnps ymm6,ymm0,YMMWORD PTR [rdx]
|
||||
vaddps ymm1,ymm1,ymm6
|
||||
vmovups YMMWORD PTR [rdx],ymm1
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
add rdx,8*4 # advance matrix C by 8 columns
|
||||
sub rax,8
|
||||
jae .LProcessColumnLoop2
|
||||
|
||||
.LProcessRemainingCountN2:
|
||||
test al,7 # test for unaligned columns
|
||||
jz .LProcessedRemainingCountN2
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm1,ymm2,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi+r9]
|
||||
vmulps ymm8,ymm3,ymm6
|
||||
vaddps ymm1,ymm1,ymm8
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rdx]
|
||||
vandnps ymm6,ymm0,ymm6
|
||||
vaddps ymm1,ymm1,ymm6
|
||||
vmaskmovps YMMWORD PTR [rdx],ymm7,ymm1
|
||||
|
||||
.LProcessedRemainingCountN2:
|
||||
test cl,1
|
||||
jz .LExitKernel
|
||||
vxorps xmm0,xmm0,xmm0 # switch to accumulate mode
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
.LProcessRowLoop1:
|
||||
vbroadcastss ymm2,DWORD PTR [rdi]
|
||||
mov rax,r8 # reload CountN
|
||||
mov rsi,r11 # reload matrix B
|
||||
mov rdx,r10 # reload matrix C
|
||||
sub rax,8
|
||||
jb .LProcessRemainingCountN1
|
||||
|
||||
.LProcessColumnLoop1:
|
||||
vmulps ymm1,ymm2,YMMWORD PTR [rsi]
|
||||
vandnps ymm6,ymm0,YMMWORD PTR [rdx]
|
||||
vaddps ymm1,ymm1,ymm6
|
||||
vmovups YMMWORD PTR [rdx],ymm1
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
add rdx,8*4 # advance matrix C by 8 columns
|
||||
sub rax,8
|
||||
jae .LProcessColumnLoop1
|
||||
|
||||
.LProcessRemainingCountN1:
|
||||
test al,7 # test for unaligned columns
|
||||
jz .LExitKernel
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm1,ymm2,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rdx]
|
||||
vandnps ymm6,ymm0,ymm6
|
||||
vaddps ymm1,ymm1,ymm6
|
||||
vmaskmovps YMMWORD PTR [rdx],ymm7,ymm1
|
||||
jmp .LExitKernel
|
||||
|
||||
.end
|
||||
@@ -0,0 +1,275 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelM1TransposeBAvx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM). This handles the special case of M=1.
|
||||
|
||||
This implementation uses AVX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine is an inner kernel to compute matrix multiplication for a
|
||||
set of rows. This handles the special case of M=1.
|
||||
|
||||
The elements in matrix B are transposed.
|
||||
|
||||
Arguments:
|
||||
|
||||
A (rdi) - Supplies the address of matrix A.
|
||||
|
||||
B (rsi) - Supplies the address of matrix B. The elements are transposed.
|
||||
|
||||
C (rdx) - Supplies the address of matrix C.
|
||||
|
||||
CountK (rcx) - Supplies the number of columns from matrix A and the number
|
||||
of columns from matrix B to iterate over.
|
||||
|
||||
CountN (r8) - Supplies the number of rows from matrix B and the number of
|
||||
columns from matrix C to iterate over.
|
||||
|
||||
ldb (r9) - Supplies the first dimension of matrix B.
|
||||
|
||||
Beta (xmm0) - Supplies the scalar beta multiplier (see SGEMM definition).
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmKernelM1TransposeBAvx
|
||||
|
||||
push rbx
|
||||
shl r9,2 # convert ldb to bytes
|
||||
mov r10,rdi
|
||||
mov r11,rsi
|
||||
|
||||
//
|
||||
// Compute the results mask for zeroing or accumulate mode.
|
||||
//
|
||||
|
||||
vxorps xmm1,xmm1,xmm1
|
||||
vcmpeqss xmm0,xmm1,xmm0
|
||||
vshufps xmm0,xmm0,xmm0,0
|
||||
|
||||
//
|
||||
// Compute the conditional load/store mask for an unaligned CountK.
|
||||
//
|
||||
|
||||
mov eax,ecx
|
||||
and eax,7
|
||||
vmovd xmm7,eax
|
||||
vshufps xmm7,xmm7,xmm7,0
|
||||
vpcmpgtd xmm6,xmm7,XMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip+16]
|
||||
vpcmpgtd xmm7,xmm7,XMMWORD PTR C_UNDERSCORE(MlasMaskMoveAvx)[rip]
|
||||
vinsertf128 ymm7,ymm7,xmm6,1
|
||||
|
||||
//
|
||||
// Process 4 rows of the matrices in a loop.
|
||||
//
|
||||
|
||||
sub r8,4
|
||||
jb .LProcessRemainingCountN
|
||||
|
||||
.LProcessRowLoop4:
|
||||
vxorps xmm2,xmm2,xmm2 # clear row accumulators
|
||||
vxorps xmm3,xmm3,xmm3
|
||||
vxorps xmm4,xmm4,xmm4
|
||||
vxorps xmm5,xmm5,xmm5
|
||||
mov rdi,r10 # reload matrix A
|
||||
mov rsi,r11 # reload matrix B
|
||||
mov rax,rcx # reload CountK
|
||||
lea r11,[rsi+r9*4] # advance matrix B by 4 rows
|
||||
sub rax,8
|
||||
jb .LProcessRemainingCountK4
|
||||
|
||||
.LProcessColumnLoop4:
|
||||
lea rbx,[rsi+r9*2] # compute matrix B plus 2 rows
|
||||
vmovups ymm1,YMMWORD PTR [rdi]
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rsi]
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rsi+r9]
|
||||
vaddps ymm3,ymm3,ymm6
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rbx]
|
||||
vaddps ymm4,ymm4,ymm6
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rbx+r9]
|
||||
vaddps ymm5,ymm5,ymm6
|
||||
add rdi,8*4 # advance matrix A by 8 columns
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
sub rax,8
|
||||
jae .LProcessColumnLoop4
|
||||
|
||||
.LProcessRemainingCountK4:
|
||||
test al,7 # test for unaligned columns
|
||||
jz .LOutput4x1Block
|
||||
lea rbx,[rsi+r9*2] # compute matrix B plus 2 rows
|
||||
vmaskmovps ymm1,ymm7,YMMWORD PTR [rdi]
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi+r9]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm3,ymm3,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rbx]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm4,ymm4,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rbx+r9]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm5,ymm5,ymm6
|
||||
|
||||
//
|
||||
// Reduce and output the row accumulators.
|
||||
//
|
||||
|
||||
.LOutput4x1Block:
|
||||
vunpcklps ymm6,ymm2,ymm3 # transpose row accumulators
|
||||
vunpckhps ymm1,ymm2,ymm3
|
||||
vunpcklps ymm2,ymm4,ymm5
|
||||
vunpckhps ymm3,ymm4,ymm5
|
||||
vunpcklpd ymm4,ymm6,ymm2
|
||||
vunpckhpd ymm5,ymm6,ymm2
|
||||
vaddps ymm4,ymm4,ymm5
|
||||
vunpcklpd ymm6,ymm1,ymm3
|
||||
vunpckhpd ymm2,ymm1,ymm3
|
||||
vaddps ymm4,ymm4,ymm6
|
||||
vaddps ymm4,ymm4,ymm2
|
||||
vextractf128 xmm5,ymm4,1
|
||||
vaddps xmm4,xmm4,xmm5
|
||||
vandnps xmm6,xmm0,XMMWORD PTR [rdx]
|
||||
vaddps xmm4,xmm4,xmm6
|
||||
vmovups XMMWORD PTR [rdx],xmm4
|
||||
add rdx,4*4 # advance matrix C by 4 columns
|
||||
sub r8,4
|
||||
jae .LProcessRowLoop4
|
||||
|
||||
.LProcessRemainingCountN:
|
||||
test r8d,2
|
||||
jnz .LProcessRowLoop2
|
||||
test r8d,1
|
||||
jnz .LProcessRowLoop1
|
||||
|
||||
.LExitKernel:
|
||||
vzeroupper
|
||||
pop rbx
|
||||
ret
|
||||
|
||||
//
|
||||
// Process 2 rows of the matrices.
|
||||
//
|
||||
|
||||
.LProcessRowLoop2:
|
||||
vxorps xmm2,xmm2,xmm2 # clear row accumulators
|
||||
vxorps xmm3,xmm3,xmm3
|
||||
mov rdi,r10 # reload matrix A
|
||||
mov rsi,r11 # reload matrix B
|
||||
mov rax,rcx # reload CountK
|
||||
lea r11,[rsi+r9*2] # advance matrix B by 2 rows
|
||||
sub rax,8
|
||||
jb .LProcessRemainingCountK2
|
||||
|
||||
.LProcessColumnLoop2:
|
||||
vmovups ymm1,YMMWORD PTR [rdi]
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rsi]
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rsi+r9]
|
||||
vaddps ymm3,ymm3,ymm6
|
||||
add rdi,8*4 # advance matrix A by 8 columns
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
sub rax,8
|
||||
jae .LProcessColumnLoop2
|
||||
|
||||
.LProcessRemainingCountK2:
|
||||
test al,7 # test for unaligned columns
|
||||
jz .LOutput2x1Block
|
||||
vmaskmovps ymm1,ymm7,YMMWORD PTR [rdi]
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi+r9]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm3,ymm3,ymm6
|
||||
|
||||
//
|
||||
// Reduce and output the row accumulators.
|
||||
//
|
||||
|
||||
.LOutput2x1Block:
|
||||
vunpcklps ymm4,ymm2,ymm3 # reduce row accumulators
|
||||
vunpckhps ymm2,ymm2,ymm3
|
||||
vaddps ymm2,ymm2,ymm4
|
||||
vextractf128 xmm4,ymm2,1
|
||||
vaddps xmm2,xmm2,xmm4
|
||||
vmovhlps xmm4,xmm2,xmm2
|
||||
vaddps xmm2,xmm2,xmm4
|
||||
vmovsd xmm3,QWORD PTR [rdx]
|
||||
vandnps xmm3,xmm0,xmm3
|
||||
vaddps xmm2,xmm2,xmm3
|
||||
vmovsd QWORD PTR [rdx],xmm2
|
||||
add rdx,2*4 # advance matrix C by 2 columns
|
||||
test r8d,1
|
||||
jz .LExitKernel
|
||||
|
||||
//
|
||||
// Process 1 row of the matrices.
|
||||
//
|
||||
|
||||
.LProcessRowLoop1:
|
||||
vxorps xmm2,xmm2,xmm2 # clear row accumulators
|
||||
mov rdi,r10 # reload matrix A
|
||||
mov rsi,r11 # reload matrix B
|
||||
mov rax,rcx # reload CountK
|
||||
sub rax,8
|
||||
jb .LProcessRemainingCountK1
|
||||
|
||||
.LProcessColumnLoop1:
|
||||
vmovups ymm1,YMMWORD PTR [rdi]
|
||||
vmulps ymm6,ymm1,YMMWORD PTR [rsi]
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
add rdi,8*4 # advance matrix A by 8 columns
|
||||
add rsi,8*4 # advance matrix B by 8 columns
|
||||
sub rax,8
|
||||
jae .LProcessColumnLoop1
|
||||
|
||||
.LProcessRemainingCountK1:
|
||||
test al,7 # test for unaligned columns
|
||||
jz .LOutput1x1Block
|
||||
vmaskmovps ymm1,ymm7,YMMWORD PTR [rdi]
|
||||
vmaskmovps ymm6,ymm7,YMMWORD PTR [rsi]
|
||||
vmulps ymm6,ymm1,ymm6
|
||||
vaddps ymm2,ymm2,ymm6
|
||||
|
||||
//
|
||||
// Reduce and output the row accumulators.
|
||||
//
|
||||
|
||||
.LOutput1x1Block:
|
||||
vhaddps ymm2,ymm2,ymm2 # reduce row accumulators
|
||||
vhaddps ymm2,ymm2,ymm2
|
||||
vextractf128 xmm4,ymm2,1
|
||||
vaddss xmm2,xmm2,xmm4
|
||||
vmovss xmm3,DWORD PTR [rdx]
|
||||
vandnps xmm3,xmm0,xmm3
|
||||
vaddss xmm2,xmm2,xmm3
|
||||
vmovss DWORD PTR [rdx],xmm2
|
||||
jmp .LExitKernel
|
||||
|
||||
.end
|
||||
+273
@@ -0,0 +1,273 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmKernelSse2.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements the kernels for the single precision matrix/matrix
|
||||
multiply operation (SGEMM).
|
||||
|
||||
This implementation uses SSE2 instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
#include "SgemmKernelCommon.h"
|
||||
#include "FgemmKernelSse2Common.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro multiplies and accumulates for a 16xN block of the output matrix.
|
||||
|
||||
Arguments:
|
||||
|
||||
RowCount - Supplies the number of rows to process.
|
||||
|
||||
VectorOffset - Supplies the byte offset from matrix B to fetch elements.
|
||||
|
||||
Shuffle - Supplies the shuffle mask to extract the element from matrix A.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
rsi - Supplies the address into the matrix B data.
|
||||
|
||||
xmm0-xmm1 - Supplies up to four elements loaded from matrix A and matrix A
|
||||
plus one row.
|
||||
|
||||
xmm8-xmm15 - Supplies the block accumulators.
|
||||
|
||||
--*/
|
||||
|
||||
.macro ComputeBlockSseBy16 RowCount, VectorOffset, Shuffle
|
||||
|
||||
movaps xmm4,XMMWORD PTR [rsi+\VectorOffset\()]
|
||||
movaps xmm5,XMMWORD PTR [rsi+\VectorOffset\()+16]
|
||||
pshufd xmm2,xmm0,\Shuffle\()
|
||||
.if \RowCount\() == 2
|
||||
pshufd xmm3,xmm1,\Shuffle\()
|
||||
movaps xmm6,xmm4
|
||||
movaps xmm7,xmm5
|
||||
.endif
|
||||
mulps xmm4,xmm2
|
||||
mulps xmm5,xmm2
|
||||
addps xmm8,xmm4
|
||||
addps xmm9,xmm5
|
||||
.if \RowCount\() == 2
|
||||
mulps xmm6,xmm3
|
||||
mulps xmm7,xmm3
|
||||
addps xmm12,xmm6
|
||||
addps xmm13,xmm7
|
||||
.endif
|
||||
movaps xmm4,XMMWORD PTR [rsi+\VectorOffset\()+32]
|
||||
movaps xmm5,XMMWORD PTR [rsi+\VectorOffset\()+48]
|
||||
.if \RowCount\() == 2
|
||||
movaps xmm6,xmm4
|
||||
movaps xmm7,xmm5
|
||||
.endif
|
||||
mulps xmm4,xmm2
|
||||
mulps xmm5,xmm2
|
||||
addps xmm10,xmm4
|
||||
addps xmm11,xmm5
|
||||
.if \RowCount\() == 2
|
||||
mulps xmm6,xmm3
|
||||
mulps xmm7,xmm3
|
||||
addps xmm14,xmm6
|
||||
addps xmm15,xmm7
|
||||
.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
|
||||
|
||||
.LProcessNextColumnLoop16xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "xorps xmm8,xmm8"
|
||||
EmitIfCountGE \RowCount\(), 1, "xorps xmm9,xmm9"
|
||||
EmitIfCountGE \RowCount\(), 1, "xorps xmm10,xmm10"
|
||||
EmitIfCountGE \RowCount\(), 1, "xorps xmm11,xmm11"
|
||||
EmitIfCountGE \RowCount\(), 2, "xorps xmm12,xmm12"
|
||||
EmitIfCountGE \RowCount\(), 2, "xorps xmm13,xmm13"
|
||||
EmitIfCountGE \RowCount\(), 2, "xorps xmm14,xmm14"
|
||||
EmitIfCountGE \RowCount\(), 2, "xorps xmm15,xmm15"
|
||||
mov rbp,rcx # reload CountK
|
||||
sub rbp,4
|
||||
jb .LProcessRemaining16xNBlocks\@
|
||||
|
||||
.LCompute16xNBlockBy4Loop\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "movups xmm0,XMMWORD PTR [rdi]"
|
||||
EmitIfCountGE \RowCount\(), 2, "movups xmm1,XMMWORD PTR [rdi+r10]"
|
||||
ComputeBlockSseBy16 2, 0, 0x00
|
||||
ComputeBlockSseBy16 2, 16*4, 0x55
|
||||
sub rsi,-32*4 # advance matrix B by 32 columns
|
||||
ComputeBlockSseBy16 2, 0, 0xAA
|
||||
ComputeBlockSseBy16 2, 16*4, 0xFF
|
||||
sub rsi,-32*4 # advance matrix B by 32 columns
|
||||
add rdi,4*4 # advance matrix A by 4 columns
|
||||
sub rbp,4
|
||||
jae .LCompute16xNBlockBy4Loop\@
|
||||
|
||||
.LProcessRemaining16xNBlocks\@:
|
||||
add rbp,4 # correct for over-subtract above
|
||||
jz .LOutput16xNBlock\@
|
||||
|
||||
.LCompute16xNBlockBy1Loop\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "movss xmm0,[rdi]"
|
||||
EmitIfCountGE \RowCount\(), 2, "movss xmm1,[rdi+r10]"
|
||||
ComputeBlockSseBy16 2, 0, 0x00
|
||||
add rsi,16*4 # advance matrix B by 16 columns
|
||||
add rdi,4 # advance matrix A by 1 column
|
||||
dec rbp
|
||||
jne .LCompute16xNBlockBy1Loop\@
|
||||
|
||||
.LOutput16xNBlock\@:
|
||||
movss xmm2,.LFgemmKernelFrame_alpha[rsp]
|
||||
shufps xmm2,xmm2,0
|
||||
EmitIfCountGE \RowCount\(), 1, "mulps xmm8,xmm2"
|
||||
# multiply by alpha
|
||||
EmitIfCountGE \RowCount\(), 1, "mulps xmm9,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 1, "mulps xmm10,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 1, "mulps xmm11,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "mulps xmm12,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "mulps xmm13,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "mulps xmm14,xmm2"
|
||||
EmitIfCountGE \RowCount\(), 2, "mulps xmm15,xmm2"
|
||||
sub r9,16
|
||||
jb .LOutputPartial16xNBlock\@
|
||||
AccumulateAndStoreBlock \RowCount\(), 4
|
||||
add rdx,16*4 # advance matrix C by 16 columns
|
||||
mov rdi,r11 # reload matrix A
|
||||
test r9,r9
|
||||
jnz .LProcessNextColumnLoop16xN\@
|
||||
jmp .LExitKernel
|
||||
|
||||
//
|
||||
// Output a partial 16xN block to the matrix.
|
||||
//
|
||||
|
||||
.LOutputPartial16xNBlock\@:
|
||||
add r9,16 # correct for over-subtract above
|
||||
cmp r9,4
|
||||
jb .LOutputPartialLessThan4xNBlock\@
|
||||
cmp r9,8
|
||||
jb .LOutputPartialLessThan8xNBlock\@
|
||||
cmp r9,12
|
||||
jb .LOutputPartialLessThan12xNBlock\@
|
||||
AccumulateAndStoreBlock \RowCount\(), 3
|
||||
and r9d,3 # check if remaining count is small
|
||||
jz .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "movaps xmm8,xmm11"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "movaps xmm12,xmm15"
|
||||
add rdx,12*4 # advance matrix C by 12 columns
|
||||
jmp .LOutputPartialLessThan4xNBlock\@
|
||||
|
||||
.LOutputPartialLessThan12xNBlock\@:
|
||||
AccumulateAndStoreBlock \RowCount\(), 2
|
||||
and r9d,3 # check if remaining count is small
|
||||
jz .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "movaps xmm8,xmm10"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "movaps xmm12,xmm14"
|
||||
add rdx,8*4 # advance matrix C by 8 columns
|
||||
jmp .LOutputPartialLessThan4xNBlock\@
|
||||
|
||||
.LOutputPartialLessThan8xNBlock\@:
|
||||
AccumulateAndStoreBlock \RowCount\(), 1
|
||||
and r9d,3 # check if remaining count is small
|
||||
jz .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "movaps xmm8,xmm9"
|
||||
# shift remaining elements down
|
||||
EmitIfCountGE \RowCount\(), 2, "movaps xmm12,xmm13"
|
||||
add rdx,4*4 # advance matrix C by 4 columns
|
||||
|
||||
.LOutputPartialLessThan4xNBlock\@:
|
||||
test r9d,2
|
||||
jz .LOutputPartial1xNBlock\@
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LSkipAccumulateOutput2xN\@
|
||||
EmitIfCountGE \RowCount\(), 1, "movsd xmm0,QWORD PTR [rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "movsd xmm1,QWORD PTR [rdx+rax]"
|
||||
EmitIfCountGE \RowCount\(), 1, "addps xmm8,xmm0"
|
||||
EmitIfCountGE \RowCount\(), 2, "addps xmm12,xmm1"
|
||||
|
||||
.LSkipAccumulateOutput2xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "movsd QWORD PTR [rdx],xmm8"
|
||||
EmitIfCountGE \RowCount\(), 2, "movsd QWORD PTR [rdx+rax],xmm12"
|
||||
test r9d,1 # check if remaining count is odd
|
||||
jz .LExitKernel
|
||||
EmitIfCountGE \RowCount\(), 1, "movhlps xmm8,xmm8"
|
||||
# shift third element down
|
||||
EmitIfCountGE \RowCount\(), 2, "movhlps xmm12,xmm12"
|
||||
add rdx,2*4 # advance matrix C by 2 columns
|
||||
|
||||
.LOutputPartial1xNBlock\@:
|
||||
test r15b,r15b # ZeroMode?
|
||||
jnz .LSkipAccumulateOutput1xN\@
|
||||
EmitIfCountGE \RowCount\(), 1, "addss xmm8,[rdx]"
|
||||
EmitIfCountGE \RowCount\(), 2, "addss xmm12,[rdx+rax]"
|
||||
|
||||
.LSkipAccumulateOutput1xN\@:
|
||||
EmitIfCountGE \RowCount\(), 1, "movss [rdx],xmm8"
|
||||
EmitIfCountGE \RowCount\(), 2, "movss [rdx+rax],xmm12"
|
||||
.ifb \Fallthrough\()
|
||||
jmp .LExitKernel
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
//
|
||||
// Generate the GEMM kernel.
|
||||
//
|
||||
|
||||
FgemmKernelSse2Function MlasGemmFloatKernelSse
|
||||
|
||||
.end
|
||||
@@ -0,0 +1,120 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmTransposePackB16x4Avx.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements routines for packing buffers for the single precision
|
||||
matrix/matrix multiply operation (SGEMM).
|
||||
|
||||
This implementation uses AVX instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
4 columns of 8 rows from the source matrix are transposed to 8 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
StoreOffset - Supplies the relative byte offset into the destination packed
|
||||
buffer.
|
||||
|
||||
Implicit Arguments:
|
||||
|
||||
rdi - Supplies the address of the destination packed buffer.
|
||||
|
||||
rsi - Supplies the address of the source matrix.
|
||||
|
||||
rdx - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
--*/
|
||||
|
||||
.macro TransposePackB8x4BlockAvx StoreOffset
|
||||
|
||||
//
|
||||
// Load 4 columns from 8 rows of the source matrix into the lower and upper
|
||||
// halves of 4 YMM registers.
|
||||
//
|
||||
|
||||
lea rax,[rsi+rdx*2]
|
||||
vmovups xmm0,XMMWORD PTR [rsi]
|
||||
vmovups xmm1,XMMWORD PTR [rsi+rdx]
|
||||
lea rsi,[rax+rdx*2]
|
||||
vmovups xmm2,XMMWORD PTR [rax]
|
||||
vmovups xmm3,XMMWORD PTR [rax+rdx]
|
||||
lea rax,[rsi+rdx*2]
|
||||
vinsertf128 ymm0,ymm0,XMMWORD PTR [rsi],1
|
||||
vinsertf128 ymm1,ymm1,XMMWORD PTR [rsi+rdx],1
|
||||
vinsertf128 ymm2,ymm2,XMMWORD PTR [rax],1
|
||||
vinsertf128 ymm3,ymm3,XMMWORD PTR [rax+rdx],1
|
||||
|
||||
//
|
||||
// Transpose the lower and upper halves of the 4 YMM registers as two 4x4
|
||||
// matrices and store the output to the destination packed buffer.
|
||||
//
|
||||
|
||||
vunpcklps ymm4,ymm0,ymm1
|
||||
vunpckhps ymm5,ymm0,ymm1
|
||||
vunpcklps ymm0,ymm2,ymm3
|
||||
vunpckhps ymm1,ymm2,ymm3
|
||||
vunpcklpd ymm2,ymm4,ymm0
|
||||
vunpckhpd ymm3,ymm4,ymm0
|
||||
vmovaps YMMWORD PTR [rdi+16*4*0+\StoreOffset\()],ymm2
|
||||
vmovaps YMMWORD PTR [rdi+16*4*1+\StoreOffset\()],ymm3
|
||||
vunpcklpd ymm0,ymm5,ymm1
|
||||
vunpckhpd ymm4,ymm5,ymm1
|
||||
vmovaps YMMWORD PTR [rdi+16*4*2+\StoreOffset\()],ymm0
|
||||
vmovaps YMMWORD PTR [rdi+16*4*3+\StoreOffset\()],ymm4
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine transposes elements from the source matrix to the destination
|
||||
packed buffer.
|
||||
|
||||
4 columns of 16 rows from the source matrix are transposed to 16 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
D (rdi) - Supplies the address of the destination packed buffer.
|
||||
|
||||
B (rsi) - Supplies the address of the source matrix.
|
||||
|
||||
ldb (rdx) - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmTransposePackB16x4Avx
|
||||
|
||||
shl rdx,2 # convert ldb to bytes
|
||||
TransposePackB8x4BlockAvx 0*4
|
||||
lea rsi,[rax+rdx*2]
|
||||
TransposePackB8x4BlockAvx 8*4
|
||||
vzeroupper
|
||||
ret
|
||||
|
||||
.end
|
||||
@@ -0,0 +1,83 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
SgemmTransposePackB16x4Sse2.s
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements routines for packing buffers for the single precision
|
||||
matrix/matrix multiply operation (SGEMM).
|
||||
|
||||
This implementation uses SSE2 instructions.
|
||||
|
||||
--*/
|
||||
|
||||
#include "asmmacro.h"
|
||||
|
||||
.intel_syntax noprefix
|
||||
|
||||
.text
|
||||
|
||||
/*++
|
||||
|
||||
Routine Description:
|
||||
|
||||
This routine transposes elements from the source matrix to the destination
|
||||
packed buffer.
|
||||
|
||||
4 columns of 16 rows from the source matrix are transposed to 16 columns of 4
|
||||
rows in the destination packed buffer.
|
||||
|
||||
Arguments:
|
||||
|
||||
D (rdi) - Supplies the address of the destination packed buffer.
|
||||
|
||||
B (rsi) - Supplies the address of the source matrix.
|
||||
|
||||
ldb (rdx) - Supplies the number of elements per row of the source matrix.
|
||||
|
||||
Return Value:
|
||||
|
||||
None.
|
||||
|
||||
--*/
|
||||
|
||||
FUNCTION_ENTRY MlasSgemmTransposePackB16x4Sse
|
||||
|
||||
shl rdx,2 # convert ldb to bytes
|
||||
mov ecx,4 # transpose four 4x4 blocks
|
||||
|
||||
.LTransposeBlockLoop:
|
||||
lea rax,[rsi+rdx*2]
|
||||
movups xmm0,XMMWORD PTR [rsi]
|
||||
movups xmm1,XMMWORD PTR [rsi+rdx]
|
||||
movups xmm2,XMMWORD PTR [rax]
|
||||
movups xmm3,XMMWORD PTR [rax+rdx]
|
||||
movaps xmm4,xmm0
|
||||
unpcklps xmm4,xmm1
|
||||
unpckhps xmm0,xmm1
|
||||
movaps xmm5,xmm2
|
||||
unpcklps xmm5,xmm3
|
||||
unpckhps xmm2,xmm3
|
||||
movaps xmm1,xmm4
|
||||
unpcklpd xmm1,xmm5
|
||||
unpckhpd xmm4,xmm5
|
||||
movaps xmm3,xmm0
|
||||
unpcklpd xmm3,xmm2
|
||||
unpckhpd xmm0,xmm2
|
||||
movaps XMMWORD PTR [rdi+16*4*0],xmm1
|
||||
movaps XMMWORD PTR [rdi+16*4*1],xmm4
|
||||
movaps XMMWORD PTR [rdi+16*4*2],xmm3
|
||||
movaps XMMWORD PTR [rdi+16*4*3],xmm0
|
||||
add rdi,4*4
|
||||
lea rsi,[rax+rdx*2]
|
||||
dec ecx
|
||||
jnz .LTransposeBlockLoop
|
||||
ret
|
||||
|
||||
.end
|
||||
Vendored
+172
@@ -0,0 +1,172 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
asmmacro.h
|
||||
|
||||
Abstract:
|
||||
|
||||
This module implements common macros for the assembly modules.
|
||||
|
||||
--*/
|
||||
|
||||
#if defined(__APPLE__)
|
||||
#define C_UNDERSCORE(symbol) _##symbol
|
||||
#else
|
||||
#define C_UNDERSCORE(symbol) symbol
|
||||
#endif
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the assembler directives to annotate a new function.
|
||||
|
||||
Arguments:
|
||||
|
||||
FunctionName - Supplies the name of the function.
|
||||
|
||||
--*/
|
||||
|
||||
.macro FUNCTION_ENTRY FunctionName
|
||||
|
||||
.p2align 4
|
||||
#if defined(__APPLE__)
|
||||
.globl _\FunctionName\()
|
||||
_\FunctionName\():
|
||||
#else
|
||||
.globl \FunctionName\()
|
||||
.type \FunctionName\(),@function
|
||||
\FunctionName\():
|
||||
#endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro generates an optimization for "add reg,128" which can instead
|
||||
be encoded as "sub reg,-128" to reduce code size by using a signed 8-bit
|
||||
value.
|
||||
|
||||
Arguments:
|
||||
|
||||
Register - Supplies the register to be added to.
|
||||
|
||||
Immediate - Supplies the immediate to add to the register.
|
||||
|
||||
--*/
|
||||
|
||||
.macro add_immed Register, Immediate
|
||||
|
||||
.if (\Immediate\() != 128)
|
||||
add \Register\(),\Immediate\()
|
||||
.else
|
||||
sub \Register\(),-\Immediate\() # smaller encoding
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
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.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCountGE Count1, Value1, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
|
||||
/*++
|
||||
Macro Description:
|
||||
This macro conditionally emits the statement if Count1 is equal to Value1
|
||||
and Count2 is equal to Value2.
|
||||
Arguments:
|
||||
Count1 - Supplies the variable used in the comparison.
|
||||
Value1 - Supplies the static used in the comparison.
|
||||
Count2 - Supplies the variable used in the comparison.
|
||||
Value2 - Supplies the static used in the comparison.
|
||||
Statement - Supplies the statement to conditionally emit.
|
||||
--*/
|
||||
|
||||
.macro EmitIfCount2EQ Count1, Value1, Count2, Value2, Statement
|
||||
|
||||
.if (\Count1\() == \Value1\()) && (\Count2\() == \Value2\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro conditionally emits the statement if Count1 is greater than or
|
||||
equal to Value1 and Count2 is greater than or equal to Value2.
|
||||
|
||||
Arguments:
|
||||
|
||||
Count1 - Supplies the variable used in the comparison.
|
||||
|
||||
Value1 - Supplies the static used in the comparison.
|
||||
|
||||
Count2 - Supplies the variable used in the comparison.
|
||||
|
||||
Value2 - Supplies the static used in the comparison.
|
||||
|
||||
Statement - Supplies the statement to conditionally emit.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitIfCount2GE Count1, Value1, Count2, Value2, Statement
|
||||
|
||||
.if (\Count1\() >= \Value1\()) && (\Count2\() >= \Value2\())
|
||||
\Statement\()
|
||||
.endif
|
||||
|
||||
.endm
|
||||
|
||||
/*++
|
||||
|
||||
Macro Description:
|
||||
|
||||
This macro emits the statement for each register listed in the register
|
||||
list. The statement can use RegItem to access the current register.
|
||||
|
||||
Arguments:
|
||||
|
||||
RegList - Supplies the list of registers.
|
||||
|
||||
Statement - Supplies the statement to emit.
|
||||
|
||||
--*/
|
||||
|
||||
.macro EmitForEachRegister RegList, Statement
|
||||
|
||||
.irp RegItem, \RegList\()
|
||||
\Statement\()
|
||||
.endr
|
||||
|
||||
.endm
|
||||
Reference in New Issue
Block a user