wehub-resource-sync
99 lines
3.6 KiB
Plaintext
99 lines
3.6 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 {
|
|
ivec2 pad;
|
|
ivec2 kernelSize;
|
|
ivec2 stride;
|
|
ivec2 dilate;
|
|
ivec4 inputSize; // w h icDiv4 n
|
|
ivec4 outputSize; // w h ocDiv4 n
|
|
ivec4 offset; // 0 0 oh
|
|
} 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()
|
|
{
|
|
// index calculation
|
|
int outputIndexW2 = int(gl_GlobalInvocationID.x);
|
|
int outputIndexH = int(gl_GlobalInvocationID.y);
|
|
|
|
int outputW2 = UP_DIV(uConstant.outputSize.x, 2);
|
|
|
|
if (outputIndexW2 >= outputW2 || outputIndexH >= uConstant.outputSize.y) {
|
|
return;
|
|
}
|
|
|
|
int outputIndexNC4 = int(gl_GlobalInvocationID.z);
|
|
int outputIndexC4 = outputIndexNC4 % uConstant.outputSize.z;
|
|
int outputIndexN = outputIndexNC4 / uConstant.outputSize.z;
|
|
|
|
FLOAT4 result0 = FLOAT4(texelFetch(uBias, ivec2(outputIndexC4, 0), 0));
|
|
FLOAT4 result1 = result0;
|
|
|
|
ivec2 inputIndexOffset = ivec2(outputIndexW2 * 2 - uConstant.pad.x, outputIndexH - uConstant.pad.y);
|
|
int inputIndexStartH = max(0, -inputIndexOffset.y);
|
|
int inputIndexEndH = min(uConstant.kernelSize.y, uConstant.inputSize.y - inputIndexOffset.y);
|
|
|
|
// accumulate result
|
|
for (int kernelIndexH = inputIndexStartH; kernelIndexH < inputIndexEndH; kernelIndexH++) {
|
|
int inputIndexH = inputIndexOffset.y + kernelIndexH;
|
|
int inputPosH = inputIndexH + outputIndexN * uConstant.inputSize.y;
|
|
|
|
FLOAT4 input0 = FLOAT4(0.0f);
|
|
|
|
int inputIndexW0 = inputIndexOffset.x;
|
|
FLOAT4 input1 = (inputIndexW0 >= 0 && inputIndexW0 < uConstant.inputSize.x) ? FLOAT4(texelFetch(uInput, ivec2(inputIndexW0 + outputIndexC4 * uConstant.inputSize.x, inputPosH), 0)) : FLOAT4(0);
|
|
|
|
for (int kernelIndexW = 0; kernelIndexW < uConstant.kernelSize.x; kernelIndexW++) {
|
|
// load input
|
|
input0 = input1;
|
|
|
|
inputIndexW0 = inputIndexOffset.x + kernelIndexW + 1;
|
|
input1 = (inputIndexW0 >= 0 && inputIndexW0 < uConstant.inputSize.x) ? FLOAT4(texelFetch(uInput, ivec2(inputIndexW0 + outputIndexC4 * uConstant.inputSize.x, inputPosH), 0)) : FLOAT4(0);
|
|
|
|
// load weight
|
|
FLOAT4 weight = FLOAT4(texelFetch(uKernel, ivec2(kernelIndexW + kernelIndexH * uConstant.kernelSize.x, outputIndexC4), 0));
|
|
result0 += input0 * weight;
|
|
result1 += input1 * weight;
|
|
}
|
|
}
|
|
|
|
#if defined(RELU_FP32) || defined(RELU_FP16)
|
|
result0 = FLOAT4(max(result0, FLOAT4(0)));
|
|
result1 = FLOAT4(max(result1, 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)));
|
|
#endif
|
|
|
|
// write output
|
|
int outputPosXBase = (outputIndexW2 * 2) + outputIndexC4 * uConstant.outputSize.x;
|
|
int outputPosY = outputIndexH + outputIndexN * uConstant.outputSize.y;
|
|
|
|
if (outputIndexW2 * 2 < uConstant.outputSize.x - 1) {
|
|
imageStore(uOutput, ivec2(outputPosXBase, outputPosY), result0);
|
|
imageStore(uOutput, ivec2(outputPosXBase + 1, outputPosY), result1);
|
|
} else {
|
|
imageStore(uOutput, ivec2(outputPosXBase, outputPosY), result0);
|
|
}
|
|
|
|
} |