vulkan-shaders/mul_mat_split_k_reduce.comp


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

#version 450

#extension GL_EXT_control_flow_attributes : enable

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

layout (binding = 0) readonly buffer A {float data_a[];};
layout (binding = 1) writeonly buffer D {float data_d[];};

layout (push_constant) uniform parameter {
    uint ne;
    uint k_num;
} p;

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

    if (idx >= p.ne) {
        return;
    }

    float result = 0.0f;

    [[unroll]] for (uint i = 0; i < p.k_num; i++) {
        result += data_a[i * p.ne + idx];
    }

    data_d[idx] = result;
}