1#version 450
 2
 3#extension GL_EXT_control_flow_attributes : enable
 4
 5layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
 6
 7layout (binding = 0) readonly buffer A {float data_a[];};
 8layout (binding = 0) readonly buffer A4 {vec4 data_a4[];};
 9layout (binding = 1) writeonly buffer D {float data_d[];};
10layout (binding = 1) writeonly buffer D4 {vec4 data_d4[];};
11
12layout (push_constant) uniform parameter {
13    uint ne;
14    uint k_num;
15} p;
16
17void main() {
18    // Each invocation handles four consecutive components
19    const uint idx = gl_GlobalInvocationID.x * 4;
20
21    if (idx >= p.ne) {
22        return;
23    }
24
25    // Check if all four components are in bounds and aligned,
26    // then use vector loads
27    if (idx + 3 < p.ne && (p.ne % 4) == 0) {
28        vec4 result = vec4(0.0f);
29
30        [[unroll]] for (uint i = 0; i < p.k_num; i++) {
31            result += data_a4[(i * p.ne + idx) / 4];
32        }
33
34        data_d4[idx / 4] = result;
35    } else {
36        [[unroll]] for (uint j = 0; j < 4; ++j) {
37            if (idx + j < p.ne) {
38                float result = 0.0f;
39
40                [[unroll]] for (uint i = 0; i < p.k_num; i++) {
41                    result += data_a[i * p.ne + idx + j];
42                }
43
44                data_d[idx + j] = result;
45            }
46        }
47    }
48}