// 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 { sfpvec4 bottom_blob_data[]; };
layout (binding = 1) readonly buffer bottom_blob_fp32 { vec4 bottom_blob_fp32_data[]; };
layout (binding = 2) writeonly buffer top_blob { sfpvec4 top_blob_data[]; };
layout (binding = 3) writeonly buffer top_blob_fp32 { vec4 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;

    if (cast_type_from == cast_type_to)
    {
        buffer_cp4(top_blob_data, gi, bottom_blob_data, gi);
        return;
    }

    const uint gi2 = gy * psc(stride) + gx;

    afpvec4 v;
    if (cast_type_from == 1)
    {
        v = afpvec4(bottom_blob_fp32_data[gi]);
    }
    else
    {
        v = buffer_ld4(bottom_blob_data, gi2);
    }

    if (cast_type_to == 1)
    {
        top_blob_fp32_data[gi] = vec4(v);
    }
    else
    {
        buffer_st4(top_blob_data, gi2, v);
    }
}
