mirrors/opencv
2
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2026-09-25 04:09:57 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #28934 from abhishek-gola:mlas_gemm

Browse Source 
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
2026-05-22 20:22:15 +03:00
committed by GitHub
parent 104d987ca2
commit bdf348c13a
67 changed files with 22056 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
3rdparty/mlas/CMakeLists.txt
Vendored
+277
View File
@@ -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)
3rdparty/mlas/LICENSE
Vendored
+21
View File
@@ -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.
3rdparty/mlas/README.md
Vendored
+116
View File
@@ -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.
3rdparty/mlas/inc/mlas.h
Vendored
+2311
View File
@@ -0,0 +1,2311 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
mlas.h
Abstract:
This module contains the public data structures and procedure prototypes
for the Microsoft Machine Learning algebra subprogram library.
--*/
#pragma once
#include <cstddef>
#include <cstdlib>
#include <cstdint>
#include <stdexcept>
//
// Define the calling convention for Windows targets.
//
#if (defined(_MSC_VER) && (_MSC_VER >= 800)) || defined(_STDCALL_SUPPORTED)
#define MLASCALL __stdcall
#else
#define MLASCALL
#endif
//
// Define the target architecture.
//
#if (defined(_M_AMD64) && !defined(_M_ARM64EC)) || defined(__x86_64__)
#define MLAS_TARGET_AMD64
#endif
#if defined(_M_IX86) || defined(__i386__)
#define MLAS_TARGET_IX86
#endif
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_IX86)
#define MLAS_TARGET_AMD64_IX86
#endif
#if defined(_M_ARM64) || defined(__aarch64__)
#define MLAS_TARGET_ARM64
#endif
#if defined(_M_ARM64EC)
#define MLAS_TARGET_ARM64EC
#endif
#if defined(_M_ARM) || defined(__arm__)
#define MLAS_TARGET_ARM
#endif
#if defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_ARM64EC) || defined(MLAS_TARGET_ARM)
#define MLAS_TARGET_ARM_ANY
#endif
#if defined(__s390x__)
#define MLAS_TARGET_S390X
#endif
#if defined(__riscv) && defined(__riscv_xlen) && (__riscv_xlen == 64)
#define MLAS_TARGET_RISCV64
#endif
#if defined(__VSX__)
#define MLAS_TARGET_POWER
#endif
#if defined(__wasm__)
#define MLAS_TARGET_WASM
#if defined(__wasm_relaxed_simd__)
#define MLAS_TARGET_WASM_RELAXED_SIMD
#define MLAS_TARGET_WASM_SIMD
#elif defined(__wasm_simd128__)
#define MLAS_TARGET_WASM_SIMD
#else
#define MLAS_TARGET_WASM_SCALAR
#endif
#endif
#if defined(__loongarch64)
#define MLAS_TARGET_LARCH64
#endif
//
// Define the support levels for the target architecture.
//
#if defined(MLAS_TARGET_AMD64) || defined (MLAS_TARGET_POWER) || defined (MLAS_TARGET_ZVECTOR)
#define MLAS_SUPPORTS_GEMM_DOUBLE
#endif
#if (!defined(_MSC_VER)) || (_MSC_VER >= 1930)
#if defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_ARM64EC)
#if !defined(__APPLE__)
// Had to temporary disable fp16 under APPLE ARM64, as compiling
// the source files require a hardware specific compilation flag.
// When building an universial binary for APPLE, this flag would
// cause trouble for x64 target.
#define MLAS_F16VEC_INTRINSICS_SUPPORTED
#endif //
#endif // ARM64
#endif // Visual Studio 16 or earlier does not support fp16 intrinsic
//
// Basic Linear Algebra Subprograms (BLAS) types.
//
#ifndef CBLAS_ENUM_DEFINED_H
#define CBLAS_ENUM_DEFINED_H
typedef enum { CblasNoTrans=111, CblasTrans=112, CblasConjTrans=113 } CBLAS_TRANSPOSE;
typedef enum { CblasUpper=121, CblasLower=122 } CBLAS_UPLO;
typedef enum { CblasNonUnit=131, CblasUnit=132 } CBLAS_DIAG;
typedef enum { CblasLeft=141, CblasRight=142} CBLAS_SIDE;
#endif
//
// Forward declare the thread pool implementation class and half precision floating point.
//
// N.B. Avoid including ONNX Runtime headers here to keep the dependencies for
// standalone MLAS test executables smaller.
//
namespace onnxruntime {
namespace concurrency {
class ThreadPool;
};
struct MLFloat16;
}; // namespace onnxruntime
using MLAS_THREADPOOL = onnxruntime::concurrency::ThreadPool;
//
// Platform routines.
//
size_t
MLASCALL
MlasGetPreferredBufferAlignment(
void
);
#ifdef MLAS_TARGET_AMD64_IX86
/**
* @brief Return whether the current CPU has over saturation problem
* when computing u8s8 matrix multiplication
* https://www.intel.com/content/www/us/en/develop/documentation/onednn-developer-guide-and-reference/top/advanced-topics/nuances-of-int8-computations.html
*/
bool
MLASCALL
MlasPlatformU8S8Overflow(
void
);
#endif
//
// Activation routines.
//
enum MLAS_ACTIVATION_KIND {
MlasIdentityActivation,
MlasReluActivation,
MlasLeakyReluActivation,
MlasTanhActivation,
MlasLogisticActivation,
MlasClipActivation,
MlasHardSigmoidActivation,
MlasActivationKindCount,
};
struct MLAS_ACTIVATION {
MLAS_ACTIVATION_KIND ActivationKind;
union {
struct {
float alpha;
} LeakyRelu;
struct {
float minimum;
float maximum;
} Clip;
struct {
float alpha;
float beta;
} HardSigmoid;
float Values[2];
} Parameters;
};
void
MLASCALL
MlasActivation(
const MLAS_ACTIVATION* Activation,
float* Buffer,
const float* Bias,
size_t M,
size_t N,
size_t ldc
);
// Struct to host backend kernel selection configuration options for MLAS
struct MLAS_BACKEND_KERNEL_SELECTOR_CONFIG {
bool use_kleidiai = true; /**< Flag to use KleidiAI backend kernels if available */
};
//
// Matrix/matrix multiply routines.
// C := alpha * op(A) * op(B) + beta * C
// op(X) = X or op(X) = transpose(X) or op(X) = conjg(transpose(X))
//
/**
* @brief Supply matrices data information to single precision gemm functions
*/
struct MLAS_SGEMM_DATA_PARAMS {
const float* A = nullptr; /**< Supplies the address of matrix A */
size_t lda = 0; /**< Supplies the first dimension of matrix A. */
const float* B = nullptr; /**< Supplies the address of matrix B */
size_t ldb = 0; /**< Supplies the first dimension of matrix B. */
float* C = nullptr; /**< Supplies the address of matrix C */
size_t ldc = 0; /**< Supplies the first dimension of matrix C. */
float alpha = 1.0f; /**< Supplies the scalar alpha multiplier (see SGEMM definition) */
float beta = 0.0f; /**< Supplies the scalar beta multiplier (see SGEMM definition) */
bool BIsPacked = false; /**< Whether B is pre-packed */
};
/**
* @brief Batched single precision matrix/matrix multiply operation (SGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param Data A array of matrices data parameters
* @param BatchSize Supplies number of multiplications in this batch
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
* @param BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
void
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,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
/**
* @brief Single precision matrix/matrix multiply operation (SGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param Data Supplies the matrices data parameters
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
* @param BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
const MLAS_SGEMM_DATA_PARAMS& Data,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
{
MlasGemmBatch(TransA, TransB, M, N, K, &Data, 1, ThreadPool, BackendKernelSelectorConfig);
}
/**
* @brief Single precision matrix/matrix multiply operation (SGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param alpha Supplies the scalar alpha multiplier (see SGEMM definition)
* @param A Supplies the address of matrix A
* @param lda Supplies the first dimension of matrix A.
* @param B Supplies the address of matrix B
* @param ldb Supplies the first dimension of matrix B.
* @param beta Supplies the scalar beta multiplier (see SGEMM definition)
* @param C Supplies the address of matrix C
* @param ldc Supplies the first dimension of matrix C.
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
* @param BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const float* B,
size_t ldb,
float beta,
float* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
{
MLAS_SGEMM_DATA_PARAMS Data;
Data.alpha = alpha;
Data.A = A;
Data.lda = lda;
Data.B = B;
Data.ldb = ldb;
Data.beta = beta;
Data.C = C;
Data.ldc = ldc;
MlasGemm(TransA, TransB, M, N, K, Data, ThreadPool, BackendKernelSelectorConfig);
}
/**
* @brief The single precision matrix/matrix multiply operation (SGEMM) with pre-packed B
The pre-packed weights `B` MUST be in accordance with the specified backend kernel selector configuration.
The caller is responsible for ensuring this.
*
* @param TransA - Supplies the transpose operation for matrix A.
* @param M - Supplies the number of rows of matrix A and matrix C.
* @param N - Supplies the number of columns of matrix B and matrix C.
* @param K - Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
* @param A - Supplies the address of matrix A.
* @param lda - Supplies the first dimension of matrix A.
* @param PackedB - Supplies the address of packed matrix B.
* @param beta - Supplies the scalar beta multiplier (see SGEMM definition).
* @param C - Supplies the address of matrix C.
* @param ldc - Supplies the first dimension of matrix C.
* @param ThreadPool - Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
* @param BackendKernelSelectorConfig - Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const void* PackedB,
float beta,
float* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
{
MLAS_SGEMM_DATA_PARAMS DataParams;
DataParams.A = A;
DataParams.lda = lda;
DataParams.B = static_cast<const float*>(PackedB);
DataParams.ldb = 0;
DataParams.C = C;
DataParams.ldc = ldc;
DataParams.alpha = alpha;
DataParams.beta = beta;
DataParams.BIsPacked = true;
MlasGemmBatch(TransA,
CblasTrans, // does not matter when B is packed
M, N, K, &DataParams, 1, ThreadPool, BackendKernelSelectorConfig);
}
/**
* @brief Supply matrices data information to double precision gemm functions
*/
struct MLAS_DGEMM_DATA_PARAMS {
const double* A = nullptr; /**< Supplies the address of matrix A */
size_t lda = 0; /**< Supplies the first dimension of matrix A. */
const double* B = nullptr; /**< Supplies the address of matrix B */
size_t ldb = 0; /**< Supplies the first dimension of matrix B. */
double* C = nullptr; /**< Supplies the address of matrix C */
size_t ldc = 0; /**< Supplies the first dimension of matrix C. */
double alpha = 1.0; /**< Supplies the scalar alpha multiplier (see SGEMM definition) */
double beta = 0.0; /**< Supplies the scalar beta multiplier (see SGEMM definition) */
};
/**
* @brief Batched double precision matrix/matrix multiply operation (DGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param Data A array of matrices data parameters
* @param BatchSize Supplies number of multiplications in this batch
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
*/
void
MLASCALL
MlasGemmBatch(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
const MLAS_DGEMM_DATA_PARAMS* Data,
size_t BatchSize,
MLAS_THREADPOOL* ThreadPool
);
/**
* @brief Double precision matrix/matrix multiply operation (DGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param Data Supplies the matrices data parameters
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
const MLAS_DGEMM_DATA_PARAMS& Data,
MLAS_THREADPOOL* ThreadPool
)
{
MlasGemmBatch(TransA, TransB, M, N, K, &Data, 1, ThreadPool);
}
/**
* @brief Double precision matrix/matrix multiply operation (DGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number
of rows of matrix B.
* @param alpha Supplies the scalar alpha multiplier (see SGEMM definition)
* @param A Supplies the address of matrix A
* @param lda Supplies the first dimension of matrix A.
* @param B Supplies the address of matrix B
* @param ldb Supplies the first dimension of matrix B.
* @param beta Supplies the scalar beta multiplier (see SGEMM definition)
* @param C Supplies the address of matrix C
* @param ldc Supplies the first dimension of matrix C.
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
double alpha,
const double* A,
size_t lda,
const double* B,
size_t ldb,
double beta,
double* C,
size_t ldc,
MLAS_THREADPOOL* ThreadPool
)
{
MLAS_DGEMM_DATA_PARAMS Data;
Data.alpha = alpha;
Data.A = A;
Data.lda = lda;
Data.B = B;
Data.ldb = ldb;
Data.beta = beta;
Data.C = C;
Data.ldc = ldc;
MlasGemmBatch(TransA, TransB, M, N, K, &Data, 1, ThreadPool);
}
enum class MLAS_QUANTIZATION_GRANULARITY {
PerMatrix,
PerColumn,
};
enum class MLAS_QGEMM_OUTPUT_MODE {
ZeroMode, // overwrite the output buffer
AccumulateMode, // accumulate to the output buffer
};
class MLAS_QGEMM_OUTPUT_PROCESSOR {
public:
virtual
void
Process(
const int32_t*, // Supplies the address of matrix to process
size_t, // Supplies the start row index of matrix
size_t, // Supplies the start col index of matrix
size_t, // Supplies the element count per row to process
size_t, // Supplies the element count per col to process
size_t // Supplies the leading dimension of matrix
) const = 0;
virtual ~MLAS_QGEMM_OUTPUT_PROCESSOR() {}
};
class MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR : public MLAS_QGEMM_OUTPUT_PROCESSOR {
public:
MLAS_QGEMM_SCALE_BIAS_OUTPUT_PROCESSOR(
float* Output,
size_t LeadingDimensionOutput,
const float* Scale,
const float* Bias,
MLAS_QGEMM_OUTPUT_MODE Mode = MLAS_QGEMM_OUTPUT_MODE::ZeroMode,
MLAS_QUANTIZATION_GRANULARITY QuantGran = MLAS_QUANTIZATION_GRANULARITY::PerMatrix) :
Output_(Output),
LeadingDimensionOutput_(LeadingDimensionOutput),
Scale_(Scale),
Bias_(Bias),
OutputMode_(Mode),
QuantGran_(QuantGran)
{
}
void
Process(
const int32_t* C,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN,
size_t ldc
) const override;
private:
template<bool HasBias, MLAS_QGEMM_OUTPUT_MODE Mode, MLAS_QUANTIZATION_GRANULARITY QuantGran>
inline
void
ProcessImpl(
const int32_t* C,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN,
size_t ldc
) const;
private:
float* Output_;
size_t LeadingDimensionOutput_;
const float* Scale_;
const float* Bias_;
MLAS_QGEMM_OUTPUT_MODE OutputMode_;
MLAS_QUANTIZATION_GRANULARITY QuantGran_;
};
/**
* @brief Supply matrices shape and data type information to quantized gemm functions
*
** NOTE: AIsSigned == true is not supported on non-ARM devices for now.
** AIsSigned == true is supported on ARM devices when BIsSigned is also true.
*
*/
struct MLAS_GEMM_QUANT_SHAPE_PARAMS {
size_t M = 0; /**< Supplies the row size of matrix A */
size_t N = 0; /**< Supplies the column size of matrix B */
size_t K = 0; /**< Supplies the column size of matrix A and row size of matrix B */
bool AIsSigned = false; /**< Indicates whether type of A is int8_t or uint8_t.*/
bool BIsSigned = false; /**< Indicates whether type of B is int8_t or uint8_t */
bool IsAccumulateMode = false; /**< Indicates whether to accumulate to matrix C or override matrix C */
};
struct MLAS_GEMM_QUANT_DATA_PARAMS {
const uint8_t* A = nullptr;
size_t lda = 0;
uint8_t ZeroPointA = 0;
const void* B = 0;
size_t ldb = 0;
const uint8_t* ZeroPointB = nullptr;
bool BIsPacked = false;
bool PerColumnZeroPoints = false;
int32_t* C = nullptr;
size_t ldc = 0;
const MLAS_QGEMM_OUTPUT_PROCESSOR* OutputProcessor = nullptr;
};
/**
* @brief Batched GEMM, for multiplying multiple pairs of matrices.
* Note: We only support uniform batching, so shapes and types of the
* input must be same: M, N, K, BIsSigned must be the
* same across all parameter blocks.
*
* @param [IN] Shape A single shape descriptor for all the multiplications
* @param [IN] DataParams Array of data descriptors for the matrices.
* @param [IN] BatchN Size of the parameters array, also number of multiplications to perform
* @param [IN] ThreadPool optional thread pool for parallel processing
*/
void
MLASCALL
MlasGemmBatch(
const MLAS_GEMM_QUANT_SHAPE_PARAMS& Shape,
const MLAS_GEMM_QUANT_DATA_PARAMS* DataParams,
const size_t BatchN,
MLAS_THREADPOOL* ThreadPool
);
inline
void
MlasGemm(
const MLAS_GEMM_QUANT_SHAPE_PARAMS &Shape,
const MLAS_GEMM_QUANT_DATA_PARAMS &DataParams,
MLAS_THREADPOOL *ThreadPool)
{
MlasGemmBatch(Shape, &DataParams, 1, ThreadPool);
}
/**
* @brief Parameters that define the shape of a dynamically quantized GEMM operation.
*
* The structure holds the dimensions of the matrices involved in the GEMM
* computation:
* C = A * B
*/
struct MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS {
size_t M = 0; /**< Row size of matrix A */
size_t N = 0; /**< Column size of matrix B */
size_t K = 0; /**< Column size of matrix A and Row size of matrix B */
};
/**
* @brief Parameters that define the data buffers and layout for a dynamic quant GEMM.
*
* This structure provides the memory pointers and strides for matrices
* involved in a dynamically quantized GEMM operation, along with the packed B format.
*/
struct MLAS_GEMM_DYN_QUANT_DATA_PARAMS {
const float* A = nullptr; /**< Pointer to input matrix A in FP32 format**/
size_t lda = 0; /**< Number of elements between adjecent rows in A*/
const void* PackedB = 0; /**< Points to packed weight matrix B */
float *C = nullptr; /**< Points to output Matric C */
size_t ldc = 0; /**< Number of elements between adjecent rows in Matrix C*/
void* Workspace = nullptr; /**< Workspace buffer for LHS Packing Allocation */
size_t WorkspaceSize = 0; /**< Workspace buffer size */
};
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,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
inline void
MlasDynamicQGemm (
const MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS& Shape,
const MLAS_GEMM_DYN_QUANT_DATA_PARAMS* DataParams,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
{
MlasDynamicQGemmBatch(Shape, DataParams, 1, ThreadPool, BackendKernelSelectorConfig);
}
/**
* @brief Determines whether a dynamic quantized GEMM implementation is available on the current platform.
*
* MlasDynamicQGemm() and MlasDynamicQGemmBatch() should only be called if this function returns true.
* @param BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
bool
MLASCALL
MlasIsDynamicQGemmAvailable(const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig);
//
// Symmetric QGEMM has limited buffer overrun.
// Currently only supported in ARM64
//
#if defined(MLAS_TARGET_ARM64)
constexpr size_t MLAS_SYMM_QGEMM_BUF_OVERRUN = 30;
#else
constexpr size_t MLAS_SYMM_QGEMM_BUF_OVERRUN = 0;
#endif
/**
* @brief Supply data parameters for symmetric quantized GEMM.
* B matrix zero point must be zero, and it must be
* pre-packed, with column sums scaled by (-ZeroPointA)
*/
struct MLAS_SYMM_QGEMM_DATA_PARAMS {
const void* A = nullptr;
size_t lda = 0;
const void* B = 0;
void* C = nullptr;
size_t ldc = 0;
// TODO!! add re-quantization parameters
};
/**
* @brief Batched QGEMM. Similar to MlasGemmBatch, but right hand side matrix
* must be symmetrically quantized and prepacked.
*
* @param [IN] Shape A single shape descriptor for all multiplicatons.
Currently A and B must be signed, and accumulation
mode not supported
* @param [IN] DataParams Array of data descriptors, one for each multiplication
* B must be prepacked
* @param [IN] BatchN Number of multiplications
* @param [IN] ThreadPool
*/
void
MLASCALL
MlasSymmQgemmBatch(
const MLAS_GEMM_QUANT_SHAPE_PARAMS& Shape,
const MLAS_SYMM_QGEMM_DATA_PARAMS* DataParams,
const size_t BatchN,
MLAS_THREADPOOL* ThreadPool
);
//
// Buffer packing routines.
//
size_t
MLASCALL
MlasGemmPackBSize(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
void
MLASCALL
MlasGemmPackB(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
size_t
MLASCALL
MlasGemmPackBSize(
size_t N,
size_t K,
bool AIsSigned,
bool BIsSigned,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
void
MLASCALL
MlasGemmPackB(
size_t N,
size_t K,
const uint8_t* B,
size_t ldb,
bool AIsSigned,
bool BIsSigned,
void* PackedB
);
/**
* @brief For symmetric quantized GEMM, returns size of the
* packing buffer needed for right hand side
* @param N Number of columns
* @param K Number of rows
* @param AIsSigned Whether left hand size is signed int8_t
* @return size of the packing buffer,
* 0 if operation not supported
*/
size_t
MLASCALL
MlasSymmQgemmPackBSize(
size_t N,
size_t K,
bool AIsSigned
);
void
MLASCALL
MlasSymmQgemmPackB(
size_t N,
size_t K,
const int8_t* B,
size_t ldb,
bool AIsSigned,
int32_t ZeroPointA,
void* PackedB
);
size_t
MLASCALL
MlasDynamicQgemmPackBSize(
size_t N,
size_t K,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
void
MLASCALL
MlasDynamicQgemmPackB(
size_t N,
size_t K,
const int8_t* B,
const float* Scales,
const float* Bias,
void* PackedB,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
//
// Convolution routines.
//
enum MLAS_CONV_ALGORITHM {
MlasConvAlgorithmGemmDirect,
MlasConvAlgorithmExpandThenGemm,
MlasConvAlgorithmExpandThenGemmSegmented,
MlasConvAlgorithmDepthwiseMultiplierGreaterThan1,
#if defined(MLAS_TARGET_WASM_SCALAR) || defined(MLAS_TARGET_ARM64)
MlasConvAlgorithmDepthwise,
#endif
};
struct MLAS_CONV_PARAMETERS {
const MLAS_ACTIVATION* Activation;
size_t Dimensions;
size_t BatchCount;
size_t GroupCount;
size_t InputChannels;
size_t InputShape[3];
size_t KernelShape[3];
size_t DilationShape[3];
size_t Padding[6];
size_t StrideShape[3];
size_t FilterCount;
size_t OutputShape[3];
size_t InputSize;
size_t OutputSize;
size_t K;
float Beta;
MLAS_CONV_ALGORITHM Algorithm;
ptrdiff_t ThreadCount;
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig = nullptr;
union {
struct {
CBLAS_TRANSPOSE TransB;
size_t ldb;
} GemmDirect;
struct {
size_t ThreadStrideN;
} ExpandThenGemmSegmented;
} u;
};
void 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);
void
MLASCALL
MlasConv(
const MLAS_CONV_PARAMETERS* Parameters,
const float* Input,
const float* Filter,
const float* Bias,
float* WorkingBuffer,
float* Output,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasConvDepthwise(
const void* const* Input,
int32_t InputZeroPoint,
bool InputIsSigned,
const void* Filter,
int32_t FilterZeroPoint,
bool FilterIsSigned,
int32_t* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
//
// Symmetric quantized integer convolution routines.
//
size_t
MlasConvSymPackWSize(
size_t GroupCount,
size_t InputChannels,
size_t OutputChannels,
size_t KernelSize,
bool InputIsSigned
);
void
MlasConvSymPackW(
size_t GroupCount,
size_t InputChannels,
size_t OutputChannels,
size_t KernelSize,
const int8_t* W,
int8_t* PackedW,
size_t PackedWSize,
bool InputIsSigned
);
int32_t
MlasConvSymFixupInputZeroPoint(
int32_t zero_point_value,
bool InputIsSigned
);
//
// Convolution operators (or maybe others in the future) need to do their
// own job partition. Since filters (right hand side B matrix) is usually
// small in size, activations are divided horizontally. We need to provide
// kernel stride units to facilitate the divide.
//
int32_t
MlasConvSymGetKernelOutputCount(
bool InputIsSigned
);
int32_t
MlasConvSymDepthwiseGetKernelOutputCnt(
bool InputIsSigned
);
/**
* @brief Returns the stride M of depthwise conv kernel
*
* Most optimized path is Symmetric conv. See
* MlasConvSymDepthwiseGetKernelOutputCnt(bool)
*
* These kernels are implemented in qdwconv.cpp using
* intrincic, all of them with stride val 1. We use
* a slightly bigger value to improve cache reuse.
*
* This needs to be changed if we optimize depthwise
* kernels.
*
* @return
*/
inline
int32_t
MlasConvDepthwiseGetKernelOutputCnt()
{
return 4;
}
int32_t
MlasSymmQgemmGetKernelOutputCnt();
int32_t
MlasQgemmGetKernelOutputCnt(
bool AIsSigned,
bool BIsSigned
);
struct MLAS_CONV_SYM_PARAMS {
const void* InputDirect;
const void* const* InputIndirection;
const void* Filter;
void* Output;
size_t InputChannels;
size_t OutputChannels;
size_t OutputCount;
size_t KernelSize;
const int32_t* Bias;
const float* Scale;
bool PerChannelScale;
int32_t OutputZeroPoint;
bool InputIsSigned;
};
void
MlasConvSym(
const MLAS_CONV_SYM_PARAMS& Params
);
void
MlasConvSymDepthwise(
const MLAS_CONV_SYM_PARAMS& Params
);
//
// Pooling routines.
//
enum MLAS_POOLING_KIND {
MlasMaximumPooling,
MlasAveragePoolingExcludePad,
MlasAveragePoolingIncludePad,
MlasPoolingKindCount,
};
void
MLASCALL
MlasPool(
MLAS_POOLING_KIND PoolingKind,
size_t Dimensions,
const int64_t* InputShape,
const int64_t* KernelShape,
const int64_t* Padding,
const int64_t* StrideShape,
const int64_t* OutputShape,
const float* Input,
float* Output,
MLAS_THREADPOOL* ThreadPool
);
template<typename T8Bits>
void
MLASCALL
MlasMaximumPool(
const T8Bits* const* Input,
T8Bits* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
//
// Miscellaneous compute routines.
//
void
MLASCALL
MlasComputeErf(
const float* Input,
float* Output,
size_t N
);
//
// Note: The Input and Output buffers for MlasComputeGeluErf must not overlap.
// In-place operation (e.g., passing the same buffer for both parameters) is unsupported.
//
void
MLASCALL
MlasComputeGeluErf(
const float* Input,
float* Output,
size_t N
);
//
// Note: The Input and Output buffers for MlasComputeSilu must not overlap.
// In-place operation (e.g., passing the same buffer for both parameters) is unsupported.
//
void
MLASCALL
MlasComputeSilu(
const float* Input,
float* Output,
size_t N
);
template <typename T>
void
MLASCALL
MlasComputeExp(
const T* Input,
T* Output,
size_t N
);
void
MLASCALL
MlasComputeLogistic(
const float* Input,
float* Output,
size_t N
);
template <typename T>
void
MLASCALL
MlasComputeSoftmax(
const T* Input,
T* Output,
size_t N,
size_t D,
bool LogSoftmax,
bool SmoothSoftmax,
float Sink,
MLAS_THREADPOOL* ThreadPool
);
template <typename T>
void
MLASCALL
MlasComputeSoftcap(
const T* Input,
T* Output,
size_t N,
T cap
);
template <typename T>
void
MLASCALL
MlasEltwiseAdd(
const T* left,
const T* right,
T* output,
size_t N
);
template <typename T>
void
MLASCALL
MlasEltwiseMul(
const T* left,
const T* right,
T* output,
size_t N
);
template<typename T>
void
MLASCALL
MlasComputeTanh(
const T* Input,
T* Output,
size_t N
);
//
// Transpose routines.
//
template<typename DataType>
void
MLASCALL
MlasTranspose(
const DataType* Input,
DataType* Output,
size_t M,
size_t N,
MLAS_THREADPOOL* ThreadPool
);
//
// Buffer reordering routines.
//
void
MLASCALL
MlasReorderInputNchw(
const float* S,
float* D,
size_t InputChannels,
size_t InputSize
);
void
MLASCALL
MlasReorderInputNhwc(
const float* S,
float* D,
size_t InputChannels,
size_t RowCount,
size_t FullRowCount
);
void
MLASCALL
MlasReorderOutputNchw(
const int64_t* OutputShape,
const float* S,
float* D,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasReorderOutputNhwc(
const int64_t* OutputShape,
const float* S,
float* D
);
void
MLASCALL
MlasReorderFilterOIHWBiBo(
const int64_t* FilterShape,
const float* S,
float* D
);
void
MLASCALL
MlasReorderFilterOIHWBo(
const int64_t* FilterShape,
const float* S,
float* D
);
//
// Single precision NCHWc routines.
//
size_t
MLASCALL
MlasNchwcGetBlockSize(
void
);
void
MLASCALL
MlasNchwcConv(
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 GroupCount,
const float* Input,
const float* Filter,
const float* Bias,
float* Output,
const MLAS_ACTIVATION* Activation,
bool ZeroMode,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig,
bool UseBf16
);
void
MLASCALL
MlasNchwcPool(
MLAS_POOLING_KIND PoolingKind,
const int64_t* InputShape,
const int64_t* KernelShape,
const int64_t* DilationShape,
const int64_t* Padding,
const int64_t* StrideShape,
const int64_t* OutputShape,
const float* Input,
float* Output,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasNchwcUpsampleNearest(
const int64_t* InputShape,
const int64_t* Scales,
const float* Input,
float* Output
);
void
MLASCALL
MlasNchwcUpsampleLinear(
size_t InputHeight,
size_t InputWidth,
size_t OutputWidth,
float InterpolationHeight,
const float* InterpolationWidth,
const float* Input,
float* Output
);
//
// Linear quantization routines.
//
template<typename OutputType>
void
MLASCALL
MlasQuantizeLinear(
const float* Input,
OutputType* Output,
size_t N,
float Scale,
OutputType ZeroPoint
);
void
MLASCALL
MlasQuantizeLinearU4(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint
);
void
MLASCALL
MlasQuantizeLinearS4(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint
);
//
// Linear dequantization routines.
//
template<typename InputType>
void
MLASCALL
MlasDequantizeLinear(
const InputType* Input,
float* Output,
size_t N,
float Scale,
InputType ZeroPoint
);
/**
* @brief Requantize a block of the intermediate buffer to the output buffer,
* optionally adding the supplied bias
*
* @param Input Input matrix
* @param InputLeadingDimension Input matrix leading dimension
* @param Output Output matrix
* @param OutputLeadingDimension Output matrix leading dimension
* @param Bias Optional bias vector, to be added
to the input before quantization
* @param Scale Quantization scale
* @param PerColumnScale true if scale is per-column
* @param ZeroPoint quantization zero point value
* @param StartM
* @param StartN
* @param CountM
* @param CountN
* @return
*/
template<typename OutputType>
void
MLASCALL
MlasRequantizeOutput(
const int32_t* Input,
size_t InputLeadingDimension,
OutputType* Output,
size_t OutputLeadingDimension,
const int32_t* Bias,
const float* Scale,
bool PerColumnScale,
OutputType ZeroPoint,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN
);
class MLAS_QGEMM_REQUANT_OUTPUT_PROCESSOR : public MLAS_QGEMM_OUTPUT_PROCESSOR
{
public:
MLAS_QGEMM_REQUANT_OUTPUT_PROCESSOR(
void* Output,
size_t OutputLeadingDimension,
const int32_t* Bias,
const float* Scale,
bool PerColumnScale,
int32_t ZeroPoint,
bool OutputIsSigned)
: Output_(Output),
OutputLeadingDimension_(OutputLeadingDimension),
Bias_(Bias),
Scale_(Scale),
PerColumnScale_(PerColumnScale),
ZeroPoint_(ZeroPoint),
OutputIsSigned_(OutputIsSigned)
{
}
void Process(const int32_t* C,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN,
size_t ldc) const override
{
if(OutputIsSigned_){
MlasRequantizeOutput(C, ldc, reinterpret_cast<int8_t*>(Output_), OutputLeadingDimension_,
Bias_, Scale_, PerColumnScale_, static_cast<int8_t>(ZeroPoint_),
StartM, StartN, CountM, CountN);
} else {
MlasRequantizeOutput(C, ldc, reinterpret_cast<uint8_t*>(Output_), OutputLeadingDimension_,
Bias_, Scale_, PerColumnScale_, static_cast<uint8_t>(ZeroPoint_),
StartM, StartN, CountM, CountN);
}
}
private:
void* Output_;
size_t OutputLeadingDimension_;
const int32_t* Bias_;
const float* Scale_;
bool PerColumnScale_;
int32_t ZeroPoint_;
bool OutputIsSigned_;
};
void
MLASCALL
MlasFindMinMaxElement(
const float* Input,
float* Min,
float* Max,
size_t N
);
size_t
MLASCALL
MlasQLinearSafePaddingElementCount(
size_t ElementSize,
size_t ElementCount
);
template<typename T8Bits>
void
MLASCALL
MlasQLinearGlobalAveragePoolNchw(
const T8Bits* Input,
float ScaleInput,
int32_t ZeroPointInput,
T8Bits* Output,
float ScaleOutput,
int32_t ZeroPointOutput,
size_t Channels,
size_t ImageSize,
int32_t* AccumulateBuffer
);
template <typename T8Bits>
void
MLASCALL
MlasQLinearGlobalAveragePoolNhwc(
const T8Bits* Input,
float ScaleInput,
int32_t ZeroPointInput,
T8Bits* Output,
float ScaleOutput,
int32_t ZeroPointOutput,
size_t Batch,
size_t ImageSize,
size_t Stride,
size_t Channels,
int32_t* AccumulateBuffer,
const T8Bits* ZeroBuffer
);
//
// InputA is of size N,
// Input B is of size 1 if IsScalarB == true, otherwise it is of size N
//
template<typename DataType>
void
MLASCALL
MlasQLinearAdd(
const DataType* InputA,
float ScaleA,
int32_t ZeroPointA,
const DataType* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
DataType* OutputC,
size_t N,
bool IsScalarB
);
template<typename DataType>
void
MLASCALL
MlasQLinearMul(
const DataType* InputA,
float ScaleA,
int32_t ZeroPointA,
const DataType* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
DataType* OutputC,
size_t N,
bool IsScalarB
);
//
// Half precision routines
//
// Any type with size=2 should work
using MLAS_FP16 = onnxruntime::MLFloat16;
constexpr size_t FP16_SIZE = sizeof(uint16_t);
//
// Half-precision floating-point routines.
//
void
MLASCALL
MlasConvertHalfToFloatBuffer(
const MLAS_FP16* Source,
float* Destination,
size_t Count
);
#define MLAS_MIN_TENSOR_SIZE_FOR_HALF_TO_FLOAT_CONVERSION_IN_PARALLEL 128000
void
MLASCALL
MlasConvertHalfToFloatBufferInParallel(
const MLAS_FP16* Source,
float* Destination,
size_t Count,
MLAS_THREADPOOL* ThreadPool
);
void
MLASCALL
MlasConvertFloatToHalfBuffer(
const float* Source,
MLAS_FP16* Destination,
size_t Count
);
void
MLASCALL
MlasConvertFloatToHalfBufferInParallel(
const float* Source,
MLAS_FP16* Destination,
size_t Count,
MLAS_THREADPOOL* ThreadPool
);
/**
* @brief rotary embedding for one hidden state vector
*
* @tparam T: data type of input, sin, cos and output. Currently only float32/16 are supported.
* @param input: input tensor, of shape [dim]
* @param sin: sin tensor, of shape [dim/2]
* @param cos: cos tensor, of shape [dim/2]
* @param dim: dimension of rotary embedding
* @param interleaved: whether the real part and imaginary parts are interleaved
* @param output: output tensor, of shape [dim]
*/
template <typename T>
void
MLASCALL
MlasRotaryEmbedOneRow(
const T* input,
const T* sin_data,
const T* cos_data,
size_t dim,
bool interleaved,
T* output
);
/**
* @brief Supply matrices data information to half precision gemm functions
*/
struct MLAS_HGEMM_DATA_PARAMS {
const MLAS_FP16* A; /**< Supplies the address of matrix A */
size_t lda; /**< Supplies the first dimension of matrix A. */
const MLAS_FP16* B; /**< Supplies the address of matrix B */
size_t ldb; /**< Supplies the first dimension of matrix B. */
MLAS_FP16* C; /**< Supplies the address of matrix C */
size_t ldc; /**< Supplies the first dimension of matrix C. */
uint16_t alpha; /**< Supplies the scalar alpha multiplier (see GEMM definition). FP16 encoding. */
uint16_t beta; /**< Supplies the scalar beta multiplier (see GEMM definition). FP16 encoding. */
};
/**
* @brief Check whether current CPU supports half precision gemm.
*/
bool
MLASCALL
MlasHGemmSupported(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB
);
/**
* @brief Check whether mlas supports GQA kernels with the type and transpose settings.
*/
template <typename T>
bool
MLASCALL
MlasGQASupported(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB
);
/**
* @brief Batched half precision matrix/matrix multiply operation (HGEMM)
*
* @param TransA Supplies the transpose operation for matrix A.
* @param TransB Supplies the transpose operation for matrix B.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number of rows of matrix B.
* @param Data A array of matrices data parameters
* @param BatchSize Supplies number of multiplications in this batch
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
*/
void
MLASCALL
MlasGemmBatch(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
const MLAS_HGEMM_DATA_PARAMS* Data,
size_t BatchSize,
MLAS_THREADPOOL* ThreadPool
);
/**
* @brief half precision matrix/matrix multiply operation (HGEMM)
* C = alpha * op(A) * op(B) + beta * C
*
* @param TransA Supplies the transpose operation for matrix A. Currently only support CblasNoTrans.
* @param TransB Supplies the transpose operation for matrix B. Currently only support CblasTrans.
* @param M Supplies the number of rows of matrix A and matrix C.
* @param N Supplies the number of columns of matrix B and matrix C.
* @param K Supplies the number of columns of matrix A and the number of rows of matrix B.
* @param A Supplies the address of matrix A
* @param lda Supplies the first dimension of matrix A.
* @param B Supplies the address of matrix B
* @param ldb Supplies the first dimension of matrix B.
* @param C Supplies the address of matrix C
* @param ldc Supplies the first dimension of matrix C.
* @param alpha Supplies the scalar alpha multiplier (see GEMM definition)
* @param beta Supplies the scalar beta multiplier (see GEMM definition)
* @param ThreadPool Supplies the thread pool object to use, else nullptr if the base library threading support
* should be used.
*/
inline
void
MlasGemm(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
const MLAS_FP16* A,
size_t lda,
const MLAS_FP16* B,
size_t ldb,
MLAS_FP16* C,
size_t ldc,
uint16_t alpha,
uint16_t beta,
MLAS_THREADPOOL* ThreadPool
) {
MLAS_HGEMM_DATA_PARAMS Data;
Data.A = A;
Data.lda = lda;
Data.B = B;
Data.ldb = ldb;
Data.C = C;
Data.ldc = ldc;
Data.alpha = alpha;
Data.beta = beta;
MlasGemmBatch(TransA, TransB, M, N, K, &Data, 1, ThreadPool);
}
/**
* @brief Whether current CPU supports FP16 acceleration.
*/
bool MLASCALL
MlasFp16AccelerationSupported();
/**
* @brief Interface for half gemm post processors.
*
* Example implementation of this interface includes activations,
* conversion from half precision to single precision, etc.
*
* Half GEMM is computed tile by tile. When a tile of result matrix
* is produced, the method Process() is called to process this tile.
* Parameters of this method describe the location and shape of the
* tile.
*/
class MLAS_HALF_GEMM_POSTPROCESSOR {
public:
virtual
void
Process(
MLAS_FP16*, /**< the address of matrix to process */
size_t, /**< the start row index of matrix */
size_t, /**< the start col index of matrix */
size_t, /**< the element count per row to process */
size_t, /**< the element count per col to process */
size_t /**< the leading dimension of matrix */
) const = 0;
virtual ~MLAS_HALF_GEMM_POSTPROCESSOR() {}
};
/**
* @brief Half precision activation functions, with optional sum tensor.
* Supplied sum tensor must be the same layout as the GEMM output tensor.
* And the supplied sum tensor will be added to the tensor before activation.
*/
class MLAS_HALF_GEMM_ACTIVATION_PROCESSOR : public MLAS_HALF_GEMM_POSTPROCESSOR
{
public:
MLAS_HALF_GEMM_ACTIVATION_PROCESSOR(
const MLAS_ACTIVATION& Activation,
const MLAS_FP16* SumBuf = nullptr)
: Activation_(Activation), SumBuf_(SumBuf)
{}
void Process(
MLAS_FP16* C,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN,
size_t ldc
) const override;
private:
const MLAS_ACTIVATION& Activation_;
const MLAS_FP16* SumBuf_;
};
inline
void
MlasFp16Activation(
const MLAS_ACTIVATION* Activation,
MLAS_FP16* Buffer,
size_t M,
size_t N,
size_t ldc
)
{
MLAS_HALF_GEMM_ACTIVATION_PROCESSOR proc(*Activation);
proc.Process(Buffer, 0, 0, M, N, ldc);
}
/**
* @brief Convert half gemm result matrix to single precision float matrix
*/
class MLAS_HALF_GEMM_2FLOAT_PROCESSOR : public MLAS_HALF_GEMM_POSTPROCESSOR {
public:
MLAS_HALF_GEMM_2FLOAT_PROCESSOR(
const MLAS_ACTIVATION& Activation,
float* Output, /**< address of the output matrix, row major */
size_t RowStride /**< row stride of the output matrix */
) : Activation_(Activation),
Output_(Output),
RowStride_(RowStride)
{}
void
Process(
MLAS_FP16* C,
size_t StartM,
size_t StartN,
size_t CountM,
size_t CountN,
size_t ldc
) const override;
private:
const MLAS_ACTIVATION& Activation_;
float* Output_;
const size_t RowStride_;
};
/**
* @brief Data parameters for half precision GEMM routine
* All except C are [in] parameters
*/
struct MLAS_HALF_GEMM_DATA_PARAMS {
const void* A = nullptr; /**< address of A */
const void* B = nullptr; /**< address of B */
const MLAS_FP16* Bias = nullptr; /**< address of Bias, vector size N */
MLAS_FP16* C = nullptr; /**< address of result matrix */
size_t lda = 0; /**< leading dimension of A */
size_t ldb = 0; /**< leading dimension of B, 0 when B is pre-packed*/
size_t ldc = 0; /**< leading dimension of C*/
const MLAS_HALF_GEMM_POSTPROCESSOR* OutputProcessor = nullptr;
bool AIsfp32 = false; /**< matrix A is fp32, needs to be casted into fp16*/
bool BIsfp32 = false; /**< matrix B is fp32, needs to be casted into fp16*/
};
/**
* @brief Half precision Batched GEMM: C = A * B + Bias
* Either A or B can be fp32 or fp16
*
* Note: We only support uniform batching, so shapes and types of the
* input must be same across all parameter blocks.
*
* @param[in] M row size of matrix A and C
* @param[in] N column size of matrix B and C
* @param[in] K column size of matrix A and row size of matrix B
* @param[in] BatchN number of batches
* @param[inout] DataParams An array (size BatchN) of parameter blocks
* @param[in] ThreadPool
* @return
*/
void
MLASCALL
MlasHalfGemmBatch(
const size_t M,
const size_t N,
const size_t K,
const size_t BatchN,
const MLAS_HALF_GEMM_DATA_PARAMS* DataParams,
MLAS_THREADPOOL* ThreadPool
);
/**
* @brief For half precision GEMM, returns size of the
* packing buffer needed for right hand side
* @param[in] N Number of columns
* @param[in] K Number of rows
* @param[in] float2half Whether the input is float that
* needs to be converted to half precision
* @return size of the packing buffer,
* 0 if operation not supported
*/
size_t
MLASCALL
MlasHalfGemmPackBSize(
size_t N,
size_t K,
bool float2half
);
/**
* @brief For half precision GEMM, pack the right hand
* side matrix B
*
* @param[in] N Number of columns
* @param[in] K Number of rows
* @param[in] B Address of matrix B
* @param[in] ldb leading dimension of input matrix B
* @param[out] PackedB Address of the packed matrix
*/
void
MLASCALL
MlasHalfGemmPackB(
size_t N,
size_t K,
const MLAS_FP16* B,
size_t ldb,
void* PackedB
);
/**
* @brief For half precision GEMM, convert the float matrix B
* to half precision and pack it into a packing buffer
*
* @param[in] N Number of columns
* @param[in] K Number of rows
* @param[in] B Address of matrix B
* @param[in] ldb leading dimension of input matrix B
* @param[out] PackedB Address of the packed matrix
*/
void
MLASCALL
MlasHalfGemmConvertPackB(
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB
);
#if defined(__aarch64__) && defined(__linux__)
/**
* @brief Whether current CPU supports Bfloat16(bf16) acceleration.
*/
bool MLASCALL
MlasBf16AccelerationSupported();
/**
* @brief Interface for bf16 gemm post processors.
*
* Example implementation of this interface includes activations,
* conversion from single precision to precision, etc.
*
* SBGEMM is computed tile by tile. When a tile of result matrix
* is produced, the method Process() is called to process this tile.
* Parameters of this method describe the location and shape of the
* tile.
*/
class MLAS_SBGEMM_POSTPROCESSOR
{
public:
virtual void Process(float*, /**< the address of matrix to process */
size_t, /**< the start row index of matrix */
size_t, /**< the start col index of matrix */
size_t, /**< the element count per row to process */
size_t, /**< the element count per col to process */
size_t /**< the leading dimension of matrix */
) const = 0;
virtual ~MLAS_SBGEMM_POSTPROCESSOR() {}
};
/**
* @brief bfloat16 precision activation functions, with optional sum tensor.
* Supplied sum tensor must be the same layout as the GEMM output tensor.
* And the supplied sum tensor will be added to the tensor before activation.
*/
class MLAS_SBGEMM_ACTIVATION_PROCESSOR : public MLAS_SBGEMM_POSTPROCESSOR
{
public:
MLAS_SBGEMM_ACTIVATION_PROCESSOR(const MLAS_ACTIVATION& Activation, const float* SumBuf = nullptr)
: Activation_(Activation), SumBuf_(SumBuf)
{
}
void Process(float* C, size_t StartM, size_t StartN, size_t CountM, size_t CountN, size_t ldc)
const override;
private:
const MLAS_ACTIVATION& Activation_;
const float* SumBuf_;
};
/**
* @brief Data parameters for bfloat16 precision GEMM routine
* All except C are [in] parameters
*/
struct MLAS_SBGEMM_DATA_PARAMS {
const void* A = nullptr; /**< address of A */
const void* B = nullptr; /**< address of B */
const float* Bias = nullptr; /**< address of Bias, vector size N */
float* C = nullptr; /**< address of result matrix */
size_t lda = 0; /**< leading dimension of A */
size_t ldb = 0; /**< leading dimension of B, 0 when B is pre-packed*/
size_t ldc = 0; /**< leading dimension of C*/
const MLAS_SBGEMM_POSTPROCESSOR* OutputProcessor = nullptr;
bool AIsfp32 = false; /**< matrix A is fp32, needs to be converted to bf16*/
bool BIsfp32 = false; /**< matrix B is fp32, needs to be converted to bf16*/
bool ZeroMode = true; /**< when true: C = A*B + Bias (if Bias != nullptr);
when false: C += A*B and Bias is ignored */
bool BIsPacked = false; /**< Whether B is pre-packed */
};
/**
* @brief Bfloat16 precision Batched GEMM: C = A * B + Bias
* Either B can be either fp32 or bf16
*
* Note: We only support uniform batching, so shapes and types of the
* input must be same across all parameter blocks.
*
* @param[in] TransA Supplies the transpose operation for matrix A.
* @param[in] TransB Supplies the transpose operation for matrix B.
* @param[in] M row size of matrix A and C
* @param[in] N column size of matrix B and C
* @param[in] K column size of matrix A and row size of matrix B
* @param[in] BatchN number of batches
* @param[inout] DataParams An array (size BatchN) of parameter blocks
* @param[in] ThreadPool
* @param[in] BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
* @return
*/
void MLASCALL
MlasSBGemmBatch(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
const size_t M,
const size_t N,
const size_t K,
const size_t BatchN,
const MLAS_SBGEMM_DATA_PARAMS* DataParams,
MLAS_THREADPOOL* ThreadPool,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
/**
* @brief For bfloat16 precision GEMM, returns size of the
* packing buffer needed for right hand side
* @param[in] TransA Supplies the transpose operation for matrix A.
* @param[in] TransB Supplies the transpose operation for matrix B.
* @param[in] BIsfp32 Is matrix B datatype FP32
* @param[in] N Number of columns
* @param[in] K Number of rows
* @param[in] BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
* @return size of the packing buffer,
* 0 if operation not supported
*/
size_t MLASCALL
MlasSBGemmPackBSize(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
bool BIsfp32,
size_t N,
size_t K,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
/**
* @brief For bfloat16 precision GEMM, convert the float matrix B
* to blfoat16 precision and pack it into a packing buffer
*
* @param[in] TransA Supplies the transpose operation for matrix A.
* @param[in] TransB Supplies the transpose operation for matrix B.
* @param[in] BIsfp32 Is matrix B datatype FP32
* @param[in] N Number of columns
* @param[in] K Number of rows
* @param[in] B Address of matrix B
* @param[in] ldb leading dimension of input matrix B
* @param[out] PackedB Address of the packed matrix
* @param[in] BackendKernelSelectorConfig Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
*/
void MLASCALL
MlasSBGemmConvertPackB(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
bool BIsfp32,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
);
#endif
/**
* @brief Indirect Depthwise convolution for fp16
* @param Input Supplies the indirect buffer for NHWC input
* @param Filter Supplies the address for filter tensor
* @param Bias Supplies the address for 1D bias tensor B, has size of M
* @param Output Supplies the address for the result tensor
* @param Channels # of input channels
* @param OutputCount # of output pixels
* @param KernelSize # kernel size
* @return
*/
void
MLASCALL
MlasConvDepthwise(
const MLAS_FP16* const* Input,
const MLAS_FP16* Filter,
const MLAS_FP16* Bias,
MLAS_FP16* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize,
MLAS_HALF_GEMM_POSTPROCESSOR* PostProc
);
inline
void
MlasTranspose(
const MLAS_FP16* Input,
MLAS_FP16* Output,
size_t M,
size_t N,
MLAS_THREADPOOL* ThreadPool
)
{
MlasTranspose(
reinterpret_cast<const uint16_t*>(Input),
reinterpret_cast<uint16_t*>(Output),
M,
N,
ThreadPool);
}
#ifdef MLAS_F16VEC_INTRINSICS_SUPPORTED
/**
* @brief Max Pooling for fp16 NHWC
* @param Input Indirect buffer to activations
* @param Output Address of the result tensor
* @param Channels C in NHWC
* @param OutputCount Number of output pixels
* @param KernelSize Size of the kernel
* @return
*/
void
MLASCALL
MlasNhwcMaxPool(
const MLAS_FP16* const* Input,
MLAS_FP16* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
/**
* @brief Avg Pooling for fp16 nhwc
* @param Input Indirect buffer to activations
* @param Output Address of the output data
* @param Channels C in NHWC
* @param OutputCount Number of output pixels
* @param KernelSize size of the kernel
* @return
*/
void
MLASCALL
MlasNhwcAvgPool(
const MLAS_FP16* const* Input,
MLAS_FP16* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
#endif
struct MlasFlashAttentionThreadedArgs {
int batch_size;
int num_heads;
int q_sequence_length;
int kv_sequence_length;
int qk_head_size;
int v_head_size;
int q_block_size;
int kv_block_size;
float scale;
int thread_count;
float* buffer;
size_t buffer_size_per_thread;
const float* query;
const float* key;
const float* value;
float* output;
};
/**
* @brief Per-thread worker function for fp32 Flash Attention
* @param thread_id Thread index
* @param args Arguments
* @return
*/
void
MLASCALL
MlasFlashAttention(
MlasFlashAttentionThreadedArgs* args,
MLAS_THREADPOOL* ThreadPool
);
/**
* @brief Enumeration of supported GELU algorithm variants.
*
* MlasGeluErf - Exact GELU implementation using the error function (erf).
* MlasGeluTanh - Approximate GELU implementation using tanh-based formulation.
*/
typedef enum MLAS_GELU_ALGORITHM {
MlasGeluErf = 0,
MlasGeluTanh = 1
} MLAS_GELU_ALGORITHM;
/**
* @brief Computes element-wise FP16 error function (erf).
*
* This routine computes:
* Output[i] = erf(Input[i])
* for N elements. Depending on platform capabilities, this may use
* vectorized FP16 intrinsics or fall back to a scalar FP32 conversion path.
*
* @param Input Pointer to input buffer of N FP16 elements.
* @param Output Pointer to output buffer of N FP16 elements.
* @param Input_tmp_fp32 Pointer to caller-allocated scratch buffer of N floats
* for FP32 input conversion (used only on fallback path).
* @param Output_tmp_fp32 Pointer to caller-allocated scratch buffer of N floats
* for FP32 output conversion (used only on fallback path).
* @param N Number of elements to process.
*/
void
MLASCALL
MlasComputeFP16Erf(
const MLAS_FP16* Input,
MLAS_FP16* Output,
float* Input_tmp_fp32,
float* Output_tmp_fp32,
size_t N
);
/**
* @brief Computes element-wise FP16 GELU activation.
*
* This routine computes:
*
* If algo == MlasGeluTanh (approximate):
* GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
*
* If algo == MlasGeluErf (exact):
* GELU(x) = 0.5 * x * (1 + erf(x / sqrt(2)))
*
* Depending on platform capabilities, this may use vectorized FP16 kernels
* (SVE/NEON) or fall back to a scalar FP32 conversion path.
*
* @param input Pointer to input buffer of FP16 elements.
* @param output Pointer to output buffer of FP16 elements.
* @param temp Temporary scratch buffer of at least 'count' FP16 elements.
* Required by certain vectorized implementations. May be unused
* in scalar fallback paths.
* @param count Number of elements to process.
* @param algo GELU algorithm variant (exact erf or tanh approximation).
*/
void
MLASCALL
MlasComputeFP16Gelu(
const MLAS_FP16* input,
MLAS_FP16* output,
MLAS_FP16* temp,
size_t count,
MLAS_GELU_ALGORITHM algo
);
3rdparty/mlas/inc/mlas_float16.h
Vendored
+115
View File
@@ -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
3rdparty/mlas/lib/aarch64/SgemmKernelNeon.S
Vendored
+482
View File
@@ -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
3rdparty/mlas/lib/aarch64/SgemvKernelNeon.S
Vendored
+303
View File
@@ -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
3rdparty/mlas/lib/aarch64/asmmacro.h
Vendored
+95
View File
@@ -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
3rdparty/mlas/lib/arm/sgemmc.cpp
Vendored
+531
View File
@@ -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);
}
3rdparty/mlas/lib/arm64/SgemmKernelNeon.asm
Vendored
+502
View File
@@ -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
3rdparty/mlas/lib/arm64ec/SgemmKernelNeon.asm
Vendored
+466
View File
@@ -0,0 +1,466 @@
;++
;
; Copyright (c) Microsoft Corporation. All rights reserved.
;
; Licensed under the MIT License.
;
; Module Name:
;
; SgemmKernelNeon.asm
;
; Abstract:
;
; This module implements the kernels for the single precision matrix/matrix
; multiply operation (SGEMM).
;
;--
#include "kxarm64.h"
TEXTAREA
;
; ClearRowAccumulators
;
; Generates the code to clear the accumulators for a single row of the output
; block.
;
MACRO
ClearRowAccumulators $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
movi $Vec1Reg..16b,#0
movi $Vec2Reg..16b,#0
IF $Columns > 8
movi $Vec3Reg..16b,#0
movi $Vec4Reg..16b,#0
ENDIF
MEND
;
; ClearBlockAccumulators
;
; Generates the code to clear the accumulators for a single row of the output
; block.
;
MACRO
ClearBlockAccumulators $Columns, $Rows
ClearRowAccumulators $Columns, v8, v9, v10, v11
IF $Rows >= 2
ClearRowAccumulators $Columns, v12, v13, v14, v15
ENDIF
MEND
;
; LoadMatrixAElementsBy4
; LoadMatrixAElementsBy1
;
; Generates the code to load 1 or 4 elements from matrix A.
;
MACRO
LoadMatrixAElementsBy4 $Rows
ldr v2,[x0],#16
IF $Rows >= 2
ldr v3,[x10],#16
ENDIF
MEND
MACRO
LoadMatrixAElementsBy1 $Rows
ldr s2,[x0],#4
IF $Rows >= 2
ldr s3,[x10],#4
ENDIF
MEND
;
; MultiplyAccumulateRow
;
; Generates the code to multiply and accumulate a single row of the output
; block.
;
MACRO
MultiplyAccumulateRow $Columns, $MatrixAReg, $Broadcast, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
fmla $Vec1Reg..4s,v4.4s,$MatrixAReg..s[$Broadcast]
fmla $Vec2Reg..4s,v5.4s,$MatrixAReg..s[$Broadcast]
IF $Columns > 8
fmla $Vec3Reg..4s,v6.4s,$MatrixAReg..s[$Broadcast]
fmla $Vec4Reg..4s,v7.4s,$MatrixAReg..s[$Broadcast]
ENDIF
MEND
;
; MultiplyAccumulateBlock
;
; Generates the code to multiply and accumulate into the output block.
;
MACRO
MultiplyAccumulateBlock $Columns, $Rows, $Broadcast
MultiplyAccumulateRow $Columns, v2, $Broadcast, v8, v9, v10, v11
IF $Rows >= 2
MultiplyAccumulateRow $Columns, v3, $Broadcast, v12, v13, v14, v15
ENDIF
MEND
;
; ComputeBlockLoop
;
; Generates the code to loop over K entries of the input matrices to produce
; the output block.
;
MACRO
ComputeBlockLoop $Mode, $Columns, $Rows
ClearBlockAccumulators $Columns, $Rows
IF $Rows >= 2
add x10,x0,x6 lsl #2 ; compute matrix A plus 1 row
ENDIF
sub x9,x3,#4 ; decrement block count to process
tbnz x9,#63,$Mode.ProcessRemaining$Columns.x$Rows.Blocks
$Mode.Compute$Columns.x$Rows.BlockBy4Loop
LoadMatrixAElementsBy4 $Rows
ldp v4,v5,[x1],#64*4
IF $Columns > 8
ldp v6,v7,[x1,#-56*4]
ENDIF
MultiplyAccumulateBlock $Columns,$Rows,0
ldp v4,v5,[x1,#-48*4]
IF $Columns > 8
ldp v6,v7,[x1,#-40*4]
ENDIF
MultiplyAccumulateBlock $Columns,$Rows,1
ldp v4,v5,[x1,#-32*4]
IF $Columns > 8
ldp v6,v7,[x1,#-24*4]
ENDIF
MultiplyAccumulateBlock $Columns,$Rows,2
ldp v4,v5,[x1,#-16*4]
IF $Columns > 8
ldp v6,v7,[x1,#-8*4]
ENDIF
MultiplyAccumulateBlock $Columns,$Rows,3
sub x9,x9,#4
tbz x9,#63,$Mode.Compute$Columns.x$Rows.BlockBy4Loop
$Mode.ProcessRemaining$Columns.x$Rows.Blocks
add x9,x9,#4 ; correct for over-subtract above
cbz x9,$Mode.Output$Columns.x$Rows.Block
$Mode.Compute$Columns.x$Rows.BlockBy1Loop
LoadMatrixAElementsBy1 $Rows
ldp v4,v5,[x1],#16*4
IF $Columns > 8
ldp v6,v7,[x1,#-8*4]
ENDIF
MultiplyAccumulateBlock $Columns,$Rows,0
sub x9,x9,#1
cbnz x9,$Mode.Compute$Columns.x$Rows.BlockBy1Loop
$Mode.Output$Columns.x$Rows.Block
MEND
;
; MultiplyAlphaRow
;
; Generates the code to multiply a single row of the output block by the alpha
; value.
;
MACRO
MultiplyAlphaRow $Columns, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF $Columns <= 4
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
ELIF $Columns <= 8
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
ELIF $Columns <= 12
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
ELSE
fmul $Vec1Reg..4s,$Vec1Reg..4s,v0.s[0]
fmul $Vec2Reg..4s,$Vec2Reg..4s,v0.s[0]
fmul $Vec3Reg..4s,$Vec3Reg..4s,v0.s[0]
fmul $Vec4Reg..4s,$Vec4Reg..4s,v0.s[0]
ENDIF
MEND
;
; MultiplyAlphaBlock
;
; Generates the code to multiply the output block by the alpha value.
;
MACRO
MultiplyAlphaBlock $Columns, $Rows
MultiplyAlphaRow $Columns, v8, v9, v10, v11
IF $Rows >= 2
MultiplyAlphaRow $Columns, v12, v13, v14, v15
ENDIF
MEND
;
; OutputRow1Element
; OutputRow2Element
; OutputRow4Element
; OutputRow8Element
; OutputRow16Element
;
; Generates the code to store elements to the output block.
;
MACRO
OutputRow1Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF "$Mode"=="Add"
ld1 {v4.s}[0],[$AddrReg]
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
st1 {v4.s}[0],[$AddrReg] ; post-increment not needed for last element
ELSE
st1 {$Vec1Reg..s}[0],[$AddrReg] ; post-increment not needed for last element
ENDIF
MEND
MACRO
OutputRow2Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF "$Mode"=="Add"
ld1 {v4.2s},[$AddrReg]
fmla v4.2s,$Vec1Reg..2s,v0.s[0]
st1 {v4.2s},[$AddrReg],#2*4
ELSE
st1 {$Vec1Reg..2s},[$AddrReg],#2*4
ENDIF
dup $Vec1Reg..4s,$Vec1Reg..s[2] ; shift remaining elements down
MEND
MACRO
OutputRow4Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF "$Mode"=="Add"
ld1 {v4.4s},[$AddrReg]
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
st1 {v4.4s},[$AddrReg],#4*4
ELSE
st1 {$Vec1Reg..4s},[$AddrReg],#4*4
ENDIF
mov $Vec1Reg..16b,$Vec2Reg..16b ; shift remaining elements down
MEND
MACRO
OutputRow8Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF "$Mode"=="Add"
ldp v4,v5,[$AddrReg]
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
stp v4,v5,[$AddrReg],#8*4
ELSE
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#8*4
ENDIF
mov $Vec1Reg..16b,$Vec3Reg..16b ; shift remaining elements down
mov $Vec2Reg..16b,$Vec4Reg..16b
MEND
MACRO
OutputRow16Element $Mode, $AddrReg, $Vec1Reg, $Vec2Reg, $Vec3Reg, $Vec4Reg
IF "$Mode"=="Add"
ldp v4,v5,[$AddrReg]
ldp v6,v7,[$AddrReg,#8*4]
fmla v4.4s,$Vec1Reg..4s,v0.s[0]
fmla v5.4s,$Vec2Reg..4s,v0.s[0]
fmla v6.4s,$Vec3Reg..4s,v0.s[0]
fmla v7.4s,$Vec4Reg..4s,v0.s[0]
stp v4,v5,[$AddrReg],#16*4
stp v6,v7,[$AddrReg,#-8*4]
ELSE
stp $Vec1Reg.,$Vec2Reg.,[$AddrReg],#16*4
stp $Vec3Reg.,$Vec4Reg.,[$AddrReg,#-8*4]
ENDIF
MEND
;
; OutputBlock
;
; Generates the code to store the output block.
;
MACRO
OutputBlock $Mode, $Columns, $Rows
OutputRow$Columns.Element $Mode, x2, v8, v9, v10, v11
IF $Rows >= 2
OutputRow$Columns.Element $Mode, x11, v12, v13, v14, v15
ENDIF
MEND
;
; ProcessRows
;
; Generates the code to process a compute and store the output block for a
; fixed number of rows.
;
MACRO
ProcessRows $Mode, $Rows
mov x4,#$Rows ; return number of rows handled
cmp x5,#8
ble $Mode.ProcessRemainingCountN$Rows
$Mode.ProcessNextColumnLoop16x$Rows
ComputeBlockLoop $Mode,16,$Rows
IF "$Mode"=="Zero"
MultiplyAlphaBlock 16,$Rows
ENDIF
sub x5,x5,#16
tbnz x5,#63,$Mode.OutputMasked16x$Rows.Block
OutputBlock $Mode,16,$Rows
mov x0,x8 ; reload matrix A
cmp x5,#8
bgt $Mode.ProcessNextColumnLoop16x$Rows
cbz x5,$Mode.ExitKernel
$Mode.ProcessRemainingCountN$Rows
ComputeBlockLoop $Mode,8,$Rows
IF "$Mode"=="Zero"
MultiplyAlphaBlock 8,$Rows
ENDIF
$Mode.OutputMasked16x$Rows.Block
tbz x5,#3,$Mode.OutputRemaining7x$Rows.Block
OutputBlock $Mode,8,$Rows
$Mode.OutputRemaining7x$Rows.Block
tbz x5,#2,$Mode.OutputRemaining3x$Rows.Block
OutputBlock $Mode,4,$Rows
$Mode.OutputRemaining3x$Rows.Block
tbz x5,#1,$Mode.OutputRemaining1x$Rows.Block
OutputBlock $Mode,2,$Rows
$Mode.OutputRemaining1x$Rows.Block
tbz x5,#0,$Mode.ExitKernel
OutputBlock $Mode,1,$Rows
MEND
SUBT "SGEMM kernel"
;++
;
; Routine Description:
;
; This routine is an inner kernel to compute matrix multiplication for a
; set of rows.
;
; Arguments:
;
; A (x0) - Supplies the address of matrix A.
;
; B (x1) - Supplies the address of matrix B. The matrix data has been packed
; using MlasSgemmCopyPackB or MlasSgemmTransposePackB.
;
; C (x2) - Supplies the address of matrix C.
;
; CountK (x3) - Supplies the number of columns from matrix A and the number
; of rows from matrix B to iterate over.
;
; CountM (x4) - Supplies the maximum number of rows that can be processed for
; matrix A and matrix C. The actual number of rows handled for this
; invocation depends on the kernel implementation.
;
; CountN (x5) - Supplies the number of columns from matrix B and matrix C to
; iterate over.
;
; lda (x6) - Supplies the first dimension of matrix A.
;
; ldc (x7) - Supplies the first dimension of matrix C.
;
; Alpha (s0) - Supplies the scalar multiplier (see SGEMM definition).
;
; Return Value:
;
; Returns the number of rows handled.
;
;--
MACRO
SgemmKernelNeonFunction $Mode
NESTED_ENTRY_COMDAT A64NAME(MlasSgemmKernel$Mode)
PROLOG_SAVE_REG_PAIR d8,d9,#-64!
PROLOG_SAVE_REG_PAIR d10,d11,#16
PROLOG_SAVE_REG_PAIR d12,d13,#32
PROLOG_SAVE_REG_PAIR d14,d15,#48
add x11,x2,x7 lsl #2 ; compute matrix C plus 1 row
mov x8,x0 ; save matrix A
;
; Process 2 rows of the matrices.
;
cmp x4,#2
blt $Mode.ProcessCountMLessThan2
ProcessRows $Mode,2
;
; Restore non-volatile registers and return.
;
$Mode.ExitKernel
mov x0,x4
EPILOG_RESTORE_REG_PAIR d14,d15,#48
EPILOG_RESTORE_REG_PAIR d12,d13,#32
EPILOG_RESTORE_REG_PAIR d10,d11,#16
EPILOG_RESTORE_REG_PAIR d8,d9,#64!
EPILOG_RETURN
;
; Process 1 row of the matrices.
;
$Mode.ProcessCountMLessThan2
ProcessRows $Mode,1
b $Mode.ExitKernel
NESTED_END
MEND
SgemmKernelNeonFunction Zero
SgemmKernelNeonFunction Add
END
3rdparty/mlas/lib/compute.cpp
Vendored
+1160
View File
@@ -0,0 +1,1160 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
compute.cpp
Abstract:
This module implements miscellaneous computation routines.
Our usage requires building platform specific versions of the algorithm to
target different instruction sets. The implementation below targets the
base instruction set (typically SSE2) while assembly implementations target
newer instruction sets (such as FMA3).
--*/
#include "mlasi.h"
#include "softmax.h"
//
// Bundles the constants for use by kernels written in assembly.
//
MLAS_INTERNAL_DATA const struct {
float LowerRange;
float UpperRange;
float LowerRangeSumExp;
float UpperRangeSumExp;
float RoundingBias;
float Log2Reciprocal;
float Log2High;
float Log2Low;
float poly_0;
float poly_1;
float poly_2;
float poly_3;
float poly_4;
float poly_56;
int32_t MinimumExponent;
int32_t MaximumExponent;
} MlasExpConstants = {
-103.9720840454f,
88.7762626647950f,
-88.3762626647949f,
88.3762626647949f,
MLAS_ROUNDING_BIAS_MAGIC,
1.44269504088896341f,
-6.93145752e-1f,
-1.42860677e-6f,
0x1.694000p-10,
0x1.125edcp-7,
0x1.555b5ap-5,
0x1.555450p-3,
0x1.fffff6p-2,
0x1.000000p+0,
int32_t(0xC1000000),
int32_t(0x3F800000),
};
MLAS_INTERNAL_DATA const float MlasMinimumF32Value = std::numeric_limits<float>::lowest();
//
// Define the parameters to execute segments of a softmax operation on worker
// threads.
//
template <typename T>
struct MLAS_SOFTMAX_WORK_BLOCK {
ptrdiff_t ThreadCountN;
bool LogSoftmax;
bool SmoothSoftmax;
float Sink;
const T* Input;
T* Output;
size_t N;
size_t D;
};
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasComputeExpVector(
MLAS_FLOAT32X4 Vector
)
/*++
Routine Description:
This routine computes the exponential function for the supplied vector.
This merges ideas from multiple vectorized expf() implementations:
1. The original polynomials of expf() are extracted from MlasComputeErf, which
was based on an answer to the following Stack Overflow post:
https://stackoverflow.com/questions/35148198/efficient-faithfully-rounded-implementation-of-error-function-erff
2. The author of the answer further refined the polynomials at:
https://forums.developer.nvidia.com/t/a-more-accurate-performance-competitive-implementation-of-expf/47528/5
Using these polynomials yields even closer results to the Microsoft
UCRT version of std::expf() than the values from the above post.
3. XNNPACK has a further useful refinement to extend the effective
range of results from [-88.376, 88.376] to [-103.972, 88.776] by
splitting the step of exponent reconstruction into two pieces. This
yields results similar to an AVX512 implementation using VSCALEFPS.
Arguments:
Vector - Supplies the values to operate on.
Return Value:
Returns the exponential function of the input.
--*/
{
Vector = MlasClampFloat32x4(Vector, MlasExpConstants.LowerRange, MlasExpConstants.UpperRange);
//
// Range reduction of the input by computing "(2 ^ m) * exp(reduced)".
//
const auto RoundingBias = MlasBroadcastFloat32x4(MlasExpConstants.RoundingBias);
auto biased = MlasMultiplyAddFloat32x4(Vector, MlasExpConstants.Log2Reciprocal, RoundingBias);
auto m = MlasSubtractFloat32x4(biased, RoundingBias);
Vector = MlasMultiplyAddFloat32x4(m, MlasExpConstants.Log2High, Vector);
Vector = MlasMultiplyAddFloat32x4(m, MlasExpConstants.Log2Low, Vector);
//
// Compute the scaling factors used to reconstruct the "(2 ^ m)" value
// from above. To cover the entire single precision floating point range,
// two scaling factors are needed to handle exponents [-150, 128].
//
const auto MinimumExponent = MlasBroadcastInt32x4(MlasExpConstants.MinimumExponent);
const auto MaximumExponent = MlasBroadcastInt32x4(MlasExpConstants.MaximumExponent);
auto overflow = MlasShiftLeftInt32x4<23>(MlasReinterpretAsInt32x4(biased));
auto normal = overflow;
#if defined(MLAS_SSE2_INTRINSICS)
// N.B. PMINSD/PMAXSD were not added until SSE 4.1, but the lower 16 bits
// are zero, so they can be ignored for this computation, so use PMINSW/PMAXSW
// instead.
normal = _mm_min_epi16(normal, MaximumExponent);
normal = _mm_max_epi16(normal, MinimumExponent);
#elif defined(MLAS_LSX_INTRINSICS)
normal = __lsx_vmin_h(normal, MaximumExponent);
normal = __lsx_vmax_h(normal, MinimumExponent);
#else
normal = MlasMinimumInt32x4(normal, MaximumExponent);
normal = MlasMaximumInt32x4(normal, MinimumExponent);
#endif
overflow = MlasSubtractInt32x4(overflow, normal);
overflow = MlasAddInt32x4(overflow, MaximumExponent);
normal = MlasAddInt32x4(normal, MaximumExponent);
//
// Compute the polynomial approximation of exp(reduced) and reconstruct
// the final result using the above scaling factors. The final term of
// the polynomial (poly_6=1.0f) is merged as the multiply/add of the
// overflow exponent (reference XNNPACK).
//
auto p = MlasBroadcastFloat32x4(MlasExpConstants.poly_0);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_1);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_2);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_3);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_4);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_56);
Vector = MlasMultiplyFloat32x4(Vector, MlasReinterpretAsFloat32x4(overflow));
p = MlasMultiplyAddFloat32x4(p, Vector, MlasReinterpretAsFloat32x4(overflow));
p = MlasMultiplyFloat32x4(p, MlasReinterpretAsFloat32x4(normal));
return p;
}
void
MLASCALL
MlasComputeExpF32Kernel(
const float* Input,
float* Output,
size_t N
)
/*++
Routine Description:
This routine implements the generic kernel for the exponential function.
Arguments:
Input - Supplies the input buffer.
Output - Supplies the output buffer.
N - Supplies the number of elements to process.
Return Value:
None.
--*/
{
while (N > 0) {
MLAS_FLOAT32X4 Vector;
if (N >= 4) {
Vector = MlasLoadFloat32x4(Input);
} else {
#if defined(MLAS_SSE2_INTRINSICS)
// N.B. SSE2 lacks a broadcast load instruction, so avoid a shuffle
// and use zeroes for the upper elements.
Vector = _mm_load_ss(Input);
#elif defined(MLAS_LSX_INTRINSICS)
Vector = (MLAS_FLOAT32X4)__lsx_vldrepl_w(Input, 0);
#else
Vector = MlasBroadcastFloat32x4(Input);
#endif
}
Vector = MlasComputeExpVector(Vector);
if (N >= 4) {
MlasStoreFloat32x4(Output, Vector);
Input += 4;
Output += 4;
N -= 4;
} else {
MlasStoreLaneFloat32x4<0>(Output, Vector);
Input += 1;
Output += 1;
N -= 1;
}
}
}
template <>
void
MLASCALL
MlasComputeExp<float>(
const float* Input,
float* Output,
size_t N
)
/*++
Routine Description:
This routine computes the exponential function.
N.B. This implementation supports in place updates of the output buffer.
Arguments:
Input - Supplies the input buffer.
Output - Supplies the output buffer.
N - Supplies the number of elements to process.
Return Value:
None.
--*/
{
#if defined(MLAS_TARGET_AMD64)
GetMlasPlatform().ComputeExpF32Kernel(Input, Output, N);
#else
MlasComputeExpF32Kernel(Input, Output, N);
#endif
}
template <>
void MLASCALL
MlasComputeExp<MLAS_FP16>(
const MLAS_FP16* Input,
MLAS_FP16* Output,
size_t N
) {
const auto* dispatch = GetMlasPlatform().SoftmaxDispatch;
if (dispatch == nullptr || dispatch->Exp_Fp16 == nullptr) {
MLAS_THROW_EX(std::runtime_error, "Exp_Fp16 is not supported.");
}
dispatch->Exp_Fp16(Input, Output, N);
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasComputeSumExpVector(
MLAS_FLOAT32X4 Vector,
MLAS_FLOAT32X4 NegativeMaximumVector
)
/*++
Routine Description:
This routine computes the exponential function for the supplied vector.
This function handles a narrower range of inputs compared to
MlasComputeExpVector in order to improve efficiency.
Arguments:
Vector - Supplies the values to operate on.
NegativeMaximumVector - Supplies the broadcasted negative maximum
value that is added to each element before computing the exponential
function.
Return Value:
Returns the exponential function of the input.
--*/
{
//
// Subtract the maximum value from every element.
//
// N.B. For each of use by the assembly kernels, this value has been negated
// so add the value instead.
//
Vector = MlasAddFloat32x4(Vector, NegativeMaximumVector);
//
// Clamp to the lower range of this function.
//
// The value should already be negative or equal to zero as every value has
// been reduced by the maximum value.
//
#if defined(MLAS_SSE2_INTRINSICS)
// N.B. MINPS and MAXPS propagates the value from the second vector if the
// value is a NaN.
#endif
Vector = MlasMaximumFloat32x4(MlasBroadcastFloat32x4(MlasExpConstants.LowerRangeSumExp), Vector);
//
// Range reduction of the input by computing "(2 ^ m) * exp(reduced)".
//
const auto RoundingBias = MlasBroadcastFloat32x4(MlasExpConstants.RoundingBias);
auto biased = MlasMultiplyAddFloat32x4(Vector, MlasExpConstants.Log2Reciprocal, RoundingBias);
auto m = MlasSubtractFloat32x4(biased, RoundingBias);
Vector = MlasMultiplyAddFloat32x4(m, MlasExpConstants.Log2High, Vector);
Vector = MlasMultiplyAddFloat32x4(m, MlasExpConstants.Log2Low, Vector);
//
// Compute the scaling factor used to reconstruct the "(2 ^ m)" value
// from above. The effective range of this function is smaller than
// MlasComputeExp to reduce the number of operations.
//
auto normal = MlasShiftLeftInt32x4<23>(MlasReinterpretAsInt32x4(biased));
normal = MlasAddInt32x4(normal, MlasBroadcastInt32x4(MlasExpConstants.MaximumExponent));
//
// Compute the polynomial approximation of exp(reduced) and reconstruct
// the final result using the above scale factor.
//
auto p = MlasBroadcastFloat32x4(MlasExpConstants.poly_0);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_1);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_2);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_3);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_4);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_56);
p = MlasMultiplyAddFloat32x4(p, Vector, MlasExpConstants.poly_56);
p = MlasMultiplyFloat32x4(p, MlasReinterpretAsFloat32x4(normal));
return p;
}
float
MLASCALL
MlasComputeSumExpF32Kernel(
const float* Input,
float* Output,
size_t N,
const float* NegativeMaximum
)
/*++
Routine Description:
This routine implements the generic kernel for the sum of exponential
functions.
Arguments:
Input - Supplies the input buffer.
Output - Optionally supplies the output buffer. When used for Softmax,
the output buffer is used to store the intermediate exp() results. When
used for LogSoftmax, the intermediate exp() results are not required.
N - Supplies the number of elements to process.
NegativeMaximum - Supplies the address of the negative maximum
value that is added to each element before computing the exponential
function.
Return Value:
Returns the sum of the exponential functions.
--*/
{
MLAS_FLOAT32X4 NegativeMaximumVector = MlasBroadcastFloat32x4(*NegativeMaximum);
float Accumulator = 0.0f;
if (N >= 4) {
MLAS_FLOAT32X4 AccumulatorVector = MlasZeroFloat32x4();
#if !defined(MLAS_SSE2_INTRINSICS)
//
// Unroll the loop for architectures that can benefit from improved
// instruction level parallelism.
//
// N.B. The extra code size is not worth the benefit for SSE2 as the
// MLAS_TARGET_AMD64 build already has specialized AVX2/AVX512F kernels
// that do this.
//
while (N >= 8) {
MLAS_FLOAT32X4 Vector0 = MlasLoadFloat32x4(Input);
MLAS_FLOAT32X4 Vector1 = MlasLoadFloat32x4(Input + 4);
Vector0 = MlasComputeSumExpVector(Vector0, NegativeMaximumVector);
Vector1 = MlasComputeSumExpVector(Vector1, NegativeMaximumVector);
AccumulatorVector = MlasAddFloat32x4(AccumulatorVector, Vector0);
AccumulatorVector = MlasAddFloat32x4(AccumulatorVector, Vector1);
if (Output != nullptr) {
MlasStoreFloat32x4(Output, Vector0);
MlasStoreFloat32x4(Output + 4, Vector1);
Output += 8;
}
Input += 8;
N -= 8;
}
#endif
while (N >= 4) {
MLAS_FLOAT32X4 Vector = MlasLoadFloat32x4(Input);
Vector = MlasComputeSumExpVector(Vector, NegativeMaximumVector);
AccumulatorVector = MlasAddFloat32x4(AccumulatorVector, Vector);
if (Output != nullptr) {
MlasStoreFloat32x4(Output, Vector);
Output += 4;
}
Input += 4;
N -= 4;
}
Accumulator = MlasReduceAddFloat32x4(AccumulatorVector);
}
while (N > 0) {
#if defined(MLAS_SSE2_INTRINSICS)
// N.B. SSE2 lacks a broadcast load instruction, so avoid a shuffle and
// use zeroes for the upper elements.
MLAS_FLOAT32X4 Vector = _mm_load_ss(Input);
#elif defined(MLAS_LSX_INTRINSICS)
MLAS_FLOAT32X4 Vector = (MLAS_FLOAT32X4)__lsx_vldrepl_w(Input, 0);
#else
MLAS_FLOAT32X4 Vector = MlasBroadcastFloat32x4(Input);
#endif
Vector = MlasComputeSumExpVector(Vector, NegativeMaximumVector);
Accumulator += MlasExtractLaneFloat32x4<0>(Vector);
if (Output != nullptr) {
MlasStoreLaneFloat32x4<0>(Output, Vector);
Output += 1;
}
Input += 1;
N -= 1;
}
return Accumulator;
}
float
MLASCALL
MlasReduceMaximumF32Kernel(
const float* Input,
size_t N
)
/*++
Routine Description:
This routine implements the generic kernel to find the maximum value of
the supplied buffer.
Arguments:
Input - Supplies the input buffer.
N - Supplies the number of elements to process.
Return Value:
Returns the maximum value of the supplied buffer.
--*/
{
float Maximum = MlasMinimumF32Value;
if (N >= 4) {
MLAS_FLOAT32X4 MaximumVector0 = MlasBroadcastFloat32x4(Maximum);
if (N >= 16) {
MLAS_FLOAT32X4 MaximumVector1 = MaximumVector0;
MLAS_FLOAT32X4 MaximumVector2 = MaximumVector0;
MLAS_FLOAT32X4 MaximumVector3 = MaximumVector0;
while (N >= 16) {
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MlasLoadFloat32x4(Input));
MaximumVector1 = MlasMaximumFloat32x4(MaximumVector1, MlasLoadFloat32x4(Input + 4));
MaximumVector2 = MlasMaximumFloat32x4(MaximumVector2, MlasLoadFloat32x4(Input + 8));
MaximumVector3 = MlasMaximumFloat32x4(MaximumVector3, MlasLoadFloat32x4(Input + 12));
Input += 16;
N -= 16;
}
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MaximumVector1);
MaximumVector2 = MlasMaximumFloat32x4(MaximumVector2, MaximumVector3);
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MaximumVector2);
}
while (N >= 4) {
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MlasLoadFloat32x4(Input));
Input += 4;
N -= 4;
}
Maximum = MlasReduceMaximumFloat32x4(MaximumVector0);
}
while (N > 0) {
Maximum = std::max(Maximum, *Input);
Input += 1;
N -= 1;
}
return Maximum;
}
void
MLASCALL
MlasReduceMinimumMaximumF32Kernel(
const float* Input,
float* Min,
float* Max,
size_t N
)
{
float tmp_min = std::numeric_limits<float>::max();
float tmp_max = std::numeric_limits<float>::lowest();
if (N >= 4) {
MLAS_FLOAT32X4 MaximumVector0 = MlasBroadcastFloat32x4(tmp_max);
MLAS_FLOAT32X4 MinimumVector0 = MlasBroadcastFloat32x4(tmp_min);
if (N >= 16) {
MLAS_FLOAT32X4 MaximumVector1 = MaximumVector0;
MLAS_FLOAT32X4 MaximumVector2 = MaximumVector0;
MLAS_FLOAT32X4 MaximumVector3 = MaximumVector0;
MLAS_FLOAT32X4 MinimumVector1 = MinimumVector0;
MLAS_FLOAT32X4 MinimumVector2 = MinimumVector0;
MLAS_FLOAT32X4 MinimumVector3 = MinimumVector0;
while (N >= 16) {
MLAS_FLOAT32X4 InputVector0 = MlasLoadFloat32x4(Input);
MLAS_FLOAT32X4 InputVector1 = MlasLoadFloat32x4(Input + 4);
MLAS_FLOAT32X4 InputVector2 = MlasLoadFloat32x4(Input + 8);
MLAS_FLOAT32X4 InputVector3 = MlasLoadFloat32x4(Input + 12);
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, InputVector0);
MaximumVector1 = MlasMaximumFloat32x4(MaximumVector1, InputVector1);
MaximumVector2 = MlasMaximumFloat32x4(MaximumVector2, InputVector2);
MaximumVector3 = MlasMaximumFloat32x4(MaximumVector3, InputVector3);
MinimumVector0 = MlasMinimumFloat32x4(MinimumVector0, InputVector0);
MinimumVector1 = MlasMinimumFloat32x4(MinimumVector1, InputVector1);
MinimumVector2 = MlasMinimumFloat32x4(MinimumVector2, InputVector2);
MinimumVector3 = MlasMinimumFloat32x4(MinimumVector3, InputVector3);
Input += 16;
N -= 16;
}
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MaximumVector1);
MaximumVector2 = MlasMaximumFloat32x4(MaximumVector2, MaximumVector3);
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, MaximumVector2);
MinimumVector0 = MlasMinimumFloat32x4(MinimumVector0, MinimumVector1);
MinimumVector2 = MlasMinimumFloat32x4(MinimumVector2, MinimumVector3);
MinimumVector0 = MlasMinimumFloat32x4(MinimumVector0, MinimumVector2);
}
while (N >= 4) {
MLAS_FLOAT32X4 InputVector0 = MlasLoadFloat32x4(Input);
MaximumVector0 = MlasMaximumFloat32x4(MaximumVector0, InputVector0);
MinimumVector0 = MlasMinimumFloat32x4(MinimumVector0, InputVector0);
Input += 4;
N -= 4;
}
tmp_min = MlasReduceMinimumFloat32x4(MinimumVector0);
tmp_max = MlasReduceMaximumFloat32x4(MaximumVector0);
}
while (N > 0) {
tmp_max = std::max(tmp_max, *Input);
tmp_min = std::min(tmp_min, *Input);
Input += 1;
N -= 1;
}
*Min = tmp_min;
*Max = tmp_max;
}
void
MLASCALL
MlasComputeSoftmaxOutputF32Kernel(
float* Output,
size_t N,
const float* Parameters
)
/*++
Routine Description:
This routine implements the generic kernel to produce the final output for
the softmax operation.
Arguments:
Output - Supplies the output buffer.
N - Supplies the number of elements to process.
Parameters - Supplies an array containing the scale value.
Return Value:
None.
--*/
{
const float Scale = Parameters[0];
const MLAS_FLOAT32X4 ScaleVector = MlasBroadcastFloat32x4(Scale);
while (N >= 16) {
MLAS_FLOAT32X4 Vector0 = MlasMultiplyFloat32x4(ScaleVector, MlasLoadFloat32x4(Output));
MLAS_FLOAT32X4 Vector1 = MlasMultiplyFloat32x4(ScaleVector, MlasLoadFloat32x4(Output + 4));
MLAS_FLOAT32X4 Vector2 = MlasMultiplyFloat32x4(ScaleVector, MlasLoadFloat32x4(Output + 8));
MLAS_FLOAT32X4 Vector3 = MlasMultiplyFloat32x4(ScaleVector, MlasLoadFloat32x4(Output + 12));
MlasStoreFloat32x4(Output, Vector0);
MlasStoreFloat32x4(Output + 4, Vector1);
MlasStoreFloat32x4(Output + 8, Vector2);
MlasStoreFloat32x4(Output + 12, Vector3);
Output += 16;
N -= 16;
}
while (N >= 4) {
MlasStoreFloat32x4(Output, MlasMultiplyFloat32x4(ScaleVector, MlasLoadFloat32x4(Output)));
Output += 4;
N -= 4;
}
while (N > 0) {
*Output *= Scale;
Output += 1;
N -= 1;
}
}
void
MLASCALL
MlasComputeLogSoftmaxOutputF32Kernel(
const float* Input,
float* Output,
size_t N,
const float* Parameters
)
/*++
Routine Description:
This routine implements the generic kernel to produce the final output for
the log softmax operation.
Arguments:
Input - Supplies the input buffer.
Output - Supplies the output buffer.
N - Supplies the number of elements to process.
Parameters - Supplies an array containing the negative maximum and
logarithm values.
Return Value:
None.
--*/
{
const float NegativeMaximum = Parameters[0];
const float Logarithm = Parameters[1];
const MLAS_FLOAT32X4 NegativeMaximumVector = MlasBroadcastFloat32x4(NegativeMaximum);
const MLAS_FLOAT32X4 LogarithmVector = MlasBroadcastFloat32x4(Logarithm);
while (N >= 16) {
MLAS_FLOAT32X4 Vector0 = MlasLoadFloat32x4(Input);
MLAS_FLOAT32X4 Vector1 = MlasLoadFloat32x4(Input + 4);
MLAS_FLOAT32X4 Vector2 = MlasLoadFloat32x4(Input + 8);
MLAS_FLOAT32X4 Vector3 = MlasLoadFloat32x4(Input + 12);
Vector0 = MlasAddFloat32x4(Vector0, NegativeMaximumVector);
Vector1 = MlasAddFloat32x4(Vector1, NegativeMaximumVector);
Vector2 = MlasAddFloat32x4(Vector2, NegativeMaximumVector);
Vector3 = MlasAddFloat32x4(Vector3, NegativeMaximumVector);
Vector0 = MlasSubtractFloat32x4(Vector0, LogarithmVector);
Vector1 = MlasSubtractFloat32x4(Vector1, LogarithmVector);
Vector2 = MlasSubtractFloat32x4(Vector2, LogarithmVector);
Vector3 = MlasSubtractFloat32x4(Vector3, LogarithmVector);
MlasStoreFloat32x4(Output, Vector0);
MlasStoreFloat32x4(Output + 4, Vector1);
MlasStoreFloat32x4(Output + 8, Vector2);
MlasStoreFloat32x4(Output + 12, Vector3);
Input += 16;
Output += 16;
N -= 16;
}
while (N >= 4) {
MLAS_FLOAT32X4 Vector = MlasLoadFloat32x4(Input);
Vector = MlasAddFloat32x4(Vector, NegativeMaximumVector);
Vector = MlasSubtractFloat32x4(Vector, LogarithmVector);
MlasStoreFloat32x4(Output, Vector);
Input += 4;
Output += 4;
N -= 4;
}
while (N > 0) {
*Output = *Input + NegativeMaximum - Logarithm;
Input += 1;
Output += 1;
N -= 1;
}
}
template <typename T>
void
MlasComputeSoftmaxThreaded(
void* Context,
ptrdiff_t Index
);
template <>
void
MlasComputeSoftmaxThreaded<float>(
void* Context,
ptrdiff_t Index
)
/*++
Routine Description:
This routine is invoked from a worker thread to execute a segment of a
softmax or log softmax operation.
Arguments:
Context - Supplies the pointer to the context for the threaded operation.
ThreadId - Supplies the current index of the threaded operation.
Return Value:
None.
--*/
{
const auto* WorkBlock = (MLAS_SOFTMAX_WORK_BLOCK<float>*)Context;
//
// Partition the operation along the N dimension.
//
size_t n;
size_t CountN;
MlasPartitionWork(Index, WorkBlock->ThreadCountN, WorkBlock->N, &n, &CountN);
//
// Compute the softmax or log softmax function.
//
const size_t D = WorkBlock->D;
const bool LogSoftmax = WorkBlock->LogSoftmax;
const bool SmoothSoftmax = WorkBlock->SmoothSoftmax;
const float Sink = WorkBlock->Sink;
const float* Input = WorkBlock->Input + n * D;
float* Output = WorkBlock->Output + n * D;
#if defined(MLAS_SSE2_INTRINSICS)
// TODO: Use std::hardware_constructive_interference_size
constexpr size_t CacheLineSize = 64;
constexpr size_t ElementsPerCacheLine = CacheLineSize / sizeof(float);
#endif
while (CountN > 0) {
#if defined(MLAS_SSE2_INTRINSICS)
//
// Prefetch the next row of the input buffer.
//
for (size_t i = 0; i * ElementsPerCacheLine < D; i++) {
_mm_prefetch((char*)(Input + D) + i * CacheLineSize, _MM_HINT_T0);
}
#endif
//
// Find the maximum value for the row.
//
float Maximum;
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64) || defined(MLAS_USE_SVE) || defined(MLAS_TARGET_RISCV64)
Maximum = GetMlasPlatform().ReduceMaximumF32Kernel(Input, D);
#else
Maximum = MlasReduceMaximumF32Kernel(Input, D);
#endif
if (SmoothSoftmax && Sink > Maximum) {
Maximum = Sink;
}
float NegativeMaximum = -Maximum;
//
// Compute the exponential function for each element of the row (save to Temp if provided) and
// compute the sum of these exponential functions.
//
float* Temp = LogSoftmax ? nullptr : Output;
float Accumulation;
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_USE_SVE) || defined(MLAS_TARGET_RISCV64)
Accumulation = GetMlasPlatform().ComputeSumExpF32Kernel(Input, Temp, D, &NegativeMaximum);
#else
Accumulation = MlasComputeSumExpF32Kernel(Input, Temp, D, &NegativeMaximum);
#endif
if (SmoothSoftmax) {
Accumulation += expf(Sink + NegativeMaximum);
}
if (LogSoftmax) {
//
// Compute the log softmax output.
//
float Parameters[] = {NegativeMaximum, std::log(Accumulation)};
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64) || defined(MLAS_USE_SVE) || defined(MLAS_TARGET_RISCV64)
GetMlasPlatform().ComputeLogSoftmaxOutputF32Kernel(Input, Output, D, Parameters);
#else
MlasComputeLogSoftmaxOutputF32Kernel(Input, Output, D, Parameters);
#endif
} else {
//
// Normalize the softmax output.
//
float Parameters[] = {1.0f / Accumulation};
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64) || defined(MLAS_USE_SVE) || defined(MLAS_TARGET_RISCV64)
GetMlasPlatform().ComputeSoftmaxOutputF32Kernel(Output, D, Parameters);
#else
MlasComputeSoftmaxOutputF32Kernel(Output, D, Parameters);
#endif
}
Input += D;
Output += D;
CountN--;
}
}
template <>
void
MlasComputeSoftmaxThreaded<MLAS_FP16>(
void* Context,
ptrdiff_t Index
)
/*++
Routine Description:
This routine is invoked from a worker thread to execute a segment of a
softmax or log softmax operation.
Arguments:
Context - Supplies the pointer to the context for the threaded operation.
ThreadId - Supplies the current index of the threaded operation.
Return Value:
None.
--*/
{
const auto* WorkBlock = (MLAS_SOFTMAX_WORK_BLOCK<MLAS_FP16>*)Context;
size_t n;
size_t CountN;
MlasPartitionWork(Index, WorkBlock->ThreadCountN, WorkBlock->N, &n, &CountN);
const size_t D = WorkBlock->D;
const bool LogSoftmax = WorkBlock->LogSoftmax;
const bool SmoothSoftmax = WorkBlock->SmoothSoftmax;
const MLAS_FP16* Input = WorkBlock->Input + n * D;
MLAS_FP16* Output = WorkBlock->Output + n * D;
const auto* dispatch = GetMlasPlatform().SoftmaxDispatch;
if (dispatch == nullptr ||
dispatch->ReduceMax_Fp16 == nullptr ||
dispatch->SumExp_Fp16 == nullptr ||
(LogSoftmax && dispatch->LogSoftmax_Fp16 == nullptr) ||
(!LogSoftmax && dispatch->Softmax_Fp16 == nullptr)) {
MLAS_THROW_EX(std::runtime_error, "Lacks kernels for fp16 softmax.");
}
while (CountN > 0) {
MLAS_FP16 Maximum = dispatch->ReduceMax_Fp16(Input, D);
MLAS_FP16 NegativeMaximum = Maximum.Negate();
if (SmoothSoftmax && !NegativeMaximum.IsNegative()) {
NegativeMaximum = MLAS_FP16::FromBits(0);
}
MLAS_FP16* Temp = LogSoftmax ? nullptr : Output;
MLAS_FP16 Accumulation = dispatch->SumExp_Fp16(Input, Temp, D, NegativeMaximum);
float accumulation_fp32 = Accumulation.ToFloat();
if (SmoothSoftmax) {
accumulation_fp32 += expf(NegativeMaximum.ToFloat());
}
if (LogSoftmax) {
dispatch->LogSoftmax_Fp16(Input, Output, D, NegativeMaximum, MLAS_FP16(std::log(accumulation_fp32)));
} else {
dispatch->Softmax_Fp16(Output, Output, D, MLAS_FP16(accumulation_fp32));
}
Input += D;
Output += D;
CountN--;
}
}
template <typename T>
void
MLASCALL
MlasComputeSoftmax(
const T* Input,
T* Output,
size_t N,
size_t D,
bool LogSoftmax,
bool SmoothSoftmax,
float Sink,
MLAS_THREADPOOL* ThreadPool
)
/*++
Routine Description:
This routine computes the softmax or log softmax function.
N.B. This implementation supports in place updates of the output buffer.
Arguments:
Input - Supplies the input buffer.
Output - Supplies the output buffer.
N - Supplies the number of rows to process.
D - Supplies the number of columns per row to process.
LogSoftmax - Supplies true if this is a log softmax operation, else false
if this is a softmax operation.
SmoothSoftmax - Supplies true if a smooth factor is used in softmax operation.
Sink - Supplies the smooth factor to use in the softmax operation.
ThreadPool - Supplies the thread pool object to use, else nullptr if the
base library threading support should be used.
Return Value:
None.
--*/
{
MLAS_SOFTMAX_WORK_BLOCK<T> WorkBlock;
//
// Capture the softmax parameters to the work block.
//
WorkBlock.LogSoftmax = LogSoftmax;
WorkBlock.SmoothSoftmax = SmoothSoftmax;
WorkBlock.Input = Input;
WorkBlock.Output = Output;
WorkBlock.N = N;
WorkBlock.D = D;
WorkBlock.Sink = Sink;
//
// Compute the number of target threads given the complexity of the softmax
// operation. Limit the number of threads to the number of rows and try to
// keep each thread processing a minimum number of elements before using
// another thread.
//
ptrdiff_t ThreadCountN = MlasGetMaximumThreadCount(ThreadPool);
if (size_t(ThreadCountN) > N) {
ThreadCountN = ptrdiff_t(N);
}
constexpr size_t MinimumElementsPerThread = 16384;
size_t BlockCount = ((N * D) / MinimumElementsPerThread) + 1;
if (size_t(ThreadCountN) > BlockCount) {
ThreadCountN = ptrdiff_t(BlockCount);
}
WorkBlock.ThreadCountN = ThreadCountN;
MlasExecuteThreaded(MlasComputeSoftmaxThreaded<T>, &WorkBlock, ThreadCountN, ThreadPool);
}
template
void
MLASCALL
MlasComputeSoftmax<float>(
const float* Input,
float* Output,
size_t N,
size_t D,
bool LogSoftmax,
bool SmoothSoftmax,
float Sink,
MLAS_THREADPOOL* ThreadPool
);
template
void
MLASCALL
MlasComputeSoftmax<MLAS_FP16>(
const MLAS_FP16* Input,
MLAS_FP16* Output,
size_t N,
size_t D,
bool LogSoftmax,
bool SmoothSoftmax,
float Sink,
MLAS_THREADPOOL* ThreadPool
);
template <>
bool
MLASCALL
MlasGQASupported<MLAS_FP16>(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB
) {
if (!MlasHGemmSupported(TransA, TransB)) {
return false;
}
const auto* softmax_dispatch = GetMlasPlatform().SoftmaxDispatch;
if (softmax_dispatch == nullptr ||
softmax_dispatch->Tanh_Fp16 == nullptr ||
softmax_dispatch->Softcap_Fp16 == nullptr ||
softmax_dispatch->SumExp_Fp16 == nullptr ||
softmax_dispatch->Softmax_Fp16 == nullptr ||
softmax_dispatch->ReduceMax_Fp16 == nullptr) {
return false;
}
return true;
}
template <>
bool
MLASCALL
MlasGQASupported<float>(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB
) {
MLAS_UNREFERENCED_PARAMETER(TransA);
MLAS_UNREFERENCED_PARAMETER(TransB);
return true;
}
3rdparty/mlas/lib/core/common/common.h
Vendored
+43
View File
@@ -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__)
3rdparty/mlas/lib/core/common/narrow.h
Vendored
+30
View File
@@ -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
3rdparty/mlas/lib/erf_neon_fp16.h
Vendored
+27
View File
@@ -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);
3rdparty/mlas/lib/flashattn.cpp
Vendored
+167
View File
@@ -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);
}
3rdparty/mlas/lib/gelu_neon_fp16.h
Vendored
+31
View File
@@ -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
);
3rdparty/mlas/lib/kleidiai/mlasi_kleidiai.h
Vendored
+236
View File
@@ -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;
}
3rdparty/mlas/lib/loongarch64/SgemmKernelLasx.S
Vendored
+33
View File
@@ -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
3rdparty/mlas/lib/loongarch64/SgemmKernelLsx.S
Vendored
+267
View File
@@ -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
3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4LSX.S
Vendored
+89
View File
@@ -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
3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4Lasx.S
Vendored
+126
View File
@@ -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
3rdparty/mlas/lib/loongarch64/asmmacro.h
Vendored
+144
View File
@@ -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
3rdparty/mlas/lib/mlasi.h
Vendored
+3102
View File
@@ -0,0 +1,3102 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
mlasi.h
Abstract:
This module contains the private data structures and procedure prototypes
for the Microsoft Machine Learning algebra subprogram library.
--*/
#pragma once
#include <algorithm>
#include <atomic>
#include <cmath>
#include <functional>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <type_traits>
#ifdef MLAS_NO_EXCEPTION
#if defined(__ANDROID__)
#include <android/log.h>
#else
#include <iostream>
#endif
#endif // MLAS_NO_EXCEPTION
// 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
#define WIN32_LEAN_AND_MEAN
#endif
#ifndef NOMINMAX
#define NOMINMAX
#endif
#include <windows.h>
#include <intrin.h>
#else
#if defined(__arm__) || defined(__aarch64__)
#include <arm_neon.h>
#endif
#if defined(__x86_64__) || defined(__i386__)
#if !defined(signature_VORTEX_ebx) && !defined(signature_NEXGEN_ebx) && !defined(signature_AMD_ebx)//workaround for Bug 96238 - [i386] cpuid.h header needs include guards
#include <cpuid.h>
#endif
#if defined(__GNUC__) && __GNUC__ >= 12
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wmaybe-uninitialized" // GCC 12 warns about uninitialized variables in immintrin.h.
#include <immintrin.h>
#pragma GCC diagnostic pop
#else
#include <immintrin.h>
#endif
#endif
#if defined(__VSX__)
#include <altivec.h>
// Undefine unwanted aliases from altivec.h.
#undef vector
#undef pixel
#undef bool
#endif
#if defined(__s390x__)
#include <vecintrin.h>
#endif
#if defined(__loongarch64)
#include <lsxintrin.h>
#endif
#if defined(MLAS_TARGET_WASM_SIMD)
#include <wasm_simd128.h>
#endif
#endif
//
// Macro to place variables at a specified alignment.
//
#ifdef _WIN32
#define MLAS_DECLSPEC_ALIGN(variable, alignment) DECLSPEC_ALIGN(alignment) variable
#else
#define MLAS_DECLSPEC_ALIGN(variable, alignment) variable __attribute__ ((aligned(alignment)))
#endif
//
// Macro to force inline expansion of a function.
//
#if defined(_MSC_VER)
#define MLAS_FORCEINLINE __forceinline
#else
#define MLAS_FORCEINLINE __attribute__ ((always_inline)) inline
#endif
//
// Macro to tag globals as internal data shared with kernels written in
// assembly. These globals are marked with having hidden visibility to avoid
// needing to access the data through the global object table.
//
#if defined(_MSC_VER)
#define MLAS_INTERNAL_DATA extern "C"
#else
#define MLAS_INTERNAL_DATA extern "C" __attribute ((visibility("hidden")))
#endif
//
// Macro to suppress unreferenced parameter warnings.
//
#define MLAS_UNREFERENCED_PARAMETER(parameter) ((void)(parameter))
#ifdef MLAS_NO_EXCEPTION
MLAS_FORCEINLINE void
MlasPrintFinalMessage(const std::string& msg)
{
#if defined(__ANDROID__)
__android_log_print(ANDROID_LOG_ERROR, "mlas", "%s", msg.c_str());
#else
// TODO, consider changing the output of the error message from std::cerr to logging when the
// exceptions are disabled, since using std::cerr might increase binary size, and std::cerr
// output might not be easily accesible on some systems such as mobile
// TODO, see if we need to change the output of the error message from std::cerr to NSLog for
// iOS
std::cerr << msg << std::endl;
#endif
}
#define MLAS_THROW_EX(ex, what) \
do { \
std::string msg = #ex; \
msg.append(what); \
MlasPrintFinalMessage(msg); \
abort(); \
} while (false)
#else
#define MLAS_THROW_EX(ex, ...) throw ex(__VA_ARGS__)
#endif // MLAS_NO_EXCEPTION
//
// Select the threading model.
//
// N.B. BUILD_MLAS_NO_ONNXRUNTIME is used to build MLAS test code outside
// of the ONNX Runtime source tree. OpenMP may or may not be enabled in this
// configuration.
//
#if !defined(BUILD_MLAS_NO_ONNXRUNTIME)
#include "core/platform/threadpool.h"
#include "core/common/cpuid_info.h"
using MLAS_CPUIDINFO = onnxruntime::CPUIDInfo;
#include "core/common/float16.h"
#else // BUILD_MLAS_NO_ONNXRUNTIME
class MLASCPUIDInfo
{
public:
static const MLASCPUIDInfo& GetCPUIDInfo()
{
static MLASCPUIDInfo cpuid_info;
return cpuid_info;
}
// ARM
bool HasArmNeonDot() const { return has_arm_neon_dot_; }
bool HasFp16VectorAcceleration() const { return has_fp16_; }
uint32_t GetCurrentCoreIdx() const { return 0xFFFFFFFF; }
int32_t GetCurrentUarch() const { return -1; }
int32_t GetCoreUarch(uint32_t coreId) const { return -1; }
bool IsCoreArmv8NarrowLd(uint32_t coreId) const { return false; }
bool IsCurrentCoreArmv8NarrowLd() const { return false; }
bool HasArmNeon_I8MM() const { return has_arm_neon_i8mm_; }
bool HasArmSVE() const { return has_arm_sve_; }
bool HasArmSVE_I8MM() const { return has_arm_sve_i8mm_; }
bool HasArmNeon_BF16() const { return has_arm_neon_bf16_; }
private:
MLASCPUIDInfo();
bool has_arm_neon_dot_{false};
bool has_fp16_{false};
bool has_arm_neon_i8mm_{false};
bool has_arm_sve_{false};
bool has_arm_sve_i8mm_{false};
bool has_arm_neon_bf16_{false};
};
using MLAS_CPUIDINFO = MLASCPUIDInfo;
#if defined(MLAS_TARGET_ARM64)
/**
* @brief IDs for cpu microarchitectures.
*
* Copied from python cpuinfo package. Can't use the definition
* from cpuinfo directly as it causes lots of compilation issues
* in many platforms that we support.
*/
enum MlasUArch {
cpuinfo_uarch_unknown = 0,
/** ARM Cortex-A32. */
cpuinfo_uarch_cortex_a32 = 0x00300332,
/** ARM Cortex-A35. */
cpuinfo_uarch_cortex_a35 = 0x00300335,
/** ARM Cortex-A53. */
cpuinfo_uarch_cortex_a53 = 0x00300353,
/** ARM Cortex-A55 revision 0 (restricted dual-issue capabilities compared to revision 1+). */
cpuinfo_uarch_cortex_a55r0 = 0x00300354,
/** ARM Cortex-A55. */
cpuinfo_uarch_cortex_a55 = 0x00300355,
/** ARM Cortex-A57. */
cpuinfo_uarch_cortex_a57 = 0x00300357,
/** ARM Cortex-A65. */
cpuinfo_uarch_cortex_a65 = 0x00300365,
/** ARM Cortex-A72. */
cpuinfo_uarch_cortex_a72 = 0x00300372,
/** ARM Cortex-A73. */
cpuinfo_uarch_cortex_a73 = 0x00300373,
/** ARM Cortex-A75. */
cpuinfo_uarch_cortex_a75 = 0x00300375,
/** ARM Cortex-A76. */
cpuinfo_uarch_cortex_a76 = 0x00300376,
/** ARM Cortex-A77. */
cpuinfo_uarch_cortex_a77 = 0x00300377,
/** ARM Cortex-A78. */
cpuinfo_uarch_cortex_a78 = 0x00300378,
};
#endif // MLAS_TARGET_ARM64
//
// Define MLAS_FP16
//
#include "mlas_float16.h"
namespace onnxruntime
{
struct MLFloat16 {
uint16_t val{0};
MLFloat16() = default;
explicit constexpr MLFloat16(uint16_t x) : val(x) {}
explicit MLFloat16(float ff) : val(MLAS_Float2Half(ff)) {}
constexpr static MLFloat16 FromBits(uint16_t x) noexcept { return MLFloat16(x); }
MLFloat16 Abs() const noexcept {
return MLFloat16(static_cast<uint16_t>(val & ~kSignMask));
}
bool IsNaN() const noexcept {
return Abs().val > kPositiveInfinityBits;
}
bool IsNegative() const noexcept {
return static_cast<int16_t>(val) < 0;
}
MLFloat16 Negate() const {
return MLFloat16(IsNaN() ? val : static_cast<uint16_t>(val ^ kSignMask));
}
static constexpr uint16_t kSignMask = 0x8000U;
static constexpr uint16_t kPositiveInfinityBits = 0x7C00U;
float ToFloat() const { return MLAS_Half2Float(val); }
operator float() const { return ToFloat(); }
MLFloat16& operator=(float ff)
{
val = MLAS_Float2Half(ff);
return *this;
}
};
inline bool
operator==(const MLFloat16& left, const MLFloat16& right)
{
return left.val == right.val;
}
inline bool
operator!=(const MLFloat16& left, const MLFloat16& right)
{
return left.val != right.val;
}
}
#endif // BUILD_MLAS_NO_ONNXRUNTIME
static_assert(sizeof(MLAS_FP16) == FP16_SIZE);
//
// Define the maximum number of threads supported by this implementation.
//
#define MLAS_MAXIMUM_THREAD_COUNT 16
//
// Define the default strides to step through slices of the input matrices.
//
#define MLAS_HGEMM_STRIDEN 128
#define MLAS_HGEMM_STRIDEK 128
#define MLAS_SGEMM_STRIDEN 128
#define MLAS_SGEMM_STRIDEK 128
#define MLAS_SGEMM_PACKED_STRIDEN 128
#define MLAS_SGEMM_PACKED_STRIDEK 256
#define MLAS_DGEMM_STRIDEN 64
#define MLAS_DGEMM_STRIDEK 128
//
// Define the alignment for segmenting a GEMM operation across multiple
// threads.
//
// All of the SGEMM kernels can efficiently handle 16 elements. AVX512F can
// efficiently handle 32 elements, but making this value dynamic is not worth
// the effort at this time.
//
#define MLAS_HGEMM_STRIDEN_THREAD_ALIGN 32
#define MLAS_SGEMM_STRIDEN_THREAD_ALIGN 16
#define MLAS_DGEMM_STRIDEN_THREAD_ALIGN 8
#define MLAS_QGEMM_STRIDEN_THREAD_ALIGN 16
//
// Define the prototypes of the platform optimized routines.
//
#if defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER) || \
defined(MLAS_TARGET_LARCH64) || defined(MLAS_TARGET_S390X) || \
defined(MLAS_TARGET_RISCV64)
typedef
size_t
(MLASCALL MLAS_GEMM_FLOAT_KERNEL)(
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
);
typedef
size_t
(MLASCALL MLAS_GEMM_DOUBLE_KERNEL)(
const double* A,
const double* B,
double* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
double alpha,
bool ZeroMode
);
#ifdef FORCE_GENERIC_ALGORITHMS
typedef
size_t
(MLASCALL MLAS_GEMM_FLOAT_KERNEL_GENERIC)(
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
);
#endif
#else
#if defined(__aarch64__) && defined(__linux__)
typedef size_t(MLASCALL MLAS_SBGEMM_FLOAT_KERNEL)(
const float* A,
const bfloat16_t* B,
float* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
const float* Bias
);
#endif
typedef
size_t
(MLASCALL MLAS_GEMM_FLOAT_KERNEL)(
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
);
typedef
size_t
(MLASCALL MLAS_GEMM_DOUBLE_KERNEL)(
const double* A,
const double* B,
double* C,
size_t CountK,
size_t CountM,
size_t CountN,
size_t lda,
size_t ldc,
double alpha
);
#endif
typedef
void
(MLASCALL MLAS_GEMV_FLOAT_KERNEL)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountN,
size_t ldb,
bool ZeroMode
);
typedef
void
(MLASCALL MLAS_SGEMM_KERNEL_M1_ROUTINE)(
const float* A,
const float* B,
float* C,
size_t CountK,
size_t CountN,
size_t ldb,
float beta
);
typedef
void
(MLASCALL MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE)(
float* D,
const float* B,
size_t ldb
);
typedef
size_t
(MLASCALL MLAS_GEMM_U8S8_KERNEL)(
const uint8_t* A,
const uint8_t* B,
int32_t* C,
size_t PackedCountK,
size_t CountM,
size_t CountN,
size_t ldc,
const int32_t* RowSumVector,
const int32_t* ColumnSumVector,
const int32_t* ZeroPointB,
bool ZeroMode
);
typedef
size_t
(MLASCALL MLAS_GEMV_U8S8_KERNEL)(
const uint8_t* A,
const uint8_t* B,
int32_t* C,
size_t CountK,
size_t CountN,
size_t ldb
);
typedef
size_t
(MLASCALL MLAS_GEMM_U8U8_KERNEL)(
const int16_t* A,
const uint8_t* B,
int32_t* C,
size_t PackedCountK,
size_t CountM,
size_t CountN,
size_t ldc,
const int32_t* RowSumVector,
const int32_t* ColumnSumVector,
const int32_t* ZeroPointB,
bool ZeroMode
);
typedef
void
(MLASCALL MLAS_CONV_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t FilterCount,
size_t InputStride,
size_t FilterStride,
size_t OutputStride,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad,
const float* Bias,
unsigned KernelFlags
);
typedef
void
(MLASCALL MLAS_CONV_DEPTHWISE_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t InputStride,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad,
const float* Bias,
unsigned KernelFlags
);
typedef
void
(MLASCALL MLAS_CONV_POINTWISE_FLOAT_KERNEL)(
const float* Input,
const float* Filter,
float* Output,
size_t StrideWidth,
size_t InputChannels,
size_t FilterCount,
size_t InputStride,
size_t FilterStride,
size_t OutputStride,
size_t OutputCount,
const float* Bias,
unsigned KernelFlags
);
typedef
void
(MLASCALL MLAS_POOL_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t StrideWidth,
size_t DilationWidth,
size_t InputStride,
size_t ActualKernelSize,
size_t KernelHeight,
size_t KernelWidth,
const float* InputBase,
size_t InputWidth,
size_t DilatedInputWidth,
size_t OutputCountLeftPad,
size_t OutputCount,
size_t OutputCountRightPad
);
typedef
void
(MLASCALL MLAS_COMPUTE_UNARY_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N
);
typedef
void
(MLASCALL MLAS_COMPUTE_ERF_FP16_KERNEL)(
const MLAS_FP16* Input,
MLAS_FP16* Output,
size_t N
);
typedef
void
(MLASCALL MLAS_COMPUTE_GELU_FP16_KERNEL)(
const MLAS_FP16* Input,
MLAS_FP16* Output,
MLAS_FP16* Temp,
size_t N,
MLAS_GELU_ALGORITHM Algo
);
typedef void
(MLASCALL MLAS_COMPUTE_TANH_FP16_KERNEL)(
const MLAS_FP16* Input,
MLAS_FP16* Output,
size_t N
);
typedef
float
(MLASCALL MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N,
const float* NegativeMaximum
);
typedef
void
(MLASCALL MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL)(
float* Output,
size_t N,
const float* Parameters
);
typedef
void
(MLASCALL MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL)(
const float* Input,
float* Output,
size_t N,
const float* Parameters
);
typedef
float
(MLASCALL MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL)(
const float* Input,
size_t N
);
typedef
void
(MLASCALL MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL)(
const float* Input,
float* Min,
float* Max,
size_t N
);
typedef
void(MLASCALL MLAS_CAST_F16_TO_F32_KERNEL)(
const unsigned short* Source,
float* Destination,
size_t Count
);
typedef void(MLASCALL MLAS_CAST_F32_TO_F16_KERNEL)(
const float* Source,
unsigned short* Destination,
size_t Count
);
typedef
void
(MLASCALL MLAS_QLINEAR_BINARY_OP_S8_KERNEL)(
const int8_t* InputA,
float ScaleA,
int32_t ZeroPointA,
const int8_t* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
int8_t* OutputC,
size_t N,
bool IsScalarB
);
typedef
void
(MLASCALL MLAS_QLINEAR_BINARY_OP_U8_KERNEL)(
const uint8_t* InputA,
float ScaleA,
int32_t ZeroPointA,
const uint8_t* InputB,
float ScaleB,
int32_t ZeroPointB,
float ScaleC,
int32_t ZeroPointC,
uint8_t* OutputC,
size_t N,
bool IsScalarB
);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_U8_KERNEL)(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
uint8_t ZeroPoint
);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_S8_KERNEL)(
const float* Input,
int8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint
);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_U16_KERNEL)(
const float* Input,
uint16_t* Output,
size_t N,
float Scale,
uint16_t ZeroPoint);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_S16_KERNEL)(
const float* Input,
int16_t* Output,
size_t N,
float Scale,
int16_t ZeroPoint);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_U4_KERNEL)(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint);
typedef
void
(MLASCALL MLAS_QUANTIZE_LINEAR_S4_KERNEL)(
const float* Input,
uint8_t* Output,
size_t N,
float Scale,
int8_t ZeroPoint);
typedef
void
(MLASCALL MLAS_DEQUANTIZE_LINEAR_U8_KERNEL)(
const uint8_t* Input,
float* Output,
size_t N,
float Scale,
uint8_t ZeroPoint);
typedef
void
(MLASCALL MLAS_DEQUANTIZE_LINEAR_S8_KERNEL)(
const int8_t* Input,
float* Output,
size_t N,
float Scale,
int8_t ZeroPoint);
template<typename InputType, typename FilterType>
struct MLAS_QUANT_KERNEL
{
typedef
void
(MLASCALL DepthwiseKernel)(
const InputType* const* Input,
InputType InputZeroPoint,
const FilterType* Filter,
FilterType FilterZeroPoint,
int32_t* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
};
typedef
void
(MLASCALL MLAS_CONV_FLOAT_FN)(
const MLAS_CONV_PARAMETERS* Parameters,
const float* Input,
const float* Filter,
const float* Bias,
float* WorkingBuffer,
float* Output,
MLAS_THREADPOOL* ThreadPool
);
typedef
bool
(MLASCALL MLAS_CONV_FLOAT_OVERRIDE)(
const MLAS_CONV_PARAMETERS* Parameters,
const float* Input,
const float* Filter,
const float* Bias,
float* WorkingBuffer,
float* Output,
MLAS_THREADPOOL* ThreadPool
);
// TODO: Investigate if overridden typedefs can be removed
typedef
void
(MLASCALL MLAS_CONV_PREPARE_FLOAT_FN)(
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
);
typedef
bool
(MLASCALL MLAS_CONV_PREPARE_FLOAT_OVERRIDE)(
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
);
typedef
bool
(MLASCALL MLAS_SGEMM_BATCH_OVERRIDE)(
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);
typedef
size_t
(MLASCALL MLAS_SGEMM_PACK_B_SIZE_OVERRIDE)(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K);
typedef
bool
(MLASCALL MLAS_SGEMM_PACK_B_OVERRIDE)(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB);
typedef
void
(MLASCALL MLAS_DYNAMIC_QGEMM_BATCH_OVERRIDE)(
const MLAS_GEMM_DYN_QUANT_SHAPE_PARAMS& Shape,
const MLAS_GEMM_DYN_QUANT_DATA_PARAMS* DataParams,
const size_t BatchN,
MLAS_THREADPOOL* ThreadPool);
typedef
size_t
(MLASCALL MLAS_DYNAMIC_QGEMM_PACK_B_SIZE_OVERRIDE)(
size_t N,
size_t K);
typedef
void
(MLASCALL MLAS_DYNAMIC_QGEMM_PACK_B_OVERRIDE)(
size_t N,
size_t K,
const int8_t* B,
const float* Scales,
const float* Bias,
void* PackedB);
#if defined(__aarch64__) && defined(__linux__)
typedef
bool
(MLASCALL MLAS_SBGEMM_BATCH_OVERRIDE)(
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);
typedef
size_t
(MLASCALL MLAS_SBGEMM_PACK_B_SIZE_OVERRIDE)(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K);
typedef
bool
(MLASCALL MLAS_SBGEMM_PACK_B_OVERRIDE)(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB);
#endif
extern "C" {
#if defined(MLAS_TARGET_AMD64_IX86)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelSse;
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelAvx;
#ifdef FORCE_GENERIC_ALGORITHMS
MLAS_GEMM_FLOAT_KERNEL_GENERIC MlasSgemmKernelZero;
MLAS_GEMM_FLOAT_KERNEL_GENERIC MlasSgemmKernelAdd;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelFma3;
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelAvx512F;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelSse;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelAvx;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelFma3;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelAvx512F;
#endif
#elif defined(MLAS_TARGET_POWER)
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernel;
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelPOWER10;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernel;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernelPOWER10;
MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8KernelVSX;
MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8KernelVSX;
#elif defined(MLAS_TARGET_S390X)
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernel;
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelZVECTOR;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernel;
MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8KernelZVECTOR;
MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8KernelZVECTOR;
#elif defined(MLAS_TARGET_LARCH64)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelLSX;
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelLasx;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelLSX;
MLAS_GEMM_DOUBLE_KERNEL MlasGemmDoubleKernelLasx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelLSX;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelLSX;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelLSX;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelLSX;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelLasx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelLasx;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelLasx;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelLasx;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelLSX;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelLSX;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelLSX;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelLasx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelLasx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelLasx;
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4LSX;
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Lasx;
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelLasx;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32KernelLasx;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelLasx;
#elif defined(MLAS_TARGET_RISCV64)
#if defined(MLAS_USE_RVV)
MLAS_GEMM_FLOAT_KERNEL MlasGemmFloatKernelRvv;
void MlasSgemmCopyPackBRvv(
float* D,
const float* B,
size_t ldb,
size_t CountX,
size_t CountY);
#endif
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);
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);
#else
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelZero;
MLAS_GEMM_FLOAT_KERNEL MlasSgemmKernelAdd;
#if defined(__aarch64__) && defined(__linux__)
MLAS_SBGEMM_FLOAT_KERNEL MlasSbgemmKernelZero;
MLAS_SBGEMM_FLOAT_KERNEL MlasSbgemmKernelAdd;
#endif
#if defined(MLAS_TARGET_ARM64) && defined(MLAS_USE_ARM_NEON_NCHWC)
// Intrinsics kernel for direct NCHW convolution
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelNeon;
#if !defined(_WIN32)
// AArch64 assembly micro-kernel for direct NCHW convolution
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelNeonAsm;
#endif
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelNeon;
#if !defined(_WIN32)
// AArch64 assembly micro-kernel for direct NCHWc convolution
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelNeonAsm;
#endif
// Intrinsics kernel for depthwise NCHWc convolution
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelNeon;
#if !defined(_WIN32)
// AArch64 assembly micro-kernel for depthwise NCHWc convolution
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelNeonAsm;
#endif
// Intrinsics kernel for pointwise NCHWc convolution
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelNeon;
#if !defined(_WIN32)
// AArch64 assembly micro-kernel for pointwise NCHWc convolution
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelNeonAsm;
#endif
#if defined(__linux__)
// AArch64 assembly fast-math micro-kernels
MLAS_CONV_FLOAT_KERNEL MlasConvNchwBf16KernelNeon;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseBf16KernelNeon;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseBf16KernelNeon;
#endif
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelNeon;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelNeon;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelNeon;
#endif
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernelZero;
MLAS_GEMM_DOUBLE_KERNEL MlasDgemmKernelAdd;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_SGEMM_KERNEL_M1_ROUTINE MlasSgemmKernelM1Avx;
MLAS_SGEMM_KERNEL_M1_ROUTINE MlasSgemmKernelM1TransposeBAvx;
#elif defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_WASM)
MLAS_GEMV_FLOAT_KERNEL MlasGemvFloatKernel;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Sse;
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE MlasSgemmTransposePackB16x4Avx;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx2;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx2;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512Core;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Core;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvx512Vnni;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvx512Vnni;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8S8KernelAvxVnni;
MLAS_GEMV_U8S8_KERNEL MlasGemvU8S8KernelAvxVnni;
MLAS_GEMM_U8S8_KERNEL MlasGemmU8U8KernelAvx2Vnni;
MLAS_GEMM_U8S8_KERNEL MlasGemmS8S8KernelAvx2Vnni;
MLAS_GEMM_U8S8_KERNEL MlasGemmS8U8KernelAvx2Vnni;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx2;
MLAS_GEMM_U8U8_KERNEL MlasGemmU8U8KernelAvx512Core;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelSse;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelSse;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelSse;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelSse;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelAvx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelAvx;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelAvx;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelAvx;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelFma3;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelFma3;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelFma3;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelFma3;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernelAvx512F;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernelAvx512F;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernelAvx512F;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernelAvx512F;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelSse;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelAvx;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernelAvx512F;
#else
MLAS_CONV_FLOAT_KERNEL MlasConvNchwFloatKernel;
MLAS_CONV_FLOAT_KERNEL MlasConvNchwcFloatKernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL MlasConvDepthwiseFloatKernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL MlasConvPointwiseFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolMaximumFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageExcludePadFloatKernel;
MLAS_POOL_FLOAT_KERNEL MlasPoolAverageIncludePadFloatKernel;
#endif
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasGeluErfKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasSiluKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32Kernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasLogisticKernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasTanhKernel;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32Kernel;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32Kernel;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32Kernel;
MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8Kernel;
MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8Kernel;
MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8Kernel;
MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8Kernel;
MLAS_QUANTIZE_LINEAR_S16_KERNEL MlasQuantizeLinearS16Kernel;
MLAS_QUANTIZE_LINEAR_U16_KERNEL MlasQuantizeLinearU16Kernel;
MLAS_QUANTIZE_LINEAR_S4_KERNEL MlasQuantizeLinearS4Kernel;
MLAS_QUANTIZE_LINEAR_U4_KERNEL MlasQuantizeLinearU4Kernel;
#if defined(MLAS_TARGET_AMD64)
MLAS_DEQUANTIZE_LINEAR_S8_KERNEL MlasDequantizeLinearS8Kernel;
MLAS_DEQUANTIZE_LINEAR_U8_KERNEL MlasDequantizeLinearU8Kernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasErfKernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeExpF32KernelAvx512F;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeLogisticF32KernelFma3;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasComputeTanhF32KernelFma3;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelFma3;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelAvx512F;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32KernelAvx;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelAvx;
MLAS_QLINEAR_BINARY_OP_S8_KERNEL MlasQLinearAddS8KernelAvx2;
MLAS_QLINEAR_BINARY_OP_U8_KERNEL MlasQLinearAddU8KernelAvx2;
MLAS_QUANTIZE_LINEAR_S8_KERNEL MlasQuantizeLinearS8KernelAvx512F;
MLAS_QUANTIZE_LINEAR_U8_KERNEL MlasQuantizeLinearU8KernelAvx512F;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasGeluErfKernelAvx512F;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL MlasSiluKernelAvx512F;
#endif
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32Kernel;
MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32Kernel;
#if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV)
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL MlasComputeSumExpF32KernelRvv;
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelRvv;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeSoftmaxOutputF32KernelRvv;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL MlasComputeLogSoftmaxOutputF32KernelRvv;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx;
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL MlasReduceMaximumF32KernelAvx512F;
MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL MlasReduceMinimumMaximumF32KernelAvx;
#endif
#if defined(MLAS_TARGET_AMD64)
MLAS_CAST_F16_TO_F32_KERNEL MlasCastF16ToF32KernelSse;
MLAS_CAST_F16_TO_F32_KERNEL MlasCastF16ToF32KernelAvx;
MLAS_CAST_F16_TO_F32_KERNEL MlasCastF16ToF32KernelAvx2;
MLAS_CAST_F32_TO_F16_KERNEL MlasCastF32ToF16KernelAvx2;
#endif
#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64)
MLAS_CAST_F16_TO_F32_KERNEL MlasCastF16ToF32KernelNeon;
MLAS_CAST_F32_TO_F16_KERNEL MlasCastF32ToF16KernelNeon;
#endif
}
//
// Define the default preferred byte alignment for buffers.
//
// MLAS_TARGET_AMD64_IX86: The typical architecture uses AVX instructions
// accessing 256-bit vectors. MLAS_TARGET_AMD64 returns a larger value if the
// platform supports 512-bit vectors to ensure that vectors are not split.
//
// MLAS_TARGET_ARM64: The kernels use "load pair" instructions to access 128-bit
// vectors, so this value keeps both vectors in the same cache line.
//
// MLAS_TARGET_ARM: Using 16 for a single 128-bit vector may be sufficient for
// this architecture, but the ONNX Runtime has historically used this larger
// value.
//
#define MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT 64
//
// Define the target number of per-thread multiplies before using another
// thread to perform additional work.
//
#define MLAS_SGEMM_THREAD_COMPLEXITY (size_t(64) * size_t(1024))
#define MLAS_DGEMM_THREAD_COMPLEXITY (size_t(64) * size_t(1024))
#define MLAS_QGEMM_THREAD_COMPLEXITY 65536
#define MLAS_HGEMM_THREAD_COMPLEXITY 65536
#if defined(__aarch64__) && defined(__linux__)
#define MLAS_SBGEMM_THREAD_COMPLEXITY (size_t(64) * size_t(1024))
#endif
//
// Single-threaded single precision matrix/matrix multiply operation.
//
void
MlasSgemmOperation(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const float* B,
size_t ldb,
float beta,
float* C,
size_t ldc
);
//
// Quantized integer matrix/matrix dispatch structure.
//
struct MLAS_GEMM_QUANT_DISPATCH;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchSse;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchLSX;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchLSX;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8U8DispatchLSX;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8S8DispatchSse41;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8S8DispatchAvx2;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8U8DispatchAvx2;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8U8DispatchAvx2Vnni;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchAvx2Vnni;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8U8DispatchAvx2Vnni;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8S8DispatchAmx;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchNeon;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmX8S8DispatchNeon;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchUdot;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchSdot;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchUmmla;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmS8S8DispatchSmmla;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchWasmSimd;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmU8X8DispatchWasmRelaxedSimd;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemmQuantDispatchDefault;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemm8X8DispatchPOWER10;
extern const MLAS_GEMM_QUANT_DISPATCH MlasGemm8X8DispatchZVECTOR;
#if defined(MLAS_TARGET_WASM_RELAXED_SIMD)
extern bool HasUSDot();
#endif
//
// Symmetric quantized qgemm dispatch structure
//
struct MLAS_SYMM_QGEMM_DISPATCH;
extern const MLAS_SYMM_QGEMM_DISPATCH MlasSymmQgemmS8DispatchNeon;
extern const MLAS_SYMM_QGEMM_DISPATCH MlasSymmQgemmS8DispatchSdot;
//
// Symmetric quantized integer convolution dispatch structure.
//
struct MLAS_CONV_SYM_DISPATCH;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymDispatchAvx2;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymDispatchAvxVnni;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymDispatchAvx512Core;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymDispatchAvx512Vnni;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymU8DispatchNeon;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymS8DispatchNeon;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymU8DispatchDot;
extern const MLAS_CONV_SYM_DISPATCH MlasConvSymS8DispatchDot;
//
// Quantized 8-bit integer/quantized 4-bit integer matrix/matrix multiply dispatch structure.
//
struct MLAS_Q8Q4GEMM_DISPATCH;
extern const MLAS_Q8Q4GEMM_DISPATCH MlasQ8Q4GemmDispatchAvx512vnni;
//
// Float/quantized 4-bit integer matrix/matrix multiply dispatch structure.
//
struct MLAS_FPQ4GEMM_DISPATCH;
extern const MLAS_FPQ4GEMM_DISPATCH MlasFpQ4GemmDispatchAvx512;
//
// Float/quantized n-bit integer matrix/matrix multiply dispatch structure.
//
struct MLAS_QNBIT_GEMM_DISPATCH;
const MLAS_QNBIT_GEMM_DISPATCH&
GetMlasQNBitGemmDispatchNeon(
bool InitializeWithDotSupport,
bool InitializeWithI8MMSupport
);
extern const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2;
extern const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx2vnni;
extern const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx512;
extern const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchAvx512vnni;
extern const MLAS_QNBIT_GEMM_DISPATCH MlasSQNBitGemmDispatchLasx;
struct MLAS_QNBIT_LUT_GEMM_DISPATCH;
extern const MLAS_QNBIT_LUT_GEMM_DISPATCH MlasLutGenKernelAvx2;
//
// Rotary embedding dispatch structure.
//
struct MLAS_ROPE_DISPATCH;
extern const MLAS_ROPE_DISPATCH MlasRopeDispatchNeon;
extern const MLAS_ROPE_DISPATCH MlasRopeDispatchAvx2;
//
// half gemm dispatch structure
//
struct MLAS_HGEMM_DISPATCH;
extern const MLAS_HGEMM_DISPATCH MlasHGemmDispatchNeon;
// softmax dispatch structure
struct MLAS_SOFTMAX_DISPATCH;
extern const MLAS_SOFTMAX_DISPATCH MlasSoftmaxDispatchNeon;
// eltwise dispatch structure
struct MLAS_ELTWISE_DISPATCH;
extern const MLAS_ELTWISE_DISPATCH MlasEltwiseDispatchNeon;
//
// Quantized depthwise convolution kernels.
//
template<typename InputType, typename FilterType>
void
MLASCALL
MlasConvDepthwiseKernel(
const InputType* const* Input,
InputType InputZeroPoint,
const FilterType* Filter,
FilterType FilterZeroPoint,
int32_t* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
template <typename InputType, typename FilterType>
void
MLASCALL
MlasConvDepthwiseKernelAvx2(
const InputType* const* Input,
InputType InputZeroPoint,
const FilterType* Filter,
FilterType FilterZeroPoint,
int32_t* Output,
size_t Channels,
size_t OutputCount,
size_t KernelSize
);
//
// Define the kernel flags for conv sym
//
#define MLAS_CONV_SYM_FLAG_INPUT_DIRECT 0x00000001
#define MLAS_CONV_SYM_FLAG_PER_CHANNEL_SCALE 0x00000002
//
// Define the post-processing parameters for conv sym: bias and re-quant params
//
struct MLAS_CONV_SYM_POST_PROCESS_PARAMS {
const int32_t* Bias;
const float* Scale;
float MinimumValue;
float MaximumValue;
int32_t OutputZeroPoint;
};
//
// Environment information class.
//
enum MlasCoreType { mlas_core_unknown = 0, mlas_core_little = 2, mlas_core_big = 3 };
struct MLAS_PLATFORM {
MLAS_PLATFORM(void);
// TODO: move to cpuinfo
bool Avx2Supported_ = false;
bool Avx512Supported_ = false;
bool ArmNeonIsQuantActivationsUnsigned = false;
// MLAS SGemm overrides
MLAS_SGEMM_BATCH_OVERRIDE* MlasSGemmBatchOverride = nullptr;
MLAS_SGEMM_PACK_B_SIZE_OVERRIDE* MlasSGemmPackBSizeOverride = nullptr;
MLAS_SGEMM_PACK_B_OVERRIDE* MlasSGemmPackBOverride = nullptr;
// MLAS Dynamic QGemm overrides
MLAS_DYNAMIC_QGEMM_BATCH_OVERRIDE* MlasDynamicQGemmBatchOverride = nullptr;
MLAS_DYNAMIC_QGEMM_PACK_B_SIZE_OVERRIDE* MlasDynamicQGemmPackBSizeOverride = nullptr;
MLAS_DYNAMIC_QGEMM_PACK_B_OVERRIDE* MlasDynamicQGemmPackBOverride = nullptr;
// MLAS Conv overrides
MLAS_CONV_PREPARE_FLOAT_OVERRIDE* MlasConvPrepareOverride = nullptr;
MLAS_CONV_FLOAT_OVERRIDE* MlasConvOverride = nullptr;
#if defined(__aarch64__) && defined(__linux__)
// SBGemm overrides
MLAS_SBGEMM_BATCH_OVERRIDE* MlasSBGemmBatchOverride = nullptr;
MLAS_SBGEMM_PACK_B_SIZE_OVERRIDE* MlasSBGemmPackBSizeOverride = nullptr;
MLAS_SBGEMM_PACK_B_OVERRIDE* MlasSBGemmPackBOverride = nullptr;
#endif
#if defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER) || defined(MLAS_TARGET_S390X) || defined(MLAS_TARGET_RISCV64)
MLAS_GEMM_FLOAT_KERNEL* GemmFloatKernel;
#endif
#if defined(MLAS_TARGET_LARCH64)
const MLAS_GEMM_QUANT_DISPATCH* GemmU8S8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmU8U8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmS8S8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmS8U8Dispatch;
MLAS_GEMM_FLOAT_KERNEL* GemmFloatKernel;
MLAS_GEMM_DOUBLE_KERNEL* GemmDoubleKernel;
MLAS_CONV_FLOAT_KERNEL* ConvNchwFloatKernel;
MLAS_CONV_FLOAT_KERNEL* ConvNchwcFloatKernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL* ConvDepthwiseFloatKernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL* ConvPointwiseFloatKernel;
MLAS_POOL_FLOAT_KERNEL* PoolFloatKernel[MlasPoolingKindCount];
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* TransposePackB16x4Routine;
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL* ReduceMaximumF32Kernel;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeSoftmaxOutputF32Kernel;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeLogSoftmaxOutputF32Kernel;
uint32_t NchwcBlockSize;
#endif
#if defined(MLAS_TARGET_AMD64_IX86)
const MLAS_GEMM_QUANT_DISPATCH* GemmU8S8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmU8U8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmS8S8Dispatch{&MlasGemmQuantDispatchDefault};
const MLAS_GEMM_QUANT_DISPATCH* GemmS8U8Dispatch{&MlasGemmQuantDispatchDefault};
#elif defined(MLAS_TARGET_ARM64)
const MLAS_GEMM_QUANT_DISPATCH* GemmU8U8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmU8S8Dispatch;
const MLAS_GEMM_QUANT_DISPATCH* GemmS8S8Dispatch;
#if defined(MLAS_USE_ARM_NEON_NCHWC)
MLAS_CONV_FLOAT_KERNEL* ConvNchwFloatKernel;
MLAS_CONV_FLOAT_KERNEL* ConvNchwcFloatKernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL* ConvDepthwiseFloatKernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL* ConvPointwiseFloatKernel;
#if defined(__linux__)
MLAS_CONV_FLOAT_KERNEL* ConvNchwBf16Kernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL* ConvDepthwiseBf16Kernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL* ConvPointwiseBf16Kernel;
#endif
MLAS_POOL_FLOAT_KERNEL* PoolFloatKernel[MlasPoolingKindCount];
uint32_t NchwcBlockSize;
#endif
#endif
const MLAS_SYMM_QGEMM_DISPATCH* SymmQgemmDispatch{nullptr};
const MLAS_CONV_SYM_DISPATCH* ConvSymU8S8Dispatch{nullptr};
const MLAS_CONV_SYM_DISPATCH* ConvSymS8S8Dispatch{nullptr};
MLAS_QUANT_KERNEL<uint8_t, int8_t>::DepthwiseKernel* ConvDepthwiseU8S8Kernel;
MLAS_QUANT_KERNEL<uint8_t, uint8_t>::DepthwiseKernel* ConvDepthwiseU8U8Kernel;
MLAS_QUANT_KERNEL<int8_t, int8_t>::DepthwiseKernel* ConvDepthwiseS8S8Kernel;
MLAS_QUANT_KERNEL<int8_t, uint8_t>::DepthwiseKernel* ConvDepthwiseS8U8Kernel;
#if defined(MLAS_TARGET_POWER) || defined(MLAS_TARGET_S390X)
MLAS_GEMM_DOUBLE_KERNEL* GemmDoubleKernel;
const MLAS_GEMM_QUANT_DISPATCH* GemmU8X8Dispatch;
MLAS_QUANTIZE_LINEAR_S8_KERNEL* QuantizeLinearS8Kernel;
MLAS_QUANTIZE_LINEAR_U8_KERNEL* QuantizeLinearU8Kernel;
MLAS_QUANTIZE_LINEAR_S16_KERNEL* QuantizeLinearS16Kernel;
MLAS_QUANTIZE_LINEAR_U16_KERNEL* QuantizeLinearU16Kernel;
MLAS_QUANTIZE_LINEAR_S4_KERNEL* QuantizeLinearS4Kernel;
MLAS_QUANTIZE_LINEAR_U4_KERNEL* QuantizeLinearU4Kernel;
#endif
#if defined(MLAS_USE_SVE) || defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_RISCV64)
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* ErfKernelRoutine;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* LogisticKernelRoutine;
MLAS_REDUCE_MAXIMUM_FLOAT_KERNEL* ReduceMaximumF32Kernel;
MLAS_COMPUTE_SUMEXP_FLOAT_KERNEL* ComputeSumExpF32Kernel;
MLAS_COMPUTE_LOGSOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeLogSoftmaxOutputF32Kernel;
MLAS_COMPUTE_SOFTMAX_OUTPUT_FLOAT_KERNEL* ComputeSoftmaxOutputF32Kernel;
#endif
MLAS_COMPUTE_ERF_FP16_KERNEL* ErfFP16KernelRoutine = nullptr;
MLAS_COMPUTE_GELU_FP16_KERNEL* GeluFP16KernelRoutine = nullptr;
MLAS_COMPUTE_TANH_FP16_KERNEL* TanhFP16KernelRoutine = nullptr;
#if defined(MLAS_TARGET_AMD64)
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* GeluErfKernelRoutine;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* SiluKernelRoutine;
MLAS_SGEMM_KERNEL_M1_ROUTINE* KernelM1Routine;
MLAS_SGEMM_KERNEL_M1_ROUTINE* KernelM1TransposeBRoutine;
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* TransposePackB16x4Routine;
MLAS_GEMM_DOUBLE_KERNEL* GemmDoubleKernel;
MLAS_GEMM_U8S8_KERNEL* GemmU8S8Kernel;
MLAS_GEMM_U8S8_KERNEL* GemmS8S8Kernel;
MLAS_GEMM_U8S8_KERNEL* GemmS8U8Kernel;
MLAS_GEMV_U8S8_KERNEL* GemvU8S8Kernel;
MLAS_GEMM_U8U8_KERNEL* GemmU8U8Kernel;
MLAS_CONV_FLOAT_KERNEL* ConvNchwFloatKernel;
MLAS_CONV_FLOAT_KERNEL* ConvNchwcFloatKernel;
MLAS_CONV_DEPTHWISE_FLOAT_KERNEL* ConvDepthwiseFloatKernel;
MLAS_CONV_POINTWISE_FLOAT_KERNEL* ConvPointwiseFloatKernel;
MLAS_POOL_FLOAT_KERNEL* PoolFloatKernel[MlasPoolingKindCount];
MLAS_QLINEAR_BINARY_OP_S8_KERNEL* QLinearAddS8Kernel;
MLAS_QLINEAR_BINARY_OP_U8_KERNEL* QLinearAddU8Kernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* ComputeExpF32Kernel;
MLAS_COMPUTE_UNARY_FLOAT_KERNEL* TanhKernelRoutine;
MLAS_REDUCE_MINIMUM_MAXIMUM_FLOAT_KERNEL* ReduceMinimumMaximumF32Kernel;
MLAS_QUANTIZE_LINEAR_S8_KERNEL* QuantizeLinearS8Kernel;
MLAS_QUANTIZE_LINEAR_U8_KERNEL* QuantizeLinearU8Kernel;
MLAS_QUANTIZE_LINEAR_S16_KERNEL* QuantizeLinearS16Kernel;
MLAS_QUANTIZE_LINEAR_U16_KERNEL* QuantizeLinearU16Kernel;
MLAS_QUANTIZE_LINEAR_S4_KERNEL* QuantizeLinearS4Kernel;
MLAS_QUANTIZE_LINEAR_U4_KERNEL* QuantizeLinearU4Kernel;
MLAS_DEQUANTIZE_LINEAR_S8_KERNEL* DequantizeLinearS8Kernel;
MLAS_DEQUANTIZE_LINEAR_U8_KERNEL* DequantizeLinearU8Kernel;
uint32_t NchwcBlockSize;
uint32_t PreferredBufferAlignment;
int32_t MaximumThreadCount;
#elif defined(MLAS_TARGET_ARM64)
static constexpr int32_t MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT * 4;
static constexpr size_t MLAS_NEON_NCHWC_BLOCK_SIZE = 16;
#else
static constexpr int32_t MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
#endif
const MLAS_FPQ4GEMM_DISPATCH* FpQ4GemmDispatch{nullptr};
const MLAS_Q8Q4GEMM_DISPATCH* Q8Q4GemmDispatch{nullptr};
const MLAS_QNBIT_GEMM_DISPATCH* QNBitGemmDispatch{nullptr};
const MLAS_QNBIT_LUT_GEMM_DISPATCH* LutGenKernel{nullptr};
MLAS_CAST_F16_TO_F32_KERNEL* CastF16ToF32Kernel;
MLAS_CAST_F32_TO_F16_KERNEL* CastF32ToF16Kernel;
const MLAS_ROPE_DISPATCH* RopeDispatch{nullptr};
const MLAS_HGEMM_DISPATCH* HGemmDispatch{nullptr};
const MLAS_SOFTMAX_DISPATCH* SoftmaxDispatch{nullptr};
const MLAS_ELTWISE_DISPATCH* EltwiseDispatch{nullptr};
};
inline
MLAS_PLATFORM& GetMlasPlatform(){
static MLAS_PLATFORM MlasPlatform;
return MlasPlatform;
}
//
// Threading support.
//
typedef
void
(MLAS_THREADED_ROUTINE)(
void* Context,
ptrdiff_t Index
);
void
MlasExecuteThreaded(
MLAS_THREADED_ROUTINE* ThreadedRoutine,
void* Context,
ptrdiff_t Iterations,
MLAS_THREADPOOL* ThreadPool
);
constexpr
size_t
MlasDivRoundup(size_t up, size_t down)
{
return (up + down - 1) / down;
}
/**
* @brief Distribute multiple iterations of work over a thread pool if supported
*
* @param ThreadPool [IN] Optional thread pool. Ignored when using OpenMP
* @param Iterations [IN] Total number of iterations
* @param Work [IN] Logic for computing a range of iterations [begin, end)
*/
void
MlasTrySimpleParallel(
MLAS_THREADPOOL* ThreadPool,
const std::ptrdiff_t Iterations,
const std::function<void(std::ptrdiff_t tid)>& Work
);
/**
* @brief Distribute many iterations of work over a thread pool if supported.
* This function is for small workloads in non-performance critical situation.
*
* @param ThreadPool [IN] Optional thread pool. Ignored when using OpenMP
* @param Iterations [IN] Total number of iterations
* @param Work [IN] Logic for computing a range of iterations [begin, end)
*/
void
MlasTryBatchParallel(
MLAS_THREADPOOL * ThreadPool,
const std::ptrdiff_t Iterations,
const std::function<void(std::ptrdiff_t tid)>& Work
);
#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(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
return onnxruntime::concurrency::ThreadPool::DegreeOfParallelism(ThreadPool);
#endif
}
inline
void
MlasPartitionWork(
ptrdiff_t ThreadId,
ptrdiff_t ThreadCount,
size_t TotalWork,
size_t* WorkIndex,
size_t* WorkRemaining
)
{
const size_t WorkPerThread = TotalWork / ThreadCount;
const size_t WorkPerThreadExtra = TotalWork % ThreadCount;
if (size_t(ThreadId) < WorkPerThreadExtra) {
*WorkIndex = (WorkPerThread + 1) * ThreadId;
*WorkRemaining = WorkPerThread + 1;
} else {
*WorkIndex = WorkPerThread * ThreadId + WorkPerThreadExtra;
*WorkRemaining = WorkPerThread;
}
}
//
// Define the minimum floating point value (and its bit value equivalent) that
// has no fractional bits. This number can be used for fast rounding of floating
// point numbers to integers.
//
#define MLAS_ROUNDING_BIAS_MAGIC 12582912.f
#define MLAS_ROUNDING_BIAS_MAGIC_BITS 0x4B400000
//
// Helpers to cast a floating point type to and from an integer bit format.
//
#if defined(_MSC_VER) && !defined(__clang__)
#pragma warning(push)
// VC++ suggests we can attempt to make 'MlasBitsOfFp32' constexpr, but it is not valid.
#pragma warning(disable:26497)
#endif
MLAS_FORCEINLINE
uint32_t
MlasBitsOfFp32(
float FloatValue
)
{
union {
uint32_t IntegerValue;
float FloatValue;
} u;
u.FloatValue = FloatValue;
return u.IntegerValue;
}
MLAS_FORCEINLINE
float
MlasFp32FromBits(
uint32_t IntegerValue
)
{
union {
uint32_t IntegerValue;
float FloatValue;
} u;
u.IntegerValue = IntegerValue;
return u.FloatValue;
}
#if defined(_MSC_VER) && !defined(__clang__)
#pragma warning(pop)
#endif
#if defined(MLAS_TARGET_WASM_SCALAR) || defined(MLAS_TARGET_ARM64)
void
MLASCALL
MlasConvDepthwiseFloat_CHW(
const MLAS_CONV_PARAMETERS* Parameters,
const float* Input,
const float* Filter,
float* Output,
const float* Zeros
);
#endif
void
MlasConvDepthwiseWithMultiplierFloat_CHW(
const MLAS_CONV_PARAMETERS* Parameters,
const float* Input,
const float* Filter,
float* Output,
const float* Zeros
);
#if defined(MLAS_TARGET_AMD64)
void
MlasConvDepthwiseMultiplier2CHWKernel7x7S2Avx512F(
const float* Input,
size_t InputHeight,
size_t InputWidth,
const float* Filter,
float* Output,
size_t OutputHeight,
size_t OutputWidth,
float Beta
);
#endif
//
// Define the missing ARM64 NEON intrinsic macros from arm64_neon.h that enable
// cross-compiler support.
//
// Also define additional standard NEON intrinsics using the MSVC aliases.
//
#if defined(_M_ARM64)
#ifndef vmaxvq_f32
#define vmaxvq_f32(src) neon_fmaxv(src)
#endif
#ifndef vminvq_f32
#define vminvq_f32(src) neon_fminv(src)
#endif
#endif
//
// Cross-platform wrappers for 32-bit vector intrinsics.
//
#if defined(MLAS_TARGET_ARM)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON32_INTRINSICS
#elif defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_ARM64EC)
#define MLAS_NEON_INTRINSICS
#define MLAS_NEON64_INTRINSICS
#elif defined(MLAS_TARGET_POWER)
#define MLAS_VSX_INTRINSICS
#elif defined(MLAS_TARGET_S390X)
#define MLAS_ZVECTOR_INTRINSICS
#elif defined(MLAS_TARGET_AMD64_IX86)
#define MLAS_SSE2_INTRINSICS
#if defined(__SSE4_1__) || (defined(_MSC_VER) && defined(__AVX__))
#define MLAS_SSE41_INTRINSICS
#endif
#if defined(__AVX__)
#define MLAS_AVX_INTRINSICS
#endif
#if defined(__AVX2__)
#define MLAS_AVX2_INTRINSICS
#endif
#if defined(__FMA__) || (defined(_MSC_VER) && defined(__AVX2__))
#define MLAS_FMA3_INTRINSICS
#endif
#elif defined(MLAS_TARGET_WASM_SIMD)
#define MLAS_WASM_SIMD_INTRINSICS
#if defined(MLAS_TARGET_WASM_RELAXED_SIMD)
#define MLAS_WASM_RELAXED_SIMD_INTRINSICS
#endif
#elif defined(MLAS_TARGET_LARCH64)
#define MLAS_LSX_INTRINSICS
#endif
#if defined(MLAS_NEON_INTRINSICS)
typedef float32x4_t MLAS_FLOAT32X4;
typedef int32x4_t MLAS_INT32X4;
#elif defined(MLAS_SSE2_INTRINSICS)
typedef __m128 MLAS_FLOAT32X4;
typedef __m128i MLAS_INT32X4;
#elif defined(MLAS_VSX_INTRINSICS)
typedef __vector float MLAS_FLOAT32X4;
typedef __vector int MLAS_INT32X4;
typedef __vector unsigned MLAS_UINT32X4;
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
typedef v128_t MLAS_FLOAT32X4;
typedef v128_t MLAS_INT32X4;
#elif defined(MLAS_LSX_INTRINSICS)
typedef __m128 MLAS_FLOAT32X4;
typedef __m128i MLAS_INT32X4;
#else
typedef float MLAS_FLOAT32X4 __attribute__ ((vector_size(16)));
typedef int32_t MLAS_INT32X4 __attribute__ ((vector_size(16)));
#endif
MLAS_FORCEINLINE
MLAS_INT32X4
MlasReinterpretAsInt32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vreinterpretq_s32_f32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_castps_si128(Vector);
#elif defined(MLAS_LSX_INTRINSICS)
return (MLAS_INT32X4)Vector;
#else
return MLAS_INT32X4(Vector);
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasCastToInt32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vcvtq_s32_f32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_cvttps_epi32(Vector);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_cts(Vector, 0);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_signed(Vector);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vftint_w_s(Vector);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return (MLAS_INT32X4)__builtin_convertvector((__f32x4)Vector, __i32x4);
#else
return MLAS_INT32X4{int32_t(Vector[0]), int32_t(Vector[1]), int32_t(Vector[2]), int32_t(Vector[3])};
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasCastToFloat32x4(MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vcvtq_f32_s32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_cvtepi32_ps(Vector);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_ctf(Vector, 0);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_float(Vector);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_convert_i32x4(Vector);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vffint_s_w(Vector);
#else
return MLAS_FLOAT32X4{float(Vector[0]), float(Vector[1]), float(Vector[2]), float(Vector[3])};
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasBroadcastInt32x4(int32_t Value)
{
#if defined(MLAS_NEON_INTRINSICS)
return vdupq_n_s32(Value);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_set1_epi32(Value);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_splat(Value);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return vec_splats(Value);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vreplgr2vr_w(Value);
#else
return MLAS_INT32X4{Value, Value, Value, Value};
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasLoadInt32x4(const int32_t* Buffer)
{
#if defined(MLAS_NEON_INTRINSICS)
return vld1q_s32(Buffer);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_loadu_si128((const __m128i*)Buffer);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vsx_ld(0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_xl(0, Buffer);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_load(Buffer);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vld((const MLAS_INT32X4*)Buffer, 0);
#else
return *((MLAS_INT32X4*)Buffer);
#endif
}
MLAS_FORCEINLINE
void
MlasStoreInt32x4(int32_t* Buffer, MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1q_s32(Buffer, Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
_mm_storeu_si128((__m128i*)Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
vec_vsx_st(Vector, 0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
vec_xst(Vector, 0, Buffer);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
wasm_v128_store(Buffer, Vector);
#elif defined(MLAS_LSX_INTRINSICS)
__lsx_vst(Vector, (MLAS_INT32X4 *)Buffer, 0);
#else
*((MLAS_INT32X4*)Buffer) = Vector;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAddInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vaddq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_add_epi32(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_add(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_add(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vadd_w(Vector1, Vector2);
#else
return Vector1 + Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasSubtractInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vsubq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_sub_epi32(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_sub(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vsub_w(Vector1, Vector2);
#else
return Vector1 - Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAndInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vandq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_and_si128(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_and(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vand_v(Vector1, Vector2);
#else
return Vector1 & Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasOrInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vorrq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_or_si128(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_or(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vor_v(Vector1, Vector2);
#else
return Vector1 | Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasAndNotInt32x4(MLAS_INT32X4 VectorNot, MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vandq_s32(vmvnq_s32(VectorNot), Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_andnot_si128(VectorNot, Vector);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_andnot(Vector, VectorNot);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vandn_v(VectorNot, Vector);
#else
return (~VectorNot) & Vector;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasXorInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return veorq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_xor_si128(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_xor(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return vec_xor(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vxor_v(Vector1, Vector2);
#else
return Vector1 ^ Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasBlendInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2, MLAS_INT32X4 Selection)
{
return MlasOrInt32x4(MlasAndInt32x4(Vector2, Selection), MlasAndNotInt32x4(Selection, Vector1));
}
template<unsigned ShiftCount>
MLAS_FORCEINLINE
MLAS_INT32X4
MlasShiftLeftInt32x4(MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vshlq_n_s32(Vector, ShiftCount);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_slli_epi32(Vector, ShiftCount);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_shl(Vector, ShiftCount);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vslli_w(Vector, ShiftCount);
#else
return Vector << ShiftCount;
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasMaximumInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vmaxq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE41_INTRINSICS)
return _mm_max_epi32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return MlasBlendInt32x4(Vector2, Vector1, _mm_cmpgt_epi32(Vector1, Vector2));
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vmaxsw(Vector1, Vector2);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_max(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_max(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vmax_w(Vector1, Vector2);
#else
return MlasBlendInt32x4(Vector2, Vector1, Vector1 > Vector2);
#endif
}
MLAS_FORCEINLINE
MLAS_INT32X4
MlasMinimumInt32x4(MLAS_INT32X4 Vector1, MLAS_INT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vminq_s32(Vector1, Vector2);
#elif defined(MLAS_SSE41_INTRINSICS)
return _mm_min_epi32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return MlasBlendInt32x4(Vector2, Vector1, _mm_cmpgt_epi32(Vector2, Vector1));
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vminsw(Vector1, Vector2);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_min(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_min(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vmin_w(Vector1, Vector2);
#else
return MlasBlendInt32x4(Vector2, Vector1, Vector2 > Vector1);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasReinterpretAsFloat32x4(MLAS_INT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vreinterpretq_f32_s32(Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_castsi128_ps(Vector);
#elif defined(MLAS_LSX_INTRINSICS)
return MLAS_FLOAT32X4(Vector);
#else
return MLAS_FLOAT32X4(Vector);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasBroadcastFloat32x4(float Value)
{
#if defined(MLAS_NEON_INTRINSICS)
return vdupq_n_f32(Value);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_set1_ps(Value);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_splat(Value);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
// Suppress wrong GCC warnings
MLAS_UNREFERENCED_PARAMETER(Value);
return vec_splats(Value);
#elif defined(MLAS_LSX_INTRINSICS)
return MLAS_FLOAT32X4{Value, Value, Value, Value};
#else
return MLAS_FLOAT32X4{Value, Value, Value, Value};
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasBroadcastFloat32x4(const float* Value)
{
#if defined(MLAS_NEON_INTRINSICS)
return vld1q_dup_f32(Value);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_load_ps1(Value);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_load32_splat(Value);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return vec_splats(*Value);
#elif defined(MLAS_LSX_INTRINSICS)
return MLAS_FLOAT32X4{*Value, *Value, *Value, *Value};
#else
return MLAS_FLOAT32X4{*Value, *Value, *Value, *Value};
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasZeroFloat32x4(void)
{
#if defined(MLAS_NEON_INTRINSICS)
return vdupq_n_f32(0.0f);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_setzero_ps();
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_const(0.0f, 0.0f, 0.0f, 0.0f);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasBroadcastFloat32x4(0.0f);
#else
return MlasBroadcastFloat32x4(0.0f);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasLoadFloat32x4(const float* Buffer)
{
#if defined(MLAS_NEON_INTRINSICS)
return vld1q_f32(Buffer);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_loadu_ps(Buffer);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vsx_ld(0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_xl(0, Buffer);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_load(Buffer);
#elif defined(MLAS_LSX_INTRINSICS)
// return MlasReinterpretAsFloat32x4(__lsx_vld((const MLAS_INT32X4 *)Buffer, 0));
return (MLAS_FLOAT32X4)__lsx_vld((const MLAS_INT32X4 *)Buffer, 0);
#else
return *((MLAS_FLOAT32X4*)Buffer);
#endif
}
MLAS_FORCEINLINE
void
MlasStoreFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1q_f32(Buffer, Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
_mm_storeu_ps(Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
vec_vsx_st(Vector, 0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
vec_xst(Vector, 0, Buffer);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
wasm_v128_store(Buffer, Vector);
#elif defined(MLAS_LSX_INTRINSICS)
__lsx_vst(MlasReinterpretAsInt32x4(Vector), Buffer, 0);
#else
*((MLAS_FLOAT32X4*)Buffer) = Vector;
#endif
}
MLAS_FORCEINLINE
void
MlasStoreAlignedFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1q_f32(Buffer, Vector);
#elif defined(MLAS_SSE2_INTRINSICS)
_mm_store_ps(Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
// Workaround for bad GCC warning that these parameters are set but not used.
MLAS_UNREFERENCED_PARAMETER(Buffer);
MLAS_UNREFERENCED_PARAMETER(Vector);
vec_st(Vector, 0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
vec_xst(Vector, 0, Buffer);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
wasm_v128_store(Buffer, Vector);
#elif defined(MLAS_LSX_INTRINSICS)
MlasStoreFloat32x4(Buffer, Vector);
#else
MlasStoreFloat32x4(Buffer, Vector);
#endif
}
template<unsigned Lane>
MLAS_FORCEINLINE
void
MlasStoreLaneFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1q_lane_f32(Buffer, Vector, Lane);
#elif defined(MLAS_SSE2_INTRINSICS)
// N.B. When building with AVX instructions, compilers optimize the following
// to a single vextractps instruction.
_mm_store_ss(Buffer, _mm_shuffle_ps(Vector, Vector, _MM_SHUFFLE(Lane, Lane, Lane, Lane)));
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
*Buffer = ((__f32x4)(Vector))[Lane];
#elif defined(MLAS_LSX_INTRINSICS)
*Buffer = Vector[Lane];
#else
*Buffer = Vector[Lane];
#endif
}
MLAS_FORCEINLINE
void
MlasStoreLowHalfFloat32x4(float* Buffer, MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
vst1_f32(Buffer, vget_low_f32(Vector));
#elif defined(MLAS_SSE2_INTRINSICS)
_mm_storel_pi((__m64*)Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
*((long long*)Buffer) = ((__vector long long)Vector)[0];
#elif defined(MLAS_LSX_INTRINSICS)
MlasStoreLaneFloat32x4<0>(&Buffer[0], Vector);
MlasStoreLaneFloat32x4<1>(&Buffer[1], Vector);
#else
MlasStoreLaneFloat32x4<0>(&Buffer[0], Vector);
MlasStoreLaneFloat32x4<1>(&Buffer[1], Vector);
#endif
}
template<unsigned Lane>
MLAS_FORCEINLINE
float
MlasExtractLaneFloat32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON_INTRINSICS)
return vgetq_lane_f32(Vector, Lane);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_cvtss_f32(_mm_shuffle_ps(Vector, Vector, _MM_SHUFFLE(Lane, Lane, Lane, Lane)));
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_extract_lane(Vector, Lane);
#elif defined(MLAS_LSX_INTRINSICS)
return Vector[Lane];
#else
return Vector[Lane];
#endif
}
#if defined(MLAS_SSE2_INTRINSICS)
template<>
MLAS_FORCEINLINE
void
MlasStoreLaneFloat32x4<0>(float* Buffer, MLAS_FLOAT32X4 Vector)
{
_mm_store_ss(Buffer, Vector);
}
template<>
MLAS_FORCEINLINE
float
MlasExtractLaneFloat32x4<0>(MLAS_FLOAT32X4 Vector)
{
return _mm_cvtss_f32(Vector);
}
template<unsigned Index0, unsigned Index1, unsigned Index2, unsigned Index3>
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasShuffleFloat32x4(MLAS_FLOAT32X4 Vector)
{
return _mm_shuffle_ps(Vector, Vector, _MM_SHUFFLE(Index3, Index2, Index1, Index0));
}
#endif
#if !defined(MLAS_SSE2_INTRINSICS) && !defined(_MSC_VER)
template<unsigned Index0, unsigned Index1, unsigned Index2, unsigned Index3>
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasShuffleFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_i32x4_shuffle(Vector1, Vector2, Index0, Index1, Index2, Index3);
#elif defined(__clang__)
return __builtin_shufflevector(Vector1, Vector2, Index0, Index1, Index2, Index3);
#elif defined(MLAS_LSX_INTRINSICS)
typedef int32_t GEN_INT32X4 __attribute__ ((vector_size(16)));
return __builtin_shuffle(Vector1, Vector2, GEN_INT32X4{Index0, Index1, Index2, Index3});
#else
return __builtin_shuffle(Vector1, Vector2, MLAS_INT32X4{Index0, Index1, Index2, Index3});
#endif
}
template<unsigned Index0, unsigned Index1, unsigned Index2, unsigned Index3>
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasShuffleFloat32x4(MLAS_FLOAT32X4 Vector)
{
return MlasShuffleFloat32x4<Index0, Index1, Index2, Index3>(Vector, Vector);
}
#endif
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasInterleaveLowFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON64_INTRINSICS)
return vzip1q_f32(Vector1, Vector2);
#elif defined(MLAS_NEON32_INTRINSICS)
float32x4x2_t zipped = vzipq_f32(Vector1, Vector2);
return zipped.val[0];
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_unpacklo_ps(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return vec_mergeh(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return (MLAS_FLOAT32X4)__lsx_vilvl_w(MlasReinterpretAsInt32x4(Vector2), MlasReinterpretAsInt32x4(Vector1));
#else
return MlasShuffleFloat32x4<0, 4, 1, 5>(Vector1, Vector2);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasInterleaveHighFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON64_INTRINSICS)
return vzip2q_f32(Vector1, Vector2);
#elif defined(MLAS_NEON32_INTRINSICS)
float32x4x2_t zipped = vzipq_f32(Vector1, Vector2);
return zipped.val[1];
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_unpackhi_ps(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return vec_mergel(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return (MLAS_FLOAT32X4)__lsx_vilvh_w(MlasReinterpretAsInt32x4(Vector2), MlasReinterpretAsInt32x4(Vector1));
#else
return MlasShuffleFloat32x4<2, 6, 3, 7>(Vector1, Vector2);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasAddFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vaddq_f32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_add_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_add(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_add(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfadd_s(Vector1, Vector2);
#else
return Vector1 + Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasSubtractFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vsubq_f32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_sub_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_sub(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_sub(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfsub_s(Vector1, Vector2);
#else
return Vector1 - Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMultiplyFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vmulq_f32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_mul_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_mul(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS)
// Suppress wrong GCC warnings
MLAS_UNREFERENCED_PARAMETER(Vector1);
MLAS_UNREFERENCED_PARAMETER(Vector2);
return vec_mul(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfmul_s(Vector1, Vector2);
#else
return Vector1 * Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMultiplyAddFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2, MLAS_FLOAT32X4 Vector3)
{
#if defined(MLAS_NEON_INTRINSICS)
#if defined(MLAS_TARGET_ARM)
// ARMv7 NEON doesn't have vfmaq_f32()
return vmlaq_f32(Vector3, Vector1, Vector2);
#else
return vfmaq_f32(Vector3, Vector1, Vector2);
#endif
#elif defined(MLAS_FMA3_INTRINSICS)
return _mm_fmadd_ps(Vector1, Vector2, Vector3);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_add_ps(_mm_mul_ps(Vector1, Vector2), Vector3);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_madd(Vector1, Vector2, Vector3);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return __builtin_s390_vfmasb(Vector1, Vector2, Vector3);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_add(wasm_f32x4_mul(Vector1, Vector2), Vector3);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfmadd_s(Vector1, Vector2, Vector3);
#else
return Vector1 * Vector2 + Vector3;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMultiplyAddFloat32x4(MLAS_FLOAT32X4 Vector1, float Scalar2, MLAS_FLOAT32X4 Vector3)
{
return MlasMultiplyAddFloat32x4(Vector1, MlasBroadcastFloat32x4(Scalar2), Vector3);
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMultiplyAddFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2, float Scalar3)
{
return MlasMultiplyAddFloat32x4(Vector1, Vector2, MlasBroadcastFloat32x4(Scalar3));
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasDivideFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON64_INTRINSICS)
return vdivq_f32(Vector1, Vector2);
#elif defined(MLAS_NEON32_INTRINSICS)
Vector1 = vsetq_lane_f32(vgetq_lane_f32(Vector1, 0) / vgetq_lane_f32(Vector2, 0), Vector1, 0);
Vector1 = vsetq_lane_f32(vgetq_lane_f32(Vector1, 1) / vgetq_lane_f32(Vector2, 1), Vector1, 1);
Vector1 = vsetq_lane_f32(vgetq_lane_f32(Vector1, 2) / vgetq_lane_f32(Vector2, 2), Vector1, 2);
Vector1 = vsetq_lane_f32(vgetq_lane_f32(Vector1, 3) / vgetq_lane_f32(Vector2, 3), Vector1, 3);
return Vector1;
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_div_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_div(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfdiv_s(Vector1, Vector2);
#else
return Vector1 / Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasGreaterThanFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vreinterpretq_f32_u32(vcgtq_f32(Vector1, Vector2));
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_cmpgt_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_gt(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return MLAS_FLOAT32X4(vec_cmpgt(Vector1, Vector2));
#elif defined(MLAS_LSX_INTRINSICS)
return (MLAS_FLOAT32X4)__lsx_vfcmp_clt_s(Vector2, Vector1);
#else
return Vector1 > Vector2;
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasAndFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_and_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_and(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasReinterpretAsFloat32x4(MlasAndInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#else
return MlasReinterpretAsFloat32x4(MlasAndInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasOrFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_or_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_or(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasReinterpretAsFloat32x4(MlasOrInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#else
return MlasReinterpretAsFloat32x4(MlasOrInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasAndNotFloat32x4(MLAS_FLOAT32X4 VectorNot, MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_andnot_ps(VectorNot, Vector);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_andnot(Vector, VectorNot);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasReinterpretAsFloat32x4(MlasAndNotInt32x4(MlasReinterpretAsInt32x4(VectorNot), MlasReinterpretAsInt32x4(Vector)));
#else
return MlasReinterpretAsFloat32x4(MlasAndNotInt32x4(MlasReinterpretAsInt32x4(VectorNot), MlasReinterpretAsInt32x4(Vector)));
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasXorFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_xor_ps(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_v128_xor(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasReinterpretAsFloat32x4(MlasXorInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#else
return MlasReinterpretAsFloat32x4(MlasXorInt32x4(MlasReinterpretAsInt32x4(Vector1), MlasReinterpretAsInt32x4(Vector2)));
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasBlendFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2, MLAS_FLOAT32X4 Selection)
{
return MlasOrFloat32x4(MlasAndFloat32x4(Vector2, Selection), MlasAndNotFloat32x4(Selection, Vector1));
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMaximumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vmaxq_f32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_max_ps(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
// Don't use vec_max to avoid undefined behavior if NAN
return vec_sel(Vector2, Vector1, vec_cmpgt(Vector1, Vector2));
#elif defined(MLAS_WASM_RELAXED_SIMD_INTRINSICS)
return wasm_f32x4_relaxed_max(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_max(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfmax_s(Vector1, Vector2);
#else
return MlasBlendFloat32x4(Vector2, Vector1, Vector1 > Vector2);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasMinimumFloat32x4(MLAS_FLOAT32X4 Vector1, MLAS_FLOAT32X4 Vector2)
{
#if defined(MLAS_NEON_INTRINSICS)
return vminq_f32(Vector1, Vector2);
#elif defined(MLAS_SSE2_INTRINSICS)
return _mm_min_ps(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
// Don't use vec_min to avoid undefined behavior if NAN
return vec_sel(Vector2, Vector1, vec_cmpgt(Vector2, Vector1));
#elif defined(MLAS_WASM_RELAXED_SIMD_INTRINSICS)
return wasm_f32x4_relaxed_min(Vector1, Vector2);
#elif defined(MLAS_WASM_SIMD_INTRINSICS)
return wasm_f32x4_min(Vector1, Vector2);
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfmin_s(Vector1, Vector2);
#else
return MlasBlendFloat32x4(Vector2, Vector1, Vector2 > Vector1);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasClampFloat32x4(MLAS_FLOAT32X4 Value, float LowerRange, float UpperRange)
{
#if defined(MLAS_SSE2_INTRINSICS)
// N.B. MINPS and MAXPS propagates the value from the second vector if the
// value is a NaN.
#endif
Value = MlasMaximumFloat32x4(MlasBroadcastFloat32x4(LowerRange), Value);
Value = MlasMinimumFloat32x4(MlasBroadcastFloat32x4(UpperRange), Value);
return Value;
}
MLAS_FORCEINLINE
float
MlasReduceAddFloat32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON64_INTRINSICS)
Vector = vpaddq_f32(Vector, Vector);
Vector = vpaddq_f32(Vector, Vector);
return vgetq_lane_f32(Vector, 0);
#elif defined(MLAS_NEON32_INTRINSICS)
float32x2_t VectorLow = vget_low_f32(Vector);
float32x2_t VectorHigh = vget_high_f32(Vector);
VectorLow = vpadd_f32(VectorLow, VectorHigh);
VectorLow = vpadd_f32(VectorLow, VectorHigh);
return vget_lane_f32(VectorLow, 0);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
Vector = MlasAddFloat32x4(Vector, MLAS_FLOAT32X4(vec_splat((__vector long long)Vector, 1)));
Vector = MlasAddFloat32x4(Vector, vec_splat(Vector, 1));
return Vector[0];
#else
Vector = MlasAddFloat32x4(Vector, MlasShuffleFloat32x4<2, 3, 2, 3>(Vector));
Vector = MlasAddFloat32x4(Vector, MlasShuffleFloat32x4<1, 1, 1, 1>(Vector));
return MlasExtractLaneFloat32x4<0>(Vector);
#endif
}
MLAS_FORCEINLINE
float
MlasReduceMaximumFloat32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON64_INTRINSICS)
return vmaxvq_f32(Vector);
#elif defined(MLAS_NEON32_INTRINSICS)
float32x2_t VectorLow = vget_low_f32(Vector);
float32x2_t VectorHigh = vget_high_f32(Vector);
VectorLow = vpmax_f32(VectorLow, VectorHigh);
VectorLow = vpmax_f32(VectorLow, VectorHigh);
return vget_lane_f32(VectorLow, 0);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
Vector = MlasMaximumFloat32x4(Vector, MLAS_FLOAT32X4(vec_splat((__vector long long)Vector, 1)));
Vector = MlasMaximumFloat32x4(Vector, vec_splat(Vector, 1));
return Vector[0];
#else
Vector = MlasMaximumFloat32x4(Vector, MlasShuffleFloat32x4<2, 3, 2, 3>(Vector));
Vector = MlasMaximumFloat32x4(Vector, MlasShuffleFloat32x4<1, 1, 1, 1>(Vector));
return MlasExtractLaneFloat32x4<0>(Vector);
#endif
}
MLAS_FORCEINLINE
float
MlasReduceMinimumFloat32x4(MLAS_FLOAT32X4 Vector)
{
#if defined(MLAS_NEON64_INTRINSICS)
return vminvq_f32(Vector);
#elif defined(MLAS_NEON32_INTRINSICS)
float32x2_t VectorLow = vget_low_f32(Vector);
float32x2_t VectorHigh = vget_high_f32(Vector);
VectorLow = vpmin_f32(VectorLow, VectorHigh);
VectorLow = vpmin_f32(VectorLow, VectorHigh);
return vget_lane_f32(VectorLow, 0);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
Vector = MlasMinimumFloat32x4(Vector, MLAS_FLOAT32X4(vec_splat((__vector long long)Vector, 1)));
Vector = MlasMinimumFloat32x4(Vector, vec_splat(Vector, 1));
return Vector[0];
#else
Vector = MlasMinimumFloat32x4(Vector, MlasShuffleFloat32x4<2, 3, 2, 3>(Vector));
Vector = MlasMinimumFloat32x4(Vector, MlasShuffleFloat32x4<1, 1, 1, 1>(Vector));
return MlasExtractLaneFloat32x4<0>(Vector);
#endif
}
// calc 2^int(N)
MLAS_FORCEINLINE
MLAS_FLOAT32X4
MlasPowerOf2Float32x4(MLAS_FLOAT32X4 Vector)
{
MLAS_INT32X4 emm0 = MlasAddInt32x4(MlasCastToInt32x4(Vector), MlasBroadcastInt32x4(127));
return MlasReinterpretAsFloat32x4(MlasShiftLeftInt32x4<23>(emm0));
}
//
// Cross-platform wrappers for 64-bit vector intrinsics.
//
#if defined(MLAS_SSE2_INTRINSICS)
typedef __m128d MLAS_FLOAT64X2;
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
typedef __vector double MLAS_FLOAT64X2;
#elif defined(MLAS_LSX_INTRINSICS)
typedef __m128d MLAS_FLOAT64X2;
#else
#define MLAS_FLOAT64X2_UNSUPPORTED
#endif
#ifndef MLAS_FLOAT64X2_UNSUPPORTED
#if defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
template<unsigned Lane>
MLAS_FORCEINLINE
double
MlasExtractLaneFloat64x2(MLAS_FLOAT64X2 Vector)
{
return Vector[Lane];
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasMultiplyAddFloat64x2(MLAS_FLOAT64X2 Vector1, MLAS_FLOAT64X2 Vector2, MLAS_FLOAT64X2 Vector3)
{
return vec_madd(Vector1, Vector2, Vector3);
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasBroadcastFloat64x2(const double *Value)
{
return MLAS_FLOAT64X2{*Value, *Value};
}
#elif defined(MLAS_LSX_INTRINSICS)
template<unsigned Lane>
MLAS_FORCEINLINE
double
MlasExtractLaneFloat64x2(MLAS_FLOAT64X2 Vector)
{
return Vector[Lane];
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasMultiplyAddFloat64x2(MLAS_FLOAT64X2 Vector1, MLAS_FLOAT64X2 Vector2, MLAS_FLOAT64X2 Vector3)
{
return __lsx_vfmadd_d(Vector1, Vector2, Vector3);
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasBroadcastFloat64x2(const double *Value)
{
return MLAS_FLOAT64X2{*Value, *Value};
}
#endif
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasBroadcastFloat64x2(double Value)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_set1_pd(Value);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return MLAS_FLOAT64X2{Value, Value};
#elif defined(MLAS_LSX_INTRINSICS)
return MLAS_FLOAT64X2{Value, Value};
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasZeroFloat64x2(void)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_setzero_pd();
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return MlasBroadcastFloat64x2(0.0f);
#elif defined(MLAS_LSX_INTRINSICS)
return MlasBroadcastFloat64x2(0.0f);
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasLoadFloat64x2(const double* Buffer)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_loadu_pd(Buffer);
#elif defined(MLAS_VSX_INTRINSICS)
return vec_vsx_ld(0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
return vec_xl(0, Buffer);
#elif defined(MLAS_LSX_INTRINSICS)
return MLAS_FLOAT64X2(__lsx_vld((const MLAS_INT32X4 *)Buffer, 0));
#endif
}
MLAS_FORCEINLINE
void
MlasStoreFloat64x2(double* Buffer, MLAS_FLOAT64X2 Vector)
{
#if defined(MLAS_SSE2_INTRINSICS)
_mm_storeu_pd(Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS)
vec_vsx_st(Vector, 0, Buffer);
#elif defined(MLAS_ZVECTOR_INTRINSICS)
vec_xst(Vector, 0, Buffer);
#elif defined(MLAS_LSX_INTRINSICS)
(__lsx_vst(MLAS_INT32X4(Vector), Buffer, 0));
#endif
}
MLAS_FORCEINLINE
void
MlasStoreAlignedFloat64x2(double* Buffer, MLAS_FLOAT64X2 Vector)
{
#if defined(MLAS_SSE2_INTRINSICS)
_mm_store_pd(Buffer, Vector);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
*((MLAS_FLOAT64X2*)Buffer) = Vector;
#elif defined(MLAS_LSX_INTRINSICS)
(__lsx_vst(MLAS_INT32X4(Vector), Buffer, 0));
#endif
}
MLAS_FORCEINLINE
MLAS_FLOAT64X2
MlasMultiplyFloat64x2(MLAS_FLOAT64X2 Vector1, MLAS_FLOAT64X2 Vector2)
{
#if defined(MLAS_SSE2_INTRINSICS)
return _mm_mul_pd(Vector1, Vector2);
#elif defined(MLAS_VSX_INTRINSICS) || defined(MLAS_ZVECTOR_INTRINSICS)
return Vector1 * Vector2;
#elif defined(MLAS_LSX_INTRINSICS)
return __lsx_vfmul_d(Vector1, Vector2);
#endif
}
#endif // !MLAS_FLOAT64X2_UNSUPPORTED
//
// Reads a platform specific time stamp counter.
//
MLAS_FORCEINLINE
uint64_t
MlasReadTimeStampCounter(void)
{
#ifdef _WIN32
#if defined(MLAS_TARGET_AMD64_IX86)
return ReadTimeStampCounter();
#else
LARGE_INTEGER PerformanceCounter;
QueryPerformanceCounter(&PerformanceCounter);
return (ULONG64)PerformanceCounter.QuadPart;
#endif
#else
#if defined(MLAS_TARGET_AMD64)
uint32_t eax, edx;
__asm__ __volatile__
(
"rdtsc"
: "=a" (eax), "=d" (edx)
);
return ((uint64_t)edx << 32) | eax;
#elif defined(MLAS_TARGET_LARCH64)
uint64_t time_cnt, id;
__asm__ __volatile__
(
"rdtime.d %0, %1\n\t"
: "=r" (time_cnt), "=r" (id)
::
);
return time_cnt;
#else
return 0;
#endif
#endif
}
//
// Aligned buffer for GEMM packing, etc.
//
constexpr size_t ThreadedBufAlignment = 64;
extern thread_local size_t ThreadedBufSize;
#ifdef _MSC_VER
extern thread_local std::unique_ptr<uint8_t, decltype(&_aligned_free)> ThreadedBufHolder;
#else
extern thread_local std::unique_ptr<uint8_t, decltype(&free)> ThreadedBufHolder;
#endif
MLAS_FORCEINLINE
constexpr size_t
UpAlignSize(size_t size)
{
size = (size + ThreadedBufAlignment - 1) / ThreadedBufAlignment;
return size * ThreadedBufAlignment;
}
MLAS_FORCEINLINE
void
MlasThreadedBufAlloc(size_t size)
{
if (size > ThreadedBufSize) {
#ifdef _MSC_VER
ThreadedBufHolder.reset(
reinterpret_cast<uint8_t*>(_aligned_malloc(size, ThreadedBufAlignment)));
#elif (__STDC_VERSION__ >= 201112L) && !defined(__APPLE__)
ThreadedBufHolder.reset(
reinterpret_cast<uint8_t*>(aligned_alloc(ThreadedBufAlignment, size)));
#else
// aligned_alloc unavailable macos 10.14 or earlier
void* ptr;
int err = posix_memalign(&ptr, ThreadedBufAlignment, size);
if (err != 0) {
ptr = nullptr;
}
ThreadedBufHolder.reset(reinterpret_cast<uint8_t*>(ptr));
#endif
ThreadedBufSize = size;
}
}
//
// Utilities for INT4 quantization.
//
template<bool Signed>
struct Int4Traits;
template<>
struct Int4Traits<true> {
using UnpackedType = int8_t;
static constexpr int8_t Min = -8;
static constexpr int8_t Max = 7;
};
template<>
struct Int4Traits<false> {
using UnpackedType = uint8_t;
static constexpr int8_t Min = 0;
static constexpr int8_t Max = 15;
};
template<typename UnpackedType>
MLAS_FORCEINLINE
void
MlasSetInt4Element(uint8_t* Output, size_t ElemIndex, UnpackedType Value)
{
static_assert(std::is_same_v<UnpackedType, uint8_t> || std::is_same_v<UnpackedType, int8_t>);
const size_t OutputIndex = ElemIndex >> 1; // which byte
const size_t NibbleIndex = ElemIndex & 0x1; // which 4-bit elem in the byte
const uint8_t Shift = static_cast<uint8_t>(NibbleIndex << 2); // Either 0 or 4
const uint8_t Mask = static_cast<uint8_t>(0xF0 >> Shift);
uint8_t* Dst = &Output[OutputIndex];
*Dst &= Mask; // Clear 4-bit lane
*Dst |= static_cast<uint8_t>((Value & 0xF) << Shift); // Set 4-bit lane
}
template<typename UnpackedType>
MLAS_FORCEINLINE
void
MlasPackInt4Elements(uint8_t* Output, UnpackedType ValueLow, UnpackedType ValueHigh)
{
static_assert(std::is_same_v<UnpackedType, uint8_t> || std::is_same_v<UnpackedType, int8_t>);
*Output = static_cast<uint8_t>(((ValueHigh & 0xF) << 4) | (ValueLow & 0xF));
}
3rdparty/mlas/lib/platform.cpp
Vendored
+1094
View File
@@ -0,0 +1,1094 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
platform.cpp
Abstract:
This module implements logic to select the best configuration for the
this platform.
--*/
#include "mlasi.h"
#ifdef MLAS_USE_SVE
#include "sve/mlasi_sve.h"
#endif
#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
#if defined(USE_KLEIDIAI)
#include "kleidiai/mlasi_kleidiai.h"
#endif
#include <cctype>
#include <cstdlib>
#include <mutex>
#include <thread>
#if defined(MLAS_TARGET_POWER)
#if defined(__linux__)
#include <sys/auxv.h>
#elif defined(_AIX)
#define POWER_10 0x40000
#define POWER_10_ANDUP (POWER_10)
#include <sys/systemcfg.h>
#define __power_10_andup() (_system_configuration.implementation & POWER_10_ANDUP)
#elif defined(__FreeBSD__)
#include <machine/cpu.h>
#include <sys/auxv.h>
#endif
#endif
#if defined(MLAS_TARGET_S390X)
#include <sys/auxv.h>
#endif
#if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV) && defined(__linux__)
#include <sys/auxv.h>
#include <asm/hwcap.h>
#ifndef COMPAT_HWCAP_ISA_V
#define COMPAT_HWCAP_ISA_V (1UL << ('V' - 'A'))
#endif
#endif
#if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV)
namespace {
bool
MlasStringEqualsIgnoreCase(
const char* value,
const char* expected
)
{
while (*value != '\0' && *expected != '\0') {
const auto lhs = static_cast<unsigned char>(*value);
const auto rhs = static_cast<unsigned char>(*expected);
if (std::tolower(lhs) != std::tolower(rhs)) {
return false;
}
++value;
++expected;
}
return *value == '\0' && *expected == '\0';
}
bool
MlasShouldForceScalarRiscv(
const char* value
)
{
if (value == nullptr || value[0] == '\0') {
return false;
}
return MlasStringEqualsIgnoreCase(value, "1") ||
MlasStringEqualsIgnoreCase(value, "true") ||
MlasStringEqualsIgnoreCase(value, "on") ||
MlasStringEqualsIgnoreCase(value, "yes");
}
} // namespace
#endif
#if defined(MLAS_TARGET_ARM64)
#if defined(_WIN32)
// N.B. Support building with downlevel versions of the Windows SDK.
#ifndef PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE
#define PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE 43
#endif
#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
MLASCPUIDInfo::MLASCPUIDInfo()
{
has_arm_neon_dot_ = (IsProcessorFeaturePresent(PF_ARM_V82_DP_INSTRUCTIONS_AVAILABLE) != 0);
// raw hack! Need CPUIDInfo implementation for more precise detection
has_fp16_ = has_arm_neon_dot_;
}
#endif
#elif defined(__linux__)
#include <sys/auxv.h>
#include <asm/hwcap.h>
// N.B. Support building with older versions of asm/hwcap.h that do not define
// this capability bit.
#ifndef HWCAP_ASIMDDP
#define HWCAP_ASIMDDP (1 << 20)
#endif
#ifndef HWCAP2_I8MM
#define HWCAP2_I8MM (1 << 13)
#endif
#ifndef HWCAP2_SVEI8MM
#define HWCAP2_SVEI8MM (1 << 9)
#endif
#ifndef HWCAP2_BF16
#define HWCAP2_BF16 (1 << 14)
#endif
#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
MLASCPUIDInfo::MLASCPUIDInfo()
{
has_arm_neon_dot_ = ((getauxval(AT_HWCAP) & HWCAP_ASIMDDP) != 0);
// raw hack! Need CPUIDInfo implementation for more precise detection
has_fp16_ = has_arm_neon_dot_;
has_arm_neon_i8mm_ = ((getauxval(AT_HWCAP2) & HWCAP2_I8MM) != 0);
has_arm_sve_i8mm_ = ((getauxval(AT_HWCAP2) & HWCAP2_SVEI8MM) != 0);
has_arm_neon_bf16_ = ((getauxval(AT_HWCAP2) & HWCAP2_BF16) != 0);
}
#endif
#else
#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
MLASCPUIDInfo::MLASCPUIDInfo() {}
#endif
#endif // Windows vs Linux vs Unknown
#else // not MLAS_TARGET_ARM64
#if defined(BUILD_MLAS_NO_ONNXRUNTIME)
MLASCPUIDInfo::MLASCPUIDInfo() {}
#endif
#endif // MLAS_TARGET_ARM64
#ifdef MLAS_TARGET_AMD64_IX86
//
// Stores a vector to build a conditional load/store mask for vmaskmovps.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint32_t MlasMaskMoveAvx[8], 32) = { 0, 1, 2, 3, 4, 5, 6, 7 };
//
// Stores a table of AVX vmaskmovps/vmaskmovpd load/store masks.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint32_t MlasMaskMoveTableAvx[16], 32) = {
0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF,
0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000,
};
//
// Stores a table of AVX512 opmask register values.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const int16_t MlasOpmask16BitTableAvx512[16], 32) = {
0x0000, 0x0001, 0x0003, 0x0007, 0x000F, 0x001F, 0x003F, 0x007F,
0x00FF, 0x01FF, 0x03FF, 0x07FF, 0x0FFF, 0x1FFF, 0x3FFF, 0x7FFF,
};
//
// Reads the processor extended control register to determine platform
// capabilities.
//
#if !defined(_XCR_XFEATURE_ENABLED_MASK)
#define _XCR_XFEATURE_ENABLED_MASK 0
#endif
#if !defined(XFEATURE_MASK_XTILE)
#define XFEATURE_XTILECFG 17
#define XFEATURE_XTILEDATA 18
#define XFEATURE_MASK_XTILECFG (1 << XFEATURE_XTILECFG)
#define XFEATURE_MASK_XTILEDATA (1 << XFEATURE_XTILEDATA)
#define XFEATURE_MASK_XTILE (XFEATURE_MASK_XTILECFG | XFEATURE_MASK_XTILEDATA)
#endif
inline
uint64_t
MlasReadExtendedControlRegister(
unsigned int ext_ctrl_reg
)
{
#if defined(_WIN32)
return _xgetbv(ext_ctrl_reg);
#else
uint32_t eax, edx;
__asm__
(
"xgetbv"
: "=a" (eax), "=d" (edx)
: "c" (ext_ctrl_reg)
);
return ((uint64_t)edx << 32) | eax;
#endif
}
#if defined(__linux__)
#include <sys/syscall.h>
#endif
bool
MlasInitAMX()
{
#if defined(__linux__)
#define ARCH_GET_XCOMP_PERM 0x1022
#define ARCH_REQ_XCOMP_PERM 0x1023
unsigned long bitmask = 0;
long rc = syscall(SYS_arch_prctl, ARCH_REQ_XCOMP_PERM, XFEATURE_XTILEDATA);
if (rc) {
return false;
}
rc = syscall(SYS_arch_prctl, ARCH_GET_XCOMP_PERM, &bitmask);
if (rc) {
return false;
}
if (bitmask & XFEATURE_MASK_XTILE) {
return true;
}
return false;
#else
return true;
#endif
}
#endif // MLAS_TARGET_AMD64_IX86
#ifdef MLAS_TARGET_LARCH64
#if defined(__linux__)
#include <sys/auxv.h>
#include <asm/hwcap.h>
#endif
//
// Stores a vector to build a conditional load/store mask for vmaskmovps.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint32_t MlasMaskMoveLasx[8], 32) = { 0, 1, 2, 3, 4, 5, 6, 7 };
//
// Stores a table of AVX vmaskmovps/vmaskmovpd load/store masks.
//
MLAS_INTERNAL_DATA MLAS_DECLSPEC_ALIGN(const uint32_t MlasMaskMoveTableLasx[16], 32) = {
0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF, 0xFFFFFFFF,
0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000, 0x00000000,
};
#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
)
/*++
Routine Description:
This routine initializes the platform support for this library.
Arguments:
None.
Return Value:
None.
--*/
{
this->ConvDepthwiseU8S8Kernel = MlasConvDepthwiseKernel<uint8_t, int8_t>;
this->ConvDepthwiseU8U8Kernel = MlasConvDepthwiseKernel<uint8_t, uint8_t>;
this->ConvDepthwiseS8S8Kernel = MlasConvDepthwiseKernel<int8_t, int8_t>;
this->ConvDepthwiseS8U8Kernel = MlasConvDepthwiseKernel<int8_t, uint8_t>;
this->CastF16ToF32Kernel = nullptr;
this->CastF32ToF16Kernel = nullptr;
#if defined(MLAS_TARGET_RISCV64)
this->GemmFloatKernel = nullptr;
this->ErfKernelRoutine = MlasErfKernel;
this->LogisticKernelRoutine = MlasLogisticKernel;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
#if defined(MLAS_USE_RVV)
bool has_rvv = true;
#if defined(__linux__)
has_rvv = (getauxval(AT_HWCAP) & COMPAT_HWCAP_ISA_V) != 0;
#endif
if (MlasShouldForceScalarRiscv(std::getenv("ORT_MLAS_RISCV_FORCE_SCALAR"))) {
has_rvv = false;
}
if (has_rvv) {
this->GemmFloatKernel = MlasGemmFloatKernelRvv;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelRvv;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelRvv;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32KernelRvv;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32KernelRvv;
}
#endif
#endif
#if defined(MLAS_TARGET_AMD64_IX86)
//
// Default to the baseline SSE2 support.
//
this->GemmFloatKernel = MlasGemmFloatKernelSse;
this->GemmU8S8Dispatch = &MlasGemmU8X8DispatchSse;
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchSse;
#if defined(MLAS_TARGET_AMD64)
this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Sse;
this->GemmDoubleKernel = MlasGemmDoubleKernelSse;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelSse;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelSse;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelSse;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelSse;
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelSse;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelSse;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelSse;
this->ComputeExpF32Kernel = MlasComputeExpF32Kernel;
this->GeluErfKernelRoutine = MlasGeluErfKernel;
this->LogisticKernelRoutine = MlasLogisticKernel;
this->SiluKernelRoutine = MlasSiluKernel;
this->TanhKernelRoutine = MlasTanhKernel;
this->ErfKernelRoutine = MlasErfKernel;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
this->ReduceMinimumMaximumF32Kernel = MlasReduceMinimumMaximumF32Kernel;
this->QLinearAddS8Kernel = MlasQLinearAddS8Kernel;
this->QLinearAddU8Kernel = MlasQLinearAddU8Kernel;
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8Kernel;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8Kernel;
this->QuantizeLinearS16Kernel = MlasQuantizeLinearS16Kernel;
this->QuantizeLinearU16Kernel = MlasQuantizeLinearU16Kernel;
this->QuantizeLinearS4Kernel = MlasQuantizeLinearS4Kernel;
this->QuantizeLinearU4Kernel = MlasQuantizeLinearU4Kernel;
this->DequantizeLinearS8Kernel = MlasDequantizeLinearS8Kernel;
this->DequantizeLinearU8Kernel = MlasDequantizeLinearU8Kernel;
#ifndef __APPLE__
#ifndef FORCE_GENERIC_ALGORITHMS
this->CastF16ToF32Kernel = &MlasCastF16ToF32KernelSse;
#else // FORCE_GENERIC_ALGORITHMS
this->CastF16ToF32Kernel = nullptr;
#endif // FORCE_GENERIC_ALGORITHMS
#endif // __APPLE__
this->NchwcBlockSize = 8;
this->PreferredBufferAlignment = MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT;
this->MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
#endif
unsigned Cpuid1[4];
#if defined(_WIN32)
__cpuid((int*)Cpuid1, 1);
#else
__cpuid(1, Cpuid1[0], Cpuid1[1], Cpuid1[2], Cpuid1[3]);
#endif
#if defined(_MSC_VER)
//
// Check if the processor supports SSE 4.1 instructions.
//
#ifndef FORCE_GENERIC_ALGORITHMS
if ((Cpuid1[2] & 0x80000) != 0) {
#else // FORCE_GENERIC_ALGORITHMS
if (false) {
#endif // FORCE_GENERIC_ALGORITHMS
this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchSse41;
}
#endif
//
// Check if the processor supports the AVX and OSXSAVE features.
//
#ifndef FORCE_GENERIC_ALGORITHMS
if ((Cpuid1[2] & 0x18000000) == 0x18000000) {
#else // FORCE_GENERIC_ALGORITHMS
if (false) {
#endif // FORCE_GENERIC_ALGORITHMS
//
// Check if the operating system supports saving SSE and AVX states.
//
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;
this->GemmDoubleKernel = MlasGemmDoubleKernelAvx;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelAvx;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelAvx;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelAvx;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelAvx;
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelAvx;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelAvx;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelAvx;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32KernelAvx;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32KernelAvx;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelAvx;
this->ReduceMinimumMaximumF32Kernel = MlasReduceMinimumMaximumF32KernelAvx;
this->GemmU8U8Kernel = nullptr;
//
// Check if the processor supports AVX2/FMA3 features.
//
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
if (((Cpuid1[2] & 0x1000) != 0) && ((Cpuid7[1] & 0x20) != 0)) {
this->Avx2Supported_ = true;
this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchAvx2;
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx2;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx2;
this->GemmU8U8Dispatch = &MlasGemmU8U8DispatchAvx2;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx2;
this->ConvSymU8S8Dispatch = &MlasConvSymDispatchAvx2;
this->GemmFloatKernel = MlasGemmFloatKernelFma3;
this->GemmDoubleKernel = MlasGemmDoubleKernelFma3;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelFma3;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelFma3;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelFma3;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelFma3;
this->ComputeExpF32Kernel = MlasComputeExpF32KernelFma3;
this->LogisticKernelRoutine = MlasComputeLogisticF32KernelFma3;
this->TanhKernelRoutine = MlasComputeTanhF32KernelFma3;
this->ErfKernelRoutine = MlasErfKernelFma3;
this->QLinearAddS8Kernel = MlasQLinearAddS8KernelAvx2;
this->QLinearAddU8Kernel = MlasQLinearAddU8KernelAvx2;
this->ConvDepthwiseU8S8Kernel = MlasConvDepthwiseKernelAvx2<uint8_t, int8_t>;
this->ConvDepthwiseU8U8Kernel = MlasConvDepthwiseKernelAvx2<uint8_t, uint8_t>;
this->ConvDepthwiseS8S8Kernel = MlasConvDepthwiseKernelAvx2<int8_t, int8_t>;
this->ConvDepthwiseS8U8Kernel = MlasConvDepthwiseKernelAvx2<int8_t, uint8_t>;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelFma3;
this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchAvx2;
this->CastF16ToF32Kernel = &MlasCastF16ToF32KernelAvx2;
this->CastF32ToF16Kernel = &MlasCastF32ToF16KernelAvx2;
this->RopeDispatch = &MlasRopeDispatchAvx2;
// TODO(vraspar): check if this really goes here or if there are other platform reqs that we need to fulfill
this->LutGenKernel = &MlasLutGenKernelAvx2;
//
// Check if the processor supports Hybrid core architecture.
//
if ((Cpuid7[3] & 0x8000) != 0) {
this->MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT * 4;
}
//
// Check if the processor supports AVXVNNI features.
//
unsigned Cpuid7_1[4];
#if defined(_WIN32)
__cpuidex((int*)Cpuid7_1, 7, 1);
#else
__cpuid_count(7, 1, Cpuid7_1[0], Cpuid7_1[1], Cpuid7_1[2], Cpuid7_1[3]);
#endif
if ((Cpuid7_1[0] & 0x10) != 0) {
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvxVnni;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvxVnni;
this->ConvSymU8S8Dispatch = &MlasConvSymDispatchAvxVnni;
this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchAvx2vnni;
}
#if !defined(ORT_MINIMAL_BUILD)
//
// Check if the processor supports AVX512F features and the
// operating system supports saving AVX512F state.
//
if (((Cpuid7[1] & 0x10000) != 0) && ((xcr0 & 0xE0) == 0xE0)) {
this->GeluErfKernelRoutine = MlasGeluErfKernelAvx512F;
this->SiluKernelRoutine = MlasSiluKernelAvx512F;
this->GemmFloatKernel = MlasGemmFloatKernelAvx512F;
this->GemmDoubleKernel = MlasGemmDoubleKernelAvx512F;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelAvx512F;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelAvx512F;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelAvx512F;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelAvx512F;
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelAvx512F;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelAvx512F;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelAvx512F;
this->ComputeExpF32Kernel = MlasComputeExpF32KernelAvx512F;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32KernelAvx512F;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelAvx512F;
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelAvx512F;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelAvx512F;
this->NchwcBlockSize = 16;
this->PreferredBufferAlignment = 64;
//
// Check if the processor supports AVX512 core features
// (AVX512BW/AVX512DQ/AVX512VL).
//
if ((Cpuid7[1] & 0xC0020000) == 0xC0020000) {
this->Avx512Supported_ = true;
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Core;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Core;
this->GemmU8U8Kernel = MlasGemmU8U8KernelAvx512Core;
this->ConvSymU8S8Dispatch = &MlasConvSymDispatchAvx512Core;
this->FpQ4GemmDispatch = &MlasFpQ4GemmDispatchAvx512;
this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchAvx512;
//
// Check if the processor supports AVX512VNNI.
//
if ((Cpuid7[2] & 0x800) != 0) {
this->GemmU8S8Kernel = MlasGemmU8S8KernelAvx512Vnni;
this->GemvU8S8Kernel = MlasGemvU8S8KernelAvx512Vnni;
this->ConvSymU8S8Dispatch = &MlasConvSymDispatchAvx512Vnni;
this->Q8Q4GemmDispatch = &MlasQ8Q4GemmDispatchAvx512vnni;
this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchAvx512vnni;
}
}
}
//
// Check if the processor supports AVX-VNNI-INT8
//
if ((Cpuid7_1[3] & 0x10) != 0) {
this->GemmU8U8Dispatch = &MlasGemmU8U8DispatchAvx2Vnni;
this->GemmS8S8Dispatch = &MlasGemmS8S8DispatchAvx2Vnni;
this->GemmS8S8Kernel = MlasGemmS8S8KernelAvx2Vnni;
this->GemmS8U8Dispatch = &MlasGemmS8U8DispatchAvx2Vnni;
this->GemmS8U8Kernel = MlasGemmS8U8KernelAvx2Vnni;
}
#ifndef __APPLE__
#if (defined(_MSC_VER) && (_MSC_VER >= 1933)) || (defined(__GNUC__) && (__GNUC__ >= 13))
//
// Check if the processor supports AVX NE CONVERT.
//
if ((Cpuid7_1[3] & (0b1 << 5)) != 0) {
this->CastF16ToF32Kernel = &MlasCastF16ToF32KernelAvx;
}
#endif // (defined(_MSC_VER) && (_MSC_VER >= 1933)) || (defined(__GNUC__) && (__GNUC__ >= 13))
//
// Check if the processor supports AMX-TILE and AMX-INT8
// features.
//
if ((Cpuid7[3] & 0b1 << 24) != 0 &&
(Cpuid7[3] & 0b1 << 25) != 0 &&
(xcr0 & XFEATURE_MASK_XTILE) == XFEATURE_MASK_XTILE) {
if (MlasInitAMX()) {
this->GemmU8S8Dispatch = &MlasGemmU8S8DispatchAmx;
}
}
#endif // __APPLE__
#endif // ORT_MINIMAL_BUILD
}
#endif // MLAS_TARGET_AMD64
}
}
#endif // MLAS_TARGET_AMD64_IX86
#if defined(MLAS_TARGET_ARM64)
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchNeon;
this->GemmU8S8Dispatch = &MlasGemmX8S8DispatchNeon;
this->GemmS8S8Dispatch = &MlasGemmX8S8DispatchNeon;
this->SymmQgemmDispatch = &MlasSymmQgemmS8DispatchNeon;
this->ConvSymU8S8Dispatch = &MlasConvSymU8DispatchNeon;
this->ConvSymS8S8Dispatch = &MlasConvSymS8DispatchNeon;
this->RopeDispatch = &MlasRopeDispatchNeon;
this->HGemmDispatch = &MlasHGemmDispatchNeon;
this->SoftmaxDispatch = &MlasSoftmaxDispatchNeon;
this->EltwiseDispatch = &MlasEltwiseDispatchNeon;
#if defined(MLAS_USE_ARM_NEON_NCHWC)
// Use the AArch64 assembly implementation on non-Windows platforms.
#if !defined(_WIN32)
// Prefer the hand written micro-kernel for the NCHW convolution path. It
// offers a tighter schedule and a specialised two-output inner loop that
// reduces pressure on the memory system compared to the generic kernel.
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelNeonAsm;
#else
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelNeon;
#endif
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelNeon;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelNeon;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelNeon;
#if defined(__linux__)
this->ConvNchwBf16Kernel = MlasConvNchwBf16KernelNeon;
this->ConvDepthwiseBf16Kernel = MlasConvDepthwiseBf16KernelNeon;
this->ConvPointwiseBf16Kernel = MlasConvPointwiseBf16KernelNeon;
#endif
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelNeon;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelNeon;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelNeon;
this->NchwcBlockSize = MLAS_NEON_NCHWC_BLOCK_SIZE;
#endif
//
// Check if the processor supports ASIMD dot product instructions.
//
// Note:
// Do NOT use ID_AA64ISAR0_EL1. It causes illegal instruction errors on Mac M1 and ARMv8-A chips
// as well as failing on other ARM chips as it is an EL1 level register that requires extra
// privileges to read.
//
// uint64_t isar0_el1;
// asm("mrs %[reg], ID_AA64ISAR0_EL1\n" : [reg] "=r"(isar0_el1) : :);
// const bool HasDotProductInstructions = ((isar0_el1 >> 44) & 0xfu) == 0x1u;
const bool HasDotProductInstructions = MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeonDot();
if (HasDotProductInstructions) {
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchUdot;
this->GemmU8S8Dispatch = &MlasGemmU8X8DispatchUdot;
this->GemmS8S8Dispatch = &MlasGemmS8S8DispatchSdot;
this->SymmQgemmDispatch = &MlasSymmQgemmS8DispatchSdot;
this->ConvSymU8S8Dispatch = &MlasConvSymU8DispatchDot;
this->ConvSymS8S8Dispatch = &MlasConvSymS8DispatchDot;
}
#if defined(USE_KLEIDIAI)
if(MLAS_CPUIDINFO::GetCPUIDInfo().HasArm_SME()){
this->MlasSGemmBatchOverride = ArmKleidiAI::MlasGemmBatch;
this->MlasSGemmPackBSizeOverride = ArmKleidiAI::MlasGemmPackBSize;
this->MlasSGemmPackBOverride = ArmKleidiAI::MlasGemmPackB;
this->MlasDynamicQGemmBatchOverride = ArmKleidiAI::MlasDynamicQGemmBatch;
this->MlasDynamicQGemmPackBSizeOverride = ArmKleidiAI::MlasDynamicQGemmPackBSize;
this->MlasDynamicQGemmPackBOverride = ArmKleidiAI::MlasDynamicQGemmPackB;
this->MlasConvPrepareOverride = ArmKleidiAI::MlasConvPrepare;
this->MlasConvOverride = ArmKleidiAI::MlasConv;
#if defined(__aarch64__) && defined(__linux__)
// Currently only an SME2 variant of SBGEMM exists
if (ArmKleidiAI::UseSME2){
this->MlasSBGemmBatchOverride = ArmKleidiAI::MlasSBGemmBatch;
this->MlasSBGemmPackBSizeOverride = ArmKleidiAI::MlasSBGemmPackBSize;
this->MlasSBGemmPackBOverride = ArmKleidiAI::MlasSBGemmPackB;
}
#endif
}
#endif
#if defined(MLAS_USE_SVE)
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
this->ErfKernelRoutine = MlasSveErfKernel;
this->LogisticKernelRoutine = MlasSveLogisticKernel;
this->ReduceMaximumF32Kernel = MlasSveReduceMaximumF32Kernel;
this->ComputeSumExpF32Kernel = MlasSveComputeSumExpF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasSveComputeLogSoftmaxOutputF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasSveComputeSoftmaxOutputF32Kernel;
}
else{
this->ErfKernelRoutine = MlasErfKernel;
this->LogisticKernelRoutine = MlasLogisticKernel;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
this->ComputeSumExpF32Kernel = MlasComputeSumExpF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
}
#endif
#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && !defined(_WIN32)
#if defined(MLAS_USE_SVE)
if (MLAS_CPUIDINFO::GetCPUIDInfo().HasArmSve()) {
this->ErfFP16KernelRoutine = MlasSveErfFP16Kernel;
this->GeluFP16KernelRoutine = MlasSveGeluFP16Kernel;
this->TanhFP16KernelRoutine = MlasSveTanhFP16Kernel;
}
else{
this->ErfFP16KernelRoutine = MlasNeonErfFP16Kernel;
this->GeluFP16KernelRoutine = MlasNeonGeluFP16Kernel;
}
#else
this->ErfFP16KernelRoutine = MlasNeonErfFP16Kernel;
this->GeluFP16KernelRoutine = MlasNeonGeluFP16Kernel;
#endif
#endif
//
// Check if the processor supports ASIMD I8MM instructions.
//
const bool HasI8MMInstructions = MLAS_CPUIDINFO::GetCPUIDInfo().HasArmNeon_I8MM();
if (HasI8MMInstructions) {
#if defined(__linux__)
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchUmmla;
this->GemmU8S8Dispatch = &MlasGemmU8X8DispatchUmmla;
this->GemmS8S8Dispatch = &MlasGemmS8S8DispatchSmmla;
#endif
}
this->ArmNeonIsQuantActivationsUnsigned = HasI8MMInstructions ? false : true;
this->QNBitGemmDispatch = &GetMlasQNBitGemmDispatchNeon(HasDotProductInstructions, HasI8MMInstructions);
#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED)
this->CastF16ToF32Kernel = &MlasCastF16ToF32KernelNeon;
this->CastF32ToF16Kernel = &MlasCastF32ToF16KernelNeon;
#endif
#endif // MLAS_TARGET_ARM64
#if defined(MLAS_TARGET_POWER)
this->GemmFloatKernel = MlasSgemmKernel;
this->GemmDoubleKernel = MlasDgemmKernel;
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8Kernel;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8Kernel;
this->QuantizeLinearS16Kernel = MlasQuantizeLinearS16Kernel;
this->QuantizeLinearU16Kernel = MlasQuantizeLinearU16Kernel;
this->QuantizeLinearS4Kernel = MlasQuantizeLinearS4Kernel;
this->QuantizeLinearU4Kernel = MlasQuantizeLinearU4Kernel;
#if defined(__linux__)
unsigned long hwcap2 = getauxval(AT_HWCAP2);
bool HasP9Instructions = hwcap2 & PPC_FEATURE2_ARCH_3_00;
#elif defined(_AIX)
bool HasP9Instructions = __power_9_andup();
#elif defined(__FreeBSD__)
unsigned long hwcap2;
elf_aux_info(AT_HWCAP2, &hwcap2, sizeof(hwcap2));
bool HasP9Instructions = hwcap2 & PPC_FEATURE2_ARCH_3_00;
#endif // __linux__
if (HasP9Instructions) {
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelVSX;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelVSX;
}
#if defined(POWER10)
#if (defined(__GNUC__) && ((__GNUC__ > 10) || (__GNUC__== 10 && __GNUC_MINOR__ >= 2))) || \
(defined(__clang__) && (__clang_major__ >= 12))
#if defined(__linux__) || defined(__FreeBSD__)
bool HasP10Instructions = ((hwcap2 & PPC_FEATURE2_MMA) && (hwcap2 & PPC_FEATURE2_ARCH_3_1));
#elif defined(_AIX)
bool HasP10Instructions = (__power_10_andup() && __power_mma_version() == MMA_V31);
#endif // __linux__
if (HasP10Instructions) {
this->GemmFloatKernel = MlasSgemmKernelPOWER10;
this->GemmDoubleKernel = MlasDgemmKernelPOWER10;
this->GemmU8X8Dispatch = &MlasGemm8X8DispatchPOWER10;
}
#endif
#endif
#endif // MLAS_TARGET_POWER
#if defined(MLAS_TARGET_S390X)
this->GemmFloatKernel = MlasSgemmKernel;
this->GemmDoubleKernel = MlasDgemmKernel;
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8Kernel;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8Kernel;
this->QuantizeLinearS16Kernel = MlasQuantizeLinearS16Kernel;
this->QuantizeLinearU16Kernel = MlasQuantizeLinearU16Kernel;
this->QuantizeLinearS4Kernel = MlasQuantizeLinearS4Kernel;
this->QuantizeLinearU4Kernel = MlasQuantizeLinearU4Kernel;
bool HasVXEInstructions = getauxval(AT_HWCAP) & HWCAP_S390_VXE;
if (HasVXEInstructions) {
this->GemmFloatKernel = MlasSgemmKernelZVECTOR;
this->GemmU8X8Dispatch = &MlasGemm8X8DispatchZVECTOR;
this->QuantizeLinearS8Kernel = MlasQuantizeLinearS8KernelZVECTOR;
this->QuantizeLinearU8Kernel = MlasQuantizeLinearU8KernelZVECTOR;
}
#endif // MLAS_TARGET_S390X
#if defined(MLAS_TARGET_LARCH64)
//
// Default to the baseline LSX support.
//
int hwcap = getauxval(AT_HWCAP);
bool cap_lasx = hwcap & HWCAP_LOONGARCH_LASX;
bool cap_lsx = hwcap & HWCAP_LOONGARCH_LSX;
if( cap_lasx ){
this->GemmFloatKernel = MlasGemmFloatKernelLasx;
this->GemmDoubleKernel = MlasGemmDoubleKernelLasx;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelLasx;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelLasx;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelLasx;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelLasx;
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelLasx;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelLasx;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelLasx;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32KernelLasx;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32KernelLasx;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32KernelLasx;
this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4Lasx;
// add new sqn-lasx kernel
this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchLasx;
this->GemmU8S8Dispatch = &MlasGemmU8X8DispatchLSX;
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchLSX;
this->GemmS8S8Dispatch = &MlasGemmS8S8DispatchLSX;
this->GemmS8U8Dispatch = &MlasGemmS8U8DispatchLSX;
}else if( cap_lsx ){
this->GemmFloatKernel = MlasGemmFloatKernelLSX;
this->GemmU8S8Dispatch = &MlasGemmU8X8DispatchLSX;
this->GemmU8U8Dispatch = &MlasGemmU8X8DispatchLSX;
this->GemmS8S8Dispatch = &MlasGemmS8S8DispatchLSX;
this->GemmS8U8Dispatch = &MlasGemmS8U8DispatchLSX;
this->TransposePackB16x4Routine = MlasSgemmTransposePackB16x4LSX;
this->GemmDoubleKernel = MlasGemmDoubleKernelLSX;
this->ConvNchwFloatKernel = MlasConvNchwFloatKernelLSX;
this->ConvNchwcFloatKernel = MlasConvNchwcFloatKernelLSX;
this->ConvDepthwiseFloatKernel = MlasConvDepthwiseFloatKernelLSX;
this->ConvPointwiseFloatKernel = MlasConvPointwiseFloatKernelLSX;
this->PoolFloatKernel[MlasMaximumPooling] = MlasPoolMaximumFloatKernelLSX;
this->PoolFloatKernel[MlasAveragePoolingExcludePad] = MlasPoolAverageExcludePadFloatKernelLSX;
this->PoolFloatKernel[MlasAveragePoolingIncludePad] = MlasPoolAverageIncludePadFloatKernelLSX;
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
}else{
this->ReduceMaximumF32Kernel = MlasReduceMaximumF32Kernel;
this->ComputeSoftmaxOutputF32Kernel = MlasComputeSoftmaxOutputF32Kernel;
this->ComputeLogSoftmaxOutputF32Kernel = MlasComputeLogSoftmaxOutputF32Kernel;
}
this->NchwcBlockSize = 8;
// this->PreferredBufferAlignment = MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT;
// this->MaximumThreadCount = MLAS_MAXIMUM_THREAD_COUNT;
#endif // MLAS_TARGET_LARCH64
}
#endif // MLAS_GEMM_ONLY
size_t
MLASCALL
MlasGetPreferredBufferAlignment(
void
)
/*++
Routine Description:
This routine returns the preferred byte alignment for buffers that are used
with this library. Buffers that are not byte aligned to this value will
function, but will not achieve best performance.
Arguments:
None.
Return Value:
Returns the preferred byte alignment for buffers.
--*/
{
#if defined(MLAS_TARGET_AMD64)
return GetMlasPlatform().PreferredBufferAlignment;
#else
return MLAS_DEFAULT_PREFERRED_BUFFER_ALIGNMENT;
#endif
}
#ifdef MLAS_TARGET_AMD64_IX86
bool
MLASCALL
MlasPlatformU8S8Overflow(
void
)
{
const auto& p = GetMlasPlatform();
return p.GemmU8U8Dispatch != p.GemmU8S8Dispatch;
}
#endif
thread_local size_t ThreadedBufSize = 0;
#ifdef _MSC_VER
thread_local std::unique_ptr<uint8_t, decltype(&_aligned_free)> ThreadedBufHolder(nullptr, &_aligned_free);
#else
thread_local std::unique_ptr<uint8_t, decltype(&free)> ThreadedBufHolder(nullptr, &free);
#endif
3rdparty/mlas/lib/power/SgemmKernelPOWER10.cpp
Vendored
+697
View File
@@ -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;
}
3rdparty/mlas/lib/power/SgemmKernelPackA.S
Vendored
+247
View File
@@ -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
3rdparty/mlas/lib/power/SgemmKernelPower.cpp
Vendored
+87
View File
@@ -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;
}
3rdparty/mlas/lib/qgemm.h
Vendored
+930
View File
@@ -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;
}
3rdparty/mlas/lib/riscv64/sgemm_kernel_rvv.cpp
Vendored
+275
View File
@@ -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)
3rdparty/mlas/lib/riscv64/sgemm_pack_b_rvv.cpp
Vendored
+115
View File
@@ -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)
3rdparty/mlas/lib/s390x/SgemmKernel.cpp
Vendored
+87
View File
@@ -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;
}
3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.cpp
Vendored
+451
View File
@@ -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;
}
3rdparty/mlas/lib/scalar/SconvDepthwiseKernelScalar.cpp
Vendored
+193
View File
@@ -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);
}
3rdparty/mlas/lib/scalar/SgemmKernelScalar.cpp
Vendored
+480
View File
@@ -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);
}
3rdparty/mlas/lib/scalar/SgemvKernelScalar.cpp
Vendored
+169
View File
@@ -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++;
}
}
3rdparty/mlas/lib/sgemm.cpp
Vendored
+1740
View File
@@ -0,0 +1,1740 @@
/*++
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.
Module Name:
sgemm.cpp
Abstract:
This module implements the single precision matrix/matrix multiply
operation (SGEMM).
--*/
#include "mlasi.h"
//
// Define the number of rows from matrix A to transpose to a local buffer.
//
// N.B. AVX processes a maximum of 4 rows, FMA3 processes a maximum of 6
// rows, and AVX512F processes a maximum of 12 rows.
//
#define MLAS_SGEMM_TRANSA_ROWS 12
//
// Define the parameters to execute segments of a SGEMM operation on worker
// threads.
//
void
MlasSgemmMultiplyBeta(
float* C,
size_t CountM,
size_t CountN,
size_t ldc,
float beta
)
/*++
Routine Description:
This routine multiplies all elements of the output matrix by the beta
scalar value.
Arguments:
C - Supplies the address of matrix C.
CountM - Supplies the number of rows from matrix C.
CountN - Supplies the number of columns from matrix C.
ldc - Supplies the first dimension of matrix C.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
Return Value:
None.
--*/
{
MLAS_FLOAT32X4 BetaBroadcast = MlasBroadcastFloat32x4(beta);
while (CountM-- > 0) {
float* c = C;
size_t n = CountN;
while (n >= 4) {
MlasStoreFloat32x4(c, MlasMultiplyFloat32x4(MlasLoadFloat32x4(c), BetaBroadcast));
c += 4;
n -= 4;
}
while (n > 0) {
#if defined(MLAS_SSE2_INTRINSICS)
_mm_store_ss(c, _mm_mul_ss(_mm_load_ss(c), BetaBroadcast));
#else
*c = *c * beta;
#endif
c += 1;
n -= 1;
}
C += ldc;
}
}
void
MlasSgemmTransposeA(
float* D,
const float* A,
size_t lda,
size_t CountY,
size_t CountX
)
/*++
Routine Description:
This routine transposes elements from the source matrix to the destination
buffer.
Arguments:
D - Supplies the address of the destination buffer.
A - Supplies the address of the source matrix.
lda - Supplies the number of elements per row of the source matrix.
CountY - Supplies the number of columns of the source matrix to transpose.
CountX - Supplies the number of rows of the source matrix to transpose.
Return Value:
None.
--*/
{
size_t ldd = CountX;
//
// Transpose elements from matrix A into the destination buffer 4 columns
// at a time.
//
while (CountX >= 4) {
float* d = D;
const float* a = A;
size_t y = CountY;
do {
float t0 = a[0];
float t1 = a[lda];
float t2 = a[lda * 2];
float t3 = a[lda * 3];
d[0] = t0;
d[1] = t1;
d[2] = t2;
d[3] = t3;
d += ldd;
a += 1;
y--;
} while (y > 0);
D += 4;
A += lda * 4;
CountX -= 4;
}
//
// Transpose elements from matrix A into the destination buffer for the
// remaining columns.
//
if (CountX >= 2) {
float* d = D;
const float* a = A;
size_t y = CountY;
do {
float t0 = a[0];
float t1 = a[lda];
d[0] = t0;
d[1] = t1;
d += ldd;
a += 1;
y--;
} while (y > 0);
D += 2;
A += lda * 2;
CountX -= 2;
}
if (CountX >= 1) {
float* d = D;
const float* a = A;
size_t y = CountY;
do {
d[0] = a[0];
d += ldd;
a += 1;
y--;
} while (y > 0);
}
}
#if !defined(MLAS_TARGET_WASM_SCALAR)
void
MlasSgemmCopyPackB(
float* D,
const float* B,
size_t ldb,
size_t CountX,
size_t CountY
)
/*++
Routine Description:
This routine copies elements from the source matrix to the destination
packed buffer.
Columns of 16 elements from the source matrix are unrolled to be physically
contiguous for better locality inside the SGEMM kernels. Any remaining
columns less than 16 elements wide are zero-padded.
Arguments:
D - Supplies the address of the destination packed buffer.
B - Supplies the address of the source matrix.
ldb - Supplies the number of elements per row of the source matrix.
CountX - Supplies the number of columns of the source matrix to copy.
CountY - Supplies the number of rows of the source matrix to copy.
Return Value:
None.
--*/
{
#if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV) && !defined(FORCE_GENERIC_ALGORITHMS)
if (GetMlasPlatform().GemmFloatKernel != nullptr) {
MlasSgemmCopyPackBRvv(D, B, ldb, CountX, CountY);
return;
}
#endif
//
// Copy data from matrix B into the destination buffer 16 columns at a
// time.
//
while (CountX >= 16) {
const float* b = B;
size_t y = CountY;
do {
#if defined(MLAS_NEON_INTRINSICS)
vst4q_f32(D, vld4q_f32(b));
#else
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
MLAS_FLOAT32X4 t1 = MlasLoadFloat32x4(&b[4]);
MLAS_FLOAT32X4 t2 = MlasLoadFloat32x4(&b[8]);
MLAS_FLOAT32X4 t3 = MlasLoadFloat32x4(&b[12]);
MlasStoreAlignedFloat32x4(&D[0], t0);
MlasStoreAlignedFloat32x4(&D[4], t1);
MlasStoreAlignedFloat32x4(&D[8], t2);
MlasStoreAlignedFloat32x4(&D[12], t3);
#endif
D += 16;
b += ldb;
y--;
} while (y > 0);
B += 16;
CountX -= 16;
}
//
// Special case the handling of the remaining columns less than 16 elements
// wide.
//
if (CountX > 0) {
MLAS_FLOAT32X4 ZeroFloat32x4 = MlasZeroFloat32x4();
#if defined(MLAS_NEON_INTRINSICS)
float32x4x4_t ZeroFloat32x4x4 = { ZeroFloat32x4, ZeroFloat32x4, ZeroFloat32x4, ZeroFloat32x4 };
#endif
size_t y = CountY;
do {
float* d = D;
const float* b = B;
#if defined(MLAS_NEON_INTRINSICS)
vst4q_f32(d, ZeroFloat32x4x4);
#else
MlasStoreAlignedFloat32x4(d, ZeroFloat32x4);
MlasStoreAlignedFloat32x4(d + 4, ZeroFloat32x4);
MlasStoreAlignedFloat32x4(d + 8, ZeroFloat32x4);
MlasStoreAlignedFloat32x4(d + 12, ZeroFloat32x4);
#endif
if ((CountX & 8) != 0) {
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(b);
MLAS_FLOAT32X4 t1 = MlasLoadFloat32x4(b + 4);
MlasStoreAlignedFloat32x4(d, t0);
MlasStoreAlignedFloat32x4(d + 4, t1);
d += 8;
b += 8;
}
if ((CountX & 4) != 0) {
MlasStoreAlignedFloat32x4(d, MlasLoadFloat32x4(b));
d += 4;
b += 4;
}
if ((CountX & 2) != 0) {
float t0 = b[0];
float t1 = b[1];
d[0] = t0;
d[1] = t1;
d += 2;
b += 2;
}
if ((CountX & 1) != 0) {
d[0] = b[0];
}
D += 16;
B += ldb;
y--;
} while (y > 0);
}
}
template<unsigned N>
inline
void
MlasSgemmTransposePackBNx4(
float* D,
const float* B,
size_t ldb
)
/*++
Routine Description:
This routine transposes elements from the source matrix to the destination
packed buffer.
4 columns of N rows from the source matrix are transposed to N columns of 4
rows in the destination packed buffer.
Arguments:
D - Supplies the address of the destination packed buffer.
B - Supplies the address of the source matrix.
ldb - Supplies the number of elements per row of the source matrix.
Return Value:
None.
--*/
{
for (unsigned n = 0; n < N / 4; n++) {
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&B[ldb * 0]);
MLAS_FLOAT32X4 t1 = MlasLoadFloat32x4(&B[ldb * 1]);
MLAS_FLOAT32X4 t2 = MlasLoadFloat32x4(&B[ldb * 2]);
MLAS_FLOAT32X4 t3 = MlasLoadFloat32x4(&B[ldb * 3]);
#if defined(MLAS_NEON_INTRINSICS)
float32x4x2_t z0 = vzipq_f32(t0, t2);
float32x4x2_t z1 = vzipq_f32(t1, t3);
float32x4x2_t o0 = vzipq_f32(z0.val[0], z1.val[0]);
float32x4x2_t o1 = vzipq_f32(z0.val[1], z1.val[1]);
t0 = o0.val[0];
t1 = o0.val[1];
t2 = o1.val[0];
t3 = o1.val[1];
#else
MLAS_FLOAT32X4 z0 = MlasInterleaveLowFloat32x4(t0, t2);
MLAS_FLOAT32X4 z1 = MlasInterleaveHighFloat32x4(t0, t2);
MLAS_FLOAT32X4 z2 = MlasInterleaveLowFloat32x4(t1, t3);
MLAS_FLOAT32X4 z3 = MlasInterleaveHighFloat32x4(t1, t3);
t0 = MlasInterleaveLowFloat32x4(z0, z2);
t1 = MlasInterleaveHighFloat32x4(z0, z2);
t2 = MlasInterleaveLowFloat32x4(z1, z3);
t3 = MlasInterleaveHighFloat32x4(z1, z3);
#endif
MlasStoreAlignedFloat32x4(&D[0], t0);
MlasStoreAlignedFloat32x4(&D[16], t1);
MlasStoreAlignedFloat32x4(&D[32], t2);
MlasStoreAlignedFloat32x4(&D[48], t3);
D += 4;
B += ldb * 4;
}
}
void
MlasSgemmTransposePackB(
float* D,
const float* B,
size_t ldb,
size_t CountY,
size_t CountX
)
/*++
Routine Description:
This routine transposes elements from the source matrix to the destination
packed buffer.
Columns of 16 elements from the source matrix are unrolled to be physically
contiguous for better locality inside the SGEMM kernels. Any remaining
columns less than 16 elements wide are zero-padded.
Arguments:
D - Supplies the address of the destination packed buffer.
B - Supplies the address of the source matrix.
ldb - Supplies the number of elements per row of the source matrix.
CountY - Supplies the number of rows of the source matrix to transpose.
CountX - Supplies the number of columns of the source matrix to transpose.
Return Value:
None.
--*/
{
//
// Transpose elements from matrix B into the packed buffer 16 rows at a
// time.
//
while (CountY >= 16) {
const float* b = B;
size_t x = CountX;
#if defined(MLAS_TARGET_AMD64) || defined(MLAS_TARGET_LARCH64)
MLAS_SGEMM_TRANSPOSE_PACKB_BLOCK_ROUTINE* SgemmTransposePackB16x4Routine =
GetMlasPlatform().TransposePackB16x4Routine;
while (x >= 4) {
SgemmTransposePackB16x4Routine(&D[0], &b[0], ldb);
D += 16 * 4;
b += 4;
x -= 4;
}
#else
while (x >= 4) {
MlasSgemmTransposePackBNx4<16>(&D[0], &b[0], ldb);
D += 16 * 4;
b += 4;
x -= 4;
}
#endif
while (x > 0) {
float t0 = b[0];
float t1 = b[ldb];
float t2 = b[ldb * 2];
float t3 = b[ldb * 3];
float t4 = b[ldb * 4];
float t5 = b[ldb * 5];
float t6 = b[ldb * 6];
float t7 = b[ldb * 7];
float t8 = b[ldb * 8];
float t9 = b[ldb * 9];
float t10 = b[ldb * 10];
float t11 = b[ldb * 11];
float t12 = b[ldb * 12];
float t13 = b[ldb * 13];
float t14 = b[ldb * 14];
float t15 = b[ldb * 15];
D[0] = t0;
D[1] = t1;
D[2] = t2;
D[3] = t3;
D[4] = t4;
D[5] = t5;
D[6] = t6;
D[7] = t7;
D[8] = t8;
D[9] = t9;
D[10] = t10;
D[11] = t11;
D[12] = t12;
D[13] = t13;
D[14] = t14;
D[15] = t15;
D += 16;
b += 1;
x--;
}
B += ldb * 16;
CountY -= 16;
}
//
// Special case the handling of the less than 16 remaining rows.
//
if (CountY > 0) {
MLAS_FLOAT32X4 ZeroFloat32x4 = MlasZeroFloat32x4();
size_t x = CountX;
//
// Transpose 4 columns at a time.
//
while (x >= 4) {
float* d = D;
const float* b = B;
if ((CountY & 8) != 0) {
MlasSgemmTransposePackBNx4<8>(&d[0], &b[0], ldb);
d += 8;
b += ldb * 8;
} else {
MlasStoreAlignedFloat32x4(&d[8], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[12], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[24], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[28], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[40], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[44], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[56], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[60], ZeroFloat32x4);
}
if ((CountY & 4) != 0) {
MlasSgemmTransposePackBNx4<4>(&d[0], &b[0], ldb);
d += 4;
b += ldb * 4;
} else {
MlasStoreAlignedFloat32x4(&d[4], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[20], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[36], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[52], ZeroFloat32x4);
}
MlasStoreAlignedFloat32x4(&d[0], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[16], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[32], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[48], ZeroFloat32x4);
if ((CountY & 2) != 0) {
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
MLAS_FLOAT32X4 t1 = MlasLoadFloat32x4(&b[ldb]);
#if defined(MLAS_SSE2_INTRINSICS)
__m128 v0 = _mm_unpacklo_ps(t0, t1);
__m128 v1 = _mm_unpackhi_ps(t0, t1);
_mm_storel_pi((__m64*)&d[0], v0);
_mm_storeh_pi((__m64*)&d[16], v0);
_mm_storel_pi((__m64*)&d[32], v1);
_mm_storeh_pi((__m64*)&d[48], v1);
#else
MlasStoreLaneFloat32x4<0>(&d[0], t0);
MlasStoreLaneFloat32x4<0>(&d[1], t1);
MlasStoreLaneFloat32x4<1>(&d[16], t0);
MlasStoreLaneFloat32x4<1>(&d[17], t1);
MlasStoreLaneFloat32x4<2>(&d[32], t0);
MlasStoreLaneFloat32x4<2>(&d[33], t1);
MlasStoreLaneFloat32x4<3>(&d[48], t0);
MlasStoreLaneFloat32x4<3>(&d[49], t1);
#endif
d += 2;
b += ldb * 2;
}
if ((CountY & 1) != 0) {
#if defined(MLAS_NEON_INTRINSICS)
MLAS_FLOAT32X4 t0 = MlasLoadFloat32x4(&b[0]);
MlasStoreLaneFloat32x4<0>(&d[0], t0);
MlasStoreLaneFloat32x4<1>(&d[16], t0);
MlasStoreLaneFloat32x4<2>(&d[32], t0);
MlasStoreLaneFloat32x4<3>(&d[48], t0);
#else
d[0] = b[0];
d[16] = b[1];
d[32] = b[2];
d[48] = b[3];
#endif
}
D += 16 * 4;
B += 4;
x -= 4;
}
//
// Transpose the remaining columns.
//
while (x > 0) {
float* d = D;
const float* b = B;
if ((CountY & 8) != 0) {
float t0 = b[0];
float t1 = b[ldb];
float t2 = b[ldb * 2];
float t3 = b[ldb * 3];
float t4 = b[ldb * 4];
float t5 = b[ldb * 5];
float t6 = b[ldb * 6];
float t7 = b[ldb * 7];
d[0] = t0;
d[1] = t1;
d[2] = t2;
d[3] = t3;
d[4] = t4;
d[5] = t5;
d[6] = t6;
d[7] = t7;
d += 8;
b += ldb * 8;
} else {
MlasStoreAlignedFloat32x4(&d[8], ZeroFloat32x4);
MlasStoreAlignedFloat32x4(&d[12], ZeroFloat32x4);
}
if ((CountY & 4) != 0) {
float t0 = b[0];
float t1 = b[ldb];
float t2 = b[ldb * 2];
float t3 = b[ldb * 3];
d[0] = t0;
d[1] = t1;
d[2] = t2;
d[3] = t3;
d += 4;
b += ldb * 4;
} else {
MlasStoreAlignedFloat32x4(&d[4], ZeroFloat32x4);
}
MlasStoreAlignedFloat32x4(d, ZeroFloat32x4);
if ((CountY & 2) != 0) {
float t0 = b[0];
float t1 = b[ldb];
d[0] = t0;
d[1] = t1;
d += 2;
b += ldb * 2;
}
if ((CountY & 1) != 0) {
d[0] = b[0];
}
D += 16;
B += 1;
x--;
}
}
}
#else //defined(MLAS_TARGET_WASM_SCALAR)
void
MlasSgemmCopyPackB(
float* D,
const float* B,
size_t ldb,
size_t CountX,
size_t CountY
)
/*++
Routine Description:
This routine copies elements from the source matrix to the destination
packed buffer.
Columns of 16 elements from the source matrix are unrolled to be physically
contiguous for better locality inside the SGEMM kernels. Any remaining
columns less than 16 elements wide are zero-padded.
Arguments:
D - Supplies the address of the destination packed buffer.
B - Supplies the address of the source matrix.
ldb - Supplies the number of elements per row of the source matrix.
CountX - Supplies the number of columns of the source matrix to copy.
CountY - Supplies the number of rows of the source matrix to copy.
Return Value:
None.
--*/
{
//
// Copy data from matrix B into the destination buffer 16 columns at a
// time.
//
while (CountX >= 16) {
const float* b = B;
size_t y = CountY;
do {
std::copy_n(b, 16, D);
D += 16;
b += ldb;
y--;
} while (y > 0);
B += 16;
CountX -= 16;
}
//
// Special case the handling of the remaining columns less than 16 elements
// wide.
//
if (CountX > 0) {
size_t y = CountY;
do {
std::fill_n(D, 16, 0.0f);
std::copy_n(B, CountX, D);
D += 16;
B += ldb;
y--;
} while (y > 0);
}
}
void
MlasSgemmTransposePackB(
float* D,
const float* B,
size_t ldb,
size_t CountY,
size_t CountX
)
/*++
Routine Description:
This routine transposes elements from the source matrix to the destination
packed buffer.
Columns of 16 elements from the source matrix are unrolled to be physically
contiguous for better locality inside the SGEMM kernels. Any remaining
columns less than 16 elements wide are zero-padded.
Arguments:
D - Supplies the address of the destination packed buffer.
B - Supplies the address of the source matrix.
ldb - Supplies the number of elements per row of the source matrix.
CountY - Supplies the number of rows of the source matrix to transpose.
CountX - Supplies the number of columns of the source matrix to transpose.
Return Value:
None.
--*/
{
//
// Transpose elements from matrix B into the packed buffer 16 rows at a
// time.
//
while (CountY >= 16) {
const float* b = B;
size_t x = CountX;
while (x >= 4) {
for (size_t row = 0; row < 16; row++) {
D[0 * 16 + row] = b[row * ldb + 0];
D[1 * 16 + row] = b[row * ldb + 1];
D[2 * 16 + row] = b[row * ldb + 2];
D[3 * 16 + row] = b[row * ldb + 3];
}
D += 16 * 4;
b += 4;
x -= 4;
}
while (x > 0) {
for (size_t row = 0; row < 16; row++) {
D[row] = b[row * ldb];
}
D += 16;
b += 1;
x--;
}
B += ldb * 16;
CountY -= 16;
}
//
// Special case the handling of the less than 16 remaining rows.
//
if (CountY > 0) {
size_t x = CountX;
//
// Transpose 4 columns at a time.
//
while (x >= 4) {
std::fill_n(D, 16 * 4, 0.0f);
for (size_t row = 0; row < CountY; row++) {
D[0 * 16 + row] = B[row * ldb + 0];
D[1 * 16 + row] = B[row * ldb + 1];
D[2 * 16 + row] = B[row * ldb + 2];
D[3 * 16 + row] = B[row * ldb + 3];
}
D += 16 * 4;
B += 4;
x -= 4;
}
//
// Transpose the remaining columns.
//
while (x > 0) {
std::fill_n(D, 16, 0.0f);
for (size_t row = 0; row < CountY; row++) {
D[row] = B[row * ldb];
}
D += 16;
B += 1;
x--;
}
}
}
#endif
MLAS_FORCEINLINE
float*
MlasSgemmKernelLoop(
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 steps through the rows of the input and output matrices calling
the kernel until all rows have been processed.
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 number of rows from matrix A and matrix C 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 alpha 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 next address of matrix C.
--*/
{
while (CountM > 0) {
size_t RowsHandled;
#if (defined(MLAS_TARGET_AMD64_IX86) || defined(MLAS_TARGET_POWER) || defined(MLAS_TARGET_S390X) || defined(MLAS_TARGET_LARCH64)) && !defined(FORCE_GENERIC_ALGORITHMS)
RowsHandled = GetMlasPlatform().GemmFloatKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode);
#elif defined(MLAS_TARGET_RISCV64) && !defined(FORCE_GENERIC_ALGORITHMS)
if (GetMlasPlatform().GemmFloatKernel != nullptr) {
RowsHandled = GetMlasPlatform().GemmFloatKernel(A, B, C, CountK, CountM, CountN, lda, ldc, alpha, ZeroMode);
} else if (ZeroMode) {
RowsHandled = MlasSgemmKernelZero(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
} else {
RowsHandled = MlasSgemmKernelAdd(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
}
#else
if (ZeroMode) {
RowsHandled = MlasSgemmKernelZero(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
} else {
RowsHandled = MlasSgemmKernelAdd(A, B, C, CountK, CountM, CountN, lda, ldc, alpha);
}
#endif
C += ldc * RowsHandled;
A += lda * RowsHandled;
CountM -= RowsHandled;
}
return C;
}
void
MlasSgemmOperation(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t M,
size_t N,
size_t K,
float alpha,
const float* A,
size_t lda,
const float* B,
size_t ldb,
float beta,
float* C,
size_t ldc
)
/*++
Routine Description:
This routine implements the single precision matrix/matrix multiply
operation (SGEMM).
Arguments:
TransA - Supplies the transpose operation for matrix A.
TransB - Supplies the transpose operation for matrix B.
M - Supplies the number of rows of matrix A and matrix C.
N - Supplies the number of columns of matrix B and matrix C.
K - Supplies the number of columns of matrix A and the number of rows of
matrix B.
alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
A - Supplies the address of matrix A.
lda - Supplies the first dimension of matrix A.
B - Supplies the address of matrix B.
ldb - Supplies the first dimension of matrix B.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
C - Supplies the address of matrix C.
ldc - Supplies the first dimension of matrix C.
Return Value:
None.
--*/
{
float PanelA[MLAS_SGEMM_TRANSA_ROWS * MLAS_SGEMM_STRIDEK];
MLAS_DECLSPEC_ALIGN(float PanelB[MLAS_SGEMM_STRIDEN * MLAS_SGEMM_STRIDEK], 16 * sizeof(float));
//
// Handle the special case of K equals zero. Apply the beta multiplier to
// the output matrix and exit.
//
if (K == 0) {
MlasSgemmMultiplyBeta(C, M, N, ldc, beta);
return;
}
//
// Handle the special case of a small M. The data from matrix B is not
// referenced multiple times, so using a local packed buffer is a wasted
// memory copy.
//
if (M == 1 && TransA == CblasNoTrans && alpha == 1.0f && (beta == 0.0f || beta == 1.0f)) {
#if !defined(FORCE_GENERIC_ALGORITHMS)
#if defined(MLAS_TARGET_AMD64)
MLAS_SGEMM_KERNEL_M1_ROUTINE* SgemmKernelM1Routine;
if (TransB == CblasNoTrans) {
SgemmKernelM1Routine = GetMlasPlatform().KernelM1Routine;
} else {
SgemmKernelM1Routine = GetMlasPlatform().KernelM1TransposeBRoutine;
}
if (SgemmKernelM1Routine != nullptr) {
SgemmKernelM1Routine(A, B, C, K, N, ldb, beta);
return;
}
#elif defined(MLAS_TARGET_ARM64) || defined(MLAS_TARGET_WASM)
if (TransB == CblasNoTrans) {
MlasGemvFloatKernel(A, B, C, K, N, ldb, (beta == 0.0f));
return;
}
#endif
#endif // !defined(FORCE_GENERIC_ALGORITHMS)
}
//
// Handle the case when both B and C are column-vectors that are contiguous in memory.
// Because transposition of such vectors doesn't change their layout, and
// Transpose(A*B) = Transpose(B) * Transpose(A), we can apply the same 'small-M'
// optimization as above, with A and B flipped.
//
if (N == 1 && ldb == 1 && ldc == 1 && alpha == 1.0f && (beta == 0.0f || beta == 1.0f)) {
#if defined(MLAS_TARGET_AMD64) && !defined(FORCE_GENERIC_ALGORITHMS)
MLAS_SGEMM_KERNEL_M1_ROUTINE* SgemmKernelM1Routine;
if (TransA == CblasNoTrans) {
SgemmKernelM1Routine = GetMlasPlatform().KernelM1TransposeBRoutine;
} else {
SgemmKernelM1Routine = GetMlasPlatform().KernelM1Routine;
}
if (SgemmKernelM1Routine != nullptr) {
SgemmKernelM1Routine(B, A, C, K, M, lda, beta);
return;
}
#endif
}
//
// Compute the strides to step through slices of the input matrices.
//
// Expand the N stride if K is small or expand the K stride if N is small
// for better utilization of the B panel. Avoid changing the K stride if
// the A panel needs to be used for transposing.
//
size_t StrideN = MLAS_SGEMM_STRIDEN;
size_t StrideK = MLAS_SGEMM_STRIDEK;
if (N >= K) {
while (StrideK / 2 >= K) {
StrideN *= 2;
StrideK /= 2;
}
} else if (TransA == CblasNoTrans) {
while (StrideN > 16 && StrideN / 2 >= N) {
StrideK *= 2;
StrideN /= 2;
}
}
//
// Step through each slice of matrix B along the N dimension.
//
size_t CountN;
for (size_t n = 0; n < N; n += CountN) {
CountN = std::min(N - n, StrideN);
//
// Multiply the output matrix by beta as needed.
//
if (beta != 0.0f && beta != 1.0f) {
MlasSgemmMultiplyBeta(C + n, M, CountN, ldc, beta);
}
//
// Step through each slice of matrix B along the K dimension.
//
size_t CountK;
bool ZeroMode = (beta == 0.0f);
for (size_t k = 0; k < K; k += CountK) {
CountK = std::min(K - k, StrideK);
//
// Copy or transpose a panel of matrix B to a local packed buffer.
//
if (TransB == CblasNoTrans) {
MlasSgemmCopyPackB(PanelB, B + n + k * ldb, ldb, CountN, CountK);
} else {
MlasSgemmTransposePackB(PanelB, B + k + n * ldb, ldb, CountN, CountK);
}
//
// Step through each slice of matrix A along the M dimension.
//
float* c = C + n;
if (TransA == CblasNoTrans) {
MlasSgemmKernelLoop(A + k, PanelB, c, CountK, M, CountN, lda, ldc, alpha, ZeroMode);
} else {
const float* a = A + k * lda;
size_t RowsRemaining = M;
while (RowsRemaining > 0) {
//
// Transpose elements from matrix A into a local buffer.
//
size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS));
MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK);
RowsRemaining -= RowsTransposed;
a += RowsTransposed;
//
// Step through the rows of the local buffer.
//
c = MlasSgemmKernelLoop(PanelA, PanelB, c, CountK, RowsTransposed, CountN, CountK, ldc, alpha, ZeroMode);
}
}
ZeroMode = false;
}
}
}
void
MlasSgemmPackedOperation(
CBLAS_TRANSPOSE TransA,
size_t M,
size_t RangeStartN,
size_t RangeCountN,
size_t K,
float alpha,
const float* A,
size_t lda,
const void* PackedB,
size_t AlignedN,
float beta,
float* C,
size_t ldc
)
/*++
Routine Description:
This routine implements the single precision matrix/matrix multiply
operation (SGEMM).
Arguments:
TransA - Supplies the transpose operation for matrix A.
M - Supplies the number of rows of matrix A and matrix C.
RangeStartN - Supplies the starting column from packed matrix B.
RangeCountN - Supplies the number of columns of matrix B and matrix C.
K - Supplies the number of columns of matrix A and the number of rows of
matrix B.
alpha - Supplies the scalar alpha multiplier (see SGEMM definition).
A - Supplies the address of matrix A.
lda - Supplies the first dimension of matrix A.
PackedB - Supplies the address of packed matrix B.
AlignedN - Supplies the total number of aligned columns for packed matrix B.
ldb - Supplies the first dimension of matrix B.
beta - Supplies the scalar beta multiplier (see SGEMM definition).
C - Supplies the address of matrix C.
ldc - Supplies the first dimension of matrix C.
Return Value:
None.
--*/
{
float PanelA[MLAS_SGEMM_TRANSA_ROWS * MLAS_SGEMM_PACKED_STRIDEK];
//
// Step through each slice of matrix B along the N dimension.
//
size_t CountN;
for (size_t n = 0; n < RangeCountN; n += CountN) {
const size_t SliceStartN = RangeStartN + n;
CountN = std::min(RangeCountN - n, size_t(MLAS_SGEMM_PACKED_STRIDEN));
//
// Multiply the output matrix by beta as needed.
//
if (beta != 0.0f && beta != 1.0f) {
MlasSgemmMultiplyBeta(C + n, M, CountN, ldc, beta);
}
//
// Step through each slice of matrix B along the K dimension.
//
size_t CountK;
bool ZeroMode = (beta == 0.0f);
for (size_t k = 0; k < K; k += CountK) {
CountK = std::min(K - k, size_t(MLAS_SGEMM_PACKED_STRIDEK));
//
// Step through each slice of matrix A along the M dimension.
//
const float* pb = (const float*)PackedB + AlignedN * k + CountK * SliceStartN;
float* c = C + n;
if (TransA == CblasNoTrans) {
MlasSgemmKernelLoop(A + k, pb, c, CountK, M, CountN, lda, ldc, alpha, ZeroMode);
} else {
const float* a = A + k * lda;
size_t RowsRemaining = M;
while (RowsRemaining > 0) {
//
// Transpose elements from matrix A into a local buffer.
//
size_t RowsTransposed = std::min(RowsRemaining, size_t(MLAS_SGEMM_TRANSA_ROWS));
MlasSgemmTransposeA(PanelA, a, lda, RowsTransposed, CountK);
RowsRemaining -= RowsTransposed;
a += RowsTransposed;
//
// Step through the rows of the local buffer.
//
c = MlasSgemmKernelLoop(PanelA, pb, c, CountK, RowsTransposed, CountN, CountK, ldc, alpha, ZeroMode);
}
}
ZeroMode = false;
}
}
}
void
MlasSgemmThreaded(
const ptrdiff_t ThreadCountM,
const ptrdiff_t ThreadCountN,
const CBLAS_TRANSPOSE TransA,
const CBLAS_TRANSPOSE TransB,
const size_t M,
const size_t N,
const size_t K,
const MLAS_SGEMM_DATA_PARAMS* DataParams,
ptrdiff_t ThreadId
)
/*++
Routine Description:
This routine is invoked from a worker thread to execute a segment of a
SGEMM operation.
Arguments:
ThreadCountM - Supplies the total thread partition on the M dimension.
ThreadCountN - Supplies the total thread partition on the N dimension.
TransA - Supplies the transpose operation on A matrix
TransB - Supplies the transpose operation on B matrix
M, N, K - Supplies the shape of the multiplication
DataParams - Supplies the data position and layout of the matrices
ThreadId - Supplies the current index of the threaded operation.
Return Value:
None.
--*/
{
const ptrdiff_t ThreadIdM = ThreadId / ThreadCountN;
const ptrdiff_t ThreadIdN = ThreadId % ThreadCountN;
//
// Partition the operation along the M dimension.
//
size_t RangeStartM;
size_t RangeCountM;
MlasPartitionWork(ThreadIdM, ThreadCountM, M, &RangeStartM, &RangeCountM);
//
// Partition the operation along the N dimension.
//
size_t RangeStartN;
size_t RangeCountN;
const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) /
MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
MlasPartitionWork(ThreadIdN, ThreadCountN, BlockedN, &RangeStartN,
&RangeCountN);
RangeStartN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
RangeCountN *= MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
RangeCountN = std::min(N - RangeStartN, RangeCountN);
//
// Dispatch the partitioned operation.
//
const size_t lda = DataParams->lda;
const size_t ldc = DataParams->ldc;
const float* A = DataParams->A + RangeStartM * ((TransA == CblasNoTrans) ? lda : 1);
float* C = DataParams->C + RangeStartM * ldc + RangeStartN;
if (DataParams->BIsPacked) {
MlasSgemmPackedOperation(TransA, RangeCountM, RangeStartN, RangeCountN,
K, DataParams->alpha, A, lda, DataParams->B,
BlockedN * MLAS_SGEMM_STRIDEN_THREAD_ALIGN, DataParams->beta, C, ldc);
} else {
const size_t ldb = DataParams->ldb;
const float* B = (const float*)DataParams->B + RangeStartN * ((TransB == CblasNoTrans) ? 1 : ldb);
MlasSgemmOperation(TransA, TransB, RangeCountM, RangeCountN, K,
DataParams->alpha, A, lda, B, ldb, DataParams->beta, C, ldc);
}
}
#if defined(_MSC_VER) && !defined(__clang__)
#pragma warning(push)
// Chance of arithmetic overflow could be reduced
#pragma warning(disable : 26451)
#endif
void
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,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
{
// Override
if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) &&
GetMlasPlatform().MlasSGemmBatchOverride != nullptr &&
// TODO: Remove once KAI supports transposing for A
TransA != CBLAS_TRANSPOSE::CblasTrans &&
GetMlasPlatform().MlasSGemmBatchOverride(TransA, TransB, M, N, K, Data, BatchSize, ThreadPool)){
return;
}
//
// Compute the number of target threads given the complexity of the SGEMM
// operation. Small requests should run using the single threaded path.
//
const double Complexity = double(M) * double(N) * double(K);
ptrdiff_t TargetThreadCount = ptrdiff_t(Complexity / double(MLAS_SGEMM_THREAD_COMPLEXITY)) + 1;
ptrdiff_t MaximumThreadCount = MlasGetMaximumThreadCount(ThreadPool);
if (TargetThreadCount >= MaximumThreadCount) {
TargetThreadCount = MaximumThreadCount;
}
//
// Segment the operation across multiple threads.
//
// N.B. Currently, the operation is segmented as a 1D partition, which
// works okay for operations involving skinny matrices.
//
ptrdiff_t ThreadsPerGemm = (TargetThreadCount + BatchSize - 1) / BatchSize;
ptrdiff_t ThreadCountM;
ptrdiff_t ThreadCountN;
if (N > M) {
const size_t BlockedN = (N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) /
MLAS_SGEMM_STRIDEN_THREAD_ALIGN;
if (size_t(ThreadsPerGemm) > BlockedN) {
ThreadsPerGemm = ptrdiff_t(BlockedN);
}
ThreadCountM = 1;
ThreadCountN = ThreadsPerGemm;
} else {
if (size_t(ThreadsPerGemm) > M) {
ThreadsPerGemm = ptrdiff_t(M);
}
ThreadCountM = ThreadsPerGemm;
ThreadCountN = 1;
}
MlasTrySimpleParallel(ThreadPool,
ThreadsPerGemm * static_cast<ptrdiff_t>(BatchSize),
[=](ptrdiff_t tid)
{
ptrdiff_t GemmIdx = tid / ThreadsPerGemm;
ptrdiff_t ThreadIdx = tid % ThreadsPerGemm;
MlasSgemmThreaded(ThreadCountM, ThreadCountN,
TransA, TransB, M, N, K, &(Data[GemmIdx]), ThreadIdx);
});
}
#if defined(_MSC_VER) && !defined(__clang__)
#pragma warning(pop)
#endif
size_t
MLASCALL
MlasGemmPackBSize(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
/*++
Routine Description:
This routine computes the length in bytes for the packed matrix B buffer.
Arguments:
N - Supplies the number of columns of matrix B.
K - Supplies the number of rows of matrix B.
BackendKernelSelectorConfig - Supplies the backend kernel selector
configuration options, else nullptr if the
default configuration should be used.
Return Value:
Returns the size in bytes for the packed matrix B buffer.
--*/
{
//
// Compute the number of bytes required to hold the packed buffer.
//
// KleidiAI or other override
#if defined(USE_KLEIDIAI)
if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) &&
GetMlasPlatform().MlasSGemmPackBSizeOverride != nullptr &&
// TODO: Remove once KAI supports transposing for A
TransA != CBLAS_TRANSPOSE::CblasTrans) {
size_t bytes_required;
//TODO pass status by reference to indicate success/fail
bytes_required = GetMlasPlatform().MlasSGemmPackBSizeOverride(TransA, TransB, N, K);
if (bytes_required != 0){// If ArmKleidiAI::MlasGemmPackBSize ran to completion
return bytes_required;
}
}
#endif
MLAS_UNREFERENCED_PARAMETER(TransA);
MLAS_UNREFERENCED_PARAMETER(TransB);
MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig);
const size_t AlignedN =
(N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1);
const size_t BytesRequired = AlignedN * K * sizeof(float);
const size_t BufferAlignment = MlasGetPreferredBufferAlignment();
const size_t AlignedBytesRequired = (BytesRequired + BufferAlignment - 1) &
~(BufferAlignment - 1);
return AlignedBytesRequired;
}
void
MLASCALL
MlasGemmPackB(
CBLAS_TRANSPOSE TransA,
CBLAS_TRANSPOSE TransB,
size_t N,
size_t K,
const float* B,
size_t ldb,
void* PackedB,
const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* BackendKernelSelectorConfig
)
/*++
Routine Description:
This routine packs the contents of matrix B to the destination buffer. The
destination buffer should be sized based on MlasGemmPackBSize(). For best
performance, the destination buffer should be aligned to the value returned
from MlasGetPreferredBufferAlignment().
Arguments:
TransB - Supplies the transpose operation for matrix B.
N - Supplies the number of columns of matrix B.
K - Supplies the number of rows of matrix B.
B - Supplies the address of matrix B.
ldb - Supplies the first dimension of matrix B.
PackedB - Supplies the address of packed matrix B.
Return Value:
None.
--*/
{
#if defined(USE_KLEIDIAI)
if ((!BackendKernelSelectorConfig || BackendKernelSelectorConfig->use_kleidiai) &&
GetMlasPlatform().MlasSGemmPackBOverride != nullptr &&
// TODO: Remove once KAI supports transposing for A
TransA != CBLAS_TRANSPOSE::CblasTrans &&
GetMlasPlatform().MlasSGemmPackBOverride(TransA, TransB, N, K, B, ldb, PackedB)){
return;
}
#endif
MLAS_UNREFERENCED_PARAMETER(TransA);
MLAS_UNREFERENCED_PARAMETER(BackendKernelSelectorConfig);
const size_t AlignedN =
(N + MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1) & ~(MLAS_SGEMM_STRIDEN_THREAD_ALIGN - 1);
//
// 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, size_t(MLAS_SGEMM_PACKED_STRIDEK));
if (TransB == CblasNoTrans) {
MlasSgemmCopyPackB((float*)PackedB, B + k * ldb, ldb, N, CountK);
} else {
MlasSgemmTransposePackB((float*)PackedB, B + k, ldb, N, CountK);
}
PackedB = (float*)PackedB + AlignedN * CountK;
}
}
3rdparty/mlas/lib/softmax.h
Vendored
+129
View File
@@ -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;
};
3rdparty/mlas/lib/sve/mlasi_sve.h
Vendored
+680
View File
@@ -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
3rdparty/mlas/lib/x86_64/FgemmKernelAvx512FCommon.h
Vendored
+529
View File
@@ -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
3rdparty/mlas/lib/x86_64/FgemmKernelAvxCommon.h
Vendored
+451
View File
@@ -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
3rdparty/mlas/lib/x86_64/FgemmKernelCommon.h
Vendored
+124
View File
@@ -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
3rdparty/mlas/lib/x86_64/FgemmKernelFma3Common.h
Vendored
+512
View File
@@ -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
3rdparty/mlas/lib/x86_64/FgemmKernelSse2Common.h
Vendored
+173
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelAvx.S
Vendored
+34
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelAvx512F.S
Vendored
+34
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelCommon.h
Vendored
+50
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelFma3.S
Vendored
+34
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelM1Avx.S
Vendored
+267
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S
Vendored
+275
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmKernelSse2.S
Vendored
+273
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S
Vendored
+120
View File
@@ -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
3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S
Vendored
+83
View File
@@ -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
3rdparty/mlas/lib/x86_64/asmmacro.h
Vendored
+172
View File
@@ -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
3rdparty/mlas/patches/0001-mlasi-opencv-threading.diff
Vendored
+46
View File
@@ -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
3rdparty/mlas/patches/0002-platform-gemm-only.diff
Vendored
+174
View File
@@ -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
3rdparty/mlas/patches/0003-mlas-h-msc-ver-guard.diff
Vendored
+19
View File
@@ -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
3rdparty/readme.txt
Vendored
+17
View File
@@ -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
CMakeLists.txt
+32
View File
@@ -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()
modules/dnn/CMakeLists.txt
+25 -1
View File
@@ -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})
modules/dnn/src/layers/attention_layer.cpp
+35 -1
View File
@@ -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;
};
modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp
+33
View File
@@ -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));
modules/dnn/src/layers/cpu_kernels/mlas_gemm.cpp
+215
View File
@@ -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
modules/dnn/src/layers/cpu_kernels/mlas_gemm.hpp
+122
View File
@@ -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
modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp
+98
View File
@@ -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
modules/dnn/src/layers/gemm_layer.cpp
+44
View File
@@ -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;
modules/dnn/src/layers/matmul_layer.cpp
+17 -3
View File
@@ -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,