请用neon c++帮我实现一个argmax函数

下面是一个使用neon指令集实现的argmax函数的示例代码：

#include <arm_neon.h>

int argmax_neon(const float* arr, int size) {
    int index = 0;
    float max_val = arr[0];

    // 使用neon指令集进行4个元素的并行比较
    float32x4_t v_max = vld1q_f32(arr);
    uint32x4_t v_index = vdupq_n_u32(0);

    for (int i = 4; i < size; i += 4) {
        float32x4_t v = vld1q_f32(arr + i);
        uint32x4_t v_mask = vcgtq_f32(v, v_max);
        v_max = vmaxq_f32(v, v_max);
        v_index = vbslq_u32(v_mask, vdupq_n_u32(i), v_index);
    }

    // 找出最大值和对应的索引
    float max_vals[4];
    uint32_t indices[4];
    vst1q_f32(max_vals, v_max);
    vst1q_u32(indices, v_index);

    for (int i = 1; i < 4; ++i) {
        if (max_vals[i] > max_val) {
            max_val = max_vals[i];
            index = indices[i];
        }
    }

    // 处理剩余的元素
    for (int i = size - size % 4; i < size; ++i) {
        if (arr[i] > max_val) {
            max_val = arr[i];
            index = i;
        }
    }

    return index;
}

int main() {
    float arr[] = {1.2, 3.4, 2.1, 5.6, 4.3, 0.9, 2.8, 3.7};
    int size = sizeof(arr) / sizeof(float);
    int index = argmax_neon(arr, size);
    printf("argmax: %d\n", index);
    return 0;
}

这个示例代码使用了neon指令集的并行比较和向量操作指令来加速argmax函数的运算。在函数中，首先加载前4个元素到向量寄存器中，然后循环处理剩余的元素，每次加载4个元素并与当前最大值进行比较，同时更新最大值和对应的索引。最后，再处理剩余的元素，找出最大值和对应的索引。最终，返回最大值的索引