llama.cpp/ggml/src/ggml-vulkan/vulkan-shaders/get_rows_back.comp

#version 450

#include "types.glsl"
#include "generic_binary_head.glsl"

layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;

void main() {
    const uint col = gl_GlobalInvocationID.x;

    if (col >= p.ne20) {
        return;
    }

    for (uint row = gl_GlobalInvocationID.y; row < p.ne21; row += gl_WorkGroupSize.y * gl_NumWorkGroups.y) {
        float sum = 0.0f;
        for (uint i = 0; i < p.ne10; ++i) {
            if (data_b[get_boffset() + i*p.nb10] == int(row)) {
                sum += data_a[get_aoffset() + i*p.nb01 + col*p.nb00];
            }
        }

        data_d[get_doffset() + row*p.nb21 + col*p.nb20] = sum;
    }
}