#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : require

layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
layout(constant_id = 3) const uint ACTIVATION = 0; // 0:none 1:relu 2:relu6

layout(std430, binding = 0) readonly buffer ActivationBuffer {
    FLOAT4 X4[];
};

layout(std430, binding = 1) readonly buffer WeightBuffer {
    uint Wp[]; // [padN, padK/8], 8x int4(raw) per uint32
};

layout(std430, binding = 2) readonly buffer MetaBuffer {
    FLOAT Meta[]; // [padN, blockStride, 2] => [scale, offset]
};

layout(std430, binding = 3) readonly buffer BiasBuffer {
    FLOAT Bias[];
};

layout(std430, binding = 4) writeonly buffer OutputBuffer {
    FLOAT Y[];
};

layout(push_constant) uniform PushConstants {
    uint K;
    uint N;
    uint blockSize;
    uint blockStride;
    uint weightStride;
} pc;

void main() {
    const uint row = gl_WorkGroupID.x;
    if (row >= pc.N) {
        return;
    }

    const uint lane = gl_LocalInvocationID.x;
    const uint stride = gl_SubgroupSize;

    const uint K8 = pc.K >> 3;
    const uint safeBlockSize = max(pc.blockSize, 1u);
    const uint block8 = max(safeBlockSize >> 3, 1u);

    const uint baseW = row * pc.weightStride;
    const uint baseM = row * pc.blockStride;

    float sum = 0.0;

    for (uint idx = lane; idx < K8; idx += stride * 2u) {
        const uint i0 = idx;
        const uint i1 = idx + stride;

        if (i0 < K8) {
            const uint kb = i0 / block8;
            const uint metaBase = (baseM + kb) * 2;
            const float scale = float(Meta[metaBase + 0]);
            const float offset = float(Meta[metaBase + 1]);

            const uint packedW = Wp[baseW + i0];
            vec4 w0 = vec4(
                float((packedW >> 0u) & 0xFu),
                float((packedW >> 4u) & 0xFu),
                float((packedW >> 8u) & 0xFu),
                float((packedW >> 12u) & 0xFu));
            vec4 w1 = vec4(
                float((packedW >> 16u) & 0xFu),
                float((packedW >> 20u) & 0xFu),
                float((packedW >> 24u) & 0xFu),
                float((packedW >> 28u) & 0xFu));

            w0 = (w0 - vec4(8.0)) * scale + vec4(offset);
            w1 = (w1 - vec4(8.0)) * scale + vec4(offset);

            const uint xBase = i0 << 1;
            vec4 x0 = vec4(X4[xBase + 0u]);
            vec4 x1 = vec4(X4[xBase + 1u]);
            sum += dot(w0, x0) + dot(w1, x1);
        }

        if (i1 < K8) {
            const uint kb = i1 / block8;
            const uint metaBase = (baseM + kb) * 2;
            const float scale = float(Meta[metaBase + 0]);
            const float offset = float(Meta[metaBase + 1]);

            const uint packedW = Wp[baseW + i1];
            vec4 w0 = vec4(
                float((packedW >> 0u) & 0xFu),
                float((packedW >> 4u) & 0xFu),
                float((packedW >> 8u) & 0xFu),
                float((packedW >> 12u) & 0xFu));
            vec4 w1 = vec4(
                float((packedW >> 16u) & 0xFu),
                float((packedW >> 20u) & 0xFu),
                float((packedW >> 24u) & 0xFu),
                float((packedW >> 28u) & 0xFu));

            w0 = (w0 - vec4(8.0)) * scale + vec4(offset);
            w1 = (w1 - vec4(8.0)) * scale + vec4(offset);

            const uint xBase = i1 << 1;
            vec4 x0 = vec4(X4[xBase + 0u]);
            vec4 x1 = vec4(X4[xBase + 1u]);
            sum += dot(w0, x0) + dot(w1, x1);
        }
    }

    // Tail K handling for K % 8 != 0.
    const uint tail = pc.K & 7u;
    if (tail > 0u && gl_SubgroupInvocationID == 0u) {
        const uint idx = K8;
        const uint packedW = Wp[baseW + idx];
        const uint kBase = idx << 3;
        vec4 x0 = vec4(X4[idx << 1]);
        vec4 x1 = vec4(0.0);
        if (tail > 4u) {
            x1 = vec4(X4[(idx << 1) + 1u]);
        }
        for (uint t = 0u; t < tail; ++t) {
            const uint k = kBase + t;
            const uint kb = k / safeBlockSize;
            const uint metaBase = (baseM + kb) * 2u;
            const float scale = float(Meta[metaBase + 0u]);
            const float offset = float(Meta[metaBase + 1u]);
            const float q = float((packedW >> (t * 4u)) & 0xFu) - 8.0;
            float x = 0.0;
            if (t < 4u) {
                x = x0[int(t)];
            } else {
                x = x1[int(t - 4u)];
            }
            sum += (q * scale + offset) * x;
        }
    }

    sum = subgroupAdd(sum);

    if (gl_SubgroupInvocationID == 0u) {
        float y = sum + float(Bias[row]);
        if (ACTIVATION == 1u) {
            y = max(y, 0.0);
        } else if (ACTIVATION == 2u) {
            y = clamp(y, 0.0, 6.0);
        }
        Y[row] = FLOAT(y);
    }
}