Speculative raymarching experiment

wave.frag

#version 460

/********************************************************************************************************************

Speculative Raymarching

Shaders run code in "waves" or "subgroups" of 32 or 64 threads (or 4 in WebGPU on my box, which is why there's no
web demo for this). The idea here is that instead of 32 pixels running the same (expensive) sdf function which is
likely to return similar results to advance their respective rays - they all band together and cooperatively march
a cone that is large enough to hold all of them.

Now, every ray running the same sdf at the same (or a very close) 3D point to march the same cone wouldn't really
help us, so instead - they speculate that the next step distance will be the same as the last one and call the sdf
function with the point they assume the previous step will lead to.

Then we loop through the sdf results and see if we can actually use them.

In theory it could be a x32 speedup. In reality not so much. Still, it kinda works. Artifacts are unfortunately a
little uglier than in standard per-pixel raymarching.

Run it with https://github.com/patriciogonzalezvivo/glslViewer/releases
$ glslViewer wave.frag

*********************************************************************************************************************/

// Enable Subgroup / Wave Intrinsics
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : require
#extension GL_KHR_shader_subgroup_shuffle : require

// glslViewer uniforms
uniform vec2 u_resolution;
uniform float u_time;
out vec4 fragColor;
float time = mod(u_time, 1000.0); // Avoid precision loss if it's been running for a while...

mat2 rot(float r){
    return mat2(cos(r), sin(r), -sin(r), cos(r));
}

float sdSphGrid(vec3 p) {
    p.z += time;
    p=fract(p)-.5;
    return length(p)-.1;
}

float sdFractal(vec3 p) {
    p.z += time*.2;
    float y, s = 1., e;
    y = p.y += 0.2;
    p.z -= round(p.z);
    p = .7-abs(p);
    for(int i = 0; i < 6; i++) {
        p = abs(p.x > p.y ? p : p.yxz) - 0.8;
        s *= e = 5. / min(dot(p, p), 0.5);
        p = abs(p) * e - 6.;
        p.z += 5.;
    }
    return min(y, length(p.yz) / s - 4e3/s);
}

float df(vec3 p) {
    return sdSphGrid(p); // Basic infinite grid of spheres test scene
    
    // return sdFractal(p); // Fractal hallway - less of an impact here :(
    
    // In a previous version I tried guessing future distances based on the distance gradient
    // derived from last_step and the step before it. It didn't work as well as linear guesses
    // for the grid-of-spheres test scene, but might be worth revisiting.
}

vec3 calc_normal(vec3 p, float eps = 0.0001) {
    vec2 k = vec2(1,-1);
    vec3 a = k.xyy * df(p + k.xyy*eps);
    vec3 b = k.yyx * df(p + k.yyx*eps);
    vec3 c = k.yxy * df(p + k.yxy*eps);
    vec3 d = k.xxx * df(p + k.xxx*eps);
    return normalize(a+b+c+d);
}

void main() {
    const float fov = 1.7;
    vec2 uv = (gl_FragCoord.xy * 2 - u_resolution) / u_resolution.y;
    vec3 ro = vec3(0);
    vec3 rd = normalize(vec3(uv, fov)); // Ray direction
    float t = 0.0;
    int max_steps = 16;
    //rd.xy *= rot(time*.1);
    //rd.xz *= rot(time*.1);
    vec3 sgrd = normalize(vec3((subgroupMin(uv) + subgroupMax(uv)) * 0.5, fov)); // Cone direction
    // Minimum dot product (maximum angle between center ray and thread ray)
    float min_dot = clamp(subgroupMin(dot(rd, sgrd)), 0.0001, 1.0); 
    float cone_tan_half_angle = sqrt(1.0 - min_dot * min_dot) / min_dot; // Subgroup cone half-angle (tan(acos(min_dot)))
    
    bool left = gl_FragCoord.x < u_resolution.x * .5;
    bool break_wave = left;
    float dist = 0.0;
    float last_step = 1.0; // Next step distance estimator, any number will do for the 1st iteration
    bool hit = false;

    for(int i = 0; i < max_steps; i++) {
        if(break_wave) { // Use standard raymarching
            last_step = dist = df(ro + rd * t);
            if(dist < .0001) {
                hit = true;
                break;
            } else {
                t += dist;
            }
        } else { // Speculative part - assumes the next steps will be the same distance as the last one (last_step)
            float guess_t = t + last_step * gl_SubgroupInvocationID; // Guess total distance (t + 0 for thread 0)
            // Get central ray distances from geometry (thread-local, relative to the tentative guess_t)
            dist = df(ro + sgrd * guess_t);
            float cone_radius = (guess_t + dist) * cone_tan_half_angle;
            dist -= cone_radius; // Reduce distance to compensate for cone size
            for(int j = 0; j < gl_SubgroupSize; j++) { // Sweep through the wave and see how far we can go
                // Thread sdf distance and ray length (accurate for thread 0, speculated ray position for the rest)
                float thread_dist = subgroupShuffle(dist, j);
                float thread_t = subgroupShuffle(guess_t, j);
                if(t > thread_t - thread_dist) { // t is inside the thread's empty sphere
                    last_step = thread_dist; // Update the next step distance estimate
                    t = thread_t + thread_dist; // Move to the end of this thread's empty sphere
                }
            }
            // If the cone has reached its limit break the wave and march individual rays 
            if(last_step < 0.01) break_wave = true;
        }
    }
    if(hit) {
        vec3 n = calc_normal(ro + rd * t);
        fragColor = vec4(n*.5+.5,1);
    } else {
        fragColor = vec4(.5);
    }
    if(abs(gl_FragCoord.x - u_resolution.x * .5) < 1.0) fragColor = vec4(0);
}

If you use it for anything and it works out, or if you come up with improvements or bug fixes - please leave a comment.

Spheres in a grid (right side is speculative marching)

Fractal (less impressive)