// Copyright 2026 pchar.cn
// SPDX-License-Identifier: BSD-3-Clause

#if NCNN_RUNTIME_CPU && NCNN_AVX512BF16 && __AVX512F__ && !__AVX512BF16__
void rotaryembed_bf16s_avx512bf16(const Mat& bottom_blob, const Mat& cos_cache, const Mat& sin_cache, Mat& top_blob, int interleaved, const Option& opt);
#endif

static void rotaryembed_bf16s(const Mat& bottom_blob, const Mat& cos_cache, const Mat& sin_cache, Mat& top_blob, int interleaved, const Option& opt)
{
#if NCNN_RUNTIME_CPU && NCNN_AVX512BF16 && __AVX512F__ && !__AVX512BF16__
    if (ncnn::cpu_support_x86_avx512_bf16())
    {
        rotaryembed_bf16s_avx512bf16(bottom_blob, cos_cache, sin_cache, top_blob, interleaved, opt);
        return;
    }
#endif

    const int embed_dim = bottom_blob.w;
    const int seqlen = bottom_blob.h;
    const int num_heads = bottom_blob.c;

    top_blob.create_like(bottom_blob, opt.blob_allocator);
    if (top_blob.empty())
        return;

    #pragma omp parallel for num_threads(opt.num_threads)
    for (int q = 0; q < num_heads; q++)
    {
        const Mat head = bottom_blob.channel(q);
        Mat out_head = top_blob.channel(q);

        for (int i = 0; i < seqlen; i++)
        {
            if (interleaved)
            {
                const unsigned short* ptr = head.row<const unsigned short>(i);
                const unsigned short* cos_ptr = cos_cache.row<const unsigned short>(i);
                const unsigned short* sin_ptr = sin_cache.row<const unsigned short>(i);
                unsigned short* outptr = out_head.row<unsigned short>(i);

                int j = 0;
#if __SSE2__
#if __AVX__
#if __AVX512F__
                const __m512i dupidx = _mm512_set_epi32(15, 15, 14, 14, 13, 13, 12, 12, 11, 11, 10, 10, 9, 9, 8, 8);
                const __m512i dupidx_lo = _mm512_set_epi32(7, 7, 6, 6, 5, 5, 4, 4, 3, 3, 2, 2, 1, 1, 0, 0);
                for (; j + 15 < embed_dim / 2; j += 16)
                {
                    __m512 a0 = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)ptr));
                    __m512 a1 = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)(ptr + 16)));

                    __m512 cs_src = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)cos_ptr));
                    __m512 ss_src = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)sin_ptr));

                    __m512 c0 = _mm512_permutexvar_ps(dupidx_lo, cs_src);
                    __m512 c1 = _mm512_permutexvar_ps(dupidx, cs_src);
                    __m512 s0 = _mm512_permutexvar_ps(dupidx_lo, ss_src);
                    __m512 s1 = _mm512_permutexvar_ps(dupidx, ss_src);

                    __m512 swap0 = _mm512_shuffle_ps(a0, a0, _MM_SHUFFLE(2, 3, 0, 1));
                    __m512 swap1 = _mm512_shuffle_ps(a1, a1, _MM_SHUFFLE(2, 3, 0, 1));

                    __m512 ss0 = _mm512_mul_ps(swap0, s0);
                    __m512 ss1 = _mm512_mul_ps(swap1, s1);

                    __m512 y0 = _mm512_fmaddsub_ps(a0, c0, ss0);
                    __m512 y1 = _mm512_fmaddsub_ps(a1, c1, ss1);

                    _mm256_storeu_si256((__m256i*)outptr, float2bfloat_avx512(y0));
                    _mm256_storeu_si256((__m256i*)(outptr + 16), float2bfloat_avx512(y1));

                    ptr += 32;
                    outptr += 32;
                    cos_ptr += 16;
                    sin_ptr += 16;
                }
#endif // __AVX512F__
#if __AVX2__
                const __m256i dupidx256 = _mm256_set_epi32(7, 7, 6, 6, 5, 5, 4, 4);
                const __m256i dupidx256_lo = _mm256_set_epi32(3, 3, 2, 2, 1, 1, 0, 0);
                for (; j + 7 < embed_dim / 2; j += 8)
                {
                    __m256 a0 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)ptr));
                    __m256 a1 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)(ptr + 8)));

                    __m256 c_src = bfloat2float_avx(_mm_loadu_si128((const __m128i*)cos_ptr));
                    __m256 s_src = bfloat2float_avx(_mm_loadu_si128((const __m128i*)sin_ptr));

                    __m256 c0 = _mm256_permutevar8x32_ps(c_src, dupidx256_lo);
                    __m256 c1 = _mm256_permutevar8x32_ps(c_src, dupidx256);
                    __m256 s0 = _mm256_permutevar8x32_ps(s_src, dupidx256_lo);
                    __m256 s1 = _mm256_permutevar8x32_ps(s_src, dupidx256);

                    __m256 swap0 = _mm256_shuffle_ps(a0, a0, _MM_SHUFFLE(2, 3, 0, 1));
                    __m256 swap1 = _mm256_shuffle_ps(a1, a1, _MM_SHUFFLE(2, 3, 0, 1));

                    __m256 ss0 = _mm256_mul_ps(swap0, s0);
                    __m256 ss1 = _mm256_mul_ps(swap1, s1);

                    __m256 y0 = _mm256_fmaddsub_ps(a0, c0, ss0);
                    __m256 y1 = _mm256_fmaddsub_ps(a1, c1, ss1);

                    _mm_storeu_si128((__m128i*)outptr, float2bfloat_avx(y0));
                    _mm_storeu_si128((__m128i*)(outptr + 8), float2bfloat_avx(y1));

                    ptr += 16;
                    outptr += 16;
                    cos_ptr += 8;
                    sin_ptr += 8;
                }
#else // __AVX2__
                for (; j + 7 < embed_dim / 2; j += 8)
                {
                    __m256 a0 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)ptr));
                    __m256 a1 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)(ptr + 8)));

                    __m128 clo4 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)cos_ptr));
                    __m128 chi4 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)(cos_ptr + 4)));
                    __m128 slo4 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)sin_ptr));
                    __m128 shi4 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)(sin_ptr + 4)));

                    __m128 clo_lo = _mm_unpacklo_ps(clo4, clo4); // [c0,c0,c1,c1]
                    __m128 clo_hi = _mm_unpackhi_ps(clo4, clo4); // [c2,c2,c3,c3]
                    __m128 chi_lo = _mm_unpacklo_ps(chi4, chi4); // [c4,c4,c5,c5]
                    __m128 chi_hi = _mm_unpackhi_ps(chi4, chi4); // [c6,c6,c7,c7]

                    __m256 c0 = combine4x2_ps(clo_lo, clo_hi);
                    __m256 c1 = combine4x2_ps(chi_lo, chi_hi);

                    __m128 slo_lo = _mm_unpacklo_ps(slo4, slo4); // [s0,s0,s1,s1]
                    __m128 slo_hi = _mm_unpackhi_ps(slo4, slo4); // [s2,s2,s3,s3]
                    __m128 shi_lo = _mm_unpacklo_ps(shi4, shi4); // [s4,s4,s5,s5]
                    __m128 shi_hi = _mm_unpackhi_ps(shi4, shi4); // [s6,s6,s7,s7]

                    __m256 s0 = combine4x2_ps(slo_lo, slo_hi);
                    __m256 s1 = combine4x2_ps(shi_lo, shi_hi);

                    __m256 swap0 = _mm256_shuffle_ps(a0, a0, _MM_SHUFFLE(2, 3, 0, 1));
                    __m256 swap1 = _mm256_shuffle_ps(a1, a1, _MM_SHUFFLE(2, 3, 0, 1));

                    __m256 ss0 = _mm256_mul_ps(swap0, s0);
                    __m256 ss1 = _mm256_mul_ps(swap1, s1);

#if __FMA__
                    __m256 y0 = _mm256_fmaddsub_ps(a0, c0, ss0);
                    __m256 y1 = _mm256_fmaddsub_ps(a1, c1, ss1);
#else
                    __m256 ac0 = _mm256_mul_ps(a0, c0);
                    __m256 ac1 = _mm256_mul_ps(a1, c1);

                    __m256 y0 = _mm256_addsub_ps(ac0, ss0);
                    __m256 y1 = _mm256_addsub_ps(ac1, ss1);
#endif
                    _mm_storeu_si128((__m128i*)outptr, float2bfloat_avx(y0));
                    _mm_storeu_si128((__m128i*)(outptr + 8), float2bfloat_avx(y1));

                    ptr += 16;
                    outptr += 16;
                    cos_ptr += 8;
                    sin_ptr += 8;
                }
#endif // __AVX2__
#endif // __AVX__
                for (; j + 3 < embed_dim / 2; j += 4)
                {
                    __m128 a0 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)ptr));
                    __m128 a1 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)(ptr + 4)));

                    __m128i c4_raw = _mm_loadl_epi64((const __m128i*)cos_ptr);
                    __m128i s4_raw = _mm_loadl_epi64((const __m128i*)sin_ptr);
                    __m128 c4 = bfloat2float_sse(c4_raw);
                    __m128 s4 = bfloat2float_sse(s4_raw);

                    __m128 clo = _mm_unpacklo_ps(c4, c4); // [c0,c0,c1,c1]
                    __m128 chi = _mm_unpackhi_ps(c4, c4); // [c2,c2,c3,c3]
                    __m128 slo = _mm_unpacklo_ps(s4, s4); // [s0,s0,s1,s1]
                    __m128 shi = _mm_unpackhi_ps(s4, s4); // [s2,s2,s3,s3]

                    __m128 swap0 = _mm_shuffle_ps(a0, a0, _MM_SHUFFLE(2, 3, 0, 1));
                    __m128 swap1 = _mm_shuffle_ps(a1, a1, _MM_SHUFFLE(2, 3, 0, 1));

                    __m128 ss0 = _mm_mul_ps(swap0, slo);
                    __m128 ss1 = _mm_mul_ps(swap1, shi);
#if __FMA__
                    __m128 y0 = _mm_fmaddsub_ps(a0, clo, ss0);
                    __m128 y1 = _mm_fmaddsub_ps(a1, chi, ss1);
#else
                    __m128 ac0 = _mm_mul_ps(a0, clo);
                    __m128 ac1 = _mm_mul_ps(a1, chi);
#if __SSE3__
                    __m128 y0 = _mm_addsub_ps(ac0, ss0);
                    __m128 y1 = _mm_addsub_ps(ac1, ss1);
#else
#if defined(__MINGW32__) && !defined(__x86_64__)
                    __attribute__((aligned(16)))
                    const float signmask128_array[4]
                        = {-0.f, 0.f, -0.f, 0.f};
                    const __m128 signmask128 = _mm_load_ps(signmask128_array);
#else
                    const __m128 signmask128 = _mm_set_ps(0.f, -0.f, 0.f, -0.f);
#endif
                    ss0 = _mm_xor_ps(ss0, signmask128);
                    ss1 = _mm_xor_ps(ss1, signmask128);
                    __m128 y0 = _mm_add_ps(ac0, ss0);
                    __m128 y1 = _mm_add_ps(ac1, ss1);
#endif
#endif
                    __m128i y01_bf16 = float2bfloat_sse(y0, y1);
                    _mm_storeu_si128((__m128i*)outptr, y01_bf16);

                    ptr += 8;
                    outptr += 8;
                    cos_ptr += 4;
                    sin_ptr += 4;
                }
                for (; j + 1 < embed_dim / 2; j += 2)
                {
                    __m128i a_raw = _mm_loadl_epi64((const __m128i*)ptr);
                    __m128 a = bfloat2float_sse(a_raw);

                    float cos0 = bfloat16_to_float32(cos_ptr[0]);
                    float cos1 = bfloat16_to_float32(cos_ptr[1]);
                    float sin0 = bfloat16_to_float32(sin_ptr[0]);
                    float sin1 = bfloat16_to_float32(sin_ptr[1]);

                    __m128 c = _mm_set_ps(cos1, cos1, cos0, cos0);
                    __m128 s = _mm_set_ps(sin1, sin1, sin0, sin0);

                    __m128 swap = _mm_shuffle_ps(a, a, _MM_SHUFFLE(2, 3, 0, 1));
                    __m128 ss = _mm_mul_ps(swap, s);

#if __FMA__
                    __m128 y = _mm_fmaddsub_ps(a, c, ss);
#else
                    __m128 ac = _mm_mul_ps(a, c);
#if __SSE3__
                    __m128 y = _mm_addsub_ps(ac, ss);
#else
#if defined(__MINGW32__) && !defined(__x86_64__)
                    __attribute__((aligned(16)))
                    const float signmask128_array[4]
                        = {-0.f, 0.f, -0.f, 0.f};
                    const __m128 signmask128 = _mm_load_ps(signmask128_array);
#else
                    const __m128 signmask128 = _mm_set_ps(0.f, -0.f, 0.f, -0.f);
#endif
                    ss = _mm_xor_ps(ss, signmask128);
                    __m128 y = _mm_add_ps(ac, ss);
#endif
#endif
                    __m128i y_bf16 = float2bfloat_sse(y, y);
                    _mm_storel_epi64((__m128i*)outptr, y_bf16);

                    ptr += 4;
                    outptr += 4;
                    cos_ptr += 2;
                    sin_ptr += 2;
                }
#endif // __SSE2__
                for (; j < embed_dim / 2; j++)
                {
                    const float x0 = bfloat16_to_float32(ptr[0]);
                    const float x1 = bfloat16_to_float32(ptr[1]);
                    const float cos_val = bfloat16_to_float32(*cos_ptr++);
                    const float sin_val = bfloat16_to_float32(*sin_ptr++);

                    outptr[0] = float32_to_bfloat16(x0 * cos_val - x1 * sin_val);
                    outptr[1] = float32_to_bfloat16(x0 * sin_val + x1 * cos_val);

                    ptr += 2;
                    outptr += 2;
                }
            }
            else
            {
                const unsigned short* ptr0 = head.row<const unsigned short>(i);
                const unsigned short* ptr1 = ptr0 + embed_dim / 2;
                const unsigned short* cos_ptr = cos_cache.row<const unsigned short>(i);
                const unsigned short* sin_ptr = sin_cache.row<const unsigned short>(i);

                unsigned short* outptr0 = out_head.row<unsigned short>(i);
                unsigned short* outptr1 = outptr0 + embed_dim / 2;

                int j = 0;
#if __SSE2__
#if __AVX__
#if __AVX512F__
                for (; j + 15 < embed_dim / 2; j += 16)
                {
                    __m512 x0 = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)ptr0));
                    __m512 x1 = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)ptr1));
                    __m512 c = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)cos_ptr));
                    __m512 s = bfloat2float_avx512(_mm256_loadu_si256((const __m256i*)sin_ptr));

                    __m512 y0 = _mm512_fnmadd_ps(x1, s, _mm512_mul_ps(x0, c));
                    __m512 y1 = _mm512_fmadd_ps(x0, s, _mm512_mul_ps(x1, c));

                    _mm256_storeu_si256((__m256i*)outptr0, float2bfloat_avx512(y0));
                    _mm256_storeu_si256((__m256i*)outptr1, float2bfloat_avx512(y1));

                    ptr0 += 16;
                    ptr1 += 16;
                    cos_ptr += 16;
                    sin_ptr += 16;
                    outptr0 += 16;
                    outptr1 += 16;
                }
#endif // __AVX512F__
                for (; j + 7 < embed_dim / 2; j += 8)
                {
                    __m256 x0 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)ptr0));
                    __m256 x1 = bfloat2float_avx(_mm_loadu_si128((const __m128i*)ptr1));
                    __m256 c = bfloat2float_avx(_mm_loadu_si128((const __m128i*)cos_ptr));
                    __m256 s = bfloat2float_avx(_mm_loadu_si128((const __m128i*)sin_ptr));

                    __m256 y0 = _mm256_comp_fnmadd_ps(x1, s, _mm256_mul_ps(x0, c));
                    __m256 y1 = _mm256_comp_fmadd_ps(x0, s, _mm256_mul_ps(x1, c));

                    _mm_storeu_si128((__m128i*)outptr0, float2bfloat_avx(y0));
                    _mm_storeu_si128((__m128i*)outptr1, float2bfloat_avx(y1));

                    ptr0 += 8;
                    ptr1 += 8;
                    cos_ptr += 8;
                    sin_ptr += 8;
                    outptr0 += 8;
                    outptr1 += 8;
                }
#endif // __AVX__
                for (; j + 3 < embed_dim / 2; j += 4)
                {
                    __m128 x0 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)ptr0));
                    __m128 x1 = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)ptr1));
                    __m128 c = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)cos_ptr));
                    __m128 s = bfloat2float_sse(_mm_loadl_epi64((const __m128i*)sin_ptr));

                    __m128 y0 = _mm_comp_fnmadd_ps(x1, s, _mm_mul_ps(x0, c));
                    __m128 y1 = _mm_comp_fmadd_ps(x0, s, _mm_mul_ps(x1, c));

                    _mm_storel_epi64((__m128i*)outptr0, float2bfloat_sse(y0, y0));
                    _mm_storel_epi64((__m128i*)outptr1, float2bfloat_sse(y1, y1));

                    ptr0 += 4;
                    ptr1 += 4;
                    cos_ptr += 4;
                    sin_ptr += 4;
                    outptr0 += 4;
                    outptr1 += 4;
                }
#endif // __SSE2__
                for (; j < embed_dim / 2; j++)
                {
                    const float x0 = bfloat16_to_float32(*ptr0++);
                    const float x1 = bfloat16_to_float32(*ptr1++);
                    const float cos_val = bfloat16_to_float32(*cos_ptr++);
                    const float sin_val = bfloat16_to_float32(*sin_ptr++);

                    *outptr0++ = float32_to_bfloat16(x0 * cos_val - x1 * sin_val);
                    *outptr1++ = float32_to_bfloat16(x0 * sin_val + x1 * cos_val);
                }
            }
        }
    }
}
