// Copyright 2025 Tencent
// SPDX-License-Identifier: BSD-3-Clause

#version 450

layout (constant_id = 0) const int cast_type_from = 0;
layout (constant_id = 1) const int cast_type_to = 1;

#define shape_constant_id_offset 2
layout (constant_id = shape_constant_id_offset + 0) const uint n = 0;
layout (constant_id = shape_constant_id_offset + 1) const uint c = 0;
layout (constant_id = shape_constant_id_offset + 2) const uint stride = 0;

layout (binding = 0) readonly buffer bottom_blob { sfpvec8 bottom_blob_data[]; };
layout (binding = 1) readonly buffer bottom_blob_fp32 { mat2x4 bottom_blob_fp32_data[]; };
layout (binding = 2) writeonly buffer top_blob { sfp top_blob_data[]; };
layout (binding = 3) writeonly buffer top_blob_fp32 { float top_blob_fp32_data[]; };

layout (push_constant) uniform parameter
{
    uint n;
    uint c;
    uint stride;
} p;

void main()
{
    const uint gx = gl_GlobalInvocationID.x;
    const uint gy = gl_GlobalInvocationID.y;

    if (gx >= psc(n) || gy >= psc(c))
        return;

    const uint gi = gy * psc(n) + gx;

    const uvec4 gi4 = (gy * 8 + uvec4(0, 1, 2, 3)) * psc(stride) + gx;
    const uvec4 gi8 = gi4 + psc(stride) * 4;

    if (cast_type_from == cast_type_to)
    {
        buffer_cp8to1(top_blob_data, gi4, gi8, bottom_blob_data, gi);
        return;
    }

    afpvec8 v;
    if (cast_type_from == 1)
    {
        mat2x4 vf = bottom_blob_fp32_data[gi];
        v[0] = afpvec4(vf[0]);
        v[1] = afpvec4(vf[1]);
    }
    else
    {
        v = buffer_ld8(bottom_blob_data, gi);
    }

    if (cast_type_to == 1)
    {
        vec4 vf0 = vec4(v[0]);
        vec4 vf1 = vec4(v[1]);
        top_blob_fp32_data[gi4.r] = vf0.r;
        top_blob_fp32_data[gi4.g] = vf0.g;
        top_blob_fp32_data[gi4.b] = vf0.b;
        top_blob_fp32_data[gi4.a] = vf0.a;
        top_blob_fp32_data[gi8.r] = vf1.r;
        top_blob_fp32_data[gi8.g] = vf1.g;
        top_blob_fp32_data[gi8.b] = vf1.b;
        top_blob_fp32_data[gi8.a] = vf1.a;
    }
    else
    {
        buffer_st1(top_blob_data, gi4.r, v[0].r);
        buffer_st1(top_blob_data, gi4.g, v[0].g);
        buffer_st1(top_blob_data, gi4.b, v[0].b);
        buffer_st1(top_blob_data, gi4.a, v[0].a);
        buffer_st1(top_blob_data, gi8.r, v[1].r);
        buffer_st1(top_blob_data, gi8.g, v[1].g);
        buffer_st1(top_blob_data, gi8.b, v[1].b);
        buffer_st1(top_blob_data, gi8.a, v[1].a);
    }
}
