#version 450

#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_control_flow_attributes : enable

layout (push_constant) uniform parameter
{
    uint KX;
    uint KY;
    float scale;
    float max_bias;
    float m0;
    float m1;
    uint n_head_log2;
    uint nrows_x;
} p;

#include "types.comp"

layout(constant_id = 0) const uint BLOCK_SIZE = 32;
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;

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

shared FLOAT_TYPE vals[BLOCK_SIZE];

// num_iters is the number of BLOCK_SIZE loop iterations we need to iterate
// over all the columns. The main function tries to pass a constant here,
// as if it were a template function, to allow unrolling.
void soft_max(uint num_iters) {
    const uint tid = gl_LocalInvocationID.x;
    const uint rowx = gl_WorkGroupID.z * 262144 + gl_WorkGroupID.y * 512 + gl_WorkGroupID.x;
    const uint rowy = (p.KY > 0) ? (rowx % p.KY) : 0;

    if (rowx >= p.nrows_x) {
        return;
    }

    float slope = 1.0f;

    // ALiBi
    if (p.max_bias > 0.0f) {
        const uint h = rowx/p.KY; // head index

        const float base = h < p.n_head_log2 ? p.m0 : p.m1;
        const uint   exp  = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1;

        slope = pow(base, exp);
    }

    // Find max
    FLOAT_TYPE max_val = uintBitsToFloat(0xFF800000);

    // Cache values while we compute the max, so we don't need to read them
    // again when we're ready to compute exp(x-max).
    const uint DATA_CACHE_SIZE = 16;
    FLOAT_TYPE data_cache[DATA_CACHE_SIZE];

    [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
        const uint col = col0 + tid;

        FLOAT_TYPE a = FLOAT_TYPE(0);
        if (col < p.KX) {
            a = data_a[rowx * p.KX + col];
        }

        FLOAT_TYPE b = FLOAT_TYPE(0);
        if (p.KY > 0 && col < p.KX) {
            b = data_b[rowy * p.KX + col];
        }

        FLOAT_TYPE v = a * p.scale + slope * b;

        if (col < p.KX) {
            max_val = max(max_val, v);
        }

        if (idx < DATA_CACHE_SIZE) {
            data_cache[idx] = v;
        }
    }

    // reduce across the workgroup
    vals[tid] = max_val;
    barrier();
    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            vals[tid] = max(vals[tid], vals[tid + s]);
        }
        barrier();
    }

    max_val = vals[0];
    barrier();

    FLOAT_TYPE sum = FLOAT_TYPE(0.0f);

    // Compute sum{exp(x - max)}
    [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
        const uint col = col0 + tid;

        if (col >= p.KX) {
            break;
        }

        // compute exp(a*scale+b*slope), add it to sum, and cache the new value
        // in data_cache if possible.
        const uint i = rowx * p.KX + col;
        FLOAT_TYPE val;
        if (idx < DATA_CACHE_SIZE) {
            val = exp(data_cache[idx] - max_val);
        } else {
            val = exp(FLOAT_TYPE(data_a[i]) * p.scale + (p.KY > 0 ? slope * FLOAT_TYPE(data_b[rowy * p.KX + col]) : FLOAT_TYPE(0.0f)) - max_val);
        }
        sum += val;
        if (idx < DATA_CACHE_SIZE) {
            data_cache[idx] = val;
        } else {
            data_d[i] = D_TYPE(val);
        }
    }

    // reduce across the workgroup
    vals[tid] = sum;
    barrier();
    [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (tid < s) {
            vals[tid] += vals[tid + s];
        }
        barrier();
    }
    sum = vals[0];

    FLOAT_TYPE rcpdivisor = 1.0/sum;

    [[unroll]] for (uint col0 = 0, idx = 0; idx < num_iters; col0 += BLOCK_SIZE, ++idx) {
        const uint col = col0 + tid;

        if (col >= p.KX) {
            continue;
        }

        if (idx < DATA_CACHE_SIZE) {
            data_d[rowx*p.KX + col] = D_TYPE(data_cache[idx] * rcpdivisor);
        } else {
            data_d[rowx*p.KX + col] *= D_TYPE(rcpdivisor);
        }
    }
}

void main() {
    // instantiate the soft_max function for several different
    // dimensions, to allow loop unrolling
    uint num_blocks = (p.KX + BLOCK_SIZE - 1) / BLOCK_SIZE;
    if (num_blocks > 32) {
        soft_max(num_blocks);
    } else if (num_blocks > 16) {
        soft_max(32);
    } else if (num_blocks > 8) {
        soft_max(16);
    } else if (num_blocks > 4) {
        soft_max(8);
    } else if (num_blocks == 4) {
        soft_max(4);
    } else if (num_blocks == 3) {
        soft_max(3);
    } else if (num_blocks == 2) {
        soft_max(2);
    } else if (num_blocks == 1) {
        soft_max(1);
    }
}