added the apis for fp16 softmax kernels

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)

onnxruntime/core/mlas/inc/mlas.h

@@ -990,11 +990,12 @@ MlasComputeErf(
     size_t N
     );
 
+template <typename T>
 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 <typename T>
 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<typename T>
 void
 MLASCALL
 MlasComputeTanh(
-    const float* Input,
-    float* Output,
+    const T* Input,
+    T* Output,
     size_t N
     );
 

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<float>:
 // threads.
 //
 
+template <typename T>
 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<float>(
     const float* Input,
     float* Output,
     size_t N
@@ -280,6 +283,20 @@ Return Value:
 #endif
 }
 
+template <>
+void MLASCALL
+MlasComputeExp<MLAS_FP16>(
+    const MLAS_FP16* Input,
+    MLAS_FP16* Output,
+    size_t N
+) {
+    const auto* dispatch = GetMlasPlatform().SoftmaxDispatch;
+    if (dispatch == nullptr || dispatch->Exp_Fp16 == nullptr) {
+        MLAS_THROW_EX(std::runtime_error, "Exp_Fp16 is not supported.");
+    }
+    dispatch->Exp_Fp16(Input, Output, N);
+}
+
 MLAS_FORCEINLINE
 MLAS_FLOAT32X4
 MlasComputeSumExpVector(
@@ -783,10 +800,18 @@ Return Value:
     }
 }
 
+template <typename T>
 void
 MlasComputeSoftmaxThreaded(
     void* Context,
     ptrdiff_t Index
+);
+
+template <>
+void
+MlasComputeSoftmaxThreaded<float>(
+    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<float>*)Context;
 
     //
     // Partition the operation along the N dimension.
@@ -906,11 +931,85 @@ Return Value:
     }
 }
 
+template <>
+void
+MlasComputeSoftmaxThreaded<MLAS_FP16>(
+    void* Context,
+    ptrdiff_t Index
+)
+/*++
+
+Routine Description:
+
+    This routine is invoked from a worker thread to execute a segment of a
+    softmax or log softmax operation.
+
+Arguments:
+
+    Context - Supplies the pointer to the context for the threaded operation.
+
+    ThreadId - Supplies the current index of the threaded operation.
+
+Return Value:
+
+    None.
+
+--*/
+{
+    const auto* WorkBlock = (MLAS_SOFTMAX_WORK_BLOCK<MLAS_FP16>*)Context;
+    size_t n;
+    size_t CountN;
+    MlasPartitionWork(Index, WorkBlock->ThreadCountN, WorkBlock->N, &n, &CountN);
+
+    const size_t D = WorkBlock->D;
+    const bool LogSoftmax = WorkBlock->LogSoftmax;
+    const bool SmoothSoftmax = WorkBlock->SmoothSoftmax;
+
+    const MLAS_FP16* Input = WorkBlock->Input + n * D;
+    MLAS_FP16* Output = WorkBlock->Output + n * D;
+
+    const auto* dispatch = GetMlasPlatform().SoftmaxDispatch;
+    if (dispatch == nullptr ||
+        dispatch->ReduceMax_Fp16 == nullptr ||
+        dispatch->SumExp_Fp16 == nullptr ||
+        (LogSoftmax && dispatch->LogSoftmax_Fp16 == nullptr) ||
+        (!LogSoftmax && dispatch->Softmax_Fp16 == nullptr)) {
+        MLAS_THROW_EX(std::runtime_error, "Lacks kernels for fp16 softmax.");
+    }
+
+    while (CountN > 0) {
+        MLAS_FP16 Maximum = dispatch->ReduceMax_Fp16(Input, D);
+        MLAS_FP16 NegativeMaximum = Maximum.Negate();
+        if (SmoothSoftmax && !NegativeMaximum.IsNegative()) {
+            NegativeMaximum = MLAS_FP16::FromBits(0);
+        }
+
+        MLAS_FP16* Temp = LogSoftmax ? nullptr : Output;
+        MLAS_FP16 Accumulation = dispatch->SumExp_Fp16(Input, Temp, D, NegativeMaximum);
+        float accumulation_fp32 = Accumulation.ToFloat();
+
+        if (SmoothSoftmax) {
+            accumulation_fp32 += expf(NegativeMaximum.ToFloat());
+        }
+
+        if (LogSoftmax) {
+            dispatch->LogSoftmax_Fp16(Input, Output, D, NegativeMaximum, MLAS_FP16(std::log(accumulation_fp32)));
+        } else {
+            dispatch->Softmax_Fp16(Output, Output, D, MLAS_FP16(1.0f / accumulation_fp32));
+        }
+
+        Input += D;
+        Output += D;
+        CountN--;
+    }
+}
+
+template <typename T>
 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<T> 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<T>, &WorkBlock, ThreadCountN, ThreadPool);
 }
+
+template
+void
+MLASCALL
+MlasComputeSoftmax<float>(
+    const float* Input,
+    float* Output,
+    size_t N,
+    size_t D,
+    bool LogSoftmax,
+    bool SmoothSoftmax,
+    MLAS_THREADPOOL* ThreadPool
+);
+
+template
+void
+MLASCALL
+MlasComputeSoftmax<MLAS_FP16>(
+    const MLAS_FP16* Input,
+    MLAS_FP16* Output,
+    size_t N,
+    size_t D,
+    bool LogSoftmax,
+    bool SmoothSoftmax,
+    MLAS_THREADPOOL* ThreadPool
+);

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

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.

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;
+};

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;
+}();

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 <arm_neon.h>
+
+#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

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 <arm_neon.h>
+#include <cassert>
+
+#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

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<float>(
     const float* Input,
     float* Output,
     size_t N
@@ -182,3 +184,18 @@ Return Value:
     MlasTanhKernel(Input, Output, N);
 #endif
 }
+
+template <>
+void
+MLASCALL
+MlasComputeTanh<MLAS_FP16>(
+    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);
+}

onnxruntime/test/mlas/unittest/test_exp.cpp

@@ -50,7 +50,435 @@ class MlasComputeExpTest : public MlasTestBase {
   }
 };
 
+class MyComputeExpTest : public MlasTestBase {
+ private:
+  MatrixGuardBuffer<float> BufferInput;
+  MatrixGuardBuffer<float> BufferOutput;
+  MatrixGuardBuffer<float> 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<int16_t*>(&x);
+    std::cout << note << std::hex << i << std::dec << std::endl;
+  }
+
+  void print_hex(std::string note, float x) {
+    int i = *reinterpret_cast<int*>(&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<int16_t*>(&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<int*>(&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<float*>(&overflow);
+    float normal_f = *reinterpret_cast<float*>(&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<int*>(&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<float*>(&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<float>(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<MlasComputeExpTest>::RegisterShortExecute() : 0;
+  if (is_short_execute) {
+    return MlasDirectShortExecuteTests<MlasComputeExpTest>::RegisterShortExecute() +
+           MlasDirectShortExecuteTests<MyComputeExpTest>::RegisterShortExecute();
+  }
+  return 0ul;
 });