// TODO: implement multi-simd softmax (llama.cpp commit e16b9fa4)

#version 450

#include "common.comp"

layout(local_size_x_id = 0) in;

layout(binding = 0) buffer restrict readonly tensorIn { float in_[]; };
layout(binding = 1) buffer restrict writeonly tensorOut { float out_[]; };

layout(push_constant) uniform PushConstants {
    uint inOff;
    uint outOff;
    int ne00;
    int ne01;
    int ne02;
} pcs;

void main() {
    if (gl_SubgroupInvocationID > 31)
        return;

    const uint i03 = gl_WorkGroupID.z;
    const uint i02 = gl_WorkGroupID.y;
    const uint i01 = gl_WorkGroupID.x;

    const uint extra_off = i03*pcs.ne02*pcs.ne01*pcs.ne00 + i02*pcs.ne01*pcs.ne00 + i01*pcs.ne00;
    const uint psrc0 = extra_off + pcs.inOff; // Based from in_
    const uint pdst = extra_off + pcs.outOff; // Based from out_

    // parallel max
    float localMax = uintBitsToFloat(0xFF800000);
    for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) {
        localMax = max(localMax, in_[psrc0 + i00]);
    }
    float max_ = subgroupMax(localMax);

    // parallel sum
    float localSum = 0.0f;
    for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) {
        const float exp_psrc0 = exp(in_[psrc0 + i00] - max_);
        localSum += exp_psrc0;
        out_[pdst + i00] = exp_psrc0;
    }

    const float sum = subgroupAdd(localSum);
    for (uint i00 = gl_SubgroupInvocationID.x; i00 < pcs.ne00; i00 += 32) {
        out_[pdst + i00] /= sum;
    }
}