wehub-resource-sync
121 lines
4.8 KiB
Plaintext
121 lines
4.8 KiB
Plaintext
#version 440 core
|
|
|
|
#if defined(FP16) || defined(RELU_FP16) || defined(RELU6_FP16)
|
|
#extension GL_AMD_gpu_shader_half_float: enable
|
|
#define FLOAT4 f16vec4
|
|
#else
|
|
#define FLOAT4 vec4
|
|
#endif
|
|
|
|
layout(set=0, binding=0) writeonly uniform image2D uOutput;
|
|
layout(set=0, binding=1) uniform sampler2D uInput;
|
|
|
|
layout(set=0, binding=2) uniform sampler2D uKernel;
|
|
|
|
layout(set=0, binding=3) uniform sampler2D uBias;
|
|
|
|
layout(set=0, binding=4) readonly uniform constBuffer {
|
|
ivec4 inputSize; // w h icDiv4 n
|
|
ivec4 outputSize; // w h ocDiv4 n
|
|
} uConstant;
|
|
|
|
#define UP_DIV(x, y) (((x)+(y)-1)/(y))
|
|
|
|
layout (local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
|
|
|
|
void main() {
|
|
// -------------
|
|
int outputIndexNHW4 = int(gl_GlobalInvocationID.x);
|
|
int outputIndexC4 = int(gl_GlobalInvocationID.y);
|
|
int outputW4 = UP_DIV(uConstant.outputSize.x, 4);
|
|
|
|
if (outputIndexNHW4 >= uConstant.outputSize.w * uConstant.outputSize.y * outputW4 || outputIndexC4 >= uConstant.outputSize.z) {
|
|
return;
|
|
}
|
|
|
|
int outputIndexW4 = outputIndexNHW4 % outputW4;
|
|
int outputIndexNH = outputIndexNHW4 / outputW4;
|
|
int outputIndexH = outputIndexNH % uConstant.outputSize.y;
|
|
int outputIndexN = outputIndexNH / uConstant.outputSize.y;
|
|
// -------------
|
|
|
|
FLOAT4 result0 = FLOAT4(texelFetch(uBias, ivec2(outputIndexC4, 0), 0));
|
|
FLOAT4 result1 = result0;
|
|
FLOAT4 result2 = result0;
|
|
FLOAT4 result3 = result0;
|
|
|
|
int outputIndexWBase = outputIndexW4 * 4;
|
|
|
|
for (int inputIndexC4 = 0; inputIndexC4 < uConstant.inputSize.z; inputIndexC4 ++) {
|
|
int inputPosWBase = outputIndexWBase + inputIndexC4 * uConstant.inputSize.x;
|
|
int inputPosH = outputIndexH + outputIndexN * uConstant.inputSize.y;
|
|
FLOAT4 inputValue0 = FLOAT4(texelFetch(uInput, ivec2(inputPosWBase, inputPosH), 0));
|
|
FLOAT4 inputValue1 = FLOAT4(texelFetch(uInput, ivec2(inputPosWBase + 1, inputPosH), 0));
|
|
FLOAT4 inputValue2 = FLOAT4(texelFetch(uInput, ivec2(inputPosWBase + 2, inputPosH), 0));
|
|
FLOAT4 inputValue3 = FLOAT4(texelFetch(uInput, ivec2(inputPosWBase + 3, inputPosH), 0));
|
|
|
|
int kernelIndexXbase = inputIndexC4 * 4;
|
|
FLOAT4 weight0 = FLOAT4(texelFetch(uKernel, ivec2(kernelIndexXbase + 0, outputIndexC4), 0));
|
|
FLOAT4 weight1 = FLOAT4(texelFetch(uKernel, ivec2(kernelIndexXbase + 1, outputIndexC4), 0));
|
|
FLOAT4 weight2 = FLOAT4(texelFetch(uKernel, ivec2(kernelIndexXbase + 2, outputIndexC4), 0));
|
|
FLOAT4 weight3 = FLOAT4(texelFetch(uKernel, ivec2(kernelIndexXbase + 3, outputIndexC4), 0));
|
|
|
|
result0 += inputValue0.x * weight0;
|
|
result0 += inputValue0.y * weight1;
|
|
result0 += inputValue0.z * weight2;
|
|
result0 += inputValue0.w * weight3;
|
|
|
|
result1 += inputValue1.x * weight0;
|
|
result1 += inputValue1.y * weight1;
|
|
result1 += inputValue1.z * weight2;
|
|
result1 += inputValue1.w * weight3;
|
|
|
|
result2 += inputValue2.x * weight0;
|
|
result2 += inputValue2.y * weight1;
|
|
result2 += inputValue2.z * weight2;
|
|
result2 += inputValue2.w * weight3;
|
|
|
|
result3 += inputValue3.x * weight0;
|
|
result3 += inputValue3.y * weight1;
|
|
result3 += inputValue3.z * weight2;
|
|
result3 += inputValue3.w * weight3;
|
|
}
|
|
|
|
#if defined(RELU_FP32) || defined(RELU_FP16)
|
|
result0 = FLOAT4(max(result0, FLOAT4(0)));
|
|
result1 = FLOAT4(max(result1, FLOAT4(0)));
|
|
result2 = FLOAT4(max(result2, FLOAT4(0)));
|
|
result3 = FLOAT4(max(result3, FLOAT4(0)));
|
|
#endif
|
|
#if defined(RELU6_FP32) || defined(RELU6_FP16)
|
|
result0 = FLOAT4(clamp(result0, FLOAT4(0), FLOAT4(6)));
|
|
result1 = FLOAT4(clamp(result1, FLOAT4(0), FLOAT4(6)));
|
|
result2 = FLOAT4(clamp(result2, FLOAT4(0), FLOAT4(6)));
|
|
result3 = FLOAT4(clamp(result3, FLOAT4(0), FLOAT4(6)));
|
|
#endif
|
|
|
|
int outputPosWBase = outputIndexWBase + outputIndexC4 * uConstant.outputSize.x;
|
|
int outputPosH = outputIndexH + outputIndexN * uConstant.outputSize.y;
|
|
int WLeft = uConstant.outputSize.x - outputIndexW4 * 4;
|
|
if (WLeft == 1) {
|
|
imageStore(uOutput, ivec2(outputPosWBase, outputPosH), result0);
|
|
return;
|
|
}
|
|
if (WLeft == 2) {
|
|
imageStore(uOutput, ivec2(outputPosWBase, outputPosH), result0);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 1, outputPosH), result1);
|
|
return;
|
|
}
|
|
if (WLeft == 3) {
|
|
imageStore(uOutput, ivec2(outputPosWBase, outputPosH), result0);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 1, outputPosH), result1);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 2, outputPosH), result2);
|
|
return;
|
|
}
|
|
imageStore(uOutput, ivec2(outputPosWBase, outputPosH), result0);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 1, outputPosH), result1);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 2, outputPosH), result2);
|
|
imageStore(uOutput, ivec2(outputPosWBase + 3, outputPosH), result3);
|
|
|
|
return;
|
|
} |