/**
 * 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"

layout(local_size_x = 8, local_size_y = 8) in;

layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; };
layout (binding = 1) readonly buffer tensorInB { float inB[]; };
layout (binding = 2) writeonly buffer tensorOut { float out_[]; };

layout (push_constant) uniform parameter {
    uint inAOff;
    uint inBOff;
    uint outOff;
    int ne00;
    int ne10;
    int ne0;
} pcs;

shared float sum[gl_WorkGroupSize.x*gl_WorkGroupSize.y];

#define UNALIGNED_INPUT inA

block_q4_1 get_unaligned_block_q4_1(uint index) {
    block_q4_1 fres;
    fres.d = u8BufToFloat16(UNALIGNED_INPUT, index);
    fres.m = u8BufToFloat16(UNALIGNED_INPUT, index+2);
    [[unroll]] for (uint it = 0; it != QK4_1 / 2; it++) {
        fres.qs[it] = UNALIGNED_INPUT[index+4+it];
    }
    return fres;
}

void main() {
    const uint nb = uint(pcs.ne00/QK4_1);

    const uint r0 = gl_WorkGroupID.x;
    const uint r1 = gl_WorkGroupID.y;

    const uint x = r0*nb; // Based from inA without base offset
    const uint y = r1*uint(pcs.ne10) + pcs.inBOff; // Based from inB

    const uint nth = gl_WorkGroupSize.x*gl_WorkGroupSize.y;
    const uint ith = gl_WorkGroupSize.y*gl_LocalInvocationID.x + gl_LocalInvocationID.y;

    const uint ix = gl_LocalInvocationID.y/4;           // 0 or 1
    const uint iy = gl_LocalInvocationID.y - 4*ix;      // 0...3

    const uint first = 4 * iy;

    float sumf = 0.0;

    for (uint i = 2*gl_LocalInvocationID.x + ix; i < nb; i += 2*gl_WorkGroupSize.x) {
        //TODO: Removing the use of pointers has been quite hairy here. If something goes wrong here, this is most likely it:

        const block_q4_1 block = get_unaligned_block_q4_1((x+i)*sizeof_block_q4_1+pcs.inAOff);

        const float d = float(block.d);
        const float m = float(block.m);

        const uint xl = first; // Based from bl->qs
        const uint yl = y + i * QK4_1 + first; // Based from inB

        vec2 acc = vec2(0.0, 0.0);

        for (int j = 0; j < 4; ++j) {
            acc.x += inB[yl+j] * (d * (block.qs[xl+j] & 0xF) + m);
            acc.y += inB[yl+j+16] * (d * (block.qs[xl+j] >> 4) + m);
        }

        sumf += d * (acc.x - acc.y);
    }

    sum[ith] = sumf;

    //
    // Accumulate the sum from all threads in the threadgroup
    //
    barrier();
    memoryBarrierShared();
    if (ith%4 == 0) {
        sum[ith] += sum[ith+1] + sum[ith+2] + sum[ith+3];
    }
    barrier();
    memoryBarrierShared();
    if (ith%16 == 0) {
        sum[ith] += sum[ith+4] + sum[ith+8] + sum[ith+12];
    }
    barrier();
    memoryBarrierShared();
    if (ith == 0) {
        for (uint i = 16; i < nth; i += 16) sum[0] += sum[i];
        out_[r1*uint(pcs.ne0) + r0 + pcs.outOff] = sum[0];
    }
}