Function kernel_softmax_f16_single

cactus/kernel/kernel_nn.cpp:416–499

}

void kernel_softmax_f16_single(const __fp16* input, __fp16* output, size_t vocab_size) {

    constexpr size_t SIMD_WIDTH = 8;
    constexpr size_t UNROLL_FACTOR = 4;
    constexpr size_t VECTORIZED_WIDTH = SIMD_WIDTH * UNROLL_FACTOR;
    const size_t vocab_vectorized = (vocab_size / VECTORIZED_WIDTH) * VECTORIZED_WIDTH;

    float32x4_t max_vec[UNROLL_FACTOR * 2];
    for (size_t u = 0; u < UNROLL_FACTOR * 2; u++) {
        max_vec[u] = vdupq_n_f32(-std::numeric_limits<float>::infinity());
    }

    for (size_t i = 0; i < vocab_vectorized; i += VECTORIZED_WIDTH) {
        for (size_t u = 0; u < UNROLL_FACTOR; u++) {
            float16x8_t x_vec_f16 = vld1q_f16(&input[i + u * SIMD_WIDTH]);
            float32x4_t x_low = vcvt_f32_f16(vget_low_f16(x_vec_f16));
            float32x4_t x_high = vcvt_f32_f16(vget_high_f16(x_vec_f16));
            max_vec[u * 2] = vmaxq_f32(max_vec[u * 2], x_low);
            max_vec[u * 2 + 1] = vmaxq_f32(max_vec[u * 2 + 1], x_high);
        }
    }

    float32x4_t final_max = max_vec[0];
    for (size_t u = 1; u < UNROLL_FACTOR * 2; u++) {
        final_max = vmaxq_f32(final_max, max_vec[u]);
    }

    float max_val = vmaxvq_f32(final_max);
    for (size_t i = vocab_vectorized; i < vocab_size; ++i) {
        max_val = std::max(max_val, static_cast<float>(input[i]));
    }

    const float32x4_t max_broadcast = vdupq_n_f32(max_val);

    float32x4_t sum_vec[UNROLL_FACTOR * 2];
    for (size_t u = 0; u < UNROLL_FACTOR * 2; u++) {
        sum_vec[u] = vdupq_n_f32(0.0f);
    }

    for (size_t i = 0; i < vocab_vectorized; i += VECTORIZED_WIDTH) {
        for (size_t u = 0; u < UNROLL_FACTOR; u++) {
            float16x8_t x_vec_f16 = vld1q_f16(&input[i + u * SIMD_WIDTH]);

            float32x4_t x_low = vcvt_f32_f16(vget_low_f16(x_vec_f16));
            float32x4_t x_high = vcvt_f32_f16(vget_high_f16(x_vec_f16));

            float32x4_t exp_low = fast_exp_f32x4(vsubq_f32(x_low, max_broadcast));
            float32x4_t exp_high = fast_exp_f32x4(vsubq_f32(x_high, max_broadcast));

            float16x8_t exp_f16 = vcombine_f16(vcvt_f16_f32(exp_low), vcvt_f16_f32(exp_high));
            vst1q_f16(&output[i + u * SIMD_WIDTH], exp_f16);

            sum_vec[u * 2] = vaddq_f32(sum_vec[u * 2], exp_low);
            sum_vec[u * 2 + 1] = vaddq_f32(sum_vec[u * 2 + 1], exp_high);
        }
    }

    float32x4_t final_sum = sum_vec[0];