-
Notifications
You must be signed in to change notification settings - Fork 21.4k
/
mm.glsl
31 lines (27 loc) · 979 Bytes
/
mm.glsl
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
30
31
#version 450 core
#define PRECISION $precision
layout(std430) buffer;
layout(std430) uniform;
/* Qualifiers: layout - storage - precision - memory */
layout(set = 0, binding = 0, rgba16f) uniform PRECISION writeonly image3D uOutput;
layout(set = 0, binding = 1) uniform PRECISION sampler3D uM1;
layout(set = 0, binding = 2) uniform PRECISION sampler3D uM2;
layout(set = 0, binding = 3) uniform restrict Block {
ivec3 WHC;
int K;
} uBlock;
layout(local_size_x_id = 1, local_size_y_id = 2, local_size_z_id = 3) in;
void main() {
const ivec3 pos = ivec3(gl_GlobalInvocationID);
if (all(lessThan(pos, uBlock.WHC))) {
const int K = uBlock.K;
vec4 mmv = vec4(0);
int ki = 0;
for (; ki < K; ++ki) {
vec4 m1ki = texelFetch(uM1, ivec3(ki, pos.y, pos.z), 0);
vec4 m2ki = texelFetch(uM2, ivec3(pos.x, ki, pos.z), 0);
mmv += m1ki * m2ki;
}
imageStore(uOutput, pos, mmv);
}
}