diff --git a/3rdparty/mlas/CMakeLists.txt b/3rdparty/mlas/CMakeLists.txt new file mode 100644 index 0000000000..381a5a162f --- /dev/null +++ b/3rdparty/mlas/CMakeLists.txt @@ -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 + #include + int main() { size_t vl = __riscv_vsetvl_e32m1(4); return static_cast(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 + "$<$:-w>" + "$<$:-include>" + "$<$:cstring>" + ) + # MLAS .S files lack .note.GNU-stack; tell the assembler the stack is non-exec. + target_compile_options(opencv_dnn_mlas PRIVATE + "$<$:-Wa,--noexecstack>" + ) +endif() +# platform.cpp's __linux__ branch calls syscall() but upstream only includes +# , not (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 $ PARENT_SCOPE) +set(MLAS_INCLUDE_DIRS ${MLAS_INC_DIR} PARENT_SCOPE) +set(OPENCV_DNN_MLAS_ENABLED 1 CACHE INTERNAL "" FORCE) diff --git a/3rdparty/mlas/LICENSE b/3rdparty/mlas/LICENSE new file mode 100644 index 0000000000..48bc6bb499 --- /dev/null +++ b/3rdparty/mlas/LICENSE @@ -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. diff --git a/3rdparty/mlas/README.md b/3rdparty/mlas/README.md new file mode 100644 index 0000000000..1d5a87130b --- /dev/null +++ b/3rdparty/mlas/README.md @@ -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//`. + +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. diff --git a/3rdparty/mlas/inc/mlas.h b/3rdparty/mlas/inc/mlas.h new file mode 100644 index 0000000000..c2868995cd --- /dev/null +++ b/3rdparty/mlas/inc/mlas.h @@ -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 +#include +#include +#include + +// +// 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(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 + 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 +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 +void +MLASCALL +MlasComputeExp( + const T* Input, + T* Output, + size_t N + ); + +void +MLASCALL +MlasComputeLogistic( + const float* Input, + float* Output, + size_t N + ); + +template +void +MLASCALL +MlasComputeSoftmax( + const T* Input, + T* Output, + size_t N, + size_t D, + bool LogSoftmax, + bool SmoothSoftmax, + float Sink, + MLAS_THREADPOOL* ThreadPool + ); + +template +void +MLASCALL +MlasComputeSoftcap( + const T* Input, + T* Output, + size_t N, + T cap + ); + +template +void +MLASCALL +MlasEltwiseAdd( + const T* left, + const T* right, + T* output, + size_t N + ); + +template +void +MLASCALL +MlasEltwiseMul( + const T* left, + const T* right, + T* output, + size_t N + ); + +template +void +MLASCALL +MlasComputeTanh( + const T* Input, + T* Output, + size_t N + ); + +// +// Transpose routines. +// + +template +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 +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 +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 +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(Output_), OutputLeadingDimension_, + Bias_, Scale_, PerColumnScale_, static_cast(ZeroPoint_), + StartM, StartN, CountM, CountN); + } else { + MlasRequantizeOutput(C, ldc, reinterpret_cast(Output_), OutputLeadingDimension_, + Bias_, Scale_, PerColumnScale_, static_cast(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 +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 +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 +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 +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 +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 +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(Input), + reinterpret_cast(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 +); diff --git a/3rdparty/mlas/inc/mlas_float16.h b/3rdparty/mlas/inc/mlas_float16.h new file mode 100644 index 0000000000..33227ea90d --- /dev/null +++ b/3rdparty/mlas/inc/mlas_float16.h @@ -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 +#include +#include + + +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(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(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(f.u >> 13); + } + } + + val |= static_cast(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 \ No newline at end of file diff --git a/3rdparty/mlas/lib/aarch64/SgemmKernelNeon.S b/3rdparty/mlas/lib/aarch64/SgemmKernelNeon.S new file mode 100644 index 0000000000..26555b5c64 --- /dev/null +++ b/3rdparty/mlas/lib/aarch64/SgemmKernelNeon.S @@ -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 diff --git a/3rdparty/mlas/lib/aarch64/SgemvKernelNeon.S b/3rdparty/mlas/lib/aarch64/SgemvKernelNeon.S new file mode 100644 index 0000000000..c935f06255 --- /dev/null +++ b/3rdparty/mlas/lib/aarch64/SgemvKernelNeon.S @@ -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 diff --git a/3rdparty/mlas/lib/aarch64/asmmacro.h b/3rdparty/mlas/lib/aarch64/asmmacro.h new file mode 100644 index 0000000000..72982db003 --- /dev/null +++ b/3rdparty/mlas/lib/aarch64/asmmacro.h @@ -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 diff --git a/3rdparty/mlas/lib/arm/sgemmc.cpp b/3rdparty/mlas/lib/arm/sgemmc.cpp new file mode 100644 index 0000000000..200ec33117 --- /dev/null +++ b/3rdparty/mlas/lib/arm/sgemmc.cpp @@ -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 +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 +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(A, B, C, CountK, CountN, lda, ldc, alpha); + } else { + RowsHandled = MlasSgemmKernel(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(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(A, B, C, CountK, CountM, CountN, lda, ldc, alpha); +} diff --git a/3rdparty/mlas/lib/arm64/SgemmKernelNeon.asm b/3rdparty/mlas/lib/arm64/SgemmKernelNeon.asm new file mode 100644 index 0000000000..33b1ed1e64 --- /dev/null +++ b/3rdparty/mlas/lib/arm64/SgemmKernelNeon.asm @@ -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 diff --git a/3rdparty/mlas/lib/arm64ec/SgemmKernelNeon.asm b/3rdparty/mlas/lib/arm64ec/SgemmKernelNeon.asm new file mode 100644 index 0000000000..3c546b9051 --- /dev/null +++ b/3rdparty/mlas/lib/arm64ec/SgemmKernelNeon.asm @@ -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 diff --git a/3rdparty/mlas/lib/compute.cpp b/3rdparty/mlas/lib/compute.cpp new file mode 100644 index 0000000000..a677ee5087 --- /dev/null +++ b/3rdparty/mlas/lib/compute.cpp @@ -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::lowest(); + +// +// Define the parameters to execute segments of a softmax operation on worker +// threads. +// + +template +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( + 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( + 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::max(); + float tmp_max = std::numeric_limits::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 +void +MlasComputeSoftmaxThreaded( + void* Context, + ptrdiff_t Index +); + +template <> +void +MlasComputeSoftmaxThreaded( + 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*)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( + 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*)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 +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 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, &WorkBlock, ThreadCountN, ThreadPool); +} + +template +void +MLASCALL +MlasComputeSoftmax( + const float* Input, + float* Output, + size_t N, + size_t D, + bool LogSoftmax, + bool SmoothSoftmax, + float Sink, + MLAS_THREADPOOL* ThreadPool +); + +template +void +MLASCALL +MlasComputeSoftmax( + 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( + 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( + CBLAS_TRANSPOSE TransA, + CBLAS_TRANSPOSE TransB +) { + MLAS_UNREFERENCED_PARAMETER(TransA); + MLAS_UNREFERENCED_PARAMETER(TransB); + return true; +} diff --git a/3rdparty/mlas/lib/core/common/common.h b/3rdparty/mlas/lib/core/common/common.h new file mode 100644 index 0000000000..f94ebbafa8 --- /dev/null +++ b/3rdparty/mlas/lib/core/common/common.h @@ -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 +#include +#include + +namespace onnxruntime { + +// Concatenate stream-like arguments into a single string. Supports the +// same `operator<<` chain that ORT_ENFORCE uses for its diagnostic. +template +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__) diff --git a/3rdparty/mlas/lib/core/common/narrow.h b/3rdparty/mlas/lib/core/common/narrow.h new file mode 100644 index 0000000000..e0ec644365 --- /dev/null +++ b/3rdparty/mlas/lib/core/common/narrow.h @@ -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(...). 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 +#include + +namespace onnxruntime { + +template +constexpr T narrow(U u) { + static_assert(std::is_arithmetic::value && std::is_arithmetic::value, + "narrow(U): T and U must be arithmetic types"); + const T t = static_cast(u); + if (static_cast(t) != u || + ((t < T{}) != (u < U{}))) { + throw std::runtime_error("onnxruntime::narrow: narrowing failed"); + } + return t; +} + +} // namespace onnxruntime diff --git a/3rdparty/mlas/lib/erf_neon_fp16.h b/3rdparty/mlas/lib/erf_neon_fp16.h new file mode 100644 index 0000000000..7918df0ea3 --- /dev/null +++ b/3rdparty/mlas/lib/erf_neon_fp16.h @@ -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 + +#include "mlasi.h" +#include "fp16_common.h" +#include "softmax_kernel_neon.h" +#include + +void MlasNeonErfFP16Kernel(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N); diff --git a/3rdparty/mlas/lib/flashattn.cpp b/3rdparty/mlas/lib/flashattn.cpp new file mode 100644 index 0000000000..fe5402ed14 --- /dev/null +++ b/3rdparty/mlas/lib/flashattn.cpp @@ -0,0 +1,167 @@ +#include + +#include "mlasi.h" + +void +MlasFlashAttentionThreaded( + void* argptr, + std::ptrdiff_t thread_id +) +{ + const MlasFlashAttentionThreadedArgs* args = reinterpret_cast(argptr); + ptrdiff_t q_block_size = static_cast(args->q_block_size); + ptrdiff_t kv_block_size = static_cast(args->kv_block_size); + ptrdiff_t batch_size = static_cast(args->batch_size); + ptrdiff_t num_heads = static_cast(args->num_heads); + ptrdiff_t q_sequence_length = static_cast(args->q_sequence_length); + ptrdiff_t kv_sequence_length = static_cast(args->kv_sequence_length); + ptrdiff_t qk_head_size = static_cast(args->qk_head_size); + ptrdiff_t v_head_size = static_cast(args->v_head_size); + float* buffer = args->buffer; + ptrdiff_t buffer_size_per_thread = static_cast(args->buffer_size_per_thread); + ptrdiff_t thread_count = static_cast(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(buffer) + thread_id * buffer_size_per_thread; + float* l = reinterpret_cast(buffer_current_thread); + float* m = l + q_block_size; + for (ptrdiff_t t = 0; t < q_block_size; ++t) { + m[t] = std::numeric_limits::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(std::min(q_block_size, q_sequence_length - q_idx)); + size_t row_size_kv_capped = static_cast(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(qk_head_size), + args->scale, + inputQ, + static_cast(qk_head_size), + inputK, + static_cast(qk_head_size), + 0.0f, + intermediate, + row_size_kv_capped); + + for (ptrdiff_t irow = 0; irow < static_cast(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(v_head_size), + row_size_kv_capped, + 1.0f, + intermediate, + row_size_kv_capped, + inputV, + static_cast(v_head_size), + ir == 0 ? 0.0f : 1.0f, + temp_output, + static_cast(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(args), + static_cast(args->thread_count), + ThreadPool); +} diff --git a/3rdparty/mlas/lib/gelu_neon_fp16.h b/3rdparty/mlas/lib/gelu_neon_fp16.h new file mode 100644 index 0000000000..faabf56113 --- /dev/null +++ b/3rdparty/mlas/lib/gelu_neon_fp16.h @@ -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 +); diff --git a/3rdparty/mlas/lib/kleidiai/mlasi_kleidiai.h b/3rdparty/mlas/lib/kleidiai/mlasi_kleidiai.h new file mode 100644 index 0000000000..a1aa241b89 --- /dev/null +++ b/3rdparty/mlas/lib/kleidiai/mlasi_kleidiai.h @@ -0,0 +1,236 @@ +// +// SPDX-FileCopyrightText: Copyright 2025 Arm Limited and/or its affiliates +// +// SPDX-License-Identifier: MIT +// + +#pragma once + +#include "../mlasi.h" +#include + +// 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; +} diff --git a/3rdparty/mlas/lib/loongarch64/SgemmKernelLasx.S b/3rdparty/mlas/lib/loongarch64/SgemmKernelLasx.S new file mode 100644 index 0000000000..d537742016 --- /dev/null +++ b/3rdparty/mlas/lib/loongarch64/SgemmKernelLasx.S @@ -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 diff --git a/3rdparty/mlas/lib/loongarch64/SgemmKernelLsx.S b/3rdparty/mlas/lib/loongarch64/SgemmKernelLsx.S new file mode 100644 index 0000000000..86b5ef8b51 --- /dev/null +++ b/3rdparty/mlas/lib/loongarch64/SgemmKernelLsx.S @@ -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 diff --git a/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4LSX.S b/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4LSX.S new file mode 100644 index 0000000000..cd1747745d --- /dev/null +++ b/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4LSX.S @@ -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 diff --git a/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4Lasx.S b/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4Lasx.S new file mode 100644 index 0000000000..e617419989 --- /dev/null +++ b/3rdparty/mlas/lib/loongarch64/SgemmTransposePackB16x4Lasx.S @@ -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 diff --git a/3rdparty/mlas/lib/loongarch64/asmmacro.h b/3rdparty/mlas/lib/loongarch64/asmmacro.h new file mode 100644 index 0000000000..837aca77dd --- /dev/null +++ b/3rdparty/mlas/lib/loongarch64/asmmacro.h @@ -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 diff --git a/3rdparty/mlas/lib/mlasi.h b/3rdparty/mlas/lib/mlasi.h new file mode 100644 index 0000000000..14af627c18 --- /dev/null +++ b/3rdparty/mlas/lib/mlasi.h @@ -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 +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef MLAS_NO_EXCEPTION +#if defined(__ANDROID__) +#include +#else +#include +#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 +#include +#else +#if defined(__arm__) || defined(__aarch64__) +#include +#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 +#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 +#pragma GCC diagnostic pop +#else +#include +#endif +#endif +#if defined(__VSX__) +#include +// Undefine unwanted aliases from altivec.h. +#undef vector +#undef pixel +#undef bool +#endif +#if defined(__s390x__) +#include +#endif +#if defined(__loongarch64) +#include +#endif +#if defined(MLAS_TARGET_WASM_SIMD) +#include +#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(val & ~kSignMask)); + } + bool IsNaN() const noexcept { + return Abs().val > kPositiveInfinityBits; + } + bool IsNegative() const noexcept { + return static_cast(val) < 0; + } + MLFloat16 Negate() const { + return MLFloat16(IsNaN() ? val : static_cast(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 +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 +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 +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::DepthwiseKernel* ConvDepthwiseU8S8Kernel; + MLAS_QUANT_KERNEL::DepthwiseKernel* ConvDepthwiseU8U8Kernel; + MLAS_QUANT_KERNEL::DepthwiseKernel* ConvDepthwiseS8S8Kernel; + MLAS_QUANT_KERNEL::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& 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& 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 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(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 +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 +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 +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 +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 +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 +MLAS_FORCEINLINE +MLAS_FLOAT32X4 +MlasShuffleFloat32x4(MLAS_FLOAT32X4 Vector) +{ + return MlasShuffleFloat32x4(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 +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 +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 ThreadedBufHolder; +#else +extern thread_local std::unique_ptr 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(_aligned_malloc(size, ThreadedBufAlignment))); +#elif (__STDC_VERSION__ >= 201112L) && !defined(__APPLE__) + ThreadedBufHolder.reset( + reinterpret_cast(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(ptr)); +#endif + + ThreadedBufSize = size; + } +} + +// +// Utilities for INT4 quantization. +// + +template +struct Int4Traits; + +template<> +struct Int4Traits { + using UnpackedType = int8_t; + static constexpr int8_t Min = -8; + static constexpr int8_t Max = 7; +}; + +template<> +struct Int4Traits { + using UnpackedType = uint8_t; + static constexpr int8_t Min = 0; + static constexpr int8_t Max = 15; +}; + +template +MLAS_FORCEINLINE +void +MlasSetInt4Element(uint8_t* Output, size_t ElemIndex, UnpackedType Value) +{ + static_assert(std::is_same_v || std::is_same_v); + + 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(NibbleIndex << 2); // Either 0 or 4 + const uint8_t Mask = static_cast(0xF0 >> Shift); + uint8_t* Dst = &Output[OutputIndex]; + + *Dst &= Mask; // Clear 4-bit lane + *Dst |= static_cast((Value & 0xF) << Shift); // Set 4-bit lane +} + +template +MLAS_FORCEINLINE +void +MlasPackInt4Elements(uint8_t* Output, UnpackedType ValueLow, UnpackedType ValueHigh) +{ + static_assert(std::is_same_v || std::is_same_v); + *Output = static_cast(((ValueHigh & 0xF) << 4) | (ValueLow & 0xF)); +} diff --git a/3rdparty/mlas/lib/platform.cpp b/3rdparty/mlas/lib/platform.cpp new file mode 100644 index 0000000000..a515413331 --- /dev/null +++ b/3rdparty/mlas/lib/platform.cpp @@ -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 +#include +#include +#include + +#if defined(MLAS_TARGET_POWER) +#if defined(__linux__) +#include +#elif defined(_AIX) +#define POWER_10 0x40000 +#define POWER_10_ANDUP (POWER_10) +#include +#define __power_10_andup() (_system_configuration.implementation & POWER_10_ANDUP) +#elif defined(__FreeBSD__) +#include +#include +#endif +#endif + + +#if defined(MLAS_TARGET_S390X) +#include +#endif + +#if defined(MLAS_TARGET_RISCV64) && defined(MLAS_USE_RVV) && defined(__linux__) +#include +#include +#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(*value); + const auto rhs = static_cast(*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 +#include +// 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 +#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 +#include +#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; + this->ConvDepthwiseU8U8Kernel = MlasConvDepthwiseKernel; + this->ConvDepthwiseS8S8Kernel = MlasConvDepthwiseKernel; + this->ConvDepthwiseS8U8Kernel = MlasConvDepthwiseKernel; + 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; + this->ConvDepthwiseU8U8Kernel = MlasConvDepthwiseKernelAvx2; + this->ConvDepthwiseS8S8Kernel = MlasConvDepthwiseKernelAvx2; + this->ConvDepthwiseS8U8Kernel = MlasConvDepthwiseKernelAvx2; + 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 ThreadedBufHolder(nullptr, &_aligned_free); +#else +thread_local std::unique_ptr ThreadedBufHolder(nullptr, &free); +#endif diff --git a/3rdparty/mlas/lib/power/SgemmKernelPOWER10.cpp b/3rdparty/mlas/lib/power/SgemmKernelPOWER10.cpp new file mode 100644 index 0000000000..9ecfa3d984 --- /dev/null +++ b/3rdparty/mlas/lib/power/SgemmKernelPOWER10.cpp @@ -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 + 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 +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 +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(ABroadcast), reinterpret_cast(BElements[0])); + __builtin_mma_xvf32gerpp (&acc[1], reinterpret_cast(ABroadcast), reinterpret_cast(BElements[1])); + __builtin_mma_xvf32gerpp (&acc[2], reinterpret_cast(ABroadcast), reinterpret_cast(BElements[2])); + __builtin_mma_xvf32gerpp (&acc[3], reinterpret_cast(ABroadcast), reinterpret_cast(BElements[3])); + if (CountM == 8) { + __builtin_mma_xvf32gerpp (&acc[4], reinterpret_cast(A2Broadcast), reinterpret_cast(BElements[0])); + __builtin_mma_xvf32gerpp (&acc[5], reinterpret_cast(A2Broadcast), reinterpret_cast(BElements[1])); + __builtin_mma_xvf32gerpp (&acc[6], reinterpret_cast(A2Broadcast), reinterpret_cast(BElements[2])); + __builtin_mma_xvf32gerpp (&acc[7], reinterpret_cast(A2Broadcast), reinterpret_cast(BElements[3])); + } +} +template +struct MlasSgemmStoreVectorMMA +{ + template + 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(&C[Row * ldc + VectorCount]); + rowC[0] = Result[Row] * AlphaBroadcast; + } else { + rowC = reinterpret_cast(&C[Row * ldc + VectorCount]); + rowC[0] += Result[Row] * AlphaBroadcast; + } + } +}; + +struct MlasSgemmMultiplyAlphaTrailingMMA +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[RowCount], + MLAS_FLOAT32X4 AlphaBroadcast + ) + { + Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast); + } +}; +template +struct MlasSgemmStoreScalarMMA +{ + template + 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 +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(&acc[0], pa1[0], pa1[4], B, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[1], pa1[5], B + 16, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[2], pa1[6], B + 32, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[3], pa1[7], B + 48, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[8], pa1[12], B + 64, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[9], pa1[13], B + 80, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[10], pa1[14], B + 96, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[11], pa1[15], B + 112, CountM); + B += 128; + pa1 += 16; + k -= 8; + } else { + MlasSgemmComputeBlockMMA(&acc[0], pa1[0], ABroadcast[0], B, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[1], ABroadcast[1], B + 16, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[2], ABroadcast[2], B + 32, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[3], ABroadcast[3], B + 48, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[4], ABroadcast[0], B + 64, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[5], ABroadcast[1], B + 80, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[6], ABroadcast[2], B + 96, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[7], ABroadcast[3], B + 112, CountM); + B += 128; + pa1 += 8; + k -= 8; + } + } + + while (k >= 4) { + if (CountM == 8) { + MlasSgemmComputeBlockMMA(&acc[0], pa1[0], pa1[4], B, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[1], pa1[5], B + 16, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[2], pa1[6], B + 32, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[3], pa1[7], B + 48, CountM); + B += 16 * 4; + pa1 += 8; + k -= 4; + } else { + MlasSgemmComputeBlockMMA(&acc[0], pa1[0], ABroadcast[0], B, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[1], ABroadcast[1], B + 16, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[2], ABroadcast[2], B + 32, CountM); + MlasSgemmComputeBlockMMA(&acc[0], pa1[3], ABroadcast[3], B + 48, CountM); + B += 16 * 4; + pa1 += 4; + k -= 4; + } + } + while (k > 0) { + if (CountM == 8) { + MlasSgemmComputeBlockMMA(&acc[0], pa1[0], pa1[1], B, CountM); + pa1 += 2; + } else { + MlasSgemmComputeBlockMMA(&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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[2]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[3]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc (Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[5]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[6]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[7]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + } + } else { + + // + // Store the partial output block. + // + + if (CountN >= 12) { + __builtin_mma_disassemble_acc (Result, &acc[0]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[2]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc (Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[5]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[6]); + MlasLoopUnroll>()(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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[1]); + MlasLoopUnroll>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc (Result, &acc[4]); + MlasLoopUnroll>()(Result, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + __builtin_mma_disassemble_acc (Result, &acc[5]); + MlasLoopUnroll>()(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>()(Result, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + __builtin_mma_disassemble_acc (Result, &acc[4]); + MlasLoopUnroll>()(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()(Accumulators[0], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll()(Accumulators[1], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + } + if (CountN >= 3) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + } + } + + break; + } + + C += 16; + CountN -= 16; + + } while (CountN > 0); + + return CountM; +} + +template +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(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; +} diff --git a/3rdparty/mlas/lib/power/SgemmKernelPackA.S b/3rdparty/mlas/lib/power/SgemmKernelPackA.S new file mode 100644 index 0000000000..fa2d79ea6d --- /dev/null +++ b/3rdparty/mlas/lib/power/SgemmKernelPackA.S @@ -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 diff --git a/3rdparty/mlas/lib/power/SgemmKernelPower.cpp b/3rdparty/mlas/lib/power/SgemmKernelPower.cpp new file mode 100644 index 0000000000..8c42348338 --- /dev/null +++ b/3rdparty/mlas/lib/power/SgemmKernelPower.cpp @@ -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; +} diff --git a/3rdparty/mlas/lib/qgemm.h b/3rdparty/mlas/lib/qgemm.h new file mode 100644 index 0000000000..b2b0be9a8b --- /dev/null +++ b/3rdparty/mlas/lib/qgemm.h @@ -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 +#include +#include + +// +// 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 +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 +MLAS_FORCEINLINE constexpr +int32_t +MlasGemmQuantFixupZeroPointA( + int32_t ZeroPointA, + bool AIsSigned) +{ + MLAS_UNREFERENCED_PARAMETER(AIsSigned); + return ZeroPointA; +} + +template +int32_t constexpr +MlasGemmQuantFixupZeroPointB( + int32_t ZeroPointB, + bool BIsSigned +) +{ + MLAS_UNREFERENCED_PARAMETER(BIsSigned); + + return ZeroPointB; +} + +template +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(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 +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 +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 +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 +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 +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 +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(); + + uint8_t* p = ThreadedBufHolder.get(); + typename KernelType::PackedAType* PanelA = + reinterpret_cast(p); + p += packASize; + typename KernelType::PackedBType* PanelB = + reinterpret_cast(p); + p += packBSize; + int32_t* RowSumBuffer = reinterpret_cast(p); + p += rowSumSize; + int32_t* ColumnSumBuffer = reinterpret_cast(p); + p += colSumSize; + int32_t* ZeroPointBBuffer = reinterpret_cast(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(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(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(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( + PackedZeroPointB + n, + ZeroPointBBuffer, + CountN, + Shape->BIsSigned); + } + + // + // Copy a panel of matrix B to a local packed buffer. + // + + MlasGemmQuantCopyPackB( + 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( + 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( + 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 +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(); + + uint8_t* p = ThreadedBufHolder.get(); + typename KernelType::PackedAType* PanelA = + reinterpret_cast(p); + p += packASize; + int32_t* RowSumBuffer = reinterpret_cast(p); + p += rowSumSize; + int32_t* ColumnSumBuffer = reinterpret_cast(p); + p += colSumSize; + int32_t* ZeroPointBBuffer = reinterpret_cast(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(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(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( + 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( + 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( + 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 +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( + 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; +} diff --git a/3rdparty/mlas/lib/riscv64/sgemm_kernel_rvv.cpp b/3rdparty/mlas/lib/riscv64/sgemm_kernel_rvv.cpp new file mode 100644 index 0000000000..c6e43e2c8b --- /dev/null +++ b/3rdparty/mlas/lib/riscv64/sgemm_kernel_rvv.cpp @@ -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 + +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 +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 +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(c, row0_block, vl, alpha); + + if constexpr (Rows >= 2) { + MlasStoreAccumulatorRvv(c + ldc, row1_block, vl, alpha); + } + if constexpr (Rows >= 3) { + MlasStoreAccumulatorRvv(c + ldc * 2, row2_block, vl, alpha); + } + if constexpr (Rows >= 4) { + MlasStoreAccumulatorRvv(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 +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(A, B, C, CountK, CountN, lda, ldc, alpha); + } + + if (CountM == 3) { + return MlasSgemmKernelRvv(A, B, C, CountK, CountN, lda, ldc, alpha); + } + + if (CountM >= 2) { + return MlasSgemmKernelRvv(A, B, C, CountK, CountN, lda, ldc, alpha); + } + + return MlasSgemmKernelRvv(A, B, C, CountK, CountN, lda, ldc, alpha); +} + +template +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( + A, B, C, CountK, CountM, CountN, lda, ldc, alpha); + } + + return MlasGemmFloatKernelRvvDispatchRows( + 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(A, B, C, CountK, CountM, CountN, lda, ldc, alpha); + } + + return MlasGemmFloatKernelRvvDispatch(A, B, C, CountK, CountM, CountN, lda, ldc, alpha); +} + +#endif // defined(MLAS_USE_RVV) diff --git a/3rdparty/mlas/lib/riscv64/sgemm_pack_b_rvv.cpp b/3rdparty/mlas/lib/riscv64/sgemm_pack_b_rvv.cpp new file mode 100644 index 0000000000..b2ec24e3fb --- /dev/null +++ b/3rdparty/mlas/lib/riscv64/sgemm_pack_b_rvv.cpp @@ -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 + +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) diff --git a/3rdparty/mlas/lib/s390x/SgemmKernel.cpp b/3rdparty/mlas/lib/s390x/SgemmKernel.cpp new file mode 100644 index 0000000000..e7d8cde71e --- /dev/null +++ b/3rdparty/mlas/lib/s390x/SgemmKernel.cpp @@ -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; +} diff --git a/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.cpp b/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.cpp new file mode 100644 index 0000000000..e87913a389 --- /dev/null +++ b/3rdparty/mlas/lib/s390x/SgemmKernelZVECTOR.cpp @@ -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 + +struct MlasSgemmBroadcastAElementsZVECTOR +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 ABroadcast[RowCount], + const float* A, + size_t lda + ) + { + ABroadcast[0][Row] = A [Row * lda]; + } +}; + +template +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 +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 +struct MlasSgemmStoreVectorZVECTOR +{ + template + 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(&C[Row * ldc + VectorCount]); + rowC[0] = Result[Row] * AlphaBroadcast; + } else { + rowC = reinterpret_cast(&C[Row * ldc + VectorCount]); + rowC[0] += Result[Row] * AlphaBroadcast; + } + } +}; + +struct MlasSgemmMultiplyAlphaTrailingZVECTOR +{ + template + MLAS_FORCEINLINE + static + void + Iteration( + MLAS_FLOAT32X4 Accumulators[RowCount], + MLAS_FLOAT32X4 AlphaBroadcast + ) + { + Accumulators[Row] = MlasMultiplyFloat32x4(Accumulators[Row], AlphaBroadcast); + } +}; +template +struct MlasSgemmStoreScalarZVECTOR +{ + template + 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 +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()(AElements, a, lda); + MlasSgemmComputeAElements(AElements, ABroadcast); + if (CountM == 8) { + MlasLoopUnroll()(AElements, a + ( lda * 4), lda); + MlasSgemmComputeAElements(AElements, A2Broadcast); + } + MlasSgemmComputeBlockZVECTOR(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM); + MlasSgemmComputeBlockZVECTOR(&acc[0], ABroadcast[1], A2Broadcast[1], B+16, CountM); + MlasSgemmComputeBlockZVECTOR(&acc[0], ABroadcast[2], A2Broadcast[2], B+32, CountM); + MlasSgemmComputeBlockZVECTOR(&acc[0], ABroadcast[3], A2Broadcast[3], B+48, CountM); + B += 16 * 4; + a += 4; + k -= 4; + } + + while (k > 0) { + MlasLoopUnroll()(ABroadcast, a, lda); + if (CountM == 8) { + MlasLoopUnroll()(A2Broadcast, a + (lda * 4), lda); + } + MlasSgemmComputeBlockZVECTOR(&acc[0], ABroadcast[0], A2Broadcast[0], B, CountM); + a += 1; + B += 16; + k -= 1; + } + if (CountN >= 16) { + + // + // Store the entire output block. + // + MlasLoopUnroll>()(acc, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 8, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 12, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 20, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 24, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 28, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + } + } else { + + // + // Store the partial output block. + // + + if (CountN >= 12) { + MlasLoopUnroll>()(acc, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 8, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 20, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(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>()(acc, C, ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(acc + 4, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(acc + 16, C + (ldc*4), ldc, AlphaBroadcast, ZeroMode); + MlasLoopUnroll>()(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>()(acc, C, ldc, AlphaBroadcast, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(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()(Accumulators[0], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll()(Accumulators[1], AlphaBroadcast); + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + if (CountN >= 2) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(Accumulators[1], C + (ldc*4), ldc, ZeroMode); + } + } + if (CountN >= 3) { + MlasLoopUnroll>()(Accumulators[0], C, ldc, ZeroMode); + if (CountM == 8) { + MlasLoopUnroll>()(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; +} diff --git a/3rdparty/mlas/lib/scalar/SconvDepthwiseKernelScalar.cpp b/3rdparty/mlas/lib/scalar/SconvDepthwiseKernelScalar.cpp new file mode 100644 index 0000000000..da1cdb9606 --- /dev/null +++ b/3rdparty/mlas/lib/scalar/SconvDepthwiseKernelScalar.cpp @@ -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); +} diff --git a/3rdparty/mlas/lib/scalar/SgemmKernelScalar.cpp b/3rdparty/mlas/lib/scalar/SgemmKernelScalar.cpp new file mode 100644 index 0000000000..496904c420 --- /dev/null +++ b/3rdparty/mlas/lib/scalar/SgemmKernelScalar.cpp @@ -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 +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 +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(A, B, C, CountK, CountN, lda, ldc, alpha); + } else { + RowsHandled = MlasSgemmKernel(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(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(A, B, C, CountK, CountM, CountN, lda, ldc, alpha); +} diff --git a/3rdparty/mlas/lib/scalar/SgemvKernelScalar.cpp b/3rdparty/mlas/lib/scalar/SgemvKernelScalar.cpp new file mode 100644 index 0000000000..609a6f251e --- /dev/null +++ b/3rdparty/mlas/lib/scalar/SgemvKernelScalar.cpp @@ -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++; + } +} diff --git a/3rdparty/mlas/lib/sgemm.cpp b/3rdparty/mlas/lib/sgemm.cpp new file mode 100644 index 0000000000..88d0308bfa --- /dev/null +++ b/3rdparty/mlas/lib/sgemm.cpp @@ -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 +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(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; + } +} diff --git a/3rdparty/mlas/lib/softmax.h b/3rdparty/mlas/lib/softmax.h new file mode 100644 index 0000000000..69fe1ae86d --- /dev/null +++ b/3rdparty/mlas/lib/softmax.h @@ -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; +}; diff --git a/3rdparty/mlas/lib/sve/mlasi_sve.h b/3rdparty/mlas/lib/sve/mlasi_sve.h new file mode 100644 index 0000000000..922945c702 --- /dev/null +++ b/3rdparty/mlas/lib/sve/mlasi_sve.h @@ -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 // 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 +MLAS_SVE_TARGET +MLAS_FORCEINLINE +MLAS_SVUINT32 +MlasSveShiftLeftUInt32(MLAS_SVBOOL Pred, MLAS_SVUINT32 Vector) +{ + return svlsl_n_u32_z(Pred, Vector, ShiftCount); +} + +template +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 +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 +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 + diff --git a/3rdparty/mlas/lib/x86_64/FgemmKernelAvx512FCommon.h b/3rdparty/mlas/lib/x86_64/FgemmKernelAvx512FCommon.h new file mode 100644 index 0000000000..9f243ee8c0 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/FgemmKernelAvx512FCommon.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/FgemmKernelAvxCommon.h b/3rdparty/mlas/lib/x86_64/FgemmKernelAvxCommon.h new file mode 100644 index 0000000000..69c8d17e27 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/FgemmKernelAvxCommon.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/FgemmKernelCommon.h b/3rdparty/mlas/lib/x86_64/FgemmKernelCommon.h new file mode 100644 index 0000000000..f3e8890fd3 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/FgemmKernelCommon.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/FgemmKernelFma3Common.h b/3rdparty/mlas/lib/x86_64/FgemmKernelFma3Common.h new file mode 100644 index 0000000000..77fc1f1d9f --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/FgemmKernelFma3Common.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/FgemmKernelSse2Common.h b/3rdparty/mlas/lib/x86_64/FgemmKernelSse2Common.h new file mode 100644 index 0000000000..2f71864f00 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/FgemmKernelSse2Common.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelAvx.S b/3rdparty/mlas/lib/x86_64/SgemmKernelAvx.S new file mode 100644 index 0000000000..a0a66f330a --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelAvx.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelAvx512F.S b/3rdparty/mlas/lib/x86_64/SgemmKernelAvx512F.S new file mode 100644 index 0000000000..c75df76030 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelAvx512F.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelCommon.h b/3rdparty/mlas/lib/x86_64/SgemmKernelCommon.h new file mode 100644 index 0000000000..5802028788 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelCommon.h @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelFma3.S b/3rdparty/mlas/lib/x86_64/SgemmKernelFma3.S new file mode 100644 index 0000000000..4725459323 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelFma3.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelM1Avx.S b/3rdparty/mlas/lib/x86_64/SgemmKernelM1Avx.S new file mode 100644 index 0000000000..5c759847e2 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelM1Avx.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S b/3rdparty/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S new file mode 100644 index 0000000000..b205c3d6d7 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelM1TransposeBAvx.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmKernelSse2.S b/3rdparty/mlas/lib/x86_64/SgemmKernelSse2.S new file mode 100644 index 0000000000..e605128537 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmKernelSse2.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S b/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S new file mode 100644 index 0000000000..644077838a --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Avx.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S b/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S new file mode 100644 index 0000000000..d3ef6f3237 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/SgemmTransposePackB16x4Sse2.S @@ -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 diff --git a/3rdparty/mlas/lib/x86_64/asmmacro.h b/3rdparty/mlas/lib/x86_64/asmmacro.h new file mode 100644 index 0000000000..7ef836c570 --- /dev/null +++ b/3rdparty/mlas/lib/x86_64/asmmacro.h @@ -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 diff --git a/3rdparty/mlas/patches/0001-mlasi-opencv-threading.diff b/3rdparty/mlas/patches/0001-mlasi-opencv-threading.diff new file mode 100644 index 0000000000..b6b0c9fc2f --- /dev/null +++ b/3rdparty/mlas/patches/0001-mlasi-opencv-threading.diff @@ -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 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(opencv_dnn_mlas_max_threads()); ++#elif defined(BUILD_MLAS_NO_ONNXRUNTIME) + MLAS_UNREFERENCED_PARAMETER(ThreadPool); + return 1; + #else diff --git a/3rdparty/mlas/patches/0002-platform-gemm-only.diff b/3rdparty/mlas/patches/0002-platform-gemm-only.diff new file mode 100644 index 0000000000..69027dad46 --- /dev/null +++ b/3rdparty/mlas/patches/0002-platform-gemm-only.diff @@ -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 diff --git a/3rdparty/mlas/patches/0003-mlas-h-msc-ver-guard.diff b/3rdparty/mlas/patches/0003-mlas-h-msc-ver-guard.diff new file mode 100644 index 0000000000..7fc1b60030 --- /dev/null +++ b/3rdparty/mlas/patches/0003-mlas-h-msc-ver-guard.diff @@ -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 diff --git a/3rdparty/readme.txt b/3rdparty/readme.txt index 0b31b26557..5aaae92ae1 100644 --- a/3rdparty/readme.txt +++ b/3rdparty/readme.txt @@ -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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 5ce37a73b7..5d4f534561 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -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() diff --git a/modules/dnn/CMakeLists.txt b/modules/dnn/CMakeLists.txt index 58c4b202e1..6b12bfddd7 100644 --- a/modules/dnn/CMakeLists.txt +++ b/modules/dnn/CMakeLists.txt @@ -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}) diff --git a/modules/dnn/src/layers/attention_layer.cpp b/modules/dnn/src/layers/attention_layer.cpp index 5a4800a8dd..42d6297562 100644 --- a/modules/dnn/src/layers/attention_layer.cpp +++ b/modules/dnn/src/layers/attention_layer.cpp @@ -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 @@ -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(), + 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(); @@ -526,6 +559,7 @@ class AttentionLayerImpl CV_FINAL : public AttentionLayer { std::vector packed_weight_q; std::vector packed_weight_k; std::vector packed_weight_v; + std::vector flash_scratch; FastGemmOpt opt; }; diff --git a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp index 9e87590d94..b7b3232052 100644 --- a/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp +++ b/modules/dnn/src/layers/cpu_kernels/fast_gemm.cpp @@ -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)); diff --git a/modules/dnn/src/layers/cpu_kernels/mlas_gemm.cpp b/modules/dnn/src/layers/cpu_kernels/mlas_gemm.cpp new file mode 100644 index 0000000000..3c48c3a038 --- /dev/null +++ b/modules/dnn/src/layers/cpu_kernels/mlas_gemm.cpp @@ -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 + +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(lda); + data.B = B; + data.ldb = static_cast(ldb); + data.C = C; + data.ldc = static_cast(ldc); + data.alpha = alpha; + data.beta = beta; + data.BIsPacked = false; + + MlasGemm(trans_a ? CblasTrans : CblasNoTrans, + trans_b ? CblasTrans : CblasNoTrans, + static_cast(M), + static_cast(N), + static_cast(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 data(batch); + for (size_t i = 0; i < batch; i++) { + data[i].A = A_base + A_offsets[i]; + data[i].lda = static_cast(lda); + data[i].B = B_base + B_offsets[i]; + data[i].ldb = static_cast(ldb); + data[i].C = C_base + C_offsets[i]; + data[i].ldc = static_cast(ldc); + data[i].alpha = alpha; + data[i].beta = beta; + data[i].BIsPacked = false; + } + + MlasGemmBatch(trans_a ? CblasTrans : CblasNoTrans, + trans_b ? CblasTrans : CblasNoTrans, + static_cast(M), + static_cast(N), + static_cast(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(N), + static_cast(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(N), + static_cast(K), + B, static_cast(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(lda); + data.B = static_cast(packed_B); + data.ldb = 0; // ignored when BIsPacked + data.C = C; + data.ldc = static_cast(ldc); + data.alpha = alpha; + data.beta = beta; + data.BIsPacked = true; + + MlasGemm(trans_a ? CblasTrans : CblasNoTrans, + trans_b ? CblasTrans : CblasNoTrans, + static_cast(M), + static_cast(N), + static_cast(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(q_block_size); + const size_t kv = static_cast(kv_block_size); + const size_t vd = static_cast(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(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 diff --git a/modules/dnn/src/layers/cpu_kernels/mlas_gemm.hpp b/modules/dnn/src/layers/cpu_kernels/mlas_gemm.hpp new file mode 100644 index 0000000000..22c749b638 --- /dev/null +++ b/modules/dnn/src/layers/cpu_kernels/mlas_gemm.hpp @@ -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 + +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 diff --git a/modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp b/modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp new file mode 100644 index 0000000000..6c1d9a5559 --- /dev/null +++ b/modules/dnn/src/layers/cpu_kernels/mlas_threading.cpp @@ -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 + +// 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(Iterations)), + [&](const cv::Range& r) { + for (int tid = r.start; tid < r.end; tid++) { + ThreadedRoutine(Context, static_cast(tid)); + } + }); +} + +void +MlasTrySimpleParallel( + MLAS_THREADPOOL* /*ThreadPool*/, + const std::ptrdiff_t Iterations, + const std::function& Work) +{ + if (Iterations <= 0) return; + if (Iterations == 1) { Work(0); return; } + + cv::parallel_for_(cv::Range(0, static_cast(Iterations)), + [&](const cv::Range& r) { + for (int tid = r.start; tid < r.end; tid++) { + Work(static_cast(tid)); + } + }); +} + +void +MlasTryBatchParallel( + MLAS_THREADPOOL* /*ThreadPool*/, + const std::ptrdiff_t Iterations, + const std::function& 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(Iterations)), + [&](const cv::Range& r) { + for (int tid = r.start; tid < r.end; tid++) { + Work(static_cast(tid)); + } + }); +} + +#endif // MLAS_OPENCV_THREADING diff --git a/modules/dnn/src/layers/gemm_layer.cpp b/modules/dnn/src/layers/gemm_layer.cpp index 0b60165220..ceac819619 100644 --- a/modules/dnn/src/layers/gemm_layer.cpp +++ b/modules/dnn/src/layers/gemm_layer.cpp @@ -17,6 +17,7 @@ using namespace cv::dnn::cuda4dnn; #include #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 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(packed_bytes), CV_8U); + if (mlasSgemmPackB(trans_a, trans_b, N, K, + Bmat.ptr(), 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(), na, + packed_B_mlas.data, + 1.f, + Y.ptr(), N)) { + return; + } + } +#endif CV_CheckGT(packed_B.size(), static_cast(0), "DNN/Gemm: constant B is not pre-packed"); if (!thin_packed_B.empty()) { fastGemmThin(rows, N, K, alpha, A.ptr(), na, 1, @@ -531,6 +570,11 @@ private: bool have_bias; std::vector packed_B; std::vector 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 broadcast_C; int real_ndims_C; FastGemmOpt opt; diff --git a/modules/dnn/src/layers/matmul_layer.cpp b/modules/dnn/src/layers/matmul_layer.cpp index 7f84cbef06..aea6901167 100644 --- a/modules/dnn/src/layers/matmul_layer.cpp +++ b/modules/dnn/src/layers/matmul_layer.cpp @@ -6,6 +6,7 @@ #include #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(); - 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,