Files
2026-07-13 13:33:03 +08:00

120 lines
2.7 KiB
Plaintext

#version 440
layout(set=0, binding=0) writeonly uniform image2D uOutput;
layout(set=0, binding=1) uniform sampler2D uInput;
layout(set=0, binding=2) uniform constBuffer{
ivec4 size; // x: limit, y: channelC4*b, z:height, w:width
} uConstant;
#define FLOAT4 vec4
//for dynamic change threads counts from outside
// from vkCreateComputePipelines->VkComputePipelineCreateInfo->VkPipelineShaderStageCreateInfo->VkSpecializationInfo->VkSpecializationMapEntry
// layout(local_size_x_id = 1, local_size_y_id = 2, local_size_z_id = 3) in;
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
void main()
{
ivec3 posTmp = ivec3(gl_GlobalInvocationID);
if (posTmp.x < uConstant.size.x) {
ivec2 pos;
pos.x = posTmp.x % uConstant.size.w;
pos.y = posTmp.x / uConstant.size.w;
vec4 value = texelFetch(uInput, pos, 0);
#ifdef NEG
value = -value;
#endif
#ifdef EXP
value = exp(value);
#endif
#ifdef SIGN
value = sign(value);
#endif
#ifdef SQRT
value = sqrt(value);
#endif
#ifdef RSQRT
value = inversesqrt(value);
#endif
#ifdef ABS
value = abs(value);
#endif
#ifdef TANH
value = tanh(value);
#endif
#ifdef SQUARE
value = value * value;
#endif
#ifdef LOG
value = log(max(value,vec4(0.0000001)));
#endif
#ifdef SIGMOID
value = 1.f / (1.f + exp(-value));
#endif
#ifdef TAN
value = tan(value);
#endif
#ifdef COS
value = cos(value);
#endif
#ifdef SIN
value = sin(value);
#endif
#ifdef CEIL
value = ceil(value);
#endif
#ifdef FLOOR
value = floor(value);
#endif
#ifdef EXPM1
value = exp(value) - FLOAT4(1);
#endif
#ifdef RECIPROCAL
value = FLOAT4(1) / value;
#endif
#ifdef SINH
value = sinh(value);
#endif
#ifdef ASINH
value = asinh(value);
#endif
#ifdef ASIN
value = asin(value);
#endif
#ifdef COSH
value = cosh(value);
#endif
#ifdef ACOSH
value = acosh(value);
#endif
#ifdef ACOS
value = acos(value);
#endif
#ifdef ATAN
value = atan(value);
#endif
#ifdef ATANH
value = atanh(value);
#endif
#ifdef LOG1P
value = log(FLOAT4(1) + value);
#endif
#ifdef ROUND
value = round(value);
#endif
#ifdef HARDSWISH
const vec4 leftMask = vec4(greaterThan(value, vec4(-3.0f)));
const vec4 rightMask = vec4(lessThan(value, vec4(3.0f)));
value = leftMask*value*(rightMask*((value+3.0f)/6.0f) + 1.0f - rightMask);
#endif
#ifdef GELU
vec4 temp = vec4(0.044715) * value * value * value;
temp = vec4(0.79788458) * (temp + value);
value = (vec4(1.0) + tanh(temp)) * value * vec4(0.5);
#endif
#ifdef SILU
value = value / (1.f + exp(-value));
#endif
imageStore(uOutput, pos, value);
}
}