diff --git a/onnxruntime/core/mlas/lib/fp16_common.h b/onnxruntime/core/mlas/lib/fp16_common.h index a76e22dbc8..59f460055a 100644 --- a/onnxruntime/core/mlas/lib/fp16_common.h +++ b/onnxruntime/core/mlas/lib/fp16_common.h @@ -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) diff --git a/onnxruntime/core/mlas/lib/pooling_fp16.cpp b/onnxruntime/core/mlas/lib/pooling_fp16.cpp index 685f3eb22a..62670a6986 100644 --- a/onnxruntime/core/mlas/lib/pooling_fp16.cpp +++ b/onnxruntime/core/mlas/lib/pooling_fp16.cpp @@ -158,14 +158,14 @@ template <> MLAS_FORCEINLINE MLAS_FLOAT16X8 PoolSummary16x8(MLAS_FLOAT16X8 agg, MLAS_FLOAT16X8 context) { - return MlasDivFloat16x8(agg, context); + return MlasDivide(agg, context); } template <> MLAS_FORCEINLINE MLAS_FLOAT16X4 PoolSummary16x4(MLAS_FLOAT16X4 agg, MLAS_FLOAT16X8 context) { - return MlasDivFloat16x4(agg, MlasToLowHalfFloat16x4(context)); + return MlasDivide(agg, MlasToLowHalfFloat16x4(context)); } diff --git a/onnxruntime/core/mlas/lib/softmax.h b/onnxruntime/core/mlas/lib/softmax.h index d1c5c39496..b91b81ff79 100644 --- a/onnxruntime/core/mlas/lib/softmax.h +++ b/onnxruntime/core/mlas/lib/softmax.h @@ -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 diff --git a/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp index 6e41fb411a..fbf20db8ca 100644 --- a/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp +++ b/onnxruntime/core/mlas/lib/softmax_kernel_neon_fp16.cpp @@ -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 +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 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 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 +MLAS_FORCEINLINE +MlasTanhConstants Get_Tanh_Constants(); -// return (p / q); -// } +template <> +MLAS_FORCEINLINE +MlasTanhConstants Get_Tanh_Constants() { + return TanhConstantsFp16x8; +} +template <> +MLAS_FORCEINLINE +MlasTanhConstants Get_Tanh_Constants() { + 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 +MLAS_FORCEINLINE +T Tanh_Vector_Fp16(T x) { + const auto constants = Get_Tanh_Constants(); + 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(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(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) { diff --git a/onnxruntime/test/mlas/unittest/test_exp.cpp b/onnxruntime/test/mlas/unittest/test_exp.cpp index d25ad071bd..8cde90d966 100644 --- a/onnxruntime/test/mlas/unittest/test_exp.cpp +++ b/onnxruntime/test/mlas/unittest/test_exp.cpp @@ -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); } } };