#version 450

#extension GL_EXT_shader_16bit_storage : require

layout (push_constant) uniform parameter
{
    uint batch_offset; uint offset_delta;
    uint IC;
    uint IW; uint IH;
    uint OW; uint OH;
    uint KW; uint KH;
    uint pelements;
    uint CHW;
    int s0; int s1;
    int p0; int p1;
    int d0; int d1;
} p;

#include "types.comp"

#define BLOCK_SIZE 256

layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;

layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
layout (binding = 1) writeonly buffer D {D_TYPE data_d[];};

void main() {
    const uint i = gl_GlobalInvocationID.x;
    if (i >= p.pelements) {
        return;
    }

    const uint ksize = p.OW * (p.KH > 1 ? p.KW : 1);
    const uint kx = i / ksize;
    const uint kd = kx * ksize;
    const uint ky = (i - kd) / p.OW;
    const uint ix = i % p.OW;

    const uint oh = gl_GlobalInvocationID.y;
    const uint batch = gl_GlobalInvocationID.z / p.IC;
    const uint ic = gl_GlobalInvocationID.z % p.IC;

    const uint iiw = ix * p.s0 + kx * p.d0 - p.p0;
    const uint iih = oh * p.s1 + ky * p.d1 - p.p1;

    const uint offset_dst =
        ((batch * p.OH + oh) * p.OW + ix) * p.CHW +
        (ic * (p.KW * p.KH) + ky * p.KW + kx);

    if (iih < 0 || iih >= p.IH || iiw < 0 || iiw >= p.IW) {
        data_d[offset_dst] = D_TYPE(0.0f);
    } else {
        const uint offset_src = ic * p.offset_delta + batch * p.batch_offset;
        data_d[offset_dst] = D_TYPE(data_a[offset_src + iih * p.IW + iiw]);
    }
}