Simple Static Raymarching in a RealityKit Post Processing Compute Shader

RaymarchPostProcessing

//
//  RaymarchPostProcessing.metal
//  RealityKit Ray Marching
//
//  Created by Arthur Schiller on 02.10.22.
//

#include <metal_stdlib>
using namespace metal;

struct Uniforms {
    float time;
    float4x4 cameraTransform;
    float4x4 inverseProjectionTransform;
    float4 topLeft;
    float4 topRight;
    float4 bottomLeft;
    float4 bottomRight;
    float4x4 viewMatrixInverse;
    float4x4 viewMatrix;
    float4 viewTranslation;
};

constant float PI = 3.1415926535897932384626433832795;
constant int maxSteps = 100;
constant int maxDist = 100;
constant float surfDist = 1e-3;

float sphereSDF(float3 position, float3 origin, float radius) {
    return length(origin-position) - radius;
}

float sdf(float3 p) {
    float3 sphereOrigin = float3(0, 0, -0.5);
    float sphere = sphereSDF(p, sphereOrigin, .10);
    return sphere;
}

float raymarch(float3 ro, float3 rd) {
    float dO = 0; // distance to origin
    float dS; // distanxe from the surface

    for (int i = 0; i < maxSteps; i++) {
        float3 p = ro + dO * rd;
        dS = sdf(p);
        dO += dS;
        if (dS<surfDist || dO>maxDist) break;
    }

    return dO;
}

float3 getNormal(float3 p) {
    float2 e = float2(1e-2, 0);
    float3 n = sdf(p) - float3(
                               sdf(p - e.xyy),
                               sdf(p - e.yxy),
                               sdf(p - e.yyx)
    );
    return normalize(n);
}

//uniform vec3 camPos;
//uniform vec2 uResolution;
//uniform sampler2D tDiffuse;
//uniform mat4 cameraWorldMatrix;
//uniform mat4 cameraProjectionMatrixInverse;
//varying vec2 vUv;

[[kernel]]
void raymarchPostProcessing(
    texture2d<half, access::read> inColor [[texture(0)]],
    texture2d<half, access::write> outColor [[texture(1)]],
    uint2 gid [[thread_position_in_grid]],
    constant Uniforms &uniforms [[buffer(0)]]
){
    // Checks to make sure that the specified thread_position_in_grid value is
    // within the bounds of the framebuffer. This ensures that non-uniform size
    // threadgroups don't trigger an error. See
    // https://developer.apple.com/documentation/metal/calculating_threadgroup_and_grid_sizes
    if (gid.x >= inColor.get_width() || gid.y >= inColor.get_height()) {
        return;
    }
    
    int width = inColor.get_width();
    int height = inColor.get_height();
    float2 resolution = float2(width, height);
    float aspectRatio = resolution.x / resolution.y;
    float2 uv = float2(gid) / resolution;
    uv.y = 1 - uv.y; // flip y coordinate to account for Metals Coordinate System

    // adjust coordinates into center
    uv.x *= aspectRatio;
    uv.y -= 0.5;
    uv.x -= 0.5 * aspectRatio;
    
    half3 currentColor = inColor.read(gid).rgb;
    
    // use static camera
    float3 ro = float3(0, 0, 1);
    float3 rd = normalize(float3(uv, -1));

    float d = raymarch(ro, rd);

    half3 col = 0;

    if (d < maxDist) {
        // Surface hit
        float3 p = ro + rd * d;
        float3 n = getNormal(p);
        col.rgb = half3(n);
    } else {
        col.rgb = currentColor;
    }
    
    half4 finalColor = half4(col.rgb, 1.0);
    outColor.write(finalColor, gid);
}