#version 450

#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_8bit_storage : enable
#extension GL_EXT_shader_16bit_storage : enable
#extension GL_EXT_shader_explicit_arithmetic_types : enable

#define Q80_BLOCK_SIZE 32

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

layout(constant_id = 0) const uint N_BATCHES = 32;

struct BatchInfo {
    uint inputOffset; // number of Q80 blocks
    uint inputSizeX; // number of Q80 blocks
    uint outputOffset;
    uint outputSizeX;
};

struct BlockQ80 {
    float16_t d;
    int8_t qs[Q80_BLOCK_SIZE];
};

layout(binding = 0) readonly buffer inputBuffer { BlockQ80 x[]; };
layout(binding = 1) buffer outputBuffer { float y[]; };
layout(binding = 2) readonly uniform batchInfosBuffer { BatchInfo infos[N_BATCHES]; };

void main() {
    const uint batchIndex = gl_WorkGroupID.y;
    const uint i = gl_WorkGroupID.x;

    const BatchInfo info = infos[batchIndex];
    const uint nBlocks = info.outputSizeX / Q80_BLOCK_SIZE; // 128
    const uint nSlices = info.inputSizeX / nBlocks;
    float16_t sums[Q80_BLOCK_SIZE];

    const uint xiOffset = info.inputOffset + i;
    const uint yiOffset = info.outputOffset + i * Q80_BLOCK_SIZE;

    [[unroll]] for (uint k = 0; k < Q80_BLOCK_SIZE; k++) {
        sums[k] = float16_t(0.0f);
    }
    for (uint n = 0; n < nSlices; n++) {
        const BlockQ80 b = x[xiOffset + n * nBlocks];
        const float16_t d = b.d;

        [[unroll]] for (uint k = 0; k < Q80_BLOCK_SIZE; k++) {
            sums[k] += float16_t(b.qs[k]) * d;
        }
    }

    [[unroll]] for (uint k = 0; k < Q80_BLOCK_SIZE; k++) {
        y[yiOffset + k] += float(sums[k]);
    }
}
