wehub-resource-sync
51 lines
1.8 KiB
Plaintext
51 lines
1.8 KiB
Plaintext
#version 440 core
|
|
|
|
layout(set=0, binding=0) uniform sampler2D uInput;
|
|
layout(set=0, binding=1) writeonly uniform image2D uOutput;
|
|
layout(set=0, binding=2) uniform sampler2D uBias;
|
|
|
|
layout(set=0, binding=3) readonly uniform constBuffer {
|
|
ivec2 pad;
|
|
ivec2 kernelSize;
|
|
ivec2 stride;
|
|
ivec2 dilate;
|
|
ivec4 inputSize;
|
|
ivec4 outputSize;
|
|
ivec4 offset;//batchOffset, hOffset, outputHeight, other
|
|
} uConstant;
|
|
|
|
layout (local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
|
|
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
|
|
|
|
void main()
|
|
{
|
|
ivec3 posTmp = ivec3(gl_GlobalInvocationID);
|
|
ivec3 pos;
|
|
pos.x = posTmp.x % uConstant.outputSize.x;
|
|
int subIndex = posTmp.x / uConstant.outputSize.x;
|
|
pos.y = subIndex % uConstant.outputSize.y;
|
|
pos.z = subIndex / uConstant.outputSize.y;
|
|
ivec3 outputSize = uConstant.outputSize.xyz;
|
|
int oz = pos.z % uConstant.outputSize.z;
|
|
int ob = pos.z / uConstant.outputSize.z;
|
|
int batchOffset = uConstant.offset.x * uConstant.outputSize.z;
|
|
if (pos.x < outputSize.x && pos.y < outputSize.y && ob < uConstant.outputSize.w)
|
|
{
|
|
int sourceXIndex = pos.x + pos.y*uConstant.outputSize.x + ob*uConstant.outputSize.x*uConstant.outputSize.y;
|
|
int sourceY = sourceXIndex / 4;
|
|
int sourceX = 4*oz + sourceXIndex % 4;
|
|
|
|
vec4 color = texelFetch(uInput, ivec2(sourceX, sourceY), 0) + texelFetch(uBias, ivec2(oz, 0), 0);
|
|
#ifdef RELU
|
|
color = max(color, vec4(0));
|
|
#endif
|
|
#ifdef RELU6
|
|
color = clamp(color, vec4(0), vec4(6));
|
|
#endif
|
|
int dzb_offset = pos.z + batchOffset;
|
|
int dz_offset = dzb_offset % uConstant.outputSize.z;
|
|
int db_offset = dzb_offset / uConstant.outputSize.z;
|
|
imageStore(uOutput, ivec2(pos.x + dz_offset * uConstant.outputSize.x, pos.y + uConstant.offset.z * db_offset + uConstant.offset.y), color);
|
|
}
|
|
}
|