Files
2022-12-29 21:56:10 +08:00

268 lines
7.3 KiB
Plaintext

#version 460
#define LOCAL_SIZE_XY 16
layout (local_size_x = LOCAL_SIZE_XY, local_size_y = LOCAL_SIZE_XY, local_size_z = 1) in;
// out
layout (set = 0, binding = 0,rgba16f) uniform image2D hdrImage;
// in
layout (set = 0, binding = 1,rgba16f) uniform image2D historyImage;
layout (set = 0, binding = 2,rgba16f) uniform image2D inTAAImage;
struct TAASharpenPushConstant
{
uint sharpenMethod;
float sharpness;
};
layout(push_constant) uniform block
{
TAASharpenPushConstant pushConstant;
};
// same with cpp.
#define SHARPEN_OFF 0
// Bloom flicking.
#define SHARPEN_RESPONSIVE 1
// Bloom stable.
#define SHARPEN_CAS 2
const int kBorderSize = 1;
const int kGroupSize = LOCAL_SIZE_XY;
const int kLdsLength = kGroupSize + kBorderSize * 2;
const int kLdsArea = kLdsLength * kLdsLength;
shared vec4 sharedColor[kLdsLength][kLdsLength];
vec4 ldsLoadColor(ivec2 GId)
{
GId += ivec2(kBorderSize);
return sharedColor[GId.x][GId.y];
}
void ldsStoreColor(ivec2 GId, vec4 color)
{
sharedColor[GId.x][GId.y] = color;
}
void storeColor(ivec2 GId, ivec2 TId, ivec2 size)
{
TId = clamp(TId, ivec2(0,0), size - ivec2(1,1));
// store color.
ldsStoreColor(GId, imageLoad(inTAAImage, TId));
}
void prepareLds(ivec2 topLeft, ivec2 workSize, int groupIndex)
{
// 4 sample per pixel.
if(groupIndex < (kLdsArea >> 2)) // 1 / 4 area pixel work.
{
// sample [0, 0]
ivec2 id0 = ivec2(
groupIndex % kLdsLength,
groupIndex / kLdsLength);
storeColor(id0, topLeft + id0, workSize);
// sample [0.25, 0.25]
ivec2 id1 = ivec2(
(groupIndex + (kLdsArea >> 2)) % kLdsLength,
(groupIndex + (kLdsArea >> 2)) / kLdsLength);
storeColor(id1, topLeft + id1, workSize);
// sample [0.5, 0.5]
ivec2 id2 = ivec2(
(groupIndex + (kLdsArea >> 1)) % kLdsLength,
(groupIndex + (kLdsArea >> 1)) / kLdsLength);
storeColor(id2, topLeft + id2, workSize);
// sample [0.75, 0.75]
ivec2 id3 = ivec2(
(groupIndex + kLdsArea * 3 / 4) % kLdsLength,
(groupIndex + kLdsArea * 3 / 4) / kLdsLength);
storeColor(id3, topLeft + id3, workSize);
}
}
float min3x(float a, float b, float c)
{
return min(min(a, b), c);
}
float max3x(float a, float b, float c)
{
return max(max(a, b), c);
}
vec3 reinhardInverse(in vec3 sdr)
{
return sdr / max(1.0f - sdr, 1e-5f);
}
vec3 RGBToYCoCg(in vec3 rgb)
{
return vec3(
0.25f * rgb.r + 0.5f * rgb.g + 0.25f * rgb.b,
0.5f * rgb.r - 0.5f * rgb.b,
-0.25f * rgb.r + 0.5f * rgb.g - 0.25f * rgb.b
);
}
vec3 YCoCgToRGB(in vec3 yCoCg)
{
return vec3(
yCoCg.x + yCoCg.y - yCoCg.z,
yCoCg.x + yCoCg.z,
yCoCg.x - yCoCg.y - yCoCg.z
);
}
vec3 ApplySharpening(ivec2 groupPos)
{
const vec3 top = ldsLoadColor(groupPos + ivec2( 0, 1)).xyz;
const vec3 left = ldsLoadColor(groupPos + ivec2( 1, 0)).xyz;
const vec3 center = ldsLoadColor(groupPos + ivec2( 0, 0)).xyz;
const vec3 right = ldsLoadColor(groupPos + ivec2(-1, 0)).xyz;
const vec3 bottom = ldsLoadColor(groupPos + ivec2( 0, -1)).xyz;
vec3 result = RGBToYCoCg(center);
float unsharpenMask = 4.0f * result.x;
unsharpenMask -= RGBToYCoCg(top).x;
unsharpenMask -= RGBToYCoCg(bottom).x;
unsharpenMask -= RGBToYCoCg(left).x;
unsharpenMask -= RGBToYCoCg(right).x;
result.x = min(result.x + 0.25f * unsharpenMask, 1.1f * result.x);
return YCoCgToRGB(result);
}
// AMD CAS Filter for sharpen.
vec3 AMDCASFilter(float sharpness, ivec2 groupPos)
{
ivec2 pixelPos = ivec2(gl_GlobalInvocationID.xy);
// a b c
// d e f
// g h i
vec3 a = ldsLoadColor(groupPos + ivec2(-1, -1)).xyz;
vec3 b = ldsLoadColor(groupPos + ivec2( 0, -1)).xyz;
vec3 c = ldsLoadColor(groupPos + ivec2( 1, -1)).xyz;
vec3 d = ldsLoadColor(groupPos + ivec2(-1, 0)).xyz;
vec3 e = ldsLoadColor(groupPos + ivec2( 0, 0)).xyz;
vec3 f = ldsLoadColor(groupPos + ivec2( 1, 0)).xyz;
vec3 g = ldsLoadColor(groupPos + ivec2(-1, 1)).xyz;
vec3 h = ldsLoadColor(groupPos + ivec2( 0, 1)).xyz;
vec3 i = ldsLoadColor(groupPos + ivec2( 1, 1)).xyz;
float mnR = min3x(min3x(d.r,e.r,f.r),b.r,h.r);
float mnG = min3x(min3x(d.g,e.g,f.g),b.g,h.g);
float mnB = min3x(min3x(d.b,e.b,f.b),b.b,h.b);
float mnR2 = min3x(min3x(mnR,a.r,c.r),g.r,i.r);
float mnG2 = min3x(min3x(mnG,a.g,c.g),g.g,i.g);
float mnB2 = min3x(min3x(mnB,a.b,c.b),g.b,i.b);
mnR = mnR + mnR2;
mnG = mnG + mnG2;
mnB = mnB + mnB2;
float mxR = max3x(max3x(d.r,e.r,f.r),b.r,h.r);
float mxG = max3x(max3x(d.g,e.g,f.g),b.g,h.g);
float mxB = max3x(max3x(d.b,e.b,f.b),b.b,h.b);
float mxR2 = max3x(max3x(mxR,a.r,c.r),g.r,i.r);
float mxG2 = max3x(max3x(mxG,a.g,c.g),g.g,i.g);
float mxB2 = max3x(max3x(mxB,a.b,c.b),g.b,i.b);
mxR = mxR + mxR2;
mxG = mxG + mxG2;
mxB = mxB + mxB2;
float rcpMR = 1.0f / mxR;
float rcpMG = 1.0f / mxG;
float rcpMB = 1.0f / mxB;
float ampR = clamp(min(mnR, 2.0f - mxR) * rcpMR, 0.0f, 1.0f);
float ampG = clamp(min(mnG, 2.0f - mxG) * rcpMG, 0.0f, 1.0f);
float ampB = clamp(min(mnB, 2.0f - mxB) * rcpMB, 0.0f, 1.0f);
// Shaping amount of sharpening.
ampR = sqrt(ampR);
ampG = sqrt(ampG);
ampB = sqrt(ampB);
// Filter shape.
// 0 w 0
// w 1 w
// 0 w 0
float peak = - 1.0f / mix(8.0f, 5.0f, clamp(sharpness, 0.0f, 1.0f));
float wR = ampR * peak;
float wG = ampG * peak;
float wB = ampB * peak;
float rcpWeightR = 1.0f / (1.0f + 4.0f * wR);
float rcpWeightG = 1.0f / (1.0f + 4.0f * wG);
float rcpWeightB = 1.0f / (1.0f + 4.0f * wB);
vec3 outColor;
outColor.r = clamp((b.r * wR + d.r * wR + f.r * wR + h.r * wR + e.r) * rcpWeightR, 0.0f, 1.0f);
outColor.g = clamp((b.g * wG + d.g * wG + f.g * wG + h.g * wG + e.g) * rcpWeightG, 0.0f, 1.0f);
outColor.b = clamp((b.b * wB + d.b * wB + f.b * wB + h.b * wB + e.b) * rcpWeightB, 0.0f, 1.0f);
return outColor;
}
void main()
{
ivec2 workSize = imageSize(inTAAImage);
ivec2 topLeft = ivec2(gl_WorkGroupID.xy) * kGroupSize - kBorderSize;
ivec2 groupPos = ivec2(gl_LocalInvocationID.xy);
int groupIndex = int(gl_LocalInvocationIndex);
prepareLds(topLeft, workSize, groupIndex);
groupMemoryBarrier();
barrier();
if (gl_GlobalInvocationID.x >= workSize.x || gl_GlobalInvocationID.y >= workSize.y)
{
return;
}
// load center color, mask valid on w channel.
vec4 colorIn = ldsLoadColor(groupPos);
// cache center color.
const vec3 center = colorIn.xyz;
// history color is after reinhard tonemapper. so here reverse.
colorIn.xyz = reinhardInverse(colorIn.xyz);
// update history image.
imageStore(historyImage, ivec2(gl_GlobalInvocationID.xy), colorIn);
// out color.
vec3 color = center;
float sharpness = pushConstant.sharpness;
if(pushConstant.sharpenMethod == SHARPEN_RESPONSIVE)
{
color = ApplySharpening(groupPos);
}
else if(pushConstant.sharpenMethod == SHARPEN_CAS)
{
color = AMDCASFilter(sharpness, groupPos);
}
// out hdr color, also tonemapper reverse.
imageStore(hdrImage, ivec2(gl_GlobalInvocationID.xy), vec4(reinhardInverse(color), 1.0f));
}