-
Notifications
You must be signed in to change notification settings - Fork 3
/
Reduce.comp
52 lines (41 loc) · 951 Bytes
/
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
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
#version 450
#extension GL_KHR_shader_subgroup_arithmetic : enable
layout(std430, binding = 0) buffer Input
{
float inputs[];
};
layout(std430, binding = 1) buffer Output
{
float outputs[];
};
layout (local_size_x_id = 1) in;
layout (constant_id = 2) const int sumSubGroupSize = 64;
layout(push_constant) uniform PushConsts
{
int n;
} consts;
shared float sdata[sumSubGroupSize];
void main()
{
float sum = 0.0;
if (gl_GlobalInvocationID.x < consts.n)
{
sum = inputs[gl_GlobalInvocationID.x];
}
sum = subgroupAdd(sum);
if (gl_SubgroupInvocationID == 0)
{
sdata[gl_SubgroupID] = sum;
}
memoryBarrierShared();
barrier();
if (gl_SubgroupID == 0)
{
sum = gl_SubgroupInvocationID < gl_NumSubgroups ? sdata[gl_SubgroupInvocationID] : 0;
sum = subgroupAdd(sum);
}
if (gl_LocalInvocationID.x == 0)
{
outputs[gl_WorkGroupID.x] = sum;
}
}