saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-13 18:08:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

added softcap api

Browse source
This commit is contained in:
Jing Fang 2025-02-04 23:17:27 +00:00
parent f27df5be80
commit 7d7b535349
6 changed files with 53 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
onnxruntime/core/mlas/inc/mlas.h
View file

@ -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<typename T>
void
MLASCALL

16
onnxruntime/core/mlas/lib/softmax.h
View file

@ -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

1
onnxruntime/core/mlas/lib/softmax_kernel_neon.cpp
View file

@ -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;

6
onnxruntime/core/mlas/lib/softmax_kernel_neon.h
View file

@ -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);

12
onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp
View file

@ -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) {
}

15
onnxruntime/core/mlas/lib/tanh.cpp
View file

@ -199,3 +199,18 @@ MlasComputeTanh<MLAS_FP16>(
}
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);
}