finished reduce max kernel

This commit is contained in:
Jing Fang 2025-02-06 01:03:43 +00:00
parent cc22f530cf
commit 30eef71dd5
2 changed files with 93 additions and 7 deletions

View file

@ -407,6 +407,25 @@ MlasReduceAdd(MLAS_FLOAT16X4 Vector)
return vget_lane_u16(vreinterpret_u16_f16(Vector), 0);
}
MLAS_FORCEINLINE
_mlas_fp16_
MlasReduceMaximum(MLAS_FLOAT16X8 Vector)
{
Vector = vpmaxq_f16(Vector, Vector);
Vector = vpmaxq_f16(Vector, Vector);
Vector = vpmaxq_f16(Vector, Vector);
return vgetq_lane_u16(vreinterpretq_u16_f16(Vector), 0);
}
MLAS_FORCEINLINE
_mlas_fp16_
MlasReduceMaximum(MLAS_FLOAT16X4 Vector)
{
Vector = vpmax_f16(Vector, Vector);
Vector = vpmax_f16(Vector, Vector);
return vget_lane_u16(vreinterpret_u16_f16(Vector), 0);
}
MLAS_FORCEINLINE
MLAS_UINT16X8
MlasCmpLessEqualFloat16x8(MLAS_FLOAT16X8 left, MLAS_FLOAT16X8 right)

View file

@ -372,7 +372,7 @@ MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N
MlasStorePartialFloat16x4(output, r0, 3);
}
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 3));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 3));
accumulator4 = MlasAdd(r0, accumulator4);
} else if (N == 2) {
auto v0 = MlasLoadPartialFloat16x4(input, 2);
@ -382,8 +382,8 @@ MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N
MlasStorePartialFloat16x4(output, r0, 2);
}
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 3));
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 2));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 3));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 2));
accumulator4 = MlasAdd(r0, accumulator4);
} else if (N == 1) {
auto v0 = MlasLoadPartialFloat16x4(input, 1);
@ -393,9 +393,9 @@ MLAS_FP16 SumExp_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N
MlasStorePartialFloat16x4(output, r0, 1);
}
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 3));
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 2));
r0 = MlasReinterpretAsFloat16(vset_lane_s16(0, MlasReinterpretAsInt16(r0), 1));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 3));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 2));
r0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0, MlasReinterpretAsInt16(r0), 1));
accumulator4 = MlasAdd(r0, accumulator4);
}
@ -678,7 +678,74 @@ void Softcap_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, co
}
MLAS_FP16 ReduceMax_Kernel_Fp16(const MLAS_FP16* Input, size_t N) {
return MLAS_FP16::FromBits(0);
const auto* input = reinterpret_cast<const _mlas_fp16_*>(Input);
auto max8 = MlasBroadcastFloat16x8((_mlas_fp16_)0xfbff);
auto max4 = MlasBroadcastFloat16x4((_mlas_fp16_)0xfbff);
while (N >= 32) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
auto v2 = MlasLoadFloat16x8(input + 16);
auto v3 = MlasLoadFloat16x8(input + 24);
v0 = MlasMaximum(v0, v1);
v2 = MlasMaximum(v2, v3);
v0 = MlasMaximum(v0, v2);
max8 = MlasMaximum(max8, v0);
input += 32;
N -= 32;
}
if (N & 16) {
auto v0 = MlasLoadFloat16x8(input);
auto v1 = MlasLoadFloat16x8(input + 8);
v0 = MlasMaximum(v0, v1);
max8 = MlasMaximum(max8, v0);
input += 16;
N -= 16;
}
if (N & 8) {
auto v0 = MlasLoadFloat16x8(input);
max8 = MlasMaximum(max8, v0);
input += 8;
N -= 8;
}
if (N & 4) {
auto v0 = MlasLoadFloat16x4(input);
max4 = MlasMaximum(max4, v0);
input += 4;
N -= 4;
}
if (N == 3) {
auto v0 = MlasLoadPartialFloat16x4(input, 3);
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 3));
max4 = MlasMaximum(max4, v0);
} else if (N == 2) {
auto v0 = MlasLoadPartialFloat16x4(input, 2);
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 3));
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 2));
max4 = MlasMaximum(max4, v0);
} else if (N == 1) {
auto v0 = MlasLoadPartialFloat16x4(input, 1);
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 3));
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 2));
v0 = MlasReinterpretAsFloat16(vset_lane_s16((int16_t)0xfbff, MlasReinterpretAsInt16(v0), 1));
max4 = MlasMaximum(max4, v0);
}
auto t = MlasMaximum(vget_low_f16(max8), vget_high_f16(max8));
t = MlasMaximum(t, max4);
_mlas_fp16_ result = MlasReduceMaximum(t);
return MLAS_FP16::FromBits(result);
}
void Softmax_Kernel_Fp16(const MLAS_FP16* Input, MLAS_FP16* Output, size_t N, const MLAS_FP16 scale) {