diff --git a/cmake/onnxruntime_mlas.cmake b/cmake/onnxruntime_mlas.cmake index ed3ad89247..cb8d3dcad5 100644 --- a/cmake/onnxruntime_mlas.cmake +++ b/cmake/onnxruntime_mlas.cmake @@ -43,6 +43,7 @@ onnxruntime_add_static_library(onnxruntime_mlas ${MLAS_SRC_DIR}/cast.cpp ${MLAS_SRC_DIR}/rotary_embedding.h ${MLAS_SRC_DIR}/rotary_embedding.cpp + ${MLAS_SRC_DIR}/softmax.h ) target_sources(onnxruntime_mlas PRIVATE @@ -97,6 +98,9 @@ function(setup_mlas_source_for_windows) ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon_fp16.cpp ${MLAS_SRC_DIR}/hgemm_kernel_neon.cpp ${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp + ${MLAS_SRC_DIR}/softmax_kernel_neon.h + ${MLAS_SRC_DIR}/softmax_kernel_neon.cpp + ${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp ) set(mlas_platform_preprocess_srcs @@ -377,6 +381,8 @@ else() ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.h ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon.cpp ${MLAS_SRC_DIR}/hgemm_kernel_neon.cpp + ${MLAS_SRC_DIR}/softmax_kernel_neon.h + ${MLAS_SRC_DIR}/softmax_kernel_neon.cpp ) set_source_files_properties(${MLAS_SRC_DIR}/sqnbitgemm_kernel_neon_int8.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+dotprod") @@ -398,6 +404,7 @@ else() ${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp ${MLAS_SRC_DIR}/rotary_embedding_kernel_neon_fp16.cpp ${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp + ${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp ) set_source_files_properties(${MLAS_SRC_DIR}/aarch64/HalfGemmKernelNeon.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") set_source_files_properties(${MLAS_SRC_DIR}/aarch64/QgemmS8S8KernelSmmla.S PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+i8mm ") @@ -411,6 +418,7 @@ else() set_source_files_properties(${MLAS_SRC_DIR}/hqnbitgemm_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") set_source_files_properties(${MLAS_SRC_DIR}/rotary_embedding_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") set_source_files_properties(${MLAS_SRC_DIR}/halfgemm_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") + set_source_files_properties(${MLAS_SRC_DIR}/softmax_kernel_neon_fp16.cpp PROPERTIES COMPILE_FLAGS " -march=armv8.2-a+fp16 ") endif() if(ONNXRUNTIME_MLAS_MULTI_ARCH) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index 7e0335cc66..d311bd92b1 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -990,11 +990,12 @@ MlasComputeErf( size_t N ); +template void MLASCALL MlasComputeExp( - const float* Input, - float* Output, + const T* Input, + T* Output, size_t N ); @@ -1006,11 +1007,12 @@ MlasComputeLogistic( size_t N ); +template void MLASCALL MlasComputeSoftmax( - const float* Input, - float* Output, + const T* Input, + T* Output, size_t N, size_t D, bool LogSoftmax, @@ -1018,11 +1020,12 @@ MlasComputeSoftmax( MLAS_THREADPOOL* ThreadPool ); +template void MLASCALL MlasComputeTanh( - const float* Input, - float* Output, + const T* Input, + T* Output, size_t N ); diff --git a/onnxruntime/core/mlas/lib/compute.cpp b/onnxruntime/core/mlas/lib/compute.cpp index 73df23e64c..97ad1ca3a5 100644 --- a/onnxruntime/core/mlas/lib/compute.cpp +++ b/onnxruntime/core/mlas/lib/compute.cpp @@ -20,6 +20,7 @@ Abstract: --*/ #include "mlasi.h" +#include "softmax.h" // // Bundles the constants for use by kernels written in assembly. @@ -68,12 +69,13 @@ MLAS_INTERNAL_DATA const float MlasMinimumF32Value = std::numeric_limits: // threads. // +template struct MLAS_SOFTMAX_WORK_BLOCK { ptrdiff_t ThreadCountN; bool LogSoftmax; bool SmoothSoftmax; - const float* Input; - float* Output; + const T* Input; + T* Output; size_t N; size_t D; }; @@ -244,9 +246,10 @@ Return Value: } } +template <> void MLASCALL -MlasComputeExp( +MlasComputeExp( const float* Input, float* Output, size_t N @@ -280,6 +283,20 @@ Return Value: #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( @@ -783,10 +800,18 @@ Return Value: } } +template void MlasComputeSoftmaxThreaded( void* Context, ptrdiff_t Index +); + +template <> +void +MlasComputeSoftmaxThreaded( + void* Context, + ptrdiff_t Index ) /*++ @@ -807,7 +832,7 @@ Return Value: --*/ { - const auto* WorkBlock = (MLAS_SOFTMAX_WORK_BLOCK*)Context; + const auto* WorkBlock = (MLAS_SOFTMAX_WORK_BLOCK*)Context; // // Partition the operation along the N dimension. @@ -906,11 +931,85 @@ Return Value: } } +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(1.0f / accumulation_fp32)); + } + + Input += D; + Output += D; + CountN--; + } +} + +template void MLASCALL MlasComputeSoftmax( - const float* Input, - float* Output, + const T* Input, + T* Output, size_t N, size_t D, bool LogSoftmax, @@ -949,7 +1048,7 @@ Return Value: --*/ { - MLAS_SOFTMAX_WORK_BLOCK WorkBlock; + MLAS_SOFTMAX_WORK_BLOCK WorkBlock; // // Capture the softmax parameters to the work block. @@ -985,5 +1084,31 @@ Return Value: WorkBlock.ThreadCountN = ThreadCountN; - MlasExecuteThreaded(MlasComputeSoftmaxThreaded, &WorkBlock, ThreadCountN, ThreadPool); + MlasExecuteThreaded(MlasComputeSoftmaxThreaded, &WorkBlock, ThreadCountN, ThreadPool); } + +template +void +MLASCALL +MlasComputeSoftmax( + const float* Input, + float* Output, + size_t N, + size_t D, + bool LogSoftmax, + bool SmoothSoftmax, + MLAS_THREADPOOL* ThreadPool +); + +template +void +MLASCALL +MlasComputeSoftmax( + const MLAS_FP16* Input, + MLAS_FP16* Output, + size_t N, + size_t D, + bool LogSoftmax, + bool SmoothSoftmax, + MLAS_THREADPOOL* ThreadPool +); diff --git a/onnxruntime/core/mlas/lib/mlasi.h b/onnxruntime/core/mlas/lib/mlasi.h index 56fad6bb34..decf44708d 100644 --- a/onnxruntime/core/mlas/lib/mlasi.h +++ b/onnxruntime/core/mlas/lib/mlasi.h @@ -1065,6 +1065,9 @@ extern const MLAS_ROPE_DISPATCH MlasRopeDispatchNeon; struct MLAS_HGEMM_DISPATCH; extern const MLAS_HGEMM_DISPATCH MlasHGemmDispatchNeon; +// softmax dispatch structure +struct MLAS_SOFTMAX_DISPATCH; +extern const MLAS_SOFTMAX_DISPATCH MlasSoftmaxDispatchNeon; // // Quantized depthwise convolution kernels. @@ -1228,6 +1231,7 @@ struct MLAS_PLATFORM { const MLAS_ROPE_DISPATCH* RopeDispatch{nullptr}; const MLAS_HGEMM_DISPATCH* HGemmDispatch{nullptr}; + const MLAS_SOFTMAX_DISPATCH* SoftmaxDispatch{nullptr}; }; inline diff --git a/onnxruntime/core/mlas/lib/platform.cpp b/onnxruntime/core/mlas/lib/platform.cpp index 026a954bbc..9db3115c1a 100644 --- a/onnxruntime/core/mlas/lib/platform.cpp +++ b/onnxruntime/core/mlas/lib/platform.cpp @@ -545,6 +545,7 @@ Return Value: this->QNBitGemmDispatch = &MlasSQNBitGemmDispatchNeon; this->RopeDispatch = &MlasRopeDispatchNeon; this->HGemmDispatch = &MlasHGemmDispatchNeon; + this->SoftmaxDispatch = &MlasSoftmaxDispatchNeon; // // Check if the processor supports ASIMD dot product instructions. diff --git a/onnxruntime/core/mlas/lib/softmax.h b/onnxruntime/core/mlas/lib/softmax.h new file mode 100644 index 0000000000..c415f79e9e --- /dev/null +++ b/onnxruntime/core/mlas/lib/softmax.h @@ -0,0 +1,113 @@ +/*++ + +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 + * @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 exponential function for each element of the input array + * @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 + */ + 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. + * @param Input Address of the input array + * @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 + * @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 scale The scale factor to apply to the output + */ + typedef void(Softmax_Fp16_Fn)( + const MLAS_FP16* Input, + MLAS_FP16* Output, + size_t N, + const MLAS_FP16 scale + ); + + Softmax_Fp16_Fn* Softmax_Fp16 = nullptr; + + /** + * @brief Compute the log softmax output for each element of the input array + * @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/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp b/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp new file mode 100644 index 0000000000..d29ef4c5ef --- /dev/null +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp @@ -0,0 +1,37 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + softmax_kernel_neon.cpp + +Abstract: + + This module implements the softmax kernels for ARM NEON. + +--*/ + +#include "softmax.h" +#include "softmax_kernel_neon.h" + +// +// Kernel dispatch structure definition. +// +const MLAS_SOFTMAX_DISPATCH MlasSoftmaxDispatchNeon = []() { + MLAS_SOFTMAX_DISPATCH d; + +#if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) + if (MlasFp16AccelerationSupported()) { + d.Tanh_Fp16 = softmax_neon::Tanh_Kernel_Fp16; + d.Exp_Fp16 = softmax_neon::Exp_Kernel_Fp16; + d.ReduceMax_Fp16 = softmax_neon::ReduceMax_Kernel_Fp16; + d.SumExp_Fp16 = softmax_neon::SumExp_Kernel_Fp16; + d.Softmax_Fp16 = softmax_neon::Softmax_Kernel_Fp16; + d.LogSoftmax_Fp16 = softmax_neon::LogSoftmax_Kernel_Fp16; + } +#endif + return d; +}(); diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon.h b/onnxruntime/core/mlas/lib/softmax_kernel_neon.h new file mode 100644 index 0000000000..4612ef397a --- /dev/null +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon.h @@ -0,0 +1,38 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + softmax_kernel_neon.h + +Abstract: + + This module includes function declarations and common helper functions for + softmax on ARM cpu. + +--*/ + +#pragma once + +#include + +#include "mlasi.h" + +namespace softmax_neon { + +void Exp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N); + +void Tanh_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N); + +MLAS_FP16 ReduceMax_Kernel_Fp16(const MLAS_FP16* Input, size_t N); + +MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 NegativeMaximum); + +void Softmax_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 scale); + +void LogSoftmax_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 NegativeMaximum, const MLAS_FP16 LogSum); + +} // namespace rope_neon diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp new file mode 100644 index 0000000000..b32920524c --- /dev/null +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp @@ -0,0 +1,51 @@ +/*++ + +Copyright (c) Microsoft Corporation. All rights reserved. + +Licensed under the MIT License. + +Module Name: + + softmax_kernel_neon_fp16.cpp + +Abstract: + + This module implements the fp16 softmax kernels for ARM NEON. + +--*/ +#include +#include + +#include "fp16_common.h" +#include "softmax.h" +#include "softmax_kernel_neon.h" + +namespace softmax_neon { + +// exp kernel for fp16. Output and input can be the same buffer. +void Exp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N) { + +} + +// tanh kernel for fp16. Output and input can be the same buffer. +void Tanh_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N) { + +} + +// reduce max kernel for fp16 +MLAS_FP16 ReduceMax_Kernel_Fp16(const MLAS_FP16* Input, size_t N) { + return MLAS_FP16::FromBits(0); +} + +MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 NegativeMaximum) { + return MLAS_FP16::FromBits(0); +} + +void Softmax_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 scale) { + +} + +void LogSoftmax_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 NegativeMaximum, const MLAS_FP16 LogSum) { +} + +} // namespace rope_neon diff --git a/onnxruntime/core/mlas/lib/tanh.cpp b/onnxruntime/core/mlas/lib/tanh.cpp index 9750337237..a65a5a4170 100644 --- a/onnxruntime/core/mlas/lib/tanh.cpp +++ b/onnxruntime/core/mlas/lib/tanh.cpp @@ -21,6 +21,7 @@ Abstract: --*/ #include "mlasi.h" +#include "softmax.h" // // Bundles the floating point constants for use by kernels written in assembly. @@ -119,9 +120,9 @@ Return Value: float Value = *Input++; - // This odd two-step process exists to ensure an input value of NaN carries through - // without modification because "std::min" and "std::max" return unreliable results - // when NaNs are involved, and it's clear from the test's reference outputs that + // This odd two-step process exists to ensure an input value of NaN carries through + // without modification because "std::min" and "std::max" return unreliable results + // when NaNs are involved, and it's clear from the test's reference outputs that // they want a NaN on output whenever the input is a NaN. float v_tmp; v_tmp = (Value < MlasTanhConstants.LowerRange) ? MlasTanhConstants.LowerRange : Value; @@ -149,9 +150,10 @@ Return Value: } } +template <> void MLASCALL -MlasComputeTanh( +MlasComputeTanh( const float* Input, float* Output, size_t N @@ -182,3 +184,18 @@ Return Value: MlasTanhKernel(Input, Output, N); #endif } + +template <> +void +MLASCALL +MlasComputeTanh( + const MLAS_FP16* Input, + MLAS_FP16* Output, + size_t N +) { + const auto* dispatch = GetMlasPlatform().SoftmaxDispatch; + if (dispatch == nullptr || dispatch->Tanh_Fp16 == nullptr) { + MLAS_THROW_EX(std::runtime_error, "Tanh_Fp16 is not supported."); + } + dispatch->Tanh_Fp16(Input, Output, N); +} diff --git a/onnxruntime/test/mlas/unittest/test_exp.cpp b/onnxruntime/test/mlas/unittest/test_exp.cpp index f9cdffef19..178e42fc3f 100644 --- a/onnxruntime/test/mlas/unittest/test_exp.cpp +++ b/onnxruntime/test/mlas/unittest/test_exp.cpp @@ -50,7 +50,435 @@ class MlasComputeExpTest : public MlasTestBase { } }; +class MyComputeExpTest : public MlasTestBase { + private: + MatrixGuardBuffer BufferInput; + MatrixGuardBuffer BufferOutput; + MatrixGuardBuffer BufferOutputReference; + +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, // -150 * ln2 + 88.7762626647950f, // 128 * ln2 + -88.3762626647949f, + 88.3762626647949f, // 127.5 * ln2 + 12582912.f, // 1.5 * 2^23 + 1.44269504088896341f, + -6.93145752e-1f, + -1.42860677e-6f, + 0x1.694000p-10, // 6! // TODO: these polynomials may be chosen by optimization, even though difference is small. Test small number errors mine vs. hers. + 0x1.125edcp-7, // 5! + 0x1.555b5ap-5, // 4! + 0x1.555450p-3, // 3! + 0x1.fffff6p-2, // 2! + 0x1.000000p+0, + int32_t(0xC1000000), // -126 + int32_t(0x3F800000), // 1.0f +}; + +const struct { + _Float16 LowerRange; + _Float16 UpperRange; + _Float16 LowerRangeSumExp; + _Float16 UpperRangeSumExp; + _Float16 RoundingBias; + _Float16 Log2Reciprocal; + _Float16 Log2High; + _Float16 Log2Low; + _Float16 Log2Lowest; + _Float16 poly_0; + _Float16 poly_1; + _Float16 poly_2; + _Float16 poly_3; + _Float16 poly_4; + _Float16 poly_56; + int16_t MinimumExponent; + int16_t MaximumExponent; +} MlasExp16Constants = { + -17.328679513f16, // -25 * ln2 + 11.090354888f16, // 16 * ln2 + -10.743781298f16, + 10.743781298f16, // 15.5 * ln2 + 1536.f16, // 1.5 * 2^10 + 1.4423828125f16, + -6.9287109375e-1f16, // 0xb98b + -2.758502960205078e-4f16, // 0x8c85 + -2.384185791015625e-7f16, // 0x8004 + 1.388888888888889e-3f16, // 1/6! 0x15b0 + 8.333333333333333e-3f16, // 1/5! 0x2044 + 4.1666666666666664e-2f16, // 1/4! 0x2955 + 1.6666666e-1f16, // 1/3! 0x3155 + 0.5f16, // 1/2! 0x3800 + 1.0f16, + int16_t(0xC800), // -14 + int16_t(0x3C00), // 15 +}; + + void print_hex(std::string note, _Float16 x) { + int16_t i = *reinterpret_cast(&x); + std::cout << note << std::hex << i << std::dec << std::endl; + } + + void print_hex(std::string note, float x) { + int i = *reinterpret_cast(&x); + std::cout << note << std::hex << i << std::dec << std::endl; + } + + void print_hex(std::string note, int x) { + std::cout << note << std::hex << x << std::dec << std::endl; + } + + _Float16 my_exp(_Float16 x) { + bool debug = false; + x = std::min(std::max(x, MlasExp16Constants.LowerRange), MlasExp16Constants.UpperRange); + + auto biased = x * MlasExp16Constants.Log2Reciprocal + MlasExp16Constants.RoundingBias; + if (debug) print_hex("biased ", biased); + auto m = biased - MlasExp16Constants.RoundingBias; + if (debug) print_hex("m ", m); + + _Float16 r = m * MlasExp16Constants.Log2High + x; + r = m * MlasExp16Constants.Log2Low + r; + r = m * MlasExp16Constants.Log2Lowest + r; + if (debug) print_hex("r ", r); + + int16_t bias_i = *reinterpret_cast(&biased); + int16_t overflow = bias_i << 10; + if (debug) print_hex("overflow ", overflow); + auto normal = overflow; + + normal = std::min(normal, MlasExp16Constants.MaximumExponent); + normal = std::max(normal, MlasExp16Constants.MinimumExponent); + if (debug) print_hex("clampped normal ", normal); + + overflow = overflow - normal; + if (debug) print_hex("lowered overflow ", overflow); + overflow = overflow + MlasExp16Constants.MaximumExponent; + if (debug) print_hex("adjusted overflow ", overflow); + normal = normal + MlasExp16Constants.MaximumExponent; + if (debug) print_hex("adjusted normal ", normal); + + auto p = (_Float16)MlasExp16Constants.poly_0; + p = p * r + (_Float16)MlasExp16Constants.poly_1; + p = p * r + (_Float16)MlasExp16Constants.poly_2; + p = p * r + (_Float16)MlasExp16Constants.poly_3; + p = p * r + (_Float16)MlasExp16Constants.poly_4; + p = p * r + (_Float16)MlasExp16Constants.poly_56; + + _Float16 overflow_f = *reinterpret_cast<_Float16*>(&overflow); + _Float16 normal_f = *reinterpret_cast<_Float16*>(&normal); + r = r * overflow_f; + p = p * r + overflow_f; + p = p * normal_f; + + return p; + } + + float my_exp(float x) { + x = std::min(std::max(x, MlasExpConstants.LowerRange), MlasExpConstants.UpperRange); + + auto biased = x * MlasExpConstants.Log2Reciprocal + MlasExpConstants.RoundingBias; + print_hex("biased ", biased); + auto m = biased - MlasExpConstants.RoundingBias; + print_hex("m ", m); + + float r = m * MlasExpConstants.Log2High + x; + r = m * MlasExpConstants.Log2Low + r; + print_hex("r ", r); + + int32_t bias_i = *reinterpret_cast(&biased); + auto overflow = bias_i << 23; + print_hex("overflow ", overflow); + auto normal = overflow; + + normal = std::min(normal, MlasExpConstants.MaximumExponent); + normal = std::max(normal, MlasExpConstants.MinimumExponent); + print_hex("clampped normal ", normal); + + overflow = overflow - normal; + print_hex("lowered overflow ", overflow); + overflow = overflow + MlasExpConstants.MaximumExponent; + print_hex("adjusted overflow ", overflow); + normal = normal + MlasExpConstants.MaximumExponent; + print_hex("adjusted normal ", normal); + + auto p = MlasExpConstants.poly_0; + p = p * r + MlasExpConstants.poly_1; + p = p * r + MlasExpConstants.poly_2; + p = p * r + MlasExpConstants.poly_3; + p = p * r + MlasExpConstants.poly_4; + p = p * r + MlasExpConstants.poly_56; + + float overflow_f = *reinterpret_cast(&overflow); + float normal_f = *reinterpret_cast(&normal); + r = r * overflow_f; + p = p * r + overflow_f; + p = p * normal_f; + + return p; + } + + float my_exp_no_overflow(float x) { + x = std::min(std::max(x, MlasExpConstants.LowerRange), MlasExpConstants.UpperRange); + + auto biased = x * MlasExpConstants.Log2Reciprocal + MlasExpConstants.RoundingBias; + print_hex("biased ", biased); + auto m = biased - MlasExpConstants.RoundingBias; + print_hex("m ", m); + + float r = m * MlasExpConstants.Log2High + x; + r = m * MlasExpConstants.Log2Low + r; + print_hex("r ", r); + + int32_t bias_i = *reinterpret_cast(&biased); + auto normal = bias_i << 23; + print_hex("clampped normal ", normal); + normal = normal + MlasExpConstants.MaximumExponent; + print_hex("adjusted normal ", normal); + + auto p = MlasExpConstants.poly_0; + p = p * r + MlasExpConstants.poly_1; + p = p * r + MlasExpConstants.poly_2; + p = p * r + MlasExpConstants.poly_3; + p = p * r + MlasExpConstants.poly_4; + p = p * r + MlasExpConstants.poly_56; + p = p * r + MlasExpConstants.poly_56; + + + float normal_f = *reinterpret_cast(&normal); + p = p * normal_f; + + return p; + } + + void Test(float x) { + float ref = std::exp(x); + float out = my_exp_no_overflow(x); + + constexpr float AbsoluteTolerance = 1e-6f; + constexpr float RelativeTolerance = 1e-6f; + + float diff = std::fabs(out - ref); + ASSERT_TRUE(diff <= AbsoluteTolerance || diff <= std::fabs(ref) * RelativeTolerance) + << " of " << 1 << ", got: " << out << ", expecting: " << ref; + std::cout << "result: " << out << ", expecting: " << ref << std::endl; + } + + void Test(_Float16 x) { + float ref = std::exp(static_cast(x)); + float out = my_exp(x); + + constexpr float AbsoluteTolerance = 1e-6f; + constexpr float RelativeTolerance = 1e-6f; + + float diff = std::abs(out - ref); + // ASSERT_TRUE(diff <= AbsoluteTolerance || diff <= std::fabs(ref) * RelativeTolerance) + // << " of " << 1 << ", got: " << out << ", expecting: " << ref << " diff " << diff / ref; + std::cout << "x " << (float)x << ", result: " << out << ", expecting: " << ref << " diff " << diff / ref << std::endl; + } + +const struct { + float LowerRange; + float UpperRange; + float alpha_13; + float alpha_11; + float alpha_9; + float alpha_7; + float alpha_5; + float alpha_3; + float alpha_1; + float beta_6; + float beta_4; + float beta_2; + float beta_0; +} MlasTanhConstants = { + -9.0f, + 9.0f, + -2.76076847742355e-16f, + 2.00018790482477e-13f, + -8.60467152213735e-11f, + 5.12229709037114e-08f, + 1.48572235717979e-05f, + 6.37261928875436e-04f, + 4.89352455891786e-03f, + 1.19825839466702e-06f, // TODO: test errors + 1.18534705686654e-04f, + 2.26843463243900e-03f, + 4.89352518554385e-03f, +}; + +const struct { + _Float16 LowerRange; + _Float16 UpperRange; + _Float16 alpha_9; + _Float16 alpha_7; + _Float16 alpha_5; + _Float16 alpha_3; + _Float16 alpha_1; + _Float16 beta_10; + _Float16 beta_8; + _Float16 beta_6; + _Float16 beta_4; + _Float16 beta_2; + _Float16 beta_0; +} MlasTanh16Constants = { + -5.0f16, + 5.0f16, + 2.755731922398589e-06f16, // 0x002e + 0.00019841269841269839f16, // 0xa80 + 0.008333333333333333f16, // 0x2044 + 0.16666666666666666f16, // 0x3155 + 1.f16, // 0x3c00 + 2.7557319223985894e-07f16, // 0x0005 + 2.48015873015873e-05f16, // 0x01a0 + 0.001388888888888889f16, // 0x15b0 + 0.041666666666666664f16, // 0x2955 + 0.5f16, // 0x3800 + 1.f16, // 0x3c00 +}; + + float my_tanh(float Value) { + float v_tmp; + v_tmp = (Value < MlasTanhConstants.LowerRange) ? MlasTanhConstants.LowerRange : Value; + Value = (v_tmp > MlasTanhConstants.UpperRange) ? MlasTanhConstants.UpperRange : v_tmp; + + float ValueSquared = Value * Value; + + float p; + p = ValueSquared * MlasTanhConstants.alpha_13 + MlasTanhConstants.alpha_11; + p = p * ValueSquared + MlasTanhConstants.alpha_9; + p = p * ValueSquared + MlasTanhConstants.alpha_7; + p = p * ValueSquared + MlasTanhConstants.alpha_5; + p = p * ValueSquared + MlasTanhConstants.alpha_3; + p = p * ValueSquared + MlasTanhConstants.alpha_1; + p = p * Value; + + float q; + q = ValueSquared * MlasTanhConstants.beta_6 + MlasTanhConstants.beta_4; + q = q * ValueSquared + MlasTanhConstants.beta_2; + q = q * ValueSquared + MlasTanhConstants.beta_0; + + return (p / q); + } + + _Float16 my_tanh(_Float16 Value) { + _Float16 v_tmp; + v_tmp = (Value < MlasTanh16Constants.LowerRange) ? MlasTanh16Constants.LowerRange : Value; + Value = (v_tmp > MlasTanh16Constants.UpperRange) ? MlasTanh16Constants.UpperRange : v_tmp; + + _Float16 ValueSquared = Value * Value; + + _Float16 p = MlasTanh16Constants.alpha_9; + p = p * ValueSquared + MlasTanh16Constants.alpha_7; + p = p * ValueSquared + MlasTanh16Constants.alpha_5; + p = p * ValueSquared + MlasTanh16Constants.alpha_3; + p = p * ValueSquared + MlasTanh16Constants.alpha_1; + p = p * Value; + + _Float16 q = MlasTanh16Constants.beta_10; + q = q * ValueSquared + MlasTanh16Constants.beta_8; + q = q * ValueSquared + MlasTanh16Constants.beta_6; + q = q * ValueSquared + MlasTanh16Constants.beta_4; + q = q * ValueSquared + MlasTanh16Constants.beta_2; + q = q * ValueSquared + MlasTanh16Constants.beta_0; + + return (p / q); + } + + _Float16 fast_tanh(_Float16 x) { + _Float16 x2 = x * x; + _Float16 a = x * (135.1350f16 + x2 * (17.3250f16 + x2 * (.3780f16 + x2))); + _Float16 b = 135.1350f16 + x2 * (62.3700f16 + x2 * (3.1500f16 + x2 * .0280f16)); + return a / b; + } + + _Float16 my_tanh_no_overflow(_Float16 Value) { + if (Value > 0.45f16) { + _Float16 exp = my_exp(2.f16 * Value); + return 1 - 2.f16 / (1 + exp); + } else { + return my_tanh(Value); + } + } + + void test_tanh(float x) { + float ref = std::tanh(x); + float out = my_tanh(x); + float diff = std::abs(out - ref); + std::cout << "result: " << out << ", expecting: " << ref << " diff " << diff / ref << std::endl; + } + + void test_tanh(_Float16 x) { + float ref = std::tanh((float)x); + float out = my_tanh(x); + float diff = std::abs(out - ref); + std::cout << "x " << (float)x << ", result: " << out << ", expecting: " << ref << " diff " << diff / ref << std::endl; + } + + void test_tanh_no_overflow(_Float16 x) { + float ref = std::tanh((float)x); + float out = fast_tanh(x); + float diff = std::abs(out - ref); + std::cout << "x " << (float)x << ", result: " << out << ", expecting: " << ref << " diff " << diff / ref << std::endl; + } + + public: + static const char* GetTestSuiteName() { + static const std::string suite_name("MyExp"); + return suite_name.c_str(); + } + + void ExecuteShort(void) override { + // print_hex("h ", (_Float16)MlasExp16Constants.Log2High); + // print_hex("l ", (_Float16)MlasExp16Constants.Log2Low); + // print_hex("ll ", MlasExp16Constants.Log2Lowest); + // print_hex("r ", (_Float16)MlasExp16Constants.Log2Reciprocal); + // print_hex("poly0 ", (_Float16)MlasExp16Constants.poly_0); + // print_hex("poly1 ", (_Float16)MlasExp16Constants.poly_1); + // print_hex("poly2 ", (_Float16)MlasExp16Constants.poly_2); + // print_hex("poly3 ", (_Float16)MlasExp16Constants.poly_3); + // print_hex("poly4 ", (_Float16)MlasExp16Constants.poly_4); + // Test(.01f16); + // print_hex("alpha_13 ", MlasTanh16Constants.alpha_13); + // print_hex("alpha_11 ", MlasTanh16Constants.alpha_11); + // print_hex("alpha_9 ", MlasTanh16Constants.alpha_9); + // print_hex("alpha_7 ", MlasTanh16Constants.alpha_7); + // print_hex("alpha_5 ", MlasTanh16Constants.alpha_5); + // print_hex("alpha_3 ", MlasTanh16Constants.alpha_3); + // print_hex("alpha_1 ", MlasTanh16Constants.alpha_1); + // print_hex("beta_10 ", MlasTanh16Constants.beta_10); + // print_hex("beta_8 ", MlasTanh16Constants.beta_8); + // print_hex("beta_6 ", MlasTanh16Constants.beta_6); + // print_hex("beta_4 ", MlasTanh16Constants.beta_4); + // print_hex("beta_2 ", MlasTanh16Constants.beta_2); + // print_hex("beta_0 ", MlasTanh16Constants.beta_0); + for (_Float16 x = 0.f16; x <= 5.f16; x += 0.005f16) { + test_tanh_no_overflow(x); + } + } +}; + static UNUSED_VARIABLE bool added_to_main = AddTestRegister([](bool is_short_execute) { // no long execute needed - return is_short_execute ? MlasDirectShortExecuteTests::RegisterShortExecute() : 0; + if (is_short_execute) { + return MlasDirectShortExecuteTests::RegisterShortExecute() + + MlasDirectShortExecuteTests::RegisterShortExecute(); + } + return 0ul; });