mirror of
https://github.com/opencv/opencv.git
synced 2026-07-21 19:33:03 +04:00
Merge pull request #28934 from abhishek-gola:mlas_gemm
Added MLAS third party module and integrated into GeMM path #28934 ### Pull Request Readiness Checklist See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request - [x] I agree to contribute to the project under Apache 2 License. - [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV - [x] The PR is proposed to the proper branch - [x] There is a reference to the original bug report and related work - [x] There is accuracy test, performance test and test data in opencv_extra repository, if applicable Patch to opencv_extra has the same branch name. - [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Abhishek Gola
GitHub
Vendored
+277
@@ -0,0 +1,277 @@
|
||||
# Vendored from microsoft/onnxruntime onnxruntime/core/mlas/. License: MIT.
|
||||
# See LICENSE and README.md in this directory for provenance.
|
||||
#
|
||||
# Vendored subset: SGEMM (sgemm.cpp + arch kernels) plus MlasFlashAttention
|
||||
# (flashattn.cpp) and the portable softmax kernels it needs from compute.cpp.
|
||||
# Non-SGEMM / non-FlashAttention dispatch lines in lib/platform.cpp are
|
||||
# `#if 0`'d out (search for "MLAS_GEMM_ONLY"). Re-vendoring upstream
|
||||
# requires re-applying that patch.
|
||||
#
|
||||
# Local patches against upstream:
|
||||
# - lib/threading.cpp -> modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp
|
||||
# (cv::parallel_for_).
|
||||
# - mlasi.h: MlasGetMaximumThreadCount() returns cv::getNumThreads().
|
||||
# - lib/core/common/{narrow,common}.h shims for ORT internals MLAS uses.
|
||||
# - lib/platform.cpp: non-SGEMM dispatch removed (`#if 0` blocks), and
|
||||
# ReduceMaximumF32Kernel / ComputeSumExpF32Kernel initialized to the
|
||||
# portable compute.cpp fallbacks.
|
||||
|
||||
set(MLAS_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/lib)
|
||||
set(MLAS_INC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/inc)
|
||||
|
||||
# Hardware-agnostic core: SGEMM dispatch + softmax/exp kernels + flash-attention.
|
||||
set(mlas_common_srcs
|
||||
${MLAS_SRC_DIR}/platform.cpp
|
||||
${CMAKE_SOURCE_DIR}/modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp
|
||||
${MLAS_SRC_DIR}/sgemm.cpp
|
||||
${MLAS_SRC_DIR}/compute.cpp
|
||||
${MLAS_SRC_DIR}/flashattn.cpp
|
||||
)
|
||||
|
||||
# Architecture detection — same as ORT's onnxruntime_mlas.cmake.
|
||||
set(MLAS_ARM FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_ARM64 FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_POWER FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_X86 FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_X86_64 FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_RISCV64 FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_LOONGARCH64 FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_S390X FALSE CACHE INTERNAL "" FORCE)
|
||||
set(MLAS_WASM FALSE CACHE INTERNAL "" FORCE)
|
||||
|
||||
if (CMAKE_SYSTEM_NAME STREQUAL "Emscripten")
|
||||
set(MLAS_WASM TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^arm64.*")
|
||||
set(MLAS_ARM64 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^aarch64.*")
|
||||
set(MLAS_ARM64 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^arm.*")
|
||||
set(MLAS_ARM TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(powerpc.*|ppc.*)")
|
||||
set(MLAS_POWER TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(i.86|x86?)$")
|
||||
set(MLAS_X86 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|amd64|AMD64)$")
|
||||
set(MLAS_X86_64 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^riscv64.*")
|
||||
set(MLAS_RISCV64 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^loongarch64.*")
|
||||
set(MLAS_LOONGARCH64 TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "^s390x$")
|
||||
set(MLAS_S390X TRUE CACHE INTERNAL "" FORCE)
|
||||
endif()
|
||||
|
||||
set(mlas_platform_srcs)
|
||||
|
||||
# State flags consumed by the top-level config summary's PNG-style chain:
|
||||
# OPENCV_DNN_MLAS_ENABLED — 1 if MLAS compiled in, 0 if skipped/unavailable
|
||||
# OPENCV_DNN_MLAS_SKIP_REASON — human-readable reason for the NO case
|
||||
# Cleared in modules/dnn/CMakeLists.txt before add_subdirectory() so stale
|
||||
# values don't survive a config where MLAS is no longer reached.
|
||||
set(OPENCV_DNN_MLAS_ENABLED 0 CACHE INTERNAL "" FORCE)
|
||||
set(OPENCV_DNN_MLAS_SKIP_REASON "" CACHE INTERNAL "" FORCE)
|
||||
|
||||
# Probe ASM language once. check_language(ASM) is a no-op when
|
||||
# CMAKE_ASM_COMPILER is already set in cache — and the Android NDK toolchain
|
||||
# pre-sets it for every ABI. So on Android the guard falls through to
|
||||
# enable_language(ASM), which then fails at generate time with
|
||||
# "CMAKE_ASM_COMPILE_OBJECT not set" on the NDK + CMake 3.22.1 combo CI uses.
|
||||
# Skip MLAS on Android for the ASM-using arches; the DNN module falls back to
|
||||
# its built-in SGEMM. Android armv7a stays enabled via the C++-only
|
||||
# sgemmc.cpp path under MLAS_ARM below.
|
||||
set(_MLAS_REQUIRES_ASM FALSE)
|
||||
if(MLAS_X86_64 OR MLAS_X86 OR MLAS_ARM64 OR MLAS_LOONGARCH64)
|
||||
set(_MLAS_REQUIRES_ASM TRUE)
|
||||
endif()
|
||||
if(ANDROID AND _MLAS_REQUIRES_ASM)
|
||||
set(OPENCV_DNN_MLAS_SKIP_REASON
|
||||
"skipped on Android ${CMAKE_ANDROID_ARCH_ABI}; enable_language(ASM) from subdir is unreliable"
|
||||
CACHE INTERNAL "" FORCE)
|
||||
message(STATUS "MLAS: skipping on Android ${CMAKE_ANDROID_ARCH_ABI} "
|
||||
"(enable_language(ASM) from subdir is unreliable here); "
|
||||
"DNN will use its built-in SGEMM")
|
||||
return()
|
||||
endif()
|
||||
include(CheckLanguage)
|
||||
set(MLAS_HAS_ASM FALSE CACHE INTERNAL "" FORCE)
|
||||
check_language(ASM)
|
||||
if(CMAKE_ASM_COMPILER)
|
||||
enable_language(ASM)
|
||||
set(MLAS_HAS_ASM TRUE CACHE INTERNAL "" FORCE)
|
||||
elseif(_MLAS_REQUIRES_ASM)
|
||||
set(OPENCV_DNN_MLAS_SKIP_REASON
|
||||
"no ASM compiler available for ${CMAKE_SYSTEM_PROCESSOR}"
|
||||
CACHE INTERNAL "" FORCE)
|
||||
message(WARNING "MLAS: ASM language unavailable on ${CMAKE_SYSTEM_PROCESSOR}; "
|
||||
"MLAS disabled (DNN will use its built-in SGEMM)")
|
||||
return()
|
||||
endif()
|
||||
|
||||
# x86_64 SGEMM kernels: MlasGemmFloatKernelSse / Avx / Fma3 / Avx512F
|
||||
# plus the M=1 fast paths and B-packing helpers.
|
||||
if(MLAS_X86_64)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelSse2.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmTransposePackB16x4Sse2.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelM1Avx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelM1TransposeBAvx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmTransposePackB16x4Avx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelFma3.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx512F.S
|
||||
)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86_64/SgemmKernelSse2.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmTransposePackB16x4Sse2.S
|
||||
PROPERTIES COMPILE_FLAGS "-msse2")
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelM1Avx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmKernelM1TransposeBAvx.S
|
||||
${MLAS_SRC_DIR}/x86_64/SgemmTransposePackB16x4Avx.S
|
||||
PROPERTIES COMPILE_FLAGS "-mavx")
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86_64/SgemmKernelFma3.S
|
||||
PROPERTIES COMPILE_FLAGS "-mavx2 -mfma -mf16c")
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86_64/SgemmKernelAvx512F.S
|
||||
PROPERTIES COMPILE_FLAGS "-mavx512f")
|
||||
endif()
|
||||
|
||||
# i386 / 32-bit x86 SGEMM kernels.
|
||||
if(MLAS_X86)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/x86/SgemmKernelSse2.S
|
||||
${MLAS_SRC_DIR}/x86/SgemmKernelAvx.S
|
||||
)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86/SgemmKernelSse2.S PROPERTIES COMPILE_FLAGS "-msse2")
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/x86/SgemmKernelAvx.S PROPERTIES COMPILE_FLAGS "-mavx")
|
||||
if(ANDROID)
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/x86/x86.get_pc_thunk.S)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# ARM 32-bit: pure C++ sgemmc (no .S kernels).
|
||||
if(MLAS_ARM)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mfpu=neon")
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/arm/sgemmc.cpp)
|
||||
endif()
|
||||
|
||||
# ARM64 / AArch64.
|
||||
if(MLAS_ARM64)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/aarch64/SgemmKernelNeon.S
|
||||
${MLAS_SRC_DIR}/aarch64/SgemvKernelNeon.S
|
||||
)
|
||||
endif()
|
||||
|
||||
# POWER (ppc64le / AIX). Always compile the base SgemmKernelPower; opt-in
|
||||
# POWER10 if the compiler supports -mcpu=power10.
|
||||
if(MLAS_POWER)
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/power/SgemmKernelPower.cpp)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/power/SgemmKernelPower.cpp
|
||||
PROPERTIES COMPILE_FLAGS "-DSINGLE")
|
||||
include(CheckCXXCompilerFlag)
|
||||
check_cxx_compiler_flag("-mcpu=power10" MLAS_HAS_POWER10)
|
||||
set(MLAS_HAS_POWER10 ${MLAS_HAS_POWER10} CACHE INTERNAL "" FORCE)
|
||||
if(MLAS_HAS_POWER10 AND MLAS_HAS_ASM)
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/power/SgemmKernelPOWER10.cpp)
|
||||
if(NOT AIX)
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/power/SgemmKernelPackA.S)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/power/SgemmKernelPackA.S
|
||||
PROPERTIES COMPILE_FLAGS "-O2 -mcpu=power10")
|
||||
endif()
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/power/SgemmKernelPOWER10.cpp
|
||||
PROPERTIES COMPILE_FLAGS "-O2 -mcpu=power10 -DSINGLE")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# LoongArch 64 (LSX + LASX).
|
||||
if(MLAS_LOONGARCH64)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/loongarch64/SgemmKernelLsx.S
|
||||
${MLAS_SRC_DIR}/loongarch64/SgemmTransposePackB16x4LSX.S
|
||||
${MLAS_SRC_DIR}/loongarch64/SgemmKernelLasx.S
|
||||
${MLAS_SRC_DIR}/loongarch64/SgemmTransposePackB16x4Lasx.S
|
||||
)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mlsx -mlasx")
|
||||
endif()
|
||||
|
||||
# IBM s390x (z/Architecture, ZVECTOR).
|
||||
if(MLAS_S390X)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/s390x/SgemmKernel.cpp
|
||||
${MLAS_SRC_DIR}/s390x/SgemmKernelZVECTOR.cpp
|
||||
)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/s390x/SgemmKernel.cpp PROPERTIES COMPILE_FLAGS "-DSINGLE")
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/s390x/SgemmKernelZVECTOR.cpp PROPERTIES COMPILE_FLAGS "-DSINGLE")
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mvx -mzvector -march=z15")
|
||||
endif()
|
||||
|
||||
# RISC-V 64. RVV vector kernels iff the compiler supports rv64gcv.
|
||||
if(MLAS_RISCV64)
|
||||
include(CheckCXXSourceCompiles)
|
||||
set(_old "${CMAKE_REQUIRED_FLAGS}")
|
||||
set(CMAKE_REQUIRED_FLAGS "-march=rv64gcv -mabi=lp64d")
|
||||
check_cxx_source_compiles("
|
||||
#include <stddef.h>
|
||||
#include <riscv_vector.h>
|
||||
int main() { size_t vl = __riscv_vsetvl_e32m1(4); return static_cast<int>(vl == 0); }"
|
||||
MLAS_HAS_RISCV64_RVV)
|
||||
set(CMAKE_REQUIRED_FLAGS "${_old}")
|
||||
set(MLAS_HAS_RISCV64_RVV ${MLAS_HAS_RISCV64_RVV} CACHE INTERNAL "" FORCE)
|
||||
if(MLAS_HAS_RISCV64_RVV)
|
||||
list(APPEND mlas_platform_srcs
|
||||
${MLAS_SRC_DIR}/riscv64/sgemm_pack_b_rvv.cpp
|
||||
${MLAS_SRC_DIR}/riscv64/sgemm_kernel_rvv.cpp
|
||||
)
|
||||
foreach(f
|
||||
${MLAS_SRC_DIR}/riscv64/sgemm_pack_b_rvv.cpp
|
||||
${MLAS_SRC_DIR}/riscv64/sgemm_kernel_rvv.cpp)
|
||||
set_source_files_properties(${f} PROPERTIES COMPILE_FLAGS "-march=rv64gcv -mabi=lp64d")
|
||||
endforeach()
|
||||
endif()
|
||||
# Scalar MlasSgemmKernelZero/Add fallback for non-RVV runtime/build.
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/scalar/SgemmKernelScalar.cpp)
|
||||
endif()
|
||||
|
||||
# WASM / unknown archs use the scalar sgemm path; link its kernel.
|
||||
if(MLAS_WASM OR NOT (MLAS_X86 OR MLAS_X86_64 OR MLAS_ARM OR MLAS_ARM64
|
||||
OR MLAS_POWER OR MLAS_LOONGARCH64 OR MLAS_S390X
|
||||
OR MLAS_RISCV64))
|
||||
list(APPEND mlas_platform_srcs ${MLAS_SRC_DIR}/scalar/SgemmKernelScalar.cpp)
|
||||
endif()
|
||||
|
||||
add_library(opencv_dnn_mlas OBJECT ${mlas_common_srcs} ${mlas_platform_srcs})
|
||||
set_target_properties(opencv_dnn_mlas PROPERTIES POSITION_INDEPENDENT_CODE ON)
|
||||
target_include_directories(opencv_dnn_mlas PUBLIC ${MLAS_INC_DIR})
|
||||
target_include_directories(opencv_dnn_mlas PRIVATE ${MLAS_SRC_DIR})
|
||||
target_include_directories(opencv_dnn_mlas PRIVATE ${CMAKE_SOURCE_DIR}/modules/core/include)
|
||||
|
||||
target_compile_definitions(opencv_dnn_mlas PRIVATE
|
||||
BUILD_MLAS_NO_ONNXRUNTIME=1
|
||||
MLAS_OPENCV_THREADING=1
|
||||
MLAS_GEMM_ONLY=1
|
||||
)
|
||||
|
||||
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
|
||||
target_compile_options(opencv_dnn_mlas PRIVATE
|
||||
"$<$<COMPILE_LANGUAGE:CXX>:-w>"
|
||||
"$<$<COMPILE_LANGUAGE:CXX>:-include>"
|
||||
"$<$<COMPILE_LANGUAGE:CXX>:cstring>"
|
||||
)
|
||||
# MLAS .S files lack .note.GNU-stack; tell the assembler the stack is non-exec.
|
||||
target_compile_options(opencv_dnn_mlas PRIVATE
|
||||
"$<$<COMPILE_LANGUAGE:ASM>:-Wa,--noexecstack>"
|
||||
)
|
||||
endif()
|
||||
# platform.cpp's __linux__ branch calls syscall() but upstream only includes
|
||||
# <sys/syscall.h>, not <unistd.h> (where syscall() is declared on glibc).
|
||||
# Force-include unistd.h on non-Windows x86_64 builds. MSVC doesn't ship
|
||||
# unistd.h and doesn't recognize the -include flag; its _WIN32 branch in
|
||||
# platform.cpp doesn't need syscall() anyway.
|
||||
if(MLAS_X86_64 AND NOT WIN32)
|
||||
set_source_files_properties(${MLAS_SRC_DIR}/platform.cpp
|
||||
PROPERTIES COMPILE_FLAGS "-include unistd.h")
|
||||
endif()
|
||||
|
||||
set(HAVE_MLAS 1 PARENT_SCOPE)
|
||||
set(MLAS_OBJECTS $<TARGET_OBJECTS:opencv_dnn_mlas> PARENT_SCOPE)
|
||||
set(MLAS_INCLUDE_DIRS ${MLAS_INC_DIR} PARENT_SCOPE)
|
||||
set(OPENCV_DNN_MLAS_ENABLED 1 CACHE INTERNAL "" FORCE)
|
||||
Vendored
+21
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) Microsoft Corporation
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
Vendored
+116
@@ -0,0 +1,116 @@
|
||||
# MLAS (Microsoft Linear Algebra Subprograms)
|
||||
|
||||
MLAS is a compute library containing processor-optimized GEMM kernels and
|
||||
platform-specific threading code. It is the default math kernel library used
|
||||
internally by ONNX Runtime.
|
||||
|
||||
## Provenance
|
||||
|
||||
- **Upstream**: https://github.com/microsoft/onnxruntime
|
||||
- **Source path**: `onnxruntime/core/mlas/`
|
||||
- **Imported**: 2026-05-04
|
||||
- **Upstream commit**: [`62f742f1aa0c3102745ed35e3d869eaee845b9ac`](https://github.com/microsoft/onnxruntime/tree/62f742f1aa0c3102745ed35e3d869eaee845b9ac/onnxruntime/core/mlas)
|
||||
(2026-04-30, last MLAS-touching commit on `main` at import time;
|
||||
released as part of ORT [v1.26.0](https://github.com/microsoft/onnxruntime/releases/tag/v1.26.0))
|
||||
- **License**: MIT (see [LICENSE](LICENSE))
|
||||
|
||||
## What is vendored
|
||||
|
||||
The SGEMM (single-precision GEMM) subset of MLAS plus `MlasFlashAttention`
|
||||
(fused multi-head attention). The rest of MLAS (quantized GEMM, conv,
|
||||
FP16-dispatch SoftMax, etc.) is excluded so the OpenCV DNN module gets the
|
||||
fast SGEMM and FlashAttention paths without dragging in the full library.
|
||||
|
||||
Source files imported verbatim from upstream:
|
||||
|
||||
- `lib/sgemm.cpp` — SGEMM dispatch and host-side glue.
|
||||
- `lib/compute.cpp` — softmax / exp / row-max / sum-exp kernels. Only the
|
||||
portable C++ fallbacks for `MlasReduceMaximumF32Kernel` and
|
||||
`MlasComputeSumExpF32Kernel` are exercised; no per-arch `.S` softmax
|
||||
kernels are vendored. The file is imported whole (FP16 / GQA template
|
||||
specializations compile but never run — `SoftmaxDispatch` stays nullptr).
|
||||
- `lib/flashattn.cpp` — the `MlasFlashAttention` / `MlasFlashAttentionThreaded`
|
||||
entry points. Depends on `MlasSgemmOperation` (in `sgemm.cpp`) and the two
|
||||
portable kernels above.
|
||||
- `lib/softmax.h` — header included by `compute.cpp`; pure FP16-dispatch
|
||||
typedefs, harmless under FP32-only builds.
|
||||
- Per-arch SGEMM kernels under `lib/<arch>/`.
|
||||
|
||||
Top-level layout:
|
||||
|
||||
- `inc/` — public MLAS headers (kept verbatim from upstream).
|
||||
- `lib/` — implementation, kept verbatim from upstream except for the local
|
||||
patches listed below. Per-architecture kernels live in subdirectories
|
||||
(`x86_64/`, `aarch64/`, `arm/`, `power/`, `riscv64/`, `loongarch64/`,
|
||||
`s390x/`, `sve/`, `kleidiai/`).
|
||||
- `CMakeLists.txt` — OpenCV-side build glue. Builds an OBJECT library
|
||||
(`opencv_dnn_mlas`) whose objects are linked directly into `opencv_dnn`.
|
||||
- `threading_opencv.cpp` — OpenCV-side replacement for `lib/threading.cpp`
|
||||
(see "Local patches" below). Carries the OpenCV license header.
|
||||
|
||||
## Copyright
|
||||
|
||||
Most files are © Microsoft Corporation and licensed MIT. Some upstream
|
||||
contributions in `lib/` carry additional MIT-licensed copyrights — they are
|
||||
preserved verbatim in the file headers:
|
||||
|
||||
- `lib/kleidiai/mlasi_kleidiai.h` — © Arm Limited 2025.
|
||||
- `lib/erf_neon_fp16.{h,cpp}`, `lib/gelu_neon_fp16.{h,cpp}` — © FUJITSU
|
||||
LIMITED 2025 (jointly with Microsoft).
|
||||
|
||||
The OpenCV-authored files in this directory (`CMakeLists.txt`,
|
||||
`threading_opencv.cpp`, this `README.md`) are licensed under OpenCV's
|
||||
top-level license (Apache 2.0).
|
||||
|
||||
## Local patches against upstream
|
||||
|
||||
These deviate from a clean upstream import and must be re-applied on every
|
||||
re-vendor. The unified-diff form of each in-tree edit lives under
|
||||
[`patches/`](patches/) (same convention as
|
||||
[`3rdparty/zlib/patches/`](../zlib/patches)); re-apply with
|
||||
`git apply --directory=3rdparty/mlas patches/*.diff` after re-importing.
|
||||
|
||||
1. `lib/threading.cpp` is dropped (not vendored). Its three threading entry
|
||||
points (`MlasExecuteThreaded`, `MlasTrySimpleParallel`,
|
||||
`MlasTryBatchParallel`) are reimplemented in
|
||||
`modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp` on top of
|
||||
`cv::parallel_for_`. No `.diff` — the file is simply absent.
|
||||
2. `lib/mlasi.h` — `MlasGetMaximumThreadCount()` returns `cv::getNumThreads()`
|
||||
when `MLAS_OPENCV_THREADING` is defined; `#include "core/mlas/inc/mlas.h"`
|
||||
is rewritten to `#include "../inc/mlas.h"` because the ORT in-tree path
|
||||
does not exist here. See `patches/0001-mlasi-opencv-threading.diff`.
|
||||
3. `lib/platform.cpp` — non-SGEMM dispatch is wrapped under `MLAS_GEMM_ONLY`
|
||||
so the SGEMM-only subset builds without the rest of the MLAS sources. The
|
||||
top-of-file `erf_neon_fp16.h` / `gelu_neon_fp16.h` includes are also gated
|
||||
on `!defined(MLAS_GEMM_ONLY)` because those headers transitively pull in
|
||||
non-vendored FP16 sources (`fp16_common.h`, `softmax_kernel_neon.h`).
|
||||
The `MLAS_GEMM_ONLY` ctor also assigns `ReduceMaximumF32Kernel` and
|
||||
`ComputeSumExpF32Kernel` to the portable `compute.cpp` fallbacks so
|
||||
`MlasFlashAttention` works without per-arch softmax kernels. See
|
||||
`patches/0002-platform-gemm-only.diff`.
|
||||
4. `inc/mlas.h` — guard `_MSC_VER` with `defined()` so `-Wundef` builds
|
||||
under GCC/Clang don't warn. See `patches/0003-mlas-h-msc-ver-guard.diff`.
|
||||
5. `lib/core/common/{narrow,common}.h` — minimal shims for ORT internals
|
||||
that MLAS calls (not present upstream as MLAS sources, only as ORT
|
||||
includes). These are new files, not edits — no `.diff` needed.
|
||||
|
||||
## Build flags
|
||||
|
||||
- `HAVE_MLAS` is set by this directory's `CMakeLists.txt` when the host
|
||||
arch/OS is wired up.
|
||||
- `BUILD_MLAS_NO_ONNXRUNTIME=1`, `MLAS_OPENCV_THREADING=1`,
|
||||
`MLAS_GEMM_ONLY=1` are set as private compile definitions on the OBJECT
|
||||
library.
|
||||
|
||||
## Caller in OpenCV
|
||||
|
||||
The thin wrapper that dispatches OpenCV GEMMs to MLAS lives at
|
||||
[modules/dnn/src/layers/cpu_kernels/mlas_gemm.{hpp,cpp}](../../modules/dnn/src/layers/cpu_kernels/).
|
||||
It only includes `mlas.h` (the public header) and falls back to the existing
|
||||
fast_gemm path when MLAS is unavailable or the requested shape is unsupported.
|
||||
|
||||
## Upstream unit tests
|
||||
|
||||
Unit tests for the SGEMM kernels live in upstream ONNX Runtime under
|
||||
`onnxruntime/test/mlas`. They are not vendored here; OpenCV's own DNN tests
|
||||
exercise the integration.
|
||||
Vendored
+2311
File diff suppressed because it is too large
Load Diff
Vendored
+115
@@ -0,0 +1,115 @@
|
||||
/*++
|
||||
|
||||
Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
|
||||
Licensed under the MIT License.
|
||||
|
||||
Module Name:
|
||||
|
||||
mlas_float16.h
|
||||
|
||||
Abstract:
|
||||
|
||||
Utilities for half precision floating type conversions. Used internally
|
||||
by MLAS on platforms without half precision support. Provided here as
|
||||
convenience for tests or other client libraries/apps.
|
||||
|
||||
--*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
|
||||
|
||||
using _mlas_fp16_ = uint16_t;
|
||||
|
||||
union fp32_bits {
|
||||
uint32_t u;
|
||||
float f;
|
||||
};
|
||||
|
||||
#if defined(_MSC_VER) && !defined(__clang__)
|
||||
#pragma warning(push)
|
||||
|
||||
/*PreFast told us to convert them to constexpr but the compiler says we can't.*/
|
||||
#pragma warning(disable : 26497)
|
||||
|
||||
/*Added whole bunch of casts, still can't get rid of these overflow warnings.*/
|
||||
#pragma warning(disable : 26450)
|
||||
#pragma warning(disable : 26451)
|
||||
#endif
|
||||
|
||||
inline
|
||||
_mlas_fp16_
|
||||
MLAS_Float2Half(float ff)
|
||||
{
|
||||
constexpr fp32_bits f32infty = {255 << 23};
|
||||
constexpr fp32_bits f16max = {(127 + 16) << 23};
|
||||
constexpr fp32_bits denorm_magic = {((127 - 15) + (23 - 10) + 1) << 23};
|
||||
constexpr uint32_t sign_mask = 0x80000000u;
|
||||
|
||||
auto val = static_cast<uint16_t>(0x0u);
|
||||
fp32_bits f;
|
||||
f.f = ff;
|
||||
|
||||
uint32_t sign = f.u & sign_mask;
|
||||
f.u ^= sign;
|
||||
|
||||
if (f.u >= f16max.u) {
|
||||
// Inf or NaN (all exponent bits set)
|
||||
val = (f.u > f32infty.u) ? 0x7e00 : 0x7c00; // NaN->qNaN and Inf->Inf
|
||||
} else {
|
||||
if (f.u < (113 << 23)) {
|
||||
// Subnormal or zero
|
||||
// use a magic value to align our 10 mantissa bits at the bottom of
|
||||
// the float. as long as FP addition is round-to-nearest-even this
|
||||
// just works.
|
||||
f.f += denorm_magic.f;
|
||||
|
||||
// and one integer subtract of the bias later, we have our final float!
|
||||
val = static_cast<uint16_t>(f.u - denorm_magic.u);
|
||||
} else {
|
||||
uint32_t mant_odd = (f.u >> 13) & 1; // resulting mantissa is odd
|
||||
|
||||
// update exponent, rounding bias part 1
|
||||
f.u += ((uint32_t)(15 - 127) << 23) + 0xfff;
|
||||
// rounding bias part 2
|
||||
f.u += mant_odd;
|
||||
// take the bits!
|
||||
val = static_cast<uint16_t>(f.u >> 13);
|
||||
}
|
||||
}
|
||||
|
||||
val |= static_cast<uint16_t>(sign >> 16);
|
||||
return val;
|
||||
}
|
||||
|
||||
inline
|
||||
float
|
||||
MLAS_Half2Float(_mlas_fp16_ val)
|
||||
{
|
||||
constexpr fp32_bits magic = {113 << 23};
|
||||
constexpr uint32_t shifted_exp = 0x7c00 << 13; // exponent mask after shift
|
||||
fp32_bits o;
|
||||
|
||||
o.u = (val & 0x7fff) << 13; // exponent/mantissa bits
|
||||
uint32_t exp = shifted_exp & o.u; // just the exponent
|
||||
o.u += (127 - 15) << 23; // exponent adjust
|
||||
|
||||
// handle exponent special cases
|
||||
if (exp == shifted_exp) { // Inf/NaN?
|
||||
o.u += (128 - 16) << 23; // extra exp adjust
|
||||
} else if (exp == 0) { // Zero/Denormal?
|
||||
o.u += 1 << 23; // extra exp adjust
|
||||
o.f -= magic.f; // renormalize
|
||||
}
|
||||
|
||||
o.u |= (val & 0x8000) << 16; // sign bit
|
||||
return o.f;
|
||||
}
|
||||
|
||||
#if defined(_MSC_VER) && !defined(__clang__)
|
||||
#pragma warning(pop)
|
||||
#endif
|
||||
+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
+1094
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
|
||||
|
||||
+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
|
||||
@@ -0,0 +1,46 @@
|
||||
Patch against upstream onnxruntime/core/mlas/lib/mlasi.h
|
||||
Base commit: 62f742f1aa0c3102745ed35e3d869eaee845b9ac (ORT v1.26.0)
|
||||
|
||||
Reroute the public-header include to the relative vendored path,
|
||||
and make MlasGetMaximumThreadCount() return cv::getNumThreads()
|
||||
when MLAS_OPENCV_THREADING is defined (so MLAS partitions according
|
||||
to OpenCV's thread budget, not ORT's).
|
||||
|
||||
--- a/3rdparty/mlas/lib/mlasi.h
|
||||
+++ b/3rdparty/mlas/lib/mlasi.h
|
||||
@@ -35,7 +35,9 @@
|
||||
#endif
|
||||
#endif // MLAS_NO_EXCEPTION
|
||||
|
||||
-#include "core/mlas/inc/mlas.h"
|
||||
+// Vendored under 3rdparty/mlas/. The ORT path "core/mlas/inc/mlas.h" only
|
||||
+// works when MLAS is part of the ORT source tree.
|
||||
+#include "../inc/mlas.h"
|
||||
|
||||
#if defined(_WIN32)
|
||||
#ifndef WIN32_LEAN_AND_MEAN
|
||||
@@ -1675,13 +1677,23 @@
|
||||
);
|
||||
|
||||
|
||||
+#if defined(MLAS_OPENCV_THREADING)
|
||||
+// Defined in 3rdparty/mlas/threading_opencv.cpp. Returns
|
||||
+// cv::getNumThreads(). Hidden behind a free function so this header doesn't
|
||||
+// need to pull <opencv2/core/utility.hpp> into every MLAS translation unit.
|
||||
+extern "C" int opencv_dnn_mlas_max_threads();
|
||||
+#endif
|
||||
+
|
||||
inline
|
||||
ptrdiff_t
|
||||
MlasGetMaximumThreadCount(
|
||||
MLAS_THREADPOOL* ThreadPool
|
||||
)
|
||||
{
|
||||
-#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
|
||||
+#if defined(MLAS_OPENCV_THREADING)
|
||||
+ MLAS_UNREFERENCED_PARAMETER(ThreadPool);
|
||||
+ return static_cast<ptrdiff_t>(opencv_dnn_mlas_max_threads());
|
||||
+#elif defined(BUILD_MLAS_NO_ONNXRUNTIME)
|
||||
MLAS_UNREFERENCED_PARAMETER(ThreadPool);
|
||||
return 1;
|
||||
#else
|
||||
+174
@@ -0,0 +1,174 @@
|
||||
Patch against upstream onnxruntime/core/mlas/lib/platform.cpp
|
||||
Base commit: 62f742f1aa0c3102745ed35e3d869eaee845b9ac (ORT v1.26.0)
|
||||
|
||||
Gate non-SGEMM dispatch behind MLAS_GEMM_ONLY so the SGEMM-only
|
||||
subset can build without the rest of the MLAS sources (quantized
|
||||
GEMM, conv, FP16 SoftMax, etc.). The original ctor is preserved
|
||||
verbatim in the #else branch for clean re-vendoring.
|
||||
|
||||
Also gate the top-of-file erf_neon_fp16.h / gelu_neon_fp16.h
|
||||
includes on !MLAS_GEMM_ONLY — they transitively pull in
|
||||
fp16_common.h / softmax_kernel_neon.h, which we don't vendor.
|
||||
|
||||
The MLAS_GEMM_ONLY ctor additionally assigns ReduceMaximumF32Kernel
|
||||
and ComputeSumExpF32Kernel to the portable compute.cpp fallbacks so
|
||||
MlasFlashAttention works without per-arch softmax kernels.
|
||||
|
||||
--- a/3rdparty/mlas/lib/platform.cpp
|
||||
+++ b/3rdparty/mlas/lib/platform.cpp
|
||||
@@ -19,7 +19,7 @@
|
||||
#ifdef MLAS_USE_SVE
|
||||
#include "sve/mlasi_sve.h"
|
||||
#endif
|
||||
-#if defined(MLAS_NEON_INTRINSICS) && defined(MLAS_F16VEC_INTRINSICS_SUPPORTED)
|
||||
+#if defined(MLAS_NEON_INTRINSICS) && defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && !defined(MLAS_GEMM_ONLY)
|
||||
#include "erf_neon_fp16.h"
|
||||
#include "gelu_neon_fp16.h"
|
||||
#endif
|
||||
@@ -288,6 +288,138 @@
|
||||
};
|
||||
|
||||
#endif
|
||||
+
|
||||
+// =============================================================================
|
||||
+// SGEMM-only constructor (vendor-local patch).
|
||||
+//
|
||||
+// When MLAS_GEMM_ONLY is defined, replace the original platform-init ctor
|
||||
+// with a stripped-down version that only assigns the four (-ish) dispatch
|
||||
+// fields read by sgemm.cpp:
|
||||
+// - GemmFloatKernel
|
||||
+// - KernelM1Routine (x86_64 only)
|
||||
+// - KernelM1TransposeBRoutine (x86_64 only)
|
||||
+// - TransposePackB16x4Routine (x86_64 / loongarch only)
|
||||
+// Plus, on the SBGemm aarch64+linux path, the SBGemm batch overrides — but
|
||||
+// those are nullptr-default and we don't enable SBGemm here.
|
||||
+//
|
||||
+// Also initializes the two softmax kernel pointers consumed by
|
||||
+// flashattn.cpp (ReduceMaximumF32Kernel, ComputeSumExpF32Kernel) to the
|
||||
+// portable fallbacks provided by compute.cpp. No SIMD-asm softmax kernels
|
||||
+// are vendored — the flash-attention path uses the portable C++ rowmax /
|
||||
+// sum-exp implementations.
|
||||
+//
|
||||
+// Every other dispatch field stays at its in-class default (most are
|
||||
+// `= nullptr`). Calling any non-SGEMM / non-FlashAttention MLAS API in this
|
||||
+// build is undefined.
|
||||
+//
|
||||
+// The original full ORT ctor is preserved unchanged below the #else for
|
||||
+// future re-vendoring — drop MLAS_GEMM_ONLY to use it.
|
||||
+// =============================================================================
|
||||
+#ifdef MLAS_GEMM_ONLY
|
||||
+MLAS_PLATFORM::MLAS_PLATFORM(void)
|
||||
+{
|
||||
+ // Portable softmax kernels (compute.cpp). flashattn.cpp dereferences these
|
||||
+ // function pointers on the AMD64 / LARCH64 path; compute.cpp's
|
||||
+ // MlasComputeSoftmax does the same on AMD64 / LARCH64 / SVE / RISCV64.
|
||||
+ // Other paths call the symbols directly. Gates mirror the MLAS_PLATFORM
|
||||
+ // member visibility in mlasi.h so we initialize the field wherever it
|
||||
+ // exists — leaving it null would crash any future code that reads it via
|
||||
+ // the struct on those targets.
|
||||
+#if defined(MLAS_TARGET_LARCH64) || defined(MLAS_USE_SVE) || \
|
||||
+ defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64)
|
||||
+ this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
|
||||
+#endif
|
||||
+#if defined(MLAS_USE_SVE) || defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64)
|
||||
+ this->ComputeSumExpF32Kernel = MlasComputeSumExpF32Kernel;
|
||||
+#endif
|
||||
+
|
||||
+ // The PreferredBufferAlignment field only exists on AMD64 (see
|
||||
+ // MLAS_PLATFORM in mlasi.h). On other targets MlasGetPreferredBufferAlignment()
|
||||
+ // returns MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT directly without
|
||||
+ // consulting the struct.
|
||||
+#if defined(MLAS_TARGET_AMD64)
|
||||
+ this->PreferredBufferAlignment = MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT;
|
||||
+#endif
|
||||
+
|
||||
+#if defined(MLAS_TARGET_AMD64_IX86)
|
||||
+ // SSE2 baseline (every x86 since 2003).
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelSse;
|
||||
+#if defined(MLAS_TARGET_AMD64)
|
||||
+ this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Sse;
|
||||
+#endif
|
||||
+
|
||||
+ unsigned Cpuid1[4];
|
||||
+#if defined(_WIN32)
|
||||
+ __cpuid((int*)Cpuid1, 1);
|
||||
+#else
|
||||
+ __cpuid(1, Cpuid1[0], Cpuid1[1], Cpuid1[2], Cpuid1[3]);
|
||||
+#endif
|
||||
+ // AVX + OSXSAVE bits (matches the original ctor's checks).
|
||||
+ if ((Cpuid1[2] & 0x18000000) == 0x18000000) {
|
||||
+ uint64_t xcr0 = MlasReadExtendedControlRegister(_XCR_XFEATURE_ENABLED_MASK);
|
||||
+ if ((xcr0 & 0x6) == 0x6) {
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelAvx;
|
||||
+#if defined(MLAS_TARGET_AMD64)
|
||||
+ this->KernelM1Routine = MlasSgemmKernelM1Avx;
|
||||
+ this->KernelM1TransposeBRoutine = MlasSgemmKernelM1TransposeBAvx;
|
||||
+ this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Avx;
|
||||
+#endif
|
||||
+ unsigned Cpuid7[4];
|
||||
+#if defined(_WIN32)
|
||||
+ __cpuidex((int*)Cpuid7, 7, 0);
|
||||
+#else
|
||||
+ __cpuid_count(7, 0, Cpuid7[0], Cpuid7[1], Cpuid7[2], Cpuid7[3]);
|
||||
+#endif
|
||||
+ // AVX2 + FMA3.
|
||||
+ if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) {
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelFma3;
|
||||
+ // AVX-512F + ZMM-state save.
|
||||
+ if (((Cpuid7[1] & 0x10000) != 0) && ((xcr0 & 0xE0) == 0xE0)) {
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelAvx512F;
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+ }
|
||||
+#endif // MLAS_TARGET_AMD64_IX86
|
||||
+
|
||||
+#if defined(MLAS_TARGET_POWER)
|
||||
+ // Default to the base SgemmKernelPower; the POWER10 detection branch in
|
||||
+ // the original ctor is omitted because the POWER10 SgemmKernel symbol
|
||||
+ // (MlasSgemmKernelPOWER10) is only present when -mcpu=power10 was
|
||||
+ // detectable at configure time. CMake conditionally compiles it; the
|
||||
+ // base kernel is always available.
|
||||
+ this->GemmFloatKernel = MlasSgemmKernel;
|
||||
+#endif
|
||||
+
|
||||
+#if defined(MLAS_TARGET_S390X)
|
||||
+ this->GemmFloatKernel = MlasSgemmKernel;
|
||||
+#endif
|
||||
+
|
||||
+#if defined(MLAS_TARGET_RISCV64)
|
||||
+ this->GemmFloatKernel = nullptr;
|
||||
+#if defined(MLAS_USE_RVV)
|
||||
+ bool has_rvv = true;
|
||||
+#if defined(__linux__)
|
||||
+ has_rvv = (getauxval(AT_HWCAP) & COMPAT_HWCAP_ISA_V) != 0;
|
||||
+#endif
|
||||
+ if (has_rvv) {
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelRvv;
|
||||
+ }
|
||||
+#endif // MLAS_USE_RVV
|
||||
+#endif // MLAS_TARGET_RISCV64
|
||||
+
|
||||
+#if defined(MLAS_TARGET_LARCH64)
|
||||
+ // No fine-grained LSX/LASX detection here — pick LASX (256-bit) since
|
||||
+ // the LoongArch64 spec requires it; LSX (128-bit) is the fallback.
|
||||
+ this->GemmFloatKernel = MlasGemmFloatKernelLasx;
|
||||
+ this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Lasx;
|
||||
+#endif
|
||||
+
|
||||
+ // ARM64 and WASM intentionally do nothing here — sgemm.cpp's #else branch
|
||||
+ // calls MlasSgemmKernelZero / MlasSgemmKernelAdd directly without going
|
||||
+ // through GetMlasPlatform().GemmFloatKernel.
|
||||
+}
|
||||
+#else // !MLAS_GEMM_ONLY
|
||||
MLAS_PLATFORM::MLAS_PLATFORM(
|
||||
void
|
||||
)
|
||||
@@ -909,6 +1041,7 @@
|
||||
#endif // MLAS_TARGET_LARCH64
|
||||
|
||||
}
|
||||
+#endif // MLAS_GEMM_ONLY
|
||||
|
||||
size_t
|
||||
MLASCALL
|
||||
@@ -0,0 +1,19 @@
|
||||
Patch against upstream onnxruntime/core/mlas/inc/mlas.h
|
||||
Base commit: 62f742f1aa0c3102745ed35e3d869eaee845b9ac (ORT v1.26.0)
|
||||
|
||||
Guard _MSC_VER with defined() so -Wundef builds on GCC/Clang
|
||||
(where _MSC_VER is not predefined) do not warn when this header is
|
||||
consumed outside the ORT build that always defines _MSC_VER through
|
||||
its toolchain wrappers.
|
||||
|
||||
--- a/3rdparty/mlas/inc/mlas.h
|
||||
+++ b/3rdparty/mlas/inc/mlas.h
|
||||
@@ -26,7 +26,7 @@
|
||||
// Define the calling convention for Windows targets.
|
||||
//
|
||||
|
||||
-#if (_MSC_VER >= 800) || defined(_STDCALL_SUPPORTED)
|
||||
+#if (defined(_MSC_VER) && (_MSC_VER >= 800)) || defined(_STDCALL_SUPPORTED)
|
||||
#define MLASCALL __stdcall
|
||||
#else
|
||||
#define MLASCALL
|
||||
Vendored
+17
@@ -65,6 +65,23 @@ jasper JasPer is a collection of software
|
||||
https://github.com/jasper-software/jasper.git
|
||||
for details and links to source code
|
||||
------------------------------------------------------------------------------------
|
||||
mlas Microsoft Linear Algebra Subprograms — processor-optimized
|
||||
GEMM kernels and platform-specific threading code.
|
||||
Vendored from ONNX Runtime (onnxruntime/core/mlas/), MIT licensed.
|
||||
|
||||
Copyright (c) Microsoft Corporation
|
||||
Additional MIT-licensed contributions in the source tree:
|
||||
Copyright 2025 Arm Limited (lib/kleidiai/)
|
||||
Copyright 2025 FUJITSU LIMITED (erf/gelu neon fp16)
|
||||
|
||||
License: see mlas/LICENSE
|
||||
Provenance and local patches: see mlas/README.md
|
||||
Upstream: https://github.com/microsoft/onnxruntime
|
||||
|
||||
Used by the dnn module's SGEMM dispatch path. Built as
|
||||
an OBJECT library and linked into opencv_dnn when the
|
||||
host arch/OS is wired up (HAVE_MLAS).
|
||||
------------------------------------------------------------------------------------
|
||||
ffmpeg FFmpeg is a complete, cross-platform solution to record,
|
||||
convert and stream audio and video. It includes libavcodec -
|
||||
the leading audio/video codec library, and also libavformat, libavutils and
|
||||
|
||||
@@ -1819,6 +1819,38 @@ if(BUILD_opencv_dnn AND OPENCV_DNN_BACKEND_DEFAULT)
|
||||
status(" Default DNN backend:" ${OPENCV_DNN_BACKEND_DEFAULT})
|
||||
endif()
|
||||
|
||||
if(BUILD_opencv_dnn AND (OPENCV_DNN_MLAS_ENABLED OR OPENCV_DNN_MLAS_SKIP_REASON))
|
||||
status(" DNN MLAS:" OPENCV_DNN_MLAS_ENABLED THEN "YES (SGEMM-only, vendored)"
|
||||
ELSE "NO (${OPENCV_DNN_MLAS_SKIP_REASON})")
|
||||
if(OPENCV_DNN_MLAS_ENABLED)
|
||||
if(MLAS_X86_64)
|
||||
status(" ASM kernels:" "YES (X86_64: SSE2, AVX, FMA3, AVX512F)")
|
||||
elseif(MLAS_X86)
|
||||
status(" ASM kernels:" "YES (X86: SSE2, AVX)")
|
||||
elseif(MLAS_ARM64)
|
||||
status(" ASM kernels:" "YES (ARM64: NEON SGEMM, NEON SGEMV)")
|
||||
elseif(MLAS_ARM)
|
||||
status(" ASM kernels:" "NO (ARM 32-bit, scalar C++ sgemmc.cpp)")
|
||||
elseif(MLAS_LOONGARCH64)
|
||||
status(" ASM kernels:" "YES (LoongArch: LSX, LASX)")
|
||||
elseif(MLAS_POWER)
|
||||
status(" ASM kernels:" MLAS_HAS_POWER10 AND MLAS_HAS_ASM
|
||||
THEN "YES (POWER10 PackA)"
|
||||
ELSE "NO (POWER base, no .S kernels)")
|
||||
elseif(MLAS_S390X)
|
||||
status(" ASM kernels:" "NO (S390X ZVECTOR via intrinsics)")
|
||||
elseif(MLAS_RISCV64)
|
||||
status(" ASM kernels:" MLAS_HAS_RISCV64_RVV
|
||||
THEN "NO (RISCV64 with RVV intrinsics)"
|
||||
ELSE "NO (RISCV64 scalar fallback)")
|
||||
elseif(MLAS_WASM)
|
||||
status(" ASM kernels:" "NO (WASM scalar fallback)")
|
||||
else()
|
||||
status(" ASM kernels:" "NO (scalar fallback)")
|
||||
endif()
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(WITH_EIGEN OR HAVE_EIGEN)
|
||||
status(" Eigen:" HAVE_EIGEN THEN "YES (ver ${EIGEN_WORLD_VERSION}.${EIGEN_MAJOR_VERSION}.${EIGEN_MINOR_VERSION})" ELSE NO)
|
||||
endif()
|
||||
|
||||
@@ -487,6 +487,30 @@ if(NOT EMSCRIPTEN)
|
||||
endif()
|
||||
endif()
|
||||
|
||||
# Vendored MLAS (Microsoft Linear Algebra Subprograms) from ONNX Runtime.
|
||||
# Sources live in 3rdparty/mlas/. Builds to an OBJECT library whose objects
|
||||
# link directly into opencv_dnn. Skipped under Emscripten: MLAS is a native
|
||||
# CPU SGEMM accelerator (asm/intrinsic kernels) and the wasm scalar fallback
|
||||
# offers no benefit for the JS bindings build that produces opencv.js.
|
||||
set(HAVE_MLAS 0)
|
||||
# Reset status flags + arch booleans
|
||||
foreach(_v OPENCV_DNN_MLAS_ENABLED OPENCV_DNN_MLAS_SKIP_REASON
|
||||
MLAS_X86_64 MLAS_X86 MLAS_ARM MLAS_ARM64 MLAS_POWER
|
||||
MLAS_LOONGARCH64 MLAS_S390X MLAS_RISCV64 MLAS_WASM
|
||||
MLAS_HAS_ASM MLAS_HAS_POWER10 MLAS_HAS_RISCV64_RVV)
|
||||
unset(${_v} CACHE)
|
||||
endforeach()
|
||||
if(NOT EMSCRIPTEN)
|
||||
add_subdirectory("${OpenCV_SOURCE_DIR}/3rdparty/mlas" "${CMAKE_BINARY_DIR}/3rdparty/mlas")
|
||||
endif()
|
||||
if(HAVE_MLAS)
|
||||
add_definitions(-DHAVE_MLAS=1)
|
||||
list(APPEND include_dirs ${MLAS_INCLUDE_DIRS})
|
||||
message(STATUS "DNN: MLAS (vendored) enabled.")
|
||||
else()
|
||||
message(STATUS "DNN: MLAS (vendored) disabled — host arch/OS not wired up.")
|
||||
endif()
|
||||
|
||||
ocv_module_include_directories(${include_dirs})
|
||||
if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
|
||||
ocv_append_source_files_cxx_compiler_options(fw_srcs "-Wno-suggest-override") # GCC
|
||||
@@ -534,7 +558,7 @@ endif()
|
||||
|
||||
ocv_install_used_external_targets(${libs} ${dnn_runtime_libs})
|
||||
|
||||
ocv_glob_module_sources(${sources_options} SOURCES ${fw_srcs} ${webnn_srcs})
|
||||
ocv_glob_module_sources(${sources_options} SOURCES ${fw_srcs} ${webnn_srcs} ${MLAS_OBJECTS})
|
||||
ocv_create_module(${libs} ${dnn_runtime_libs})
|
||||
ocv_add_samples()
|
||||
ocv_add_accuracy_tests(${dnn_runtime_libs})
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
#include "../precomp.hpp"
|
||||
#include "cpu_kernels/fast_gemm.hpp"
|
||||
#include "cpu_kernels/softmax.hpp"
|
||||
#include "cpu_kernels/mlas_gemm.hpp"
|
||||
|
||||
#include <opencv2/dnn/shape_utils.hpp>
|
||||
|
||||
@@ -366,7 +367,39 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
|
||||
parallel_for_(Range(0, loops), fn, nstripes);
|
||||
}
|
||||
|
||||
// attention_prob = softmax(scale * Q @ K^T + mask)
|
||||
const int num_non_blob_inputs = (int)inputs.size();
|
||||
const bool has_mask = (!blobs.empty() && num_non_blob_inputs >= 2) ||
|
||||
( blobs.empty() && num_non_blob_inputs >= 4);
|
||||
|
||||
if (mlasAvailable() && !has_mask &&
|
||||
batch_size > 0 && num_heads > 0 && seq_len > 0 &&
|
||||
qkv_head_sizes[0] > 0 && qkv_head_sizes[2] > 0)
|
||||
{
|
||||
const int B = (int)batch_size;
|
||||
const int H = (int)num_heads;
|
||||
const int S = (int)seq_len;
|
||||
const int Dqk = (int)qkv_head_sizes[0];
|
||||
const int Dv = (int)qkv_head_sizes[2];
|
||||
// ORT-style default tiling, clamped to the actual sequence length.
|
||||
const int q_block = std::min(256, S);
|
||||
const int kv_block = std::min(256, S);
|
||||
const int threads = std::max(1, cv::getNumThreads());
|
||||
|
||||
const size_t per_thread =
|
||||
mlasFlashAttentionBufferBytesPerThread(q_block, kv_block, Dv);
|
||||
flash_scratch.resize((size_t)threads * per_thread);
|
||||
|
||||
if (mlasFlashAttention(Q, K, V,
|
||||
outputs[0].ptr<float>(),
|
||||
B, H, S, S, Dqk, Dv, scale,
|
||||
q_block, kv_block,
|
||||
flash_scratch.data(), threads))
|
||||
{
|
||||
return; // fast path done; outputs[0] is fully written
|
||||
}
|
||||
}
|
||||
|
||||
// Compute Softmax(scale * MatMul(Q, K))
|
||||
auto &attention_prob = internals[1];
|
||||
{
|
||||
auto *output = attention_prob.ptr<float>();
|
||||
@@ -526,6 +559,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer {
|
||||
std::vector<float> packed_weight_q;
|
||||
std::vector<float> packed_weight_k;
|
||||
std::vector<float> packed_weight_v;
|
||||
std::vector<unsigned char> flash_scratch;
|
||||
|
||||
FastGemmOpt opt;
|
||||
};
|
||||
|
||||
@@ -11,6 +11,7 @@
|
||||
|
||||
#include "../../precomp.hpp"
|
||||
#include "fast_gemm.hpp"
|
||||
#include "mlas_gemm.hpp"
|
||||
|
||||
#define CV_CPU_OPTIMIZATION_DECLARATIONS_ONLY
|
||||
#include "fast_gemm_kernels.simd.hpp"
|
||||
@@ -500,6 +501,20 @@ void fastGemm(bool trans_a, bool trans_b, int ma, int na, int mb, int nb,
|
||||
return fast_gemm_thin(alpha, beta, M, N, K, a, lda0, lda1, b, ldb0, c, ldc, opt.multi_thread);
|
||||
}
|
||||
|
||||
#ifdef HAVE_MLAS
|
||||
const bool a_row_major = (lda0 == 1 || lda1 == 1);
|
||||
const bool b_row_major = (ldb0 == 1 || ldb1 == 1);
|
||||
if (a_row_major && b_row_major) {
|
||||
const int phys_lda = std::max(lda0, lda1);
|
||||
const int phys_ldb = std::max(ldb0, ldb1);
|
||||
if (mlasSgemm(trans_a, trans_b, M, N, K,
|
||||
alpha, A, phys_lda, B, phys_ldb,
|
||||
beta, C, ldc)) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
opt_NEON::fastGemmKernel(M, N, K, alpha, a, lda0, lda1,
|
||||
@@ -594,6 +609,24 @@ void fastGemmBatch(size_t batch, const size_t *A_offsets, const size_t *B_offset
|
||||
return;
|
||||
}
|
||||
|
||||
#ifdef HAVE_MLAS
|
||||
bool a_ok = false, b_ok = false;
|
||||
bool mlas_trans_a = false, mlas_trans_b = false;
|
||||
int mlas_lda = 0, mlas_ldb = 0;
|
||||
if (lda1 == 1) { a_ok = true; mlas_trans_a = false; mlas_lda = lda0; }
|
||||
else if (lda0 == 1) { a_ok = true; mlas_trans_a = true; mlas_lda = lda1; }
|
||||
if (ldb1 == 1) { b_ok = true; mlas_trans_b = false; mlas_ldb = ldb0; }
|
||||
else if (ldb0 == 1) { b_ok = true; mlas_trans_b = true; mlas_ldb = ldb1; }
|
||||
if (a_ok && b_ok) {
|
||||
if (mlasSgemmBatch(batch, A_offsets, B_offsets, C_offsets,
|
||||
mlas_trans_a, mlas_trans_b, M, N, K,
|
||||
alpha, A, mlas_lda, B, mlas_ldb,
|
||||
beta, C, ldc)) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
#if CV_TRY_NEON
|
||||
if (opt.use_neon) {
|
||||
opt_NEON::fastGemmBatchKernel(batch, A_offsets, B_offsets, C_offsets, M, N, K, alpha, a, lda0, lda1, b, ldb0, ldb1, beta, c, ldc, sizeof(float));
|
||||
|
||||
@@ -0,0 +1,215 @@
|
||||
// This file is part of OpenCV project.
|
||||
// It is subject to the license terms in the LICENSE file found in the top-level directory
|
||||
// of this distribution and at http://opencv.org/license.html.
|
||||
// Copyright (C) 2026, BigVision LLC, all rights reserved.
|
||||
// Third party copyrights are property of their respective owners.
|
||||
|
||||
#include "../../precomp.hpp"
|
||||
#include "mlas_gemm.hpp"
|
||||
|
||||
#ifdef HAVE_MLAS
|
||||
|
||||
#include "mlas.h"
|
||||
#include <vector>
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
bool mlasAvailable() {
|
||||
static const bool ok = []() {
|
||||
const size_t a = MlasGetPreferredBufferAlignment();
|
||||
return a > 0 && a <= 256;
|
||||
}();
|
||||
return ok;
|
||||
}
|
||||
|
||||
bool mlasSgemm(bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A, int lda,
|
||||
const float* B, int ldb,
|
||||
float beta,
|
||||
float* C, int ldc)
|
||||
{
|
||||
if (!mlasAvailable()) return false;
|
||||
if (M <= 0 || N <= 0 || K <= 0) return false;
|
||||
|
||||
MLAS_SGEMM_DATA_PARAMS data;
|
||||
data.A = A;
|
||||
data.lda = static_cast<size_t>(lda);
|
||||
data.B = B;
|
||||
data.ldb = static_cast<size_t>(ldb);
|
||||
data.C = C;
|
||||
data.ldc = static_cast<size_t>(ldc);
|
||||
data.alpha = alpha;
|
||||
data.beta = beta;
|
||||
data.BIsPacked = false;
|
||||
|
||||
MlasGemm(trans_a ? CblasTrans : CblasNoTrans,
|
||||
trans_b ? CblasTrans : CblasNoTrans,
|
||||
static_cast<size_t>(M),
|
||||
static_cast<size_t>(N),
|
||||
static_cast<size_t>(K),
|
||||
data,
|
||||
/*ThreadPool=*/nullptr,
|
||||
/*BackendKernelSelectorConfig=*/nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool mlasSgemmBatch(size_t batch,
|
||||
const size_t* A_offsets,
|
||||
const size_t* B_offsets,
|
||||
const size_t* C_offsets,
|
||||
bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A_base, int lda,
|
||||
const float* B_base, int ldb,
|
||||
float beta,
|
||||
float* C_base, int ldc)
|
||||
{
|
||||
if (!mlasAvailable()) return false;
|
||||
if (batch == 0 || M <= 0 || N <= 0 || K <= 0) return false;
|
||||
|
||||
std::vector<MLAS_SGEMM_DATA_PARAMS> data(batch);
|
||||
for (size_t i = 0; i < batch; i++) {
|
||||
data[i].A = A_base + A_offsets[i];
|
||||
data[i].lda = static_cast<size_t>(lda);
|
||||
data[i].B = B_base + B_offsets[i];
|
||||
data[i].ldb = static_cast<size_t>(ldb);
|
||||
data[i].C = C_base + C_offsets[i];
|
||||
data[i].ldc = static_cast<size_t>(ldc);
|
||||
data[i].alpha = alpha;
|
||||
data[i].beta = beta;
|
||||
data[i].BIsPacked = false;
|
||||
}
|
||||
|
||||
MlasGemmBatch(trans_a ? CblasTrans : CblasNoTrans,
|
||||
trans_b ? CblasTrans : CblasNoTrans,
|
||||
static_cast<size_t>(M),
|
||||
static_cast<size_t>(N),
|
||||
static_cast<size_t>(K),
|
||||
data.data(),
|
||||
batch,
|
||||
/*ThreadPool=*/nullptr,
|
||||
/*BackendKernelSelectorConfig=*/nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
size_t mlasSgemmPackBSize(bool trans_a, bool trans_b, int N, int K)
|
||||
{
|
||||
if (!mlasAvailable()) return 0;
|
||||
if (N <= 0 || K <= 0) return 0;
|
||||
return MlasGemmPackBSize(trans_a ? CblasTrans : CblasNoTrans,
|
||||
trans_b ? CblasTrans : CblasNoTrans,
|
||||
static_cast<size_t>(N),
|
||||
static_cast<size_t>(K),
|
||||
/*BackendKernelSelectorConfig=*/nullptr);
|
||||
}
|
||||
|
||||
bool mlasSgemmPackB(bool trans_a, bool trans_b, int N, int K,
|
||||
const float* B, int ldb, void* packed_B)
|
||||
{
|
||||
if (!mlasAvailable()) return false;
|
||||
if (N <= 0 || K <= 0 || B == nullptr || packed_B == nullptr) return false;
|
||||
MlasGemmPackB(trans_a ? CblasTrans : CblasNoTrans,
|
||||
trans_b ? CblasTrans : CblasNoTrans,
|
||||
static_cast<size_t>(N),
|
||||
static_cast<size_t>(K),
|
||||
B, static_cast<size_t>(ldb),
|
||||
packed_B,
|
||||
/*BackendKernelSelectorConfig=*/nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool mlasSgemmPacked(bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A, int lda,
|
||||
const void* packed_B,
|
||||
float beta,
|
||||
float* C, int ldc)
|
||||
{
|
||||
if (!mlasAvailable()) return false;
|
||||
if (M <= 0 || N <= 0 || K <= 0) return false;
|
||||
|
||||
MLAS_SGEMM_DATA_PARAMS data;
|
||||
data.A = A;
|
||||
data.lda = static_cast<size_t>(lda);
|
||||
data.B = static_cast<const float*>(packed_B);
|
||||
data.ldb = 0; // ignored when BIsPacked
|
||||
data.C = C;
|
||||
data.ldc = static_cast<size_t>(ldc);
|
||||
data.alpha = alpha;
|
||||
data.beta = beta;
|
||||
data.BIsPacked = true;
|
||||
|
||||
MlasGemm(trans_a ? CblasTrans : CblasNoTrans,
|
||||
trans_b ? CblasTrans : CblasNoTrans,
|
||||
static_cast<size_t>(M),
|
||||
static_cast<size_t>(N),
|
||||
static_cast<size_t>(K),
|
||||
data,
|
||||
/*ThreadPool=*/nullptr,
|
||||
/*BackendKernelSelectorConfig=*/nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
size_t mlasFlashAttentionBufferBytesPerThread(int q_block_size,
|
||||
int kv_block_size,
|
||||
int v_head_size)
|
||||
{
|
||||
if (q_block_size <= 0 || kv_block_size <= 0 || v_head_size <= 0) return 0;
|
||||
// flashattn.cpp lays out the per-thread scratch as:
|
||||
// l[q_block_size] + m[q_block_size]
|
||||
// + intermediate[q_block_size * kv_block_size]
|
||||
// + temp_output[q_block_size * v_head_size]
|
||||
const size_t q = static_cast<size_t>(q_block_size);
|
||||
const size_t kv = static_cast<size_t>(kv_block_size);
|
||||
const size_t vd = static_cast<size_t>(v_head_size);
|
||||
return (q * (2 + kv + vd)) * sizeof(float);
|
||||
}
|
||||
|
||||
bool mlasFlashAttention(const float* query, const float* key, const float* value,
|
||||
float* output,
|
||||
int batch_size, int num_heads,
|
||||
int q_seq_len, int kv_seq_len,
|
||||
int qk_head_size, int v_head_size,
|
||||
float scale,
|
||||
int q_block_size, int kv_block_size,
|
||||
void* scratch, int thread_count)
|
||||
{
|
||||
if (!mlasAvailable()) return false;
|
||||
if (batch_size <= 0 || num_heads <= 0) return false;
|
||||
if (q_seq_len <= 0 || kv_seq_len <= 0) return false;
|
||||
if (qk_head_size <= 0 || v_head_size <= 0) return false;
|
||||
if (q_block_size <= 0 || kv_block_size <= 0) return false;
|
||||
if (thread_count <= 0 || scratch == nullptr) return false;
|
||||
if (query == nullptr || key == nullptr || value == nullptr || output == nullptr)
|
||||
return false;
|
||||
|
||||
MlasFlashAttentionThreadedArgs args;
|
||||
args.batch_size = batch_size;
|
||||
args.num_heads = num_heads;
|
||||
args.q_sequence_length = q_seq_len;
|
||||
args.kv_sequence_length = kv_seq_len;
|
||||
args.qk_head_size = qk_head_size;
|
||||
args.v_head_size = v_head_size;
|
||||
args.q_block_size = q_block_size;
|
||||
args.kv_block_size = kv_block_size;
|
||||
args.scale = scale;
|
||||
args.thread_count = thread_count;
|
||||
args.buffer = static_cast<float*>(scratch);
|
||||
args.buffer_size_per_thread = mlasFlashAttentionBufferBytesPerThread(
|
||||
q_block_size, kv_block_size, v_head_size);
|
||||
args.query = query;
|
||||
args.key = key;
|
||||
args.value = value;
|
||||
args.output = output;
|
||||
|
||||
MlasFlashAttention(&args, /*ThreadPool=*/nullptr);
|
||||
return true;
|
||||
}
|
||||
|
||||
}} // cv::dnn
|
||||
|
||||
#endif // HAVE_MLAS
|
||||
@@ -0,0 +1,122 @@
|
||||
// This file is part of OpenCV project.
|
||||
// It is subject to the license terms in the LICENSE file found in the top-level directory
|
||||
// of this distribution and at http://opencv.org/license.html.
|
||||
// Copyright (C) 2026, BigVision LLC, all rights reserved.
|
||||
// Third party copyrights are property of their respective owners.
|
||||
|
||||
#ifndef OPENCV_DNN_MLAS_GEMM_HPP
|
||||
#define OPENCV_DNN_MLAS_GEMM_HPP
|
||||
|
||||
#include <cstddef>
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
#ifdef HAVE_MLAS
|
||||
|
||||
// True if MLAS is usable on this host. False signals callers to fall back.
|
||||
bool mlasAvailable();
|
||||
|
||||
// Row-major SGEMM: C := alpha * op(A) * op(B) + beta * C, op(X) = X or X^T.
|
||||
// Returns false if MLAS is unavailable or M/N/K <= 0.
|
||||
bool mlasSgemm(bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A, int lda,
|
||||
const float* B, int ldb,
|
||||
float beta,
|
||||
float* C, int ldc);
|
||||
|
||||
// Batched SGEMM with per-batch element offsets into A_base/B_base/C_base.
|
||||
// M/N/K and leading dims are shared across the batch.
|
||||
bool mlasSgemmBatch(size_t batch,
|
||||
const size_t* A_offsets,
|
||||
const size_t* B_offsets,
|
||||
const size_t* C_offsets,
|
||||
bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A_base, int lda,
|
||||
const float* B_base, int ldb,
|
||||
float beta,
|
||||
float* C_base, int ldc);
|
||||
|
||||
// Pack B once, reuse across many mlasSgemmPacked() calls. Returns the
|
||||
// required buffer size in bytes; caller allocates and passes to mlasSgemmPackB.
|
||||
size_t mlasSgemmPackBSize(bool trans_a, bool trans_b, int N, int K);
|
||||
|
||||
bool mlasSgemmPackB(bool trans_a, bool trans_b, int N, int K,
|
||||
const float* B, int ldb, void* packed_B);
|
||||
|
||||
// mlasSgemm with a pre-packed B from mlasSgemmPackB.
|
||||
bool mlasSgemmPacked(bool trans_a, bool trans_b,
|
||||
int M, int N, int K,
|
||||
float alpha,
|
||||
const float* A, int lda,
|
||||
const void* packed_B,
|
||||
float beta,
|
||||
float* C, int ldc);
|
||||
|
||||
// Scratch-buffer size (in bytes) per worker thread for mlasFlashAttention.
|
||||
// The caller must allocate `thread_count * this` bytes for the scratch
|
||||
// pointer. Returns 0 if any argument is non-positive.
|
||||
size_t mlasFlashAttentionBufferBytesPerThread(int q_block_size,
|
||||
int kv_block_size,
|
||||
int v_head_size);
|
||||
|
||||
// Multi-head attention via MLAS flash-attention. Computes
|
||||
// output[b, i, h, :] = softmax(scale * Q[b,h,i,:] @ K[b,h,:,:]^T) @ V[b,h,:,:]
|
||||
// fused into one tiled kernel without materializing the q_seq x kv_seq
|
||||
// attention matrix.
|
||||
//
|
||||
// Layouts (row-major contiguous, FP32):
|
||||
// query : [batch, num_heads, q_seq_len, qk_head_size]
|
||||
// key : [batch, num_heads, kv_seq_len, qk_head_size]
|
||||
// value : [batch, num_heads, kv_seq_len, v_head_size]
|
||||
// output: [batch, q_seq_len, num_heads, v_head_size] (heads *after* seq)
|
||||
//
|
||||
// scale - usually 1 / sqrt(qk_head_size).
|
||||
// q_block_size - tile size along the q sequence (e.g. 256).
|
||||
// kv_block_size - tile size along the kv sequence (e.g. 256).
|
||||
// scratch - caller-owned buffer of at least
|
||||
// thread_count * mlasFlashAttentionBufferBytesPerThread(...).
|
||||
// thread_count - number of MLAS workers to fan out across (typically
|
||||
// cv::getNumThreads()).
|
||||
//
|
||||
// Returns false if MLAS is unavailable or arguments are invalid.
|
||||
bool mlasFlashAttention(const float* query, const float* key, const float* value,
|
||||
float* output,
|
||||
int batch_size, int num_heads,
|
||||
int q_seq_len, int kv_seq_len,
|
||||
int qk_head_size, int v_head_size,
|
||||
float scale,
|
||||
int q_block_size, int kv_block_size,
|
||||
void* scratch, int thread_count);
|
||||
|
||||
#else // HAVE_MLAS
|
||||
|
||||
inline bool mlasAvailable() { return false; }
|
||||
|
||||
inline bool mlasSgemm(bool, bool, int, int, int, float,
|
||||
const float*, int, const float*, int,
|
||||
float, float*, int) { return false; }
|
||||
|
||||
inline bool mlasSgemmBatch(size_t, const size_t*, const size_t*, const size_t*,
|
||||
bool, bool, int, int, int, float,
|
||||
const float*, int, const float*, int,
|
||||
float, float*, int) { return false; }
|
||||
|
||||
inline size_t mlasSgemmPackBSize(bool, bool, int, int) { return 0; }
|
||||
inline bool mlasSgemmPackB(bool, bool, int, int, const float*, int, void*) { return false; }
|
||||
inline bool mlasSgemmPacked(bool, bool, int, int, int, float,
|
||||
const float*, int, const void*, float, float*, int) { return false; }
|
||||
|
||||
inline size_t mlasFlashAttentionBufferBytesPerThread(int, int, int) { return 0; }
|
||||
inline bool mlasFlashAttention(const float*, const float*, const float*, float*,
|
||||
int, int, int, int, int, int, float,
|
||||
int, int, void*, int) { return false; }
|
||||
|
||||
#endif // HAVE_MLAS
|
||||
|
||||
}} // cv::dnn
|
||||
|
||||
#endif // OPENCV_DNN_MLAS_GEMM_HPP
|
||||
@@ -0,0 +1,98 @@
|
||||
// This file is part of OpenCV project.
|
||||
// It is subject to the license terms in the LICENSE file found in the top-level directory
|
||||
// of this distribution and at http://opencv.org/license.html.
|
||||
|
||||
// OpenCV-backed implementation of MLAS's threading primitives.
|
||||
//
|
||||
// Replaces lib/threading.cpp from upstream MLAS (which, when built with
|
||||
// BUILD_MLAS_NO_ONNXRUNTIME and a nullptr ThreadPool, runs everything in a
|
||||
// serial for-loop). MLAS internally calls MlasGetMaximumThreadCount() to
|
||||
// pick a partition count; for the standalone build that returns 1, so even
|
||||
// with a parallel `MlasTrySimpleParallel` MLAS would still emit a single
|
||||
// iteration. We patch both halves:
|
||||
//
|
||||
// 1. MlasGetMaximumThreadCount() in mlasi.h returns cv::getNumThreads()
|
||||
// when MLAS_OPENCV_THREADING is defined (see the small patch in
|
||||
// mlasi.h flagged with that define).
|
||||
// 2. The three threaded entry points below dispatch to cv::parallel_for_.
|
||||
//
|
||||
// Compiled into the opencv_dnn_mlas object library only — the
|
||||
// MLAS_OPENCV_THREADING guard makes the file a no-op when picked up by
|
||||
// opencv_dnn's recursive src glob, since MLAS internal headers and the
|
||||
// MLAS_GEMM_ONLY / BUILD_MLAS_NO_ONNXRUNTIME defines are not in scope there.
|
||||
|
||||
#if defined(MLAS_OPENCV_THREADING)
|
||||
|
||||
#include "mlasi.h"
|
||||
#include "qgemm.h" // for MLAS_GEMM_QUANT_DISPATCH definition (stub below)
|
||||
|
||||
#include <opencv2/core/utility.hpp>
|
||||
|
||||
// MLAS_GEMM_ONLY stub: mlasi.h's MLAS_PLATFORM struct uses
|
||||
// GemmS8S8Dispatch{&MlasGemmQuantDispatchDefault}
|
||||
// as in-class initializers. The real definition lives in qgemm_kernel_default.cpp
|
||||
// which we don't compile in the SGEMM-only build. Provide a zero-initialized
|
||||
// instance so mlasi.h links — it's never read because we never call MlasQgemm.
|
||||
extern "C++" const MLAS_GEMM_QUANT_DISPATCH MlasGemmQuantDispatchDefault{};
|
||||
|
||||
extern "C" int opencv_dnn_mlas_max_threads()
|
||||
{
|
||||
int n = cv::getNumThreads();
|
||||
return n > 0 ? n : 1;
|
||||
}
|
||||
|
||||
void
|
||||
MlasExecuteThreaded(
|
||||
MLAS_THREADED_ROUTINE* ThreadedRoutine,
|
||||
void* Context,
|
||||
ptrdiff_t Iterations,
|
||||
MLAS_THREADPOOL* /*ThreadPool*/)
|
||||
{
|
||||
if (Iterations <= 0) return;
|
||||
if (Iterations == 1) { ThreadedRoutine(Context, 0); return; }
|
||||
|
||||
cv::parallel_for_(cv::Range(0, static_cast<int>(Iterations)),
|
||||
[&](const cv::Range& r) {
|
||||
for (int tid = r.start; tid < r.end; tid++) {
|
||||
ThreadedRoutine(Context, static_cast<ptrdiff_t>(tid));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void
|
||||
MlasTrySimpleParallel(
|
||||
MLAS_THREADPOOL* /*ThreadPool*/,
|
||||
const std::ptrdiff_t Iterations,
|
||||
const std::function<void(std::ptrdiff_t tid)>& Work)
|
||||
{
|
||||
if (Iterations <= 0) return;
|
||||
if (Iterations == 1) { Work(0); return; }
|
||||
|
||||
cv::parallel_for_(cv::Range(0, static_cast<int>(Iterations)),
|
||||
[&](const cv::Range& r) {
|
||||
for (int tid = r.start; tid < r.end; tid++) {
|
||||
Work(static_cast<std::ptrdiff_t>(tid));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
void
|
||||
MlasTryBatchParallel(
|
||||
MLAS_THREADPOOL* /*ThreadPool*/,
|
||||
const std::ptrdiff_t Iterations,
|
||||
const std::function<void(std::ptrdiff_t tid)>& Work)
|
||||
{
|
||||
// MLAS only calls this for "non-performance-critical" small batches, but
|
||||
// there is no reason to serialize it on a 24-core host either.
|
||||
if (Iterations <= 0) return;
|
||||
if (Iterations == 1) { Work(0); return; }
|
||||
|
||||
cv::parallel_for_(cv::Range(0, static_cast<int>(Iterations)),
|
||||
[&](const cv::Range& r) {
|
||||
for (int tid = r.start; tid < r.end; tid++) {
|
||||
Work(static_cast<std::ptrdiff_t>(tid));
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#endif // MLAS_OPENCV_THREADING
|
||||
@@ -17,6 +17,7 @@ using namespace cv::dnn::cuda4dnn;
|
||||
|
||||
#include <opencv2/dnn/shape_utils.hpp>
|
||||
#include "cpu_kernels/fast_gemm.hpp"
|
||||
#include "cpu_kernels/mlas_gemm.hpp"
|
||||
|
||||
namespace cv { namespace dnn {
|
||||
|
||||
@@ -279,6 +280,30 @@ public:
|
||||
}
|
||||
}
|
||||
}
|
||||
#ifdef HAVE_MLAS
|
||||
std::vector<Mat> outputs;
|
||||
outputs_arr.getMatVector(outputs);
|
||||
const auto shape_A = shape(inputs[0]);
|
||||
const auto shape_Y = shape(outputs[0]);
|
||||
const int na = shape_A[shape_A.size() - 1];
|
||||
const int ma = shape_A[shape_A.size() - 2];
|
||||
const int N = shape_Y[shape_Y.size() - 1];
|
||||
const int K = trans_a ? ma : na;
|
||||
const Mat& Bmat = blobs[0];
|
||||
const int ldb = Bmat.size[Bmat.dims - 1];
|
||||
const size_t packed_bytes = mlasSgemmPackBSize(trans_a, trans_b, N, K);
|
||||
if (packed_bytes > 0) {
|
||||
packed_B_mlas.create(1, static_cast<int>(packed_bytes), CV_8U);
|
||||
if (mlasSgemmPackB(trans_a, trans_b, N, K,
|
||||
Bmat.ptr<const float>(), ldb,
|
||||
packed_B_mlas.data)) {
|
||||
packed_B_mlas_N = N;
|
||||
packed_B_mlas_K = K;
|
||||
} else {
|
||||
packed_B_mlas.release();
|
||||
}
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
if (constC(mode) && flatten_a) {
|
||||
@@ -361,6 +386,20 @@ public:
|
||||
}
|
||||
|
||||
if (constB(mode)) {
|
||||
#ifdef HAVE_MLAS
|
||||
if (!packed_B_mlas.empty() &&
|
||||
packed_B_mlas_N == N && packed_B_mlas_K == K)
|
||||
{
|
||||
if (mlasSgemmPacked(trans_a, trans_b, rows, N, K,
|
||||
alpha,
|
||||
A.ptr<const float>(), na,
|
||||
packed_B_mlas.data,
|
||||
1.f,
|
||||
Y.ptr<float>(), N)) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
CV_CheckGT(packed_B.size(), static_cast<size_t>(0), "DNN/Gemm: constant B is not pre-packed");
|
||||
if (!thin_packed_B.empty()) {
|
||||
fastGemmThin(rows, N, K, alpha, A.ptr<const float>(), na, 1,
|
||||
@@ -531,6 +570,11 @@ private:
|
||||
bool have_bias;
|
||||
std::vector<float> packed_B;
|
||||
std::vector<float> thin_packed_B;
|
||||
#ifdef HAVE_MLAS
|
||||
cv::Mat packed_B_mlas;
|
||||
int packed_B_mlas_N = 0;
|
||||
int packed_B_mlas_K = 0;
|
||||
#endif
|
||||
std::vector<float> broadcast_C;
|
||||
int real_ndims_C;
|
||||
FastGemmOpt opt;
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
#include <opencv2/dnn/shape_utils.hpp>
|
||||
#include "cpu_kernels/fast_gemm.hpp"
|
||||
#include "cpu_kernels/mlas_gemm.hpp"
|
||||
|
||||
// OpenVINO backend
|
||||
#include "../op_inf_engine.hpp"
|
||||
@@ -267,9 +268,22 @@ class MatMulLayerImpl CV_FINAL : public MatMulLayer {
|
||||
if (blobs.empty()) {
|
||||
const auto &B = inputs[1];
|
||||
const auto *b = B.ptr<const float>();
|
||||
fastGemmBatch(helper.batch, helper.A_offsets.data(), helper.B_offsets.data(), helper.C_offsets.data(),
|
||||
helper.M, helper.N, helper.K, alpha, a, helper.lda0, helper.lda1,
|
||||
b, helper.ldb0, helper.ldb1, beta, y, helper.ldc, opt);
|
||||
bool done = false;
|
||||
if (mlasAvailable() && helper.M > 0 && helper.N > 0 && helper.K > 0) {
|
||||
const auto A_shape = shape(A);
|
||||
const auto B_shape = shape(B);
|
||||
const int lda_mem = A_shape.back();
|
||||
const int ldb_mem = B_shape.back();
|
||||
done = mlasSgemmBatch(helper.batch,
|
||||
helper.A_offsets.data(), helper.B_offsets.data(), helper.C_offsets.data(),
|
||||
trans_a, trans_b, helper.M, helper.N, helper.K,
|
||||
alpha, a, lda_mem, b, ldb_mem, beta, y, helper.ldc);
|
||||
}
|
||||
if (!done) {
|
||||
fastGemmBatch(helper.batch, helper.A_offsets.data(), helper.B_offsets.data(), helper.C_offsets.data(),
|
||||
helper.M, helper.N, helper.K, alpha, a, helper.lda0, helper.lda1,
|
||||
b, helper.ldb0, helper.ldb1, beta, y, helper.ldc, opt);
|
||||
}
|
||||
} else if (!thin_packed_B.empty()) {
|
||||
fastGemmThin(helper.M, helper.N, helper.K, alpha,
|
||||
a, helper.lda0, helper.lda1,
|
||||
|
||||
Reference in New Issue
Block a user