Add GEVulkanHiZDepth for screen space reflection

Author: deveee

data/shaders/ge_shaders/displace_mask.frag

@@ -19,8 +19,179 @@ layout(push_constant) uniform Constants
 #include "../utils/screen_space_reflection.frag"
 
 layout (set = 2, binding = 2) uniform samplerCube u_skybox_texture;
-layout (set = 3, binding = 2) uniform sampler2D u_displace_color;
-layout (set = 3, binding = 3) uniform sampler2DShadow u_depth;
+layout (set = 3, binding = 0) uniform sampler2D u_displace_color;
+layout (set = 3, binding = 1) uniform sampler2DShadow u_depth;
+layout (set = 3, binding = 2) uniform sampler2D u_hiz_depth;
+
+#ifdef PBR_ENABLED
+
+// Start tracing in this level.
+#define HIZ_START_LEVEL      0
+// Stop tracing if current level is higher than this. (higher level means lower value)
+#define HIZ_STOP_LEVEL       0
+#define HIZ_MAX_LEVEL        6
+
+// Set to 1 to disable HiZ and perform naive linear search.
+#define DEBUG_LINEAR_SEARCH  0
+#define MAX_THICKNESS        0.001
+
+vec3 intersectDepthPlane(vec3 o, vec3 d, float z)
+{
+    return o + d * z;
+}
+
+// Index of the cell that contains the given 2D position.
+ivec2 getCell(vec2 screenUV, ivec2 cellCount)
+{
+    return ivec2(screenUV * cellCount);
+}
+
+// The number of cells in the quad tree at the given level.
+ivec2 getCellCount(int level)
+{
+    return textureSize(u_hiz_depth, level);
+}
+
+// Returns screen space position of the intersection
+// between o + d*t and the closest cell boundary at current HiZ level.
+vec3 intersectCellBoundary(
+    vec3 pos, vec3 dir,
+    ivec2 cell, ivec2 cellCount,
+    vec2 crossStep, vec2 crossOffset)
+{
+    vec3 intersection = vec3(0.0);
+
+    vec2 index = cell + crossStep;
+    vec2 boundary = index / vec2(cellCount); // Screen space position of the boundary
+    boundary += crossOffset;
+
+    vec2 delta = boundary - pos.xy;
+    delta /= dir.xy;
+    float t = min(delta.x, delta.y);
+
+    intersection = intersectDepthPlane(pos, dir, t);
+    return intersection;
+}
+
+bool crossedCellBoundary(ivec2 oldCellIx, ivec2 newCellIx)
+{
+    return any(notEqual(oldCellIx, newCellIx));
+}
+
+// Minimum depth of the current cell in the current HiZ level.
+float getMinDepthPlane(ivec2 cellIx, int level)
+{
+    return texelFetch(u_hiz_depth, cellIx, level).x;
+}
+
+float getMaxTraceDistance(vec3 p, vec3 v)
+{
+    vec3 traceDistances;
+    if (v.x < 0.0)
+        traceDistances.x = p.x / (-v.x);
+    else
+        traceDistances.x = (1.0 - p.x) / v.x;
+
+    if (v.y < 0.0)
+        traceDistances.y = p.y / (-v.y);
+    else
+        traceDistances.y = (1.0 - p.y) / v.y;
+
+    if (v.z < 0.0)
+        traceDistances.z = p.z / (-v.z);
+    else
+        traceDistances.z = (1.0 - p.z) / v.z;
+
+    return min(traceDistances.x, min(traceDistances.y, traceDistances.z));
+}
+
+// p            : Screen space position
+// v            : Screen space reflection direction
+// hitPointSS   : Returns screen space hit point
+// Return value : Whether RT actually hit a surface
+bool traceHiZ(vec3 p, vec3 v, out vec2 hitPointSS)
+{
+    const int maxLevel = min(HIZ_MAX_LEVEL, textureQueryLevels(u_hiz_depth) - 1); // Last mip level
+    float maxTraceDistance = getMaxTraceDistance(p, v);
+
+    // Get the cell cross direction and a small offset to enter
+    // the next cell when doing cell crossing.
+    vec2 crossStep = vec2(v.x >= 0 ? 1 : -1, v.y >= 0 ? 1 : -1);
+    vec2 crossOffset = crossStep / u_camera.m_viewport.zw / 128.;
+    crossStep = clamp(crossStep, 0.0, 1.0);
+
+    // Set current ray to the original screen coordinate and depth.
+    vec3 ray = p;
+    float minZ = ray.z;
+    float maxZ = ray.z + v.z * maxTraceDistance;
+    float deltaZ = maxZ - minZ;
+
+    vec3 o = ray;
+    vec3 d = v * maxTraceDistance;
+
+    int level = HIZ_START_LEVEL;
+    int deepestLevel = level;
+#if DEBUG_LINEAR_SEARCH
+    level = 0;
+#endif
+    uint iterations = 0;
+    bool isBackwardRay = v.z < 0;
+    float rayDir = isBackwardRay ? -1.0 : 1.0;
+
+    // Cross to next cell s.t. we don't get a self-intersection immediately.
+    ivec2 startCellCount = getCellCount(level);
+    ivec2 rayCell = getCell(ray.xy, startCellCount);
+    ray = intersectCellBoundary(o, d, rayCell, startCellCount, crossStep, crossOffset * 64.);
+
+    while (level >= HIZ_STOP_LEVEL && ray.z * rayDir <= maxZ * rayDir &&
+        iterations < u_hiz_iterations)
+    {
+        // Get the cell number of our current ray.
+        ivec2 cellCount = getCellCount(level);
+        ivec2 oldCellIx = getCell(ray.xy, cellCount);
+
+        // Get the minimum depth plane in which the current ray resides.
+        float cellMinZ = getMinDepthPlane(oldCellIx, level);
+
+        // Intersect only if ray depth is below the minimum depth plane.
+        vec3 tempRay;
+        if (cellMinZ > ray.z && !isBackwardRay)
+            tempRay = intersectDepthPlane(o, d, (cellMinZ - minZ) / deltaZ);
+        else
+            tempRay = ray;
+
+        ivec2 newCellIx = getCell(tempRay.xy, cellCount);
+        float thickness = level == 0 ? (ray.z - cellMinZ) : 0;
+
+        bool crossed = (isBackwardRay && (cellMinZ > ray.z))
+                    || (thickness > MAX_THICKNESS) || crossedCellBoundary(oldCellIx, newCellIx);
+
+        if (crossed)
+        {
+            ray = intersectCellBoundary(o, d, oldCellIx, cellCount, crossStep, crossOffset);
+            level = min(maxLevel, level + 1);
+            deepestLevel = max(deepestLevel, level);
+#if DEBUG_LINEAR_SEARCH
+            level = 0;
+#endif
+        }
+        else
+        {
+            ray = tempRay;
+            level = level - 1;
+        }
+
+        iterations += 1;
+    }
+
+    // Results
+    //debugDeepestLevel = deepestLevel;
+    //debugIterations = iterations;
+    hitPointSS = ray.xy;
+    return level < HIZ_STOP_LEVEL && iterations < u_hiz_iterations;
+}
+
+#endif
 
 void main()
 {
@@ -62,13 +233,42 @@ void main()
 
         // fallback to skybox
         vec4 fallback = texture(u_skybox_texture, world_reflection);
+
+        // early exit if normal is facing camera too directly (no meaningful reflection)
+        if (normal.z < -0.75)
+        {
+            o_displace_ssr = fallback;
+            return;
+        }
+
         vec4 result;
         vec2 viewport_scale = u_camera.m_viewport.zw / u_camera.m_screensize;
         vec2 viewport_offset = u_camera.m_viewport.xy / u_camera.m_screensize;
-        vec2 coords = RayCast(reflected, xpos, u_camera.m_projection_matrix,
-            viewport_scale, viewport_offset, u_depth);
+        bool hit = true;
+        vec2 coords;
+        if (u_hiz_iterations == 0)
+        {
+            coords = RayCast(reflected, xpos, u_camera.m_projection_matrix,
+                viewport_scale, viewport_offset, u_depth);
+        }
+        else
+        {
+            vec3 positionSS = CalcCoordFromPosition(xpos,
+                u_camera.m_projection_matrix, vec2(1.0), vec2(0.0));
+            vec3 positionCS = positionSS;
+            positionCS.xy = 2.0 * positionCS.xy - 1.0;
+            vec3 position2VS = xpos + 1000.0 * reflected;
+            vec4 position2CS = u_camera.m_projection_matrix * vec4(position2VS, 1.0);
+            position2CS /= position2CS.w;
+            vec3 position2SS = position2CS.xyz;
+            position2SS.xy = vec2(0.5) + 0.5 * position2SS.xy;
+            vec3 reflectionDirSS = normalize(position2SS - positionSS);
+            // Trace HiZ to find the hit point.
+            hit = traceHiZ(positionSS, reflectionDirSS, coords);
+            coords = coords * viewport_scale + viewport_offset;
+        }
         vec2 viewport_coords = (coords - viewport_offset) / viewport_scale;
-        if (viewport_coords.x < 0. || viewport_coords.x > 1. ||
+        if (!hit || viewport_coords.x < 0. || viewport_coords.x > 1. ||
             viewport_coords.y < 0. || viewport_coords.y > 1.)
         {
             result = fallback;

data/shaders/ge_shaders/hiz_depth.comp

@@ -0,0 +1,62 @@
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(binding = 0) uniform sampler2D u_depth;
+layout(binding = 1, r32f) uniform writeonly image2D u_hiz_depth;
+
+layout(push_constant) uniform PushConstants
+{
+    ivec3 u_offset_miplevel;
+} pc;
+
+void main()
+{
+    ivec2 dst = ivec2(gl_GlobalInvocationID.xy);
+    ivec2 current_size = imageSize(u_hiz_depth);
+
+    if (dst.x >= current_size.x || dst.y >= current_size.y)
+        return;
+
+    if (pc.u_offset_miplevel.z == 0)
+    {
+        // at level 0, do a 1:1 copy
+        float d = texelFetch(u_depth, dst + pc.u_offset_miplevel.xy, 0).r;
+        imageStore(u_hiz_depth, dst, vec4(d));
+    }
+    else
+    {
+        // at higher mip levels, read a 2×2 block from previous level
+        ivec2 src = dst * 2;
+        int prev_level = pc.u_offset_miplevel.z - 1;
+        ivec2 prev_size = textureSize(u_depth, prev_level);
+        float d0 = texelFetch(u_depth, src + ivec2(0, 0), prev_level).r;
+        float d1 = texelFetch(u_depth, src + ivec2(1, 0), prev_level).r;
+        float d2 = texelFetch(u_depth, src + ivec2(0, 1), prev_level).r;
+        float d3 = texelFetch(u_depth, src + ivec2(1, 1), prev_level).r;
+        float min_depth = min(min(d0, d1), min(d2, d3));
+        //float max_depth = max(max(d0, d1), max(d2, d3));
+        bool extra_sample_x = (current_size.x * 2) < prev_size.x;
+        bool extra_sample_y = (current_size.y * 2) < prev_size.y;
+        if (extra_sample_x)
+        {
+            float d4 = texelFetch(u_depth, src + ivec2(2, 0), prev_level).r;
+            float d5 = texelFetch(u_depth, src + ivec2(2, 1), prev_level).r;
+            min_depth = min(min_depth, min(d4, d5));
+            //max_depth = max(max_depth, max(d4, d5));
+        }
+        if (extra_sample_y)
+        {
+            float d6 = texelFetch(u_depth, src + ivec2(0, 2), prev_level).r;
+            float d7 = texelFetch(u_depth, src + ivec2(1, 2), prev_level).r;
+            min_depth = min(min_depth, min(d6, d7));
+            //max_depth = max(max_depth, max(d6, d7));
+        }
+        if (extra_sample_x && extra_sample_y)
+        {
+            float d8 = texelFetch(u_depth, src + ivec2(2, 2), prev_level).r;
+            min_depth = min(min_depth, d8);
+            //max_depth = max(max_depth, d8);
+        }
+        imageStore(u_hiz_depth, dst, vec4(min_depth));
+        //imageStore(u_hiz_depth, dst, vec4(min_depth, max_depth, 0.0, 0.0));
+    }
+}

data/shaders/ge_shaders/utils/constants_utils.glsl

@@ -3,6 +3,7 @@ layout (constant_id = 1) const float u_specular_levels_minus_one = 0.0;
 layout (constant_id = 2) const bool u_deferred = false;
 layout (constant_id = 3) const bool u_has_skybox = true;
 layout (constant_id = 4) const bool u_ssr = false;
+layout (constant_id = 5) const uint u_hiz_iterations = 0;
 
 vec3 convertColor(vec3 input_color)
 {

data/shaders/sp_displace_ssr.frag

@@ -49,6 +49,14 @@ void main()
 
     // fallback to skybox
     vec4 fallback = texture(u_skybox_texture, world_reflection);
+
+    // early exit if normal is facing camera too directly (no meaningful reflection)
+    if (normal.z < -0.75)
+    {
+        o_displace_ssr = fallback;
+        return;
+    }
+
     vec4 result;
     vec2 viewport_scale = vec2(1.0);
     vec2 viewport_offset = vec2(0.0);

lib/graphics_engine/CMakeLists.txt

@@ -88,6 +88,7 @@ set(GE_SOURCES
     src/ge_vulkan_environment_map.cpp
     src/ge_vulkan_fbo_texture.cpp
     src/ge_vulkan_features.cpp
+    src/ge_vulkan_hiz_depth.cpp
     src/ge_vulkan_light_handler.cpp
     src/ge_vulkan_mesh_cache.cpp
     src/ge_vulkan_mesh_scene_node.cpp

lib/graphics_engine/include/ge_main.hpp

@@ -33,6 +33,11 @@ enum GEScreenSpaceReflectionType : unsigned
 {
     GSSRT_DISABLED = 0,
     GSSRT_FAST,
+    GSSRT_HIZ,
+    GSSRT_HIZ100 = GSSRT_HIZ,
+    GSSRT_HIZ200,
+    GSSRT_HIZ400,
+    GSSRT_COUNT,
 };
 
 struct GEConfig