SHADE_Y3/Assets/Shaders/SSAO_CS.glsl

#version 450

const uint NUM_SAMPLES = 64;
const uint NUM_ROTATIONS = 16;
const int ROTATION_KERNEL_W = 4;
const int ROTATION_KERNEL_H = 4;

// can perhaps pass in as push constant.
const float RADIUS = 0.5f;
const float BIAS = 0.025f;

layout(local_size_x = 16, local_size_y = 16) in;
layout(set = 4, binding = 0, rgba32f) uniform image2D positions;
layout(set = 4, binding = 1, rgba32f) uniform image2D normals;
layout(set = 4, binding = 2, rgba32f) uniform image2D outputImage;


// SSAO data
layout(std430, set = 5, binding = 0) buffer SSAOData
{
  vec4 samples[NUM_SAMPLES];
  vec4 rotations[NUM_ROTATIONS];

} ssaoData;

layout(set = 2, binding = 0) uniform CameraData
{
  vec4 position;
  mat4 vpMat;
  mat4 viewMat;
  mat4 projMat;

} cameraData;

shared vec4 sharedRotations[NUM_ROTATIONS];

void main()
{ 
  ivec2 size = imageSize (outputImage);

  // load rotations into shared memory
  uint localThreadIndex = gl_LocalInvocationIndex;
  if (localThreadIndex < NUM_ROTATIONS)
    sharedRotations[localThreadIndex] = ssaoData.rotations[localThreadIndex];

  barrier();

  ivec2 globalThread = ivec2 (gl_GlobalInvocationID.xy);

  // load all the necessary variables
  vec3 viewSpacePos     = imageLoad (positions, globalThread).rgb;
  vec3 viewSpaceNormal  = normalize (imageLoad (normals, globalThread).rgb);

  // Get random vector
  uint randomVecXIndex = globalThread.x % ROTATION_KERNEL_W;
  uint randomVecYIndex = globalThread.y % ROTATION_KERNEL_H;
  vec3 randomVec        = (sharedRotations[randomVecYIndex * ROTATION_KERNEL_W + randomVecXIndex].rgb) * 2.0f - 1.0f;

  // Gram schmidt
  vec3 tangent = normalize (randomVec - (viewSpaceNormal * dot(viewSpaceNormal, randomVec)));
  vec3 bitangent = cross (tangent, viewSpaceNormal);
  mat3 TBN = mat3 (tangent, bitangent, viewSpaceNormal);

  float occlusion = 0.0f;
  for (int i = 0; i < NUM_SAMPLES; ++i)
  {
    // We want to get a position at an offset from the view space position. Offset scaled by radius.
    vec3 samplePos = TBN * ssaoData.samples[i].rgb;
    samplePos = viewSpacePos + samplePos * RADIUS;

    // Now we take that offset position and bring it to clip space
    vec4 offsetPos = vec4 (samplePos, 1.0f);
    offsetPos = cameraData.projMat * offsetPos;

    // then we do perspective division
    offsetPos.xyz /= offsetPos.w;

    // and bring it from [-1, 1] to screen coordinates
    offsetPos.xyz = ((offsetPos.xyz * 0.5f) + 0.5f);
    offsetPos.xy *= vec2(size.xy);

    float sampleDepth = imageLoad (positions, ivec2 (offsetPos.xy)).z;

    float rangeCheck = smoothstep (0.0f, 1.0f, RADIUS / abs (viewSpacePos.z - sampleDepth));
    occlusion += (sampleDepth >= samplePos.z + BIAS ? 1.0f : 0.0f) * rangeCheck;
  }

  occlusion =  (occlusion / float(NUM_SAMPLES));

  // store result into result image
  imageStore(outputImage, ivec2(gl_GlobalInvocationID.xy), vec4(occlusion.rrr, 0.0f));
}