#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;
layout(constant_id = 1) const uint N_Z = 1;

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

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

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

void main() {
    const uint batchIndex = gl_WorkGroupID.y;
    const uint zIndex = gl_WorkGroupID.z;
    const uint b = zIndex * N_BATCHES + batchIndex;

    const uint d = gl_WorkGroupID.x;

    const BatchInfo info = infos[b];
    const uint xiOffset = info.inputOffset + d * Q80_BLOCK_SIZE;
    const uint yiOffset = info.outputOffset + d;

    float amax = 0.0;
    [[unroll]] for (uint j = 0; j < Q80_BLOCK_SIZE; ++j) {
        amax = max(amax, abs(x[xiOffset + j]));
    }

    const float dd = amax / 127.0f;
    const float id = dd != 0.0f ? 1.0f / dd : 0.0f;

    y[yiOffset].d = float16_t(dd);

    [[unroll]] for (uint j = 0; j < Q80_BLOCK_SIZE; ++j) {
        const float v = x[xiOffset + j];
        y[yiOffset].qs[j] = int8_t(clamp(round(v * id), -127.0f, 127.0f));
    }
}
