From 7d7b535349077c4e5851d0bcf7b83d1557794d77 Mon Sep 17 00:00:00 2001 From: Jing Fang Date: Tue, 4 Feb 2025 23:17:27 +0000 Subject: [PATCH] added softcap api --- onnxruntime/core/mlas/inc/mlas.h | 9 +++++++++ onnxruntime/core/mlas/lib/softmax.h | 16 ++++++++++++++++ .../core/mlas/lib/softmax_kernel_neon.cpp | 1 + onnxruntime/core/mlas/lib/softmax_kernel_neon.h | 6 ++++-- .../core/mlas/lib/softmax_kernel_neon_fp16.cpp | 12 ++++++++---- onnxruntime/core/mlas/lib/tanh.cpp | 15 +++++++++++++++ 6 files changed, 53 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/mlas/inc/mlas.h b/onnxruntime/core/mlas/inc/mlas.h index d311bd92b1..329f32fbbd 100644 --- a/onnxruntime/core/mlas/inc/mlas.h +++ b/onnxruntime/core/mlas/inc/mlas.h @@ -1020,6 +1020,15 @@ MlasComputeSoftmax( MLAS_THREADPOOL* ThreadPool ); +void +MLASCALL +MlasComputeSoftcap( + const MLAS_FP16* Input, + MLAS_FP16* Output, + size_t N, + MLAS_FP16 cap + ); + template void MLASCALL diff --git a/onnxruntime/core/mlas/lib/softmax.h b/onnxruntime/core/mlas/lib/softmax.h index c415f79e9e..6c9bbd706b 100644 --- a/onnxruntime/core/mlas/lib/softmax.h +++ b/onnxruntime/core/mlas/lib/softmax.h @@ -34,6 +34,22 @@ struct MLAS_SOFTMAX_DISPATCH { 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 + * @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 diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp b/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp index d29ef4c5ef..77ad4b99a0 100644 --- a/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp @@ -26,6 +26,7 @@ const MLAS_SOFTMAX_DISPATCH MlasSoftmaxDispatchNeon = []() { #if defined(MLAS_F16VEC_INTRINSICS_SUPPORTED) && defined(MLAS_TARGET_ARM64) if (MlasFp16AccelerationSupported()) { d.Tanh_Fp16 = softmax_neon::Tanh_Kernel_Fp16; + d.Softcap_Fp16 = softmax_neon::Softcap_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; diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon.h b/onnxruntime/core/mlas/lib/softmax_kernel_neon.h index 4612ef397a..53207a0448 100644 --- a/onnxruntime/core/mlas/lib/softmax_kernel_neon.h +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon.h @@ -23,10 +23,12 @@ Abstract: 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); +void Softcap_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 Softcap); + +void Exp_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); diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp index b32920524c..113ec53e91 100644 --- a/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp @@ -22,13 +22,17 @@ Abstract: 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) { } -// 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) { +void Softcap_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 Softcap) { + +} + +// 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) { } diff --git a/onnxruntime/core/mlas/lib/tanh.cpp b/onnxruntime/core/mlas/lib/tanh.cpp index a65a5a4170..0a13b88918 100644 --- a/onnxruntime/core/mlas/lib/tanh.cpp +++ b/onnxruntime/core/mlas/lib/tanh.cpp @@ -199,3 +199,18 @@ MlasComputeTanh( } dispatch->Tanh_Fp16(Input, Output, N); } + +void +MLASCALL +MlasComputeSoftcap( + const MLAS_FP16* Input, + MLAS_FP16* Output, + size_t N, + MLAS_FP16 cap +) { + const auto* dispatch = GetMlasPlatform().SoftmaxDispatch; + if (dispatch == nullptr || dispatch->Softcap_Fp16 == nullptr) { + MLAS_THROW_EX(std::runtime_error, "Softcap_Fp16 is not supported."); + } + dispatch->Softcap_Fp16(Input, Output, N, cap); +}