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

finished tanh and softcap

Browse source
This commit is contained in:
Jing Fang 2025-02-06 00:39:30 +00:00
parent 7dd6ceede1
commit cc22f530cf
5 changed files with 287 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/core/mlas/lib/fp16_common.h
View file

@ -230,14 +230,14 @@ MlasMultiply(MLAS_FLOAT16X4 Vector1, MLAS_FLOAT16X4 Vector2)
MLAS_FORCEINLINE
MLAS_FLOAT16X8
MlasDivFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2)
MlasDivide(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2)
{
return vdivq_f16(Vector1, Vector2);
}
MLAS_FORCEINLINE
MLAS_FLOAT16X4
MlasDivFloat16x4(MLAS_FLOAT16X4 Vector1, MLAS_FLOAT16X4 Vector2)
MlasDivide(MLAS_FLOAT16X4 Vector1, MLAS_FLOAT16X4 Vector2)
{
return vdiv_f16(Vector1, Vector2);
}
@ -270,13 +270,6 @@ MlasMultiplyAddFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2, _mlas_f
MlasMultiplyAdd(Vector1, Vector2, MlasBroadcastFloat16x8(Scalar3));
}
MLAS_FORCEINLINE
MLAS_FLOAT16X8
MlasDivideFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2)
{
return vdivq_f16(Vector1, Vector2);
}
MLAS_FORCEINLINE
MLAS_FLOAT16X8
MlasGreaterThanFloat16x8(MLAS_FLOAT16X8 Vector1, MLAS_FLOAT16X8 Vector2)

4
onnxruntime/core/mlas/lib/pooling_fp16.cpp
View file

@ -158,14 +158,14 @@ template <>
MLAS_FORCEINLINE MLAS_FLOAT16X8
PoolSummary16x8<AveragePoolAggregation>(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 context)
{
return MlasDivFloat16x8(agg, context);
return MlasDivide(agg, context);
}
template <>
MLAS_FORCEINLINE MLAS_FLOAT16X4
PoolSummary16x4<AveragePoolAggregation>(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X8 context)
{
return MlasDivFloat16x4(agg, MlasToLowHalfFloat16x4(context));
return MlasDivide(agg, MlasToLowHalfFloat16x4(context));
}

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

@ -22,7 +22,7 @@ Abstract:
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 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
*/
@ -36,7 +36,7 @@ struct MLAS_SOFTMAX_DISPATCH {
/**
* @brief Compute the softcap function for each element of the input array. Use tanh activation.
* @param Input Address of the input array
* @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
@ -51,8 +51,8 @@ struct MLAS_SOFTMAX_DISPATCH {
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
* @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
*/
@ -79,7 +79,7 @@ struct MLAS_SOFTMAX_DISPATCH {
/**
* @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
* @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
@ -95,7 +95,7 @@ struct MLAS_SOFTMAX_DISPATCH {
/**
* @brief Compute the softmax output for each element of the input array
* @param Input Address of the input array
* @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.
* @param N Number of elements in the input array
* @param scale The scale factor to apply to the output

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

@ -405,76 +405,276 @@ MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N
return MLAS_FP16::FromBits(result);
}
const struct {
_mlas_fp16_ LowerRange;
_mlas_fp16_ UpperRange;
_mlas_fp16_ alpha_9;
_mlas_fp16_ alpha_7;
_mlas_fp16_ alpha_5;
_mlas_fp16_ alpha_3;
_mlas_fp16_ alpha_1;
_mlas_fp16_ beta_10;
_mlas_fp16_ beta_8;
_mlas_fp16_ beta_6;
_mlas_fp16_ beta_4;
_mlas_fp16_ beta_2;
_mlas_fp16_ beta_0;
} MlasTanh16Constants = {
0xc500, // -5.0f16
0x4500, // 5.0f16
0x002e, // 1/9!
0x0a80, // 1/7!
0x2044, // 1/5!
0x3155, // 1/3!
0x3c00, // 1
0x0005, // 1/10!
0x01a0, // 1/8!
0x15b0, // 1/6!
0x2955, // 1/4!
0x3800, // 1/2!
0x3c00, // 1
template <typename T>
struct MlasTanhConstants {
T LowerRange;
T UpperRange;
T alpha_7;
T alpha_5;
T alpha_3;
T alpha_1;
T beta_6;
T beta_4;
T beta_2;
T beta_0;
};
// _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;
const MlasTanhConstants<_mlas_fp16_> TanhConstantsFp16 = {
0xc308, // -3.51562
0x4308, // 3.51562
0x0001,
0x00f9,
0x1138,
0x1d03,
0x0014,
0x07c5,
0x18a5,
0x1d03,
};
// _Float16 ValueSquared = Value * Value;
const MlasTanhConstants<float16x8_t> TanhConstantsFp16x8 = {
MlasBroadcastFloat16x8(TanhConstantsFp16.LowerRange),
MlasBroadcastFloat16x8(TanhConstantsFp16.UpperRange),
MlasBroadcastFloat16x8(TanhConstantsFp16.alpha_7),
MlasBroadcastFloat16x8(TanhConstantsFp16.alpha_5),
MlasBroadcastFloat16x8(TanhConstantsFp16.alpha_3),
MlasBroadcastFloat16x8(TanhConstantsFp16.alpha_1),
MlasBroadcastFloat16x8(TanhConstantsFp16.beta_6),
MlasBroadcastFloat16x8(TanhConstantsFp16.beta_4),
MlasBroadcastFloat16x8(TanhConstantsFp16.beta_2),
MlasBroadcastFloat16x8(TanhConstantsFp16.beta_0),
};
// _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;
const MlasTanhConstants<float16x4_t> TanhConstantsFp16x4 = {
MlasBroadcastFloat16x4(TanhConstantsFp16.LowerRange),
MlasBroadcastFloat16x4(TanhConstantsFp16.UpperRange),
MlasBroadcastFloat16x4(TanhConstantsFp16.alpha_7),
MlasBroadcastFloat16x4(TanhConstantsFp16.alpha_5),
MlasBroadcastFloat16x4(TanhConstantsFp16.alpha_3),
MlasBroadcastFloat16x4(TanhConstantsFp16.alpha_1),
MlasBroadcastFloat16x4(TanhConstantsFp16.beta_6),
MlasBroadcastFloat16x4(TanhConstantsFp16.beta_4),
MlasBroadcastFloat16x4(TanhConstantsFp16.beta_2),
MlasBroadcastFloat16x4(TanhConstantsFp16.beta_0),
};
// _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;
template <typename T>
MLAS_FORCEINLINE
MlasTanhConstants<T> Get_Tanh_Constants();
// return (p / q);
// }
template <>
MLAS_FORCEINLINE
MlasTanhConstants<float16x8_t> Get_Tanh_Constants<float16x8_t>() {
return TanhConstantsFp16x8;
}
template <>
MLAS_FORCEINLINE
MlasTanhConstants<float16x4_t> Get_Tanh_Constants<float16x4_t>() {
return TanhConstantsFp16x4;
}
// _Float16 my_tanh_no_overflow(_Float16 Value) {
// if (Value > 0.5f16) {
// _Float16 exp = my_exp(Value);
// return (exp - 1.0f16/exp) / (exp + 1.0f16/exp);
// } else {
// return my_tanh(Value);
// }
// }
// TODO(fajin): optimize polynomial coefficients
template <typename T>
MLAS_FORCEINLINE
T Tanh_Vector_Fp16(T x) {
const auto constants = Get_Tanh_Constants<T>();
x = MlasClamp(x, constants.LowerRange, constants.UpperRange);
T x_2 = MlasMultiply(x, x);
T p = MlasMultiplyAdd(constants.alpha_7, x_2, constants.alpha_5);
p = MlasMultiplyAdd(p, x_2, constants.alpha_3);
p = MlasMultiplyAdd(p, x_2, constants.alpha_1);
p = MlasMultiply(p, x);
T q = MlasMultiplyAdd(constants.beta_6, x_2, constants.beta_4);
q = MlasMultiplyAdd(q, x_2, constants.beta_2);
q = MlasMultiplyAdd(q, x_2, constants.beta_0);
return MlasDivide(p / q);
}
void Tanh_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N) {
const auto* input = reinterpret_cast<const _mlas_fp16_*>(Input);
auto* output = reinterpret_cast<_mlas_fp16_*>(Output);
while (N >= 32) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
auto v2 = MlasLoadFloat16x8(input + 16);
auto v3 = MlasLoadFloat16x8(input + 24);
auto r0 = Tanh_Vector_Fp16(v0);
auto r1 = Tanh_Vector_Fp16(v1);
auto r2 = Tanh_Vector_Fp16(v2);
auto r3 = Tanh_Vector_Fp16(v3);
MlasStoreFloat16x8(output, r0);
MlasStoreFloat16x8(output + 8, r1);
MlasStoreFloat16x8(output + 16, r2);
MlasStoreFloat16x8(output + 24, r3);
input += 32;
output += 32;
N -= 32;
}
if (N & 16) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
auto r0 = Tanh_Vector_Fp16(v0);
auto r1 = Tanh_Vector_Fp16(v1);
MlasStoreFloat16x8(output, r0);
MlasStoreFloat16x8(output + 8, r1);
input += 16;
output += 16;
N -= 16;
}
if (N & 8) {
auto v0 = MlasLoadFloat16x8(input);
auto r0 = Tanh_Vector_Fp16(v0);
MlasStoreFloat16x8(output, r0);
input += 8;
output += 8;
N -= 8;
}
if (N & 4) {
auto v0 = MlasLoadFloat16x4(input);
auto r0 = Tanh_Vector_Fp16(v0);
MlasStoreFloat16x4(output, r0);
input += 4;
output += 4;
N -= 4;
}
if (N == 3) {
auto v0 = MlasLoadPartialFloat16x4(input, 3);
auto r0 = Tanh_Vector_Fp16(v0);
MlasStorePartialFloat16x4(output, r0, 3);
} else if (N == 2) {
auto v0 = MlasLoadPartialFloat16x4(input, 2);
auto r0 = Tanh_Vector_Fp16(v0);
MlasStorePartialFloat16x4(output, r0, 2);
} else if (N == 1) {
auto v0 = MlasLoadPartialFloat16x4(input, 1);
auto r0 = Tanh_Vector_Fp16(v0);
MlasStorePartialFloat16x4(output, r0, 1);
}
}
void Softcap_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 Softcap) {
const auto* input = reinterpret_cast<const _mlas_fp16_*>(Input);
auto* output = reinterpret_cast<_mlas_fp16_*>(Output);
auto softcap8 = MlasBroadcastFloat16x8(Softcap.val);
auto softcap4 = MlasBroadcastFloat16x4(Softcap.val);
auto one8 = MlasBroadcastFloat16x8((_mlas_fp16_)0x3c00);
auto one4 = MlasBroadcastFloat16x4((_mlas_fp16_)0x3c00);
auto softcap_reciprocal8 = MlasDivide(one8, softcap8);
auto softcap_reciprocal4 = MlasDivide(one4, softcap4);
while (N >= 32) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
auto v2 = MlasLoadFloat16x8(input + 16);
auto v3 = MlasLoadFloat16x8(input + 24);
v0 = MlasMultiply(v0, softcap_reciprocal8);
v1 = MlasMultiply(v1, softcap_reciprocal8);
v2 = MlasMultiply(v2, softcap_reciprocal8);
v3 = MlasMultiply(v3, softcap_reciprocal8);
v0 = Tanh_Vector_Fp16(v0);
v1 = Tanh_Vector_Fp16(v1);
v2 = Tanh_Vector_Fp16(v2);
v3 = Tanh_Vector_Fp16(v3);
v0 = MlasMultiply(v0, softcap8);
v1 = MlasMultiply(v1, softcap8);
v2 = MlasMultiply(v2, softcap8);
v3 = MlasMultiply(v3, softcap8);
MlasStoreFloat16x8(output, v0);
MlasStoreFloat16x8(output + 8, v1);
MlasStoreFloat16x8(output + 16, v2);
MlasStoreFloat16x8(output + 24, v3);
input += 32;
output += 32;
N -= 32;
}
if (N & 16) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
v0 = MlasMultiply(v0, softcap_reciprocal8);
v1 = MlasMultiply(v1, softcap_reciprocal8);
v0 = Tanh_Vector_Fp16(v0);
v1 = Tanh_Vector_Fp16(v1);
v0 = MlasMultiply(v0, softcap8);
v1 = MlasMultiply(v1, softcap8);
MlasStoreFloat16x8(output, v0);
MlasStoreFloat16x8(output + 8, v1);
input += 16;
output += 16;
N -= 16;
}
if (N & 8) {
auto v0 = MlasLoadFloat16x8(input);
v0 = MlasMultiply(v0, softcap_reciprocal8);
v0 = Tanh_Vector_Fp16(v0);
v0 = MlasMultiply(v0, softcap8);
MlasStoreFloat16x8(output, v0);
input += 8;
output += 8;
N -= 8;
}
if (N & 4) {
auto v0 = MlasLoadFloat16x4(input);
v0 = MlasMultiply(v0, softcap_reciprocal4);
v0 = Tanh_Vector_Fp16(v0);
v0 = MlasMultiply(v0, softcap4);
MlasStoreFloat16x4(output, v0);
input += 4;
output += 4;
N -= 4;
}
if (N == 3) {
auto v0 = MlasLoadPartialFloat16x4(input, 3);
v0 = MlasMultiply(v0, softcap_reciprocal4);
v0 = Tanh_Vector_Fp16(v0);
v0 = MlasMultiply(v0, softcap4);
MlasStorePartialFloat16x4(output, v0, 3);
} else if (N == 2) {
auto v0 = MlasLoadPartialFloat16x4(input, 2);
v0 = MlasMultiply(v0, softcap_reciprocal4);
v0 = Tanh_Vector_Fp16(v0);
v0 = MlasMultiply(v0, softcap4);
MlasStorePartialFloat16x4(output, v0, 2);
} else if (N == 1) {
auto v0 = MlasLoadPartialFloat16x4(input, 1);
v0 = MlasMultiply(v0, softcap_reciprocal4);
v0 = Tanh_Vector_Fp16(v0);
v0 = MlasMultiply(v0, softcap4);
MlasStorePartialFloat16x4(output, v0, 1);
}
}
MLAS_FP16 ReduceMax_Kernel_Fp16(const MLAS_FP16* Input, size_t N) {

44
onnxruntime/test/mlas/unittest/test_exp.cpp
View file

@ -326,31 +326,31 @@ const struct {
const struct {
_Float16 LowerRange;
_Float16 UpperRange;
_Float16 alpha_13;
_Float16 alpha_11;
_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, // c500
5.0f16, // 4500
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
-3.51562f16,
3.51562f16,
-2.76076847742355e-16f16,
2.00018790482477e-13f16,
-8.60467152213735e-11f16,
5.12229709037114e-08f16,
1.48572235717979e-05f16,
6.37261928875436e-04f16,
4.89352455891786e-03f16,
1.19825839466702e-06f16, // TODO: test errors
1.18534705686654e-04f16,
2.26843463243900e-03f16,
4.89352518554385e-03f16,
};
float my_tanh(float Value) {
@ -384,16 +384,16 @@ const struct {
_Float16 ValueSquared = Value * Value;
_Float16 p = MlasTanh16Constants.alpha_9;
_Float16 p = MlasTanh16Constants.alpha_13;
p = p * ValueSquared + MlasTanh16Constants.alpha_11;
p = p * ValueSquared + 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;
_Float16 q = MlasTanh16Constants.beta_6;
q = q * ValueSquared + MlasTanh16Constants.beta_4;
q = q * ValueSquared + MlasTanh16Constants.beta_2;
q = q * ValueSquared + MlasTanh16Constants.beta_0;
@ -463,19 +463,19 @@ const struct {
// Test(.01f16);
print_hex("lower range ", MlasTanh16Constants.LowerRange);
print_hex("upper range ", MlasTanh16Constants.UpperRange);
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 <= 9.f16; x += 0.005f16) {
test_tanh_no_overflow(x);
test_tanh(x);
}
}
};