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

Optimize ARM64EC build (#8515)

Browse source 
Add sgemm and qgemm optimized kernels for ARM64EC configuration.
This commit is contained in:
Tracy Sharpe 2021-07-27 23:46:39 -07:00 committed by GitHub
parent 1798698545
commit 539d1d44c1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
10 changed files with 1154 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
cmake/onnxruntime_mlas.cmake
View file

@ -36,24 +36,43 @@ if (onnxruntime_BUILD_WEBASSEMBLY)
)
endif()
elseif(MSVC)
if(onnxruntime_target_platform STREQUAL "ARM64")
set(mlas_platform_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/qgemm_kernel_neon.cpp
${ONNXRUNTIME_ROOT}/core/mlas/lib/qgemm_kernel_udot.cpp
)
if((onnxruntime_target_platform STREQUAL "ARM64") OR (onnxruntime_target_platform STREQUAL "ARM64EC"))
set(PREPROCESS_ARMASM_FLAGS "")
set(ARMASM_FLAGS "")
set(mlas_platform_preprocess_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/QgemmU8X8KernelNeon.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/QgemmU8X8KernelUdot.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/SgemmKernelNeon.asm
)
if(onnxruntime_target_platform STREQUAL "ARM64")
set(mlas_platform_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/qgemm_kernel_neon.cpp
${ONNXRUNTIME_ROOT}/core/mlas/lib/qgemm_kernel_udot.cpp
)
set(mlas_platform_preprocess_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/QgemmU8X8KernelNeon.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/QgemmU8X8KernelUdot.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64/SgemmKernelNeon.asm
)
else()
set(mlas_platform_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/qgemm_kernel_neon.cpp
)
set(mlas_platform_preprocess_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64ec/QgemmU8X8KernelNeon.asm
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm64ec/SgemmKernelNeon.asm
)
string(APPEND PREPROCESS_ARMASM_FLAGS " /arm64EC")
string(APPEND ARMASM_FLAGS " -machine ARM64EC")
endif()
if(CMAKE_BUILD_TYPE STREQUAL "Debug")
set(ARMASM_FLAGS "-g")
else()
set(ARMASM_FLAGS "")
string(APPEND ARMASM_FLAGS " -g")
endif()
# Remove double quotes from flag strings.
separate_arguments(PREPROCESS_ARMASM_FLAGS NATIVE_COMMAND "${PREPROCESS_ARMASM_FLAGS}")
separate_arguments(ARMASM_FLAGS NATIVE_COMMAND "${ARMASM_FLAGS}")
# Run the C precompiler on each input before the assembler.
foreach(asm_filename ${mlas_platform_preprocess_srcs})
get_filename_component(asm_filename_base ${asm_filename} NAME_WLE)
@ -62,7 +81,7 @@ elseif(MSVC)
add_custom_command(
OUTPUT ${obj_filename}
COMMAND
cl.exe /P ${asm_filename} /Fi${preprocess_filename}
cl.exe ${PREPROCESS_ARMASM_FLAGS} /P ${asm_filename} /Fi${preprocess_filename}
COMMAND
armasm64.exe ${ARMASM_FLAGS} ${preprocess_filename} ${obj_filename}
DEPENDS ${asm_filename}
@ -70,7 +89,7 @@ elseif(MSVC)
)
list(APPEND mlas_platform_srcs ${obj_filename})
endforeach()
elseif((onnxruntime_target_platform STREQUAL "ARM") OR (onnxruntime_target_platform STREQUAL "ARM64EC"))
elseif(onnxruntime_target_platform STREQUAL "ARM")
set(mlas_platform_srcs
${ONNXRUNTIME_ROOT}/core/mlas/lib/arm/sgemmc.cpp
)

4
onnxruntime/core/common/cpuid_info.cc
View file

@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#if defined(_M_IX86) || defined(_M_X64) || defined(__i386__) || defined(__x86_64__)
#if defined(_M_IX86) || (defined(_M_X64) && !defined(_M_ARM64EC)) || defined(__i386__) || defined(__x86_64__)
#define CPUIDINFO_ARCH_X86
#endif
@ -26,7 +26,7 @@
#if defined(_MSC_VER) && defined(CPUIDINFO_ARCH_ARM)
// pytorch cpu info does not work for Windows ARM
// 1. msvc report syntax error in file src/arm/api.h
// 2. features reporting micro-arch in Windows is missing
// 2. features reporting micro-arch in Windows is missing
#else
#define CPUINFO_INCLUDED

9
onnxruntime/core/mlas/inc/mlas.h
View file

@ -47,7 +47,10 @@ Abstract:
#if defined(_M_ARM64) || defined(__aarch64__)
#define MLAS_TARGET_ARM64
#endif
#if defined(_M_ARM) || defined(_M_ARM64EC) || defined(__arm__)
#if defined(_M_ARM64EC)
#define MLAS_TARGET_ARM64EC
#endif
#if defined(_M_ARM) || defined(__arm__)
#define MLAS_TARGET_ARM
#endif
#if defined(__VSX__)
@ -967,8 +970,8 @@ MlasQuantizeLinear(
/**
* @brief Requantize a block of the intermediate buffer to the output buffer,
* optionally adding the supplied bias
*
* @param Input Input matrix
*
* @param Input Input matrix
* @param InputLeadingDimension Input matrix leading dimension
* @param Output Output matrix
* @param OutputLeadingDimension Output matrix leading dimension

8
onnxruntime/core/mlas/lib/arm/sgemmc.cpp
View file

@ -55,7 +55,7 @@ Arguments:
ldc - Supplies the first dimension of matrix C.
alpha - Supplies the scaler multiplier (see SGEMM definition).
alpha - Supplies the scalar multiplier (see SGEMM definition).
Return Value:
@ -403,7 +403,7 @@ Arguments:
ldc - Supplies the first dimension of matrix C.
alpha - Supplies the scaler multiplier (see SGEMM definition).
alpha - Supplies the scalar multiplier (see SGEMM definition).
Return Value:
@ -465,7 +465,7 @@ Arguments:
ldc - Supplies the first dimension of matrix C.
alpha - Supplies the scaler multiplier (see SGEMM definition).
alpha - Supplies the scalar multiplier (see SGEMM definition).
Return Value:
@ -519,7 +519,7 @@ Arguments:
ldc - Supplies the first dimension of matrix C.
alpha - Supplies the scaler multiplier (see SGEMM definition).
alpha - Supplies the scalar multiplier (see SGEMM definition).
Return Value:

16
onnxruntime/core/mlas/lib/arm64/QgemmU8X8KernelUdot.asm
View file

@ -24,11 +24,11 @@ Abstract:
// Stack frame layout for the U8X8 kernel.
//
#define GemmU8XKernelFrame_SavedNeonRegisters (4 * 8)
#define GemmU8XKernelFrame_SavedRegisters GemmU8XKernelFrame_SavedNeonRegisters
#define GemmU8XKernelFrame_ColumnSumBuffer (0 + GemmU8XKernelFrame_SavedRegisters)
#define GemmU8XKernelFrame_ZeroPointB (8 + GemmU8XKernelFrame_SavedRegisters)
#define GemmU8XKernelFrame_ZeroMode (16 + GemmU8XKernelFrame_SavedRegisters)
#define GemmU8X8KernelFrame_SavedNeonRegisters (4 * 8)
#define GemmU8X8KernelFrame_SavedRegisters GemmU8X8KernelFrame_SavedNeonRegisters
#define GemmU8X8KernelFrame_ColumnSumBuffer (0 + GemmU8X8KernelFrame_SavedRegisters)
#define GemmU8X8KernelFrame_ZeroPointB (8 + GemmU8X8KernelFrame_SavedRegisters)
#define GemmU8X8KernelFrame_ZeroMode (16 + GemmU8X8KernelFrame_SavedRegisters)
TEXTAREA
@ -87,9 +87,9 @@ Return Value:
PROLOG_SAVE_REG_PAIR d8,d9,#-32!
PROLOG_SAVE_REG_PAIR d10,d11,#16
ldr x8,[sp,#GemmU8XKernelFrame_ColumnSumBuffer]
ldr x9,[sp,#GemmU8XKernelFrame_ZeroPointB]
ldrb w13,[sp,#GemmU8XKernelFrame_ZeroMode]
ldr x8,[sp,#GemmU8X8KernelFrame_ColumnSumBuffer]
ldr x9,[sp,#GemmU8X8KernelFrame_ZeroPointB]
ldrb w13,[sp,#GemmU8X8KernelFrame_ZeroMode]
mov x14,x0
ld1 {v11.4s},[x7]
mov x15,x3

2
onnxruntime/core/mlas/lib/arm64/SgemmKernelNeon.asm
View file

@ -435,7 +435,7 @@ $Mode.OutputRemaining1x$Rows.Block
;
; ldc (x7) - Supplies the first dimension of matrix C.
;
; Alpha (s0) - Supplies the scaler multiplier (see SGEMM definition).
; Alpha (s0) - Supplies the scalar multiplier (see SGEMM definition).
;
; Return Value:
;

629
onnxruntime/core/mlas/lib/arm64ec/QgemmU8X8KernelNeon.asm Normal file
View file

@ -0,0 +1,629 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
QgemmU8X8KernelNeon.asm
Abstract:
This module implements the kernels for the quantized integer matrix/matrix
multiply operation (QGEMM).
--*/
#include "kxarm64.h"
//
// Stack frame layout for the U8X8 kernel.
//
#define GemmU8X8KernelFrame_SavedNeonRegisters (8 * 8)
#define GemmU8X8KernelFrame_SavedRegisters GemmU8X8KernelFrame_SavedNeonRegisters
#define GemmU8X8KernelFrame_ColumnSumBuffer (0 + GemmU8X8KernelFrame_SavedRegisters)
#define GemmU8X8KernelFrame_ZeroPointB (8 + GemmU8X8KernelFrame_SavedRegisters)
#define GemmU8X8KernelFrame_ZeroMode (16 + GemmU8X8KernelFrame_SavedRegisters)
//
// Define instruction aliases not implemented by ARMASM64.
//
MACRO
uxtl $DestReg, $SrcReg
ushll $DestReg.,$SrcReg.,#0
MEND
TEXTAREA
/*++
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. The matrix data has been packed
using MlasGemmU8X8CopyPackANeon.
B (x1) - Supplies the address of matrix B. The matrix data has been packed
using MlasGemmU8X8CopyPackBNeon.
C (x2) - Supplies the address of matrix C.
PackedCountK (x3) - Supplies the number of packed columns from matrix A and
the number of packed 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.
ldc (x6) - Supplies the first dimension of matrix C.
RowSumBuffer (x7) - Supplies the sum of each row from matrix A multiplied by
the zero point offset of matrix B. These values are accumulated into every
row of matrix C.
ColumnSumBuffer - Supplies the sum of each column from matrix B multiplied
by the zero point offset of matrix A. These values are accumulated into
every column of matrix C.
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.
--*/
NESTED_ENTRY_COMDAT A64NAME(MlasGemmU8X8KernelNeon)
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
ldr x8,[sp,#GemmU8X8KernelFrame_ColumnSumBuffer]
ldr x9,[sp,#GemmU8X8KernelFrame_ZeroPointB]
ldrb w15,[sp,#GemmU8X8KernelFrame_ZeroMode]
mov x16,x0
ld1 {v7.4s},[x7] // load RowSumBuffer
mov x17,x3
cmp x4,#1 // CountM == 1?
beq ProcessNextColumnLoopM1
cmp x4,#4 // CountM < 4?
blo ProcessNextColumnLoopM2
//
// Process 4 rows of the matrices.
//
ProcessNextColumnLoopM4
ld1 {v0.8b},[x1],#8 // load packed B0
mov x0,x16 // reload matrix A
ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer0
mov x3,x17 // reload PackedCountK
ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer1
uxtl v0.8h,v0.8b
dup v9.4s,v7.s[0]
dup v11.4s,v7.s[1]
dup v13.4s,v7.s[2]
dup v15.4s,v7.s[3]
cbz x9,SkipScaleByZeroPointBM4
ld1 {v6.4s},[x9],#16 // load ZeroPointB0
mul v8.4s,v9.4s,v6.4s
mul v10.4s,v11.4s,v6.4s
mul v12.4s,v13.4s,v6.4s
mul v14.4s,v15.4s,v6.4s
ld1 {v6.4s},[x9],#16 // load ZeroPointB1
mul v9.4s,v9.4s,v6.4s
mul v11.4s,v11.4s,v6.4s
mul v13.4s,v13.4s,v6.4s
mul v15.4s,v15.4s,v6.4s
ld1 {v4.8b},[x0],#8 // load first packed A0
add v8.4s,v8.4s,v2.4s
add v9.4s,v9.4s,v3.4s
add v10.4s,v10.4s,v2.4s
add v11.4s,v11.4s,v3.4s
ld1 {v5.8b},[x0],#8 // load first packed A1
add v12.4s,v12.4s,v2.4s
add v13.4s,v13.4s,v3.4s
add v14.4s,v14.4s,v2.4s
add v15.4s,v15.4s,v3.4s
b ComputeBlockLoopM4
SkipScaleByZeroPointBM4
ld1 {v4.8b},[x0],#8 // load first packed A0
add v8.4s,v9.4s,v2.4s
add v9.4s,v9.4s,v3.4s
add v10.4s,v11.4s,v2.4s
add v11.4s,v11.4s,v3.4s
ld1 {v5.8b},[x0],#8 // load first packed A1
add v12.4s,v13.4s,v2.4s
add v13.4s,v13.4s,v3.4s
add v14.4s,v15.4s,v2.4s
add v15.4s,v15.4s,v3.4s
ComputeBlockLoopM4
uxtl v2.8h,v4.8b
uxtl v3.8h,v5.8b
ld1 {v1.8b},[x1],#8 // load packed B1
umlal v8.4s,v0.4h,v2.h[0]
umlal2 v9.4s,v0.8h,v2.h[0]
umlal v10.4s,v0.4h,v2.h[4]
umlal2 v11.4s,v0.8h,v2.h[4]
uxtl v1.8h,v1.8b
umlal v12.4s,v0.4h,v3.h[0]
umlal2 v13.4s,v0.8h,v3.h[0]
umlal v14.4s,v0.4h,v3.h[4]
umlal2 v15.4s,v0.8h,v3.h[4]
ld1 {v0.8b},[x1],#8 // load packed B2
umlal v8.4s,v1.4h,v2.h[1]
umlal2 v9.4s,v1.8h,v2.h[1]
umlal v10.4s,v1.4h,v2.h[5]
umlal2 v11.4s,v1.8h,v2.h[5]
uxtl v0.8h,v0.8b
umlal v12.4s,v1.4h,v3.h[1]
umlal2 v13.4s,v1.8h,v3.h[1]
umlal v14.4s,v1.4h,v3.h[5]
umlal2 v15.4s,v1.8h,v3.h[5]
ld1 {v1.8b},[x1],#8 // load packed B3
sub x3,x3,#1
cbz x3,ComputeBlockLoopFinishM4
umlal v8.4s,v0.4h,v2.h[2]
umlal2 v9.4s,v0.8h,v2.h[2]
umlal v10.4s,v0.4h,v2.h[6]
umlal2 v11.4s,v0.8h,v2.h[6]
uxtl v1.8h,v1.8b
ld1 {v4.8b},[x0],#8 // load next packed A0
umlal v12.4s,v0.4h,v3.h[2]
umlal2 v13.4s,v0.8h,v3.h[2]
umlal v14.4s,v0.4h,v3.h[6]
umlal2 v15.4s,v0.8h,v3.h[6]
ld1 {v0.8b},[x1],#8 // load packed B0
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
umlal v10.4s,v1.4h,v2.h[7]
umlal2 v11.4s,v1.8h,v2.h[7]
uxtl v0.8h,v0.8b
ld1 {v5.8b},[x0],#8 // load next packed A1
umlal v12.4s,v1.4h,v3.h[3]
umlal2 v13.4s,v1.8h,v3.h[3]
umlal v14.4s,v1.4h,v3.h[7]
umlal2 v15.4s,v1.8h,v3.h[7]
b ComputeBlockLoopM4
ComputeBlockLoopFinishM4
umlal v8.4s,v0.4h,v2.h[2] // finish computing tail vectors
umlal2 v9.4s,v0.8h,v2.h[2]
add x10,x2,x6,lsl #2 // compute output row 2
umlal v10.4s,v0.4h,v2.h[6]
umlal2 v11.4s,v0.8h,v2.h[6]
uxtl v1.8h,v1.8b
umlal v12.4s,v0.4h,v3.h[2]
umlal2 v13.4s,v0.8h,v3.h[2]
umlal v14.4s,v0.4h,v3.h[6]
umlal2 v15.4s,v0.8h,v3.h[6]
add x11,x10,x6,lsl #2 // compute output row 3
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
umlal v10.4s,v1.4h,v2.h[7]
umlal2 v11.4s,v1.8h,v2.h[7]
umlal v12.4s,v1.4h,v3.h[3]
umlal2 v13.4s,v1.8h,v3.h[3]
add x12,x11,x6,lsl #2 // compute output row 4
umlal v14.4s,v1.4h,v3.h[7]
umlal2 v15.4s,v1.8h,v3.h[7]
subs x5,x5,#8 // adjust CountN remaining
blo StoreOutputPartialM4
cbnz x15,SkipAccumulateOutputM4
ldp q0,q1,[x2]
ldp q2,q3,[x10]
add v8.4s,v8.4s,v0.4s
add v9.4s,v9.4s,v1.4s
ldp q0,q1,[x11]
add v10.4s,v10.4s,v2.4s
add v11.4s,v11.4s,v3.4s
ldp q2,q3,[x12]
add v12.4s,v12.4s,v0.4s
add v13.4s,v13.4s,v1.4s
add v14.4s,v14.4s,v2.4s
add v15.4s,v15.4s,v3.4s
SkipAccumulateOutputM4
stp q8,q9,[x2],#32
stp q10,q11,[x10]
stp q12,q13,[x11]
stp q14,q15,[x12]
cbnz x5,ProcessNextColumnLoopM4
ExitKernelM4
mov x0,#4 // return number of rows handled
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
//
// Store the partial 1 to 7 columns either overwriting the output matrix or
// accumulating into the existing contents of the output matrix.
//
StoreOutputPartialM4
cbz x15,StoreOutputPartialAddModeM4
StoreOutputPartialZeroModeM4
tbz x5,#2,StoreOutputPartial2ZeroModeM4
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
st1 {v10.4s},[x10],#16
mov v10.16b,v11.16b
st1 {v12.4s},[x11],#16
mov v12.16b,v13.16b
st1 {v14.4s},[x12],#16
mov v14.16b,v15.16b
StoreOutputPartial2ZeroModeM4
tbz x5,#1,StoreOutputPartial1ZeroModeM4
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
st1 {v10.2s},[x10],#8
dup v10.4s,v10.s[2]
st1 {v12.2s},[x11],#8
dup v12.4s,v12.s[2]
st1 {v14.2s},[x12],#8
dup v14.4s,v14.s[2]
StoreOutputPartial1ZeroModeM4
tbz x5,#0,ExitKernelM4
st1 {v8.s}[0],[x2]
st1 {v10.s}[0],[x10]
st1 {v12.s}[0],[x11]
st1 {v14.s}[0],[x12]
b ExitKernelM4
StoreOutputPartialAddModeM4
tbz x5,#2,StoreOutputPartial2AddModeM4
ld1 {v0.4s},[x2]
ld1 {v1.4s},[x10]
ld1 {v2.4s},[x11]
ld1 {v3.4s},[x12]
add v8.4s,v8.4s,v0.4s
add v10.4s,v10.4s,v1.4s
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
st1 {v10.4s},[x10],#16
mov v10.16b,v11.16b
add v12.4s,v12.4s,v2.4s
add v14.4s,v14.4s,v3.4s
st1 {v12.4s},[x11],#16
mov v12.16b,v13.16b
st1 {v14.4s},[x12],#16
mov v14.16b,v15.16b
StoreOutputPartial2AddModeM4
tbz x5,#1,StoreOutputPartial1AddModeM4
ld1 {v0.2s},[x2]
ld1 {v1.2s},[x10]
ld1 {v2.2s},[x11]
ld1 {v3.2s},[x12]
add v8.4s,v8.4s,v0.4s
add v10.4s,v10.4s,v1.4s
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
st1 {v10.2s},[x10],#8
dup v10.4s,v10.s[2]
add v12.4s,v12.4s,v2.4s
add v14.4s,v14.4s,v3.4s
st1 {v12.2s},[x11],#8
dup v12.4s,v12.s[2]
st1 {v14.2s},[x12],#8
dup v14.4s,v14.s[2]
StoreOutputPartial1AddModeM4
tbz x5,#0,ExitKernelM4
ld1 {v0.s}[0],[x2]
ld1 {v1.s}[0],[x10]
add v8.4s,v8.4s,v0.4s
ld1 {v2.s}[0],[x11]
add v10.4s,v10.4s,v1.4s
ld1 {v3.s}[0],[x12]
add v12.4s,v12.4s,v2.4s
st1 {v8.s}[0],[x2]
st1 {v10.s}[0],[x10]
add v14.4s,v14.4s,v3.4s
st1 {v12.s}[0],[x11]
st1 {v14.s}[0],[x12]
b ExitKernelM4
//
// Process 2 rows of the matrices.
//
ProcessNextColumnLoopM2
ld1 {v0.8b},[x1],#8 // load packed B0
mov x0,x16 // reload matrix A
ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer0
mov x3,x17 // reload PackedCountK
ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer1
uxtl v0.8h,v0.8b
dup v9.4s,v7.s[0]
dup v11.4s,v7.s[1]
cbz x9,SkipScaleByZeroPointBM2
ld1 {v14.4s},[x9],#16 // load ZeroPointB0
ld1 {v15.4s},[x9],#16 // load ZeroPointB1
mul v8.4s,v9.4s,v14.4s
mul v10.4s,v11.4s,v14.4s
mul v9.4s,v9.4s,v15.4s
mul v11.4s,v11.4s,v15.4s
ld1 {v4.8b},[x0],#8 // load first packed A0
add v8.4s,v8.4s,v2.4s
add v9.4s,v9.4s,v3.4s
add v10.4s,v10.4s,v2.4s
add v11.4s,v11.4s,v3.4s
b ComputeBlockLoopM2
SkipScaleByZeroPointBM2
ld1 {v4.8b},[x0],#8 // load first packed A0
add v8.4s,v9.4s,v2.4s
add v9.4s,v9.4s,v3.4s
add v10.4s,v11.4s,v2.4s
add v11.4s,v11.4s,v3.4s
ComputeBlockLoopM2
uxtl v2.8h,v4.8b
ld1 {v1.8b},[x1],#8 // load packed B1
umlal v8.4s,v0.4h,v2.h[0]
umlal2 v9.4s,v0.8h,v2.h[0]
umlal v10.4s,v0.4h,v2.h[4]
umlal2 v11.4s,v0.8h,v2.h[4]
uxtl v1.8h,v1.8b
ld1 {v0.8b},[x1],#8 // load packed B2
umlal v8.4s,v1.4h,v2.h[1]
umlal2 v9.4s,v1.8h,v2.h[1]
umlal v10.4s,v1.4h,v2.h[5]
umlal2 v11.4s,v1.8h,v2.h[5]
uxtl v0.8h,v0.8b
ld1 {v1.8b},[x1],#8 // load packed B3
sub x3,x3,#1
cbz x3,ComputeBlockLoopFinishM2
umlal v8.4s,v0.4h,v2.h[2]
umlal2 v9.4s,v0.8h,v2.h[2]
umlal v10.4s,v0.4h,v2.h[6]
umlal2 v11.4s,v0.8h,v2.h[6]
uxtl v1.8h,v1.8b
ld1 {v4.8b},[x0],#8 // load next packed A0
ld1 {v0.8b},[x1],#8 // load packed B0
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
umlal v10.4s,v1.4h,v2.h[7]
umlal2 v11.4s,v1.8h,v2.h[7]
uxtl v0.8h,v0.8b
b ComputeBlockLoopM2
ComputeBlockLoopFinishM2
umlal v8.4s,v0.4h,v2.h[2] // finish computing tail vectors
umlal2 v9.4s,v0.8h,v2.h[2]
add x10,x2,x6,lsl #2 // compute output row 2
umlal v10.4s,v0.4h,v2.h[6]
umlal2 v11.4s,v0.8h,v2.h[6]
uxtl v1.8h,v1.8b
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
umlal v10.4s,v1.4h,v2.h[7]
umlal2 v11.4s,v1.8h,v2.h[7]
subs x5,x5,#8 // adjust CountN remaining
blo StoreOutputPartialM2
cbnz x15,SkipAccumulateOutputM2
ldp q0,q1,[x2]
ldp q2,q3,[x10]
add v8.4s,v8.4s,v0.4s
add v9.4s,v9.4s,v1.4s
add v10.4s,v10.4s,v2.4s
add v11.4s,v11.4s,v3.4s
SkipAccumulateOutputM2
stp q8,q9,[x2],#32
stp q10,q11,[x10]
cbnz x5,ProcessNextColumnLoopM2
ExitKernelM2
mov x0,#2 // return number of rows handled
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
//
// Store the partial 1 to 7 columns either overwriting the output matrix or
// accumulating into the existing contents of the output matrix.
//
StoreOutputPartialM2
cbz x15,StoreOutputPartialAddModeM2
StoreOutputPartialZeroModeM2
tbz x5,#2,StoreOutputPartial2ZeroModeM2
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
st1 {v10.4s},[x10],#16
mov v10.16b,v11.16b
StoreOutputPartial2ZeroModeM2
tbz x5,#1,StoreOutputPartial1ZeroModeM2
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
st1 {v10.2s},[x10],#8
dup v10.4s,v10.s[2]
StoreOutputPartial1ZeroModeM2
tbz x5,#0,ExitKernelM2
st1 {v8.s}[0],[x2]
st1 {v10.s}[0],[x10]
b ExitKernelM2
StoreOutputPartialAddModeM2
tbz x5,#2,StoreOutputPartial2AddModeM2
ld1 {v0.4s},[x2]
ld1 {v1.4s},[x10]
add v8.4s,v8.4s,v0.4s
add v10.4s,v10.4s,v1.4s
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
st1 {v10.4s},[x10],#16
mov v10.16b,v11.16b
StoreOutputPartial2AddModeM2
tbz x5,#1,StoreOutputPartial1AddModeM2
ld1 {v0.2s},[x2]
ld1 {v1.2s},[x10]
add v8.4s,v8.4s,v0.4s
add v10.4s,v10.4s,v1.4s
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
st1 {v10.2s},[x10],#8
dup v10.4s,v10.s[2]
StoreOutputPartial1AddModeM2
tbz x5,#0,ExitKernelM2
ld1 {v0.s}[0],[x2]
ld1 {v1.s}[0],[x10]
add v8.4s,v8.4s,v0.4s
add v10.4s,v10.4s,v1.4s
st1 {v8.s}[0],[x2]
st1 {v10.s}[0],[x10]
b ExitKernelM2
//
// Process 1 row of the matrices.
//
ProcessNextColumnLoopM1
ld1 {v0.8b},[x1],#8 // load packed B0
mov x0,x16 // reload matrix A
ld1 {v2.4s},[x8],#16 // load ColumnSumBuffer0
mov x3,x17 // reload PackedCountK
ld1 {v3.4s},[x8],#16 // load ColumnSumBuffer1
uxtl v0.8h,v0.8b
dup v9.4s,v7.s[0]
cbz x9,SkipScaleByZeroPointBM1
ld1 {v14.4s},[x9],#16 // load ZeroPointB0
ld1 {v15.4s},[x9],#16 // load ZeroPointB1
mul v8.4s,v9.4s,v14.4s
mul v9.4s,v9.4s,v15.4s
ldr s4,[x0],#4 // load first packed A0
add v8.4s,v8.4s,v2.4s
add v9.4s,v9.4s,v3.4s
b ComputeBlockLoopM1
SkipScaleByZeroPointBM1
ldr s4,[x0],#4 // load first packed A0
add v8.4s,v9.4s,v2.4s
add v9.4s,v9.4s,v3.4s
ComputeBlockLoopM1
uxtl v2.8h,v4.8b
ld1 {v1.8b},[x1],#8 // load packed B1
umlal v8.4s,v0.4h,v2.h[0]
umlal2 v9.4s,v0.8h,v2.h[0]
uxtl v1.8h,v1.8b
ld1 {v0.8b},[x1],#8 // load packed B2
umlal v8.4s,v1.4h,v2.h[1]
umlal2 v9.4s,v1.8h,v2.h[1]
uxtl v0.8h,v0.8b
ld1 {v1.8b},[x1],#8 // load packed B3
sub x3,x3,#1
cbz x3,ComputeBlockLoopFinishM1
umlal v8.4s,v0.4h,v2.h[2]
umlal2 v9.4s,v0.8h,v2.h[2]
uxtl v1.8h,v1.8b
ldr s4,[x0],#4 // load first packed A0
ld1 {v0.8b},[x1],#8 // load packed B0
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
uxtl v0.8h,v0.8b
b ComputeBlockLoopM1
ComputeBlockLoopFinishM1
umlal v8.4s,v0.4h,v2.h[2] // finish computing tail vectors
umlal2 v9.4s,v0.8h,v2.h[2]
uxtl v1.8h,v1.8b
umlal v8.4s,v1.4h,v2.h[3]
umlal2 v9.4s,v1.8h,v2.h[3]
subs x5,x5,#8 // adjust CountN remaining
blo StoreOutputPartialM1
cbnz x15,SkipAccumulateOutputM1
ldp q0,q1,[x2]
add v8.4s,v8.4s,v0.4s
add v9.4s,v9.4s,v1.4s
SkipAccumulateOutputM1
stp q8,q9,[x2],#32
cbnz x5,ProcessNextColumnLoopM1
ExitKernelM1
mov x0,#1 // return number of rows handled
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
//
// Store the partial 1 to 7 columns either overwriting the output matrix or
// accumulating into the existing contents of the output matrix.
//
StoreOutputPartialM1
cbz x15,StoreOutputPartialAddModeM1
StoreOutputPartialZeroModeM1
tbz x5,#2,StoreOutputPartial2ZeroModeM1
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
StoreOutputPartial2ZeroModeM1
tbz x5,#1,StoreOutputPartial1ZeroModeM1
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
StoreOutputPartial1ZeroModeM1
tbz x5,#0,ExitKernelM1
st1 {v8.s}[0],[x2]
b ExitKernelM1
StoreOutputPartialAddModeM1
tbz x5,#2,StoreOutputPartial2AddModeM1
ld1 {v0.4s},[x2]
add v8.4s,v8.4s,v0.4s
st1 {v8.4s},[x2],#16
mov v8.16b,v9.16b // shift remaining elements down
StoreOutputPartial2AddModeM1
tbz x5,#1,StoreOutputPartial1AddModeM1
ld1 {v0.2s},[x2]
add v8.4s,v8.4s,v0.4s
st1 {v8.2s},[x2],#8
dup v8.4s,v8.s[2] // shift remaining elements down
StoreOutputPartial1AddModeM1
tbz x5,#0,ExitKernelM1
ld1 {v0.s}[0],[x2]
add v8.4s,v8.4s,v0.4s
st1 {v8.s}[0],[x2]
b ExitKernelM1
NESTED_END MlasGemmU8X8KernelNeon
END

466
onnxruntime/core/mlas/lib/arm64ec/SgemmKernelNeon.asm Normal file
View file

@ -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

2
onnxruntime/core/mlas/lib/mlasi.h
View file

@ -905,7 +905,7 @@ MlasConvDepthwiseFloat_CHW(
#if defined(MLAS_TARGET_ARM)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON32_INTRINSICS
#elif defined(MLAS_TARGET_ARM64)
#elif defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_ARM64EC)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON64_INTRINSICS
#elif defined(MLAS_TARGET_POWER)

6
onnxruntime/core/mlas/lib/qgemm.h
View file

@ -10,7 +10,7 @@ Module Name:
Abstract:
This module defines the set of template functions to implement a kernel of
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, there needs to specialize template functions below:
@ -676,9 +676,9 @@ MlasGemmU8X8GetDispatch(
else {
GemmU8X8Dispatch = MlasPlatform.GemmU8U8Dispatch;
}
#elif defined(MLAS_NEON64_INTRINSICS)
#elif defined(MLAS_TARGET_ARM64)
GemmU8X8Dispatch = MlasPlatform.GemmU8X8Dispatch;
#elif defined(MLAS_NEON32_INTRINSICS) && !defined(_MSC_VER)
#elif defined(MLAS_TARGET_ARM64EC) || (defined(MLAS_TARGET_ARM) && !defined(_MSC_VER))
GemmU8X8Dispatch = &MlasGemmU8X8DispatchNeon;
#else
GemmU8X8Dispatch = &MlasGemmU8X8DispatchDefault;