/**
 * Copyright (c) 2023 Nomic, Inc. All rights reserved.
 *
 * This software is licensed under the terms of the Software for Open Models License (SOM),
 * version 1.0, as detailed in the LICENSE_SOM.txt file. A copy of this license should accompany
 * this software. Except as expressly granted in the SOM license, all rights are reserved by Nomic, Inc.
 */

#version 450

#include "common.comp"

#define nth 32

layout(local_size_x = nth) 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;

shared float buf[nth];

void main() {
    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
    buf[gl_LocalInvocationID.x] = uintBitsToFloat(0xFF800000);
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += nth) {
        buf[gl_LocalInvocationID.x] = max(buf[gl_LocalInvocationID.x], in_[psrc0 + i00]);
    }

    // reduce
    barrier();
    memoryBarrierShared();
    [[unroll]] for (uint i = nth/2; i > 0; i /= 2) {
        if (gl_LocalInvocationID.x < i) {
            buf[gl_LocalInvocationID.x] = max(buf[gl_LocalInvocationID.x], buf[gl_LocalInvocationID.x + i]);
        }
        barrier();
        memoryBarrierShared();
    }

    // broadcast
    const float max_ = buf[0];

    // parallel sum
    buf[gl_LocalInvocationID.x] = 0.0;
    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += nth) {
        buf[gl_LocalInvocationID.x] += exp(in_[psrc0 + i00] - max_);
    }

    // reduce
    barrier();
    memoryBarrierShared();
    [[unroll]] for (uint i = nth/2; i > 0; i /= 2) {
        if (gl_LocalInvocationID.x < i) {
            buf[gl_LocalInvocationID.x] += buf[gl_LocalInvocationID.x + i];
        }
        barrier();
        memoryBarrierShared();
    }

    // broadcast
    const float sum = buf[0];

    for (uint i00 = gl_LocalInvocationID.x; i00 < pcs.ne00; i00 += nth) {
        out_[pdst + i00] = exp(in_[psrc0 + i00] - max_) / sum;
    }
}