wehub-resource-sync
73 lines
2.1 KiB
Plaintext
73 lines
2.1 KiB
Plaintext
layout(set=0, binding=0) buffer destbuffer{
|
|
int data[];
|
|
}uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer sourceBuffer{
|
|
FLOAT data[];
|
|
}uInput;
|
|
|
|
layout(set = 0, binding = 2) uniform constBuffer {
|
|
ivec4 size;//inside, axis, outside, reduceAxis
|
|
}uConst;
|
|
|
|
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
|
|
shared FLOAT local_buffer[256];
|
|
shared int local_index[256];
|
|
|
|
void main()
|
|
{
|
|
int index = int(gl_GlobalInvocationID.x) / uConst.size.w;
|
|
int lidIndex = int(gl_LocalInvocationID.x);
|
|
int lidx = lidIndex / uConst.size.w;
|
|
int lid = lidIndex % uConst.size.w;
|
|
|
|
// y: index in outside, x: index in inside
|
|
int x = index % uConst.size.x;
|
|
int y = index / uConst.size.x;
|
|
int W = uConst.size.x;
|
|
int H = uConst.size.y;
|
|
int C = uConst.size.z;
|
|
|
|
if(y < uConst.size.z && lid < uConst.size.y) {
|
|
int offset = y * H * W + x;
|
|
FLOAT maxValue = uInput.data[offset + lid * W];
|
|
int maxIndex = lid;
|
|
for(int i = lid+uConst.size.w; i < uConst.size.y; i+=uConst.size.w) {
|
|
FLOAT value = uInput.data[offset + i * W];
|
|
#ifndef ARGMIN
|
|
if (value > maxValue) {
|
|
maxValue = value;
|
|
maxIndex = i;
|
|
}
|
|
#else
|
|
if (value < maxValue) {
|
|
maxValue = value;
|
|
maxIndex = i;
|
|
}
|
|
#endif
|
|
}
|
|
local_buffer[lid + lidx * uConst.size.w] = maxValue;
|
|
local_index[lid + lidx * uConst.size.w] = maxIndex;
|
|
}
|
|
barrier();
|
|
if(y < uConst.size.z && lid == 0) {
|
|
FLOAT maxValue = local_buffer[lidx * uConst.size.w];
|
|
int maxIndex = local_index[lidx * uConst.size.w];
|
|
for (int t=1; t<uConst.size.w && t < uConst.size.y; ++t) {
|
|
FLOAT next = local_buffer[t + lidx * uConst.size.w];
|
|
#ifndef ARGMIN
|
|
if (next > maxValue) {
|
|
maxValue = next;
|
|
maxIndex = local_index[t + lidx * uConst.size.w];
|
|
}
|
|
#else
|
|
if (next < maxValue) {
|
|
maxValue = next;
|
|
maxIndex = local_index[t + lidx * uConst.size.w];
|
|
}
|
|
#endif
|
|
}
|
|
uOutput.data[index] = maxIndex;
|
|
}
|
|
}
|