wehub-resource-sync
78 lines
2.5 KiB
Plaintext
78 lines
2.5 KiB
Plaintext
layout(set=0, binding=0) writeonly buffer destBuffer{
|
|
FLOAT4 data[];
|
|
}uOutput;
|
|
|
|
layout(set=0, binding=1) readonly buffer sourceBuffer0{
|
|
FLOAT4 data[];
|
|
} uInput;
|
|
|
|
layout(set=0, binding=2) readonly buffer sourceBuffer1{
|
|
FLOAT4 data[];
|
|
} uKernel;
|
|
|
|
layout(set=0, binding=3) readonly buffer sourceBuffer2{
|
|
FLOAT4 data[];
|
|
} uBias;
|
|
|
|
layout(set=0, binding=4) readonly uniform constBuffer {
|
|
ivec2 pad;
|
|
ivec2 kernelSize;
|
|
ivec2 stride;
|
|
ivec2 dilate;
|
|
ivec4 inputSize;
|
|
ivec4 outputSize;
|
|
int batch;
|
|
int group;
|
|
} uConstant;
|
|
|
|
layout (local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
|
|
|
|
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
|
|
|
|
void main()
|
|
{
|
|
int workIndex = int(gl_GlobalInvocationID.x);
|
|
if (workIndex < uConstant.outputSize.x * uConstant.outputSize.y * uConstant.outputSize.z * uConstant.outputSize.w)
|
|
{
|
|
ivec3 pos;
|
|
pos.x = workIndex % uConstant.outputSize.x;
|
|
int temp = workIndex / uConstant.outputSize.x;
|
|
pos.y = temp % uConstant.outputSize.y;
|
|
pos.z = temp / uConstant.outputSize.y;
|
|
ivec3 inputSize = uConstant.inputSize.xyz;
|
|
ivec3 outputSize = uConstant.outputSize.xyz;
|
|
|
|
ivec2 oxy = pos.xy + uConstant.pad;
|
|
int fz = pos.z / uConstant.outputSize.w;
|
|
int sb = pos.z % uConstant.outputSize.w;
|
|
|
|
FLOAT4 color = uBias.data[fz];
|
|
for (int fy=0; fy<uConstant.kernelSize.y; ++fy)
|
|
{
|
|
int sy = oxy.y - fy*uConstant.dilate.y;
|
|
int y = sy / uConstant.stride.y;
|
|
if (sy % uConstant.stride.y == 0 && y == clamp(y, 0, inputSize.y-1))
|
|
{
|
|
for (int fx=0; fx<uConstant.kernelSize.x; ++fx)
|
|
{
|
|
int sx = oxy.x - fx*uConstant.dilate.x;
|
|
int x = sx / uConstant.stride.x;
|
|
if (sx % uConstant.stride.x == 0 && x == clamp(x, 0, inputSize.x-1))
|
|
{
|
|
FLOAT4 inputColor = uInput.data[x + y * uConstant.inputSize.x + pos.z * uConstant.inputSize.x * uConstant.inputSize.y];
|
|
FLOAT4 kernelColor = uKernel.data[fx+fy*uConstant.kernelSize.x + fz * uConstant.kernelSize.x * uConstant.kernelSize.y];
|
|
color += inputColor*kernelColor;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
#ifdef RELU
|
|
color = max(color, FLOAT4(0));
|
|
#endif
|
|
#ifdef RELU6
|
|
color = clamp(color, FLOAT4(0), FLOAT4(6));
|
|
#endif
|
|
uOutput.data[pos.x + pos.y * uConstant.outputSize.x + pos.z * uConstant.outputSize.x * uConstant.outputSize.y] = color;
|
|
}
|
|
}
|