wehub-resource-sync
96 lines
3.4 KiB
Plaintext
96 lines
3.4 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 oz = pos.z / uConstant.outputSize.w;
|
|
int sb = pos.z % uConstant.outputSize.w;
|
|
|
|
FLOAT4 color = uBias.data[oz];
|
|
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))
|
|
{
|
|
int sBegin = x
|
|
+ y * uConstant.inputSize.x
|
|
+ sb * uConstant.inputSize.x * uConstant.inputSize.y;
|
|
int kBegin = fx
|
|
+ fy * uConstant.kernelSize.x
|
|
+ oz * uConstant.kernelSize.x * uConstant.kernelSize.y;
|
|
for (int fz=0; fz<uConstant.inputSize.z; ++fz)
|
|
{
|
|
FLOAT4 S = uInput.data[sBegin
|
|
+ fz * uConstant.inputSize.x * uConstant.inputSize.y * uConstant.inputSize.w
|
|
];
|
|
int kOffst = 4 * (kBegin * uConstant.inputSize.z + fz);
|
|
FLOAT4 K0 = uKernel.data[kOffst + 0];
|
|
FLOAT4 K1 = uKernel.data[kOffst + 1];
|
|
FLOAT4 K2 = uKernel.data[kOffst + 2];
|
|
FLOAT4 K3 = uKernel.data[kOffst + 3];
|
|
color += K0 * S.x;
|
|
color += K1 * S.y;
|
|
color += K2 * S.z;
|
|
color += K3 * S.w;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
#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;
|
|
}
|
|
}
|