Discover ray tracing with Metal

[[kernel]]
void rtKernel(primitive_acceleration_structure accelerationStructure [[buffer(0)]],
              /* ... */)
{
    // Generate ray
    ray r = generateCameraRay(tid);

    // Create an intersector
    intersector<triangle_data> intersector;

    // Intersect with scene
    intersection_result<triangle_data> intersection;

    intersection = intersector.intersect(r, accelerationStructure);

    // shading...
}

let accelerationStructureDescriptor = MTLPrimitiveAccelerationStructureDescriptor()

// Create geometry descriptor(s)
let geometryDescriptor = MTLAccelerationStructureTriangleGeometryDescriptor()

geometryDescriptor.vertexBuffer = vertexBuffer
geometryDescriptor.triangleCount = triangleCount

accelerationStructureDescriptor.geometryDescriptors = [ geometryDescriptor ]

// Query for acceleration structure sizes
let sizes = device.accelerationStructureSizes(descriptor: accelerationStructureDescriptor)

// Allocate acceleration structure
let accelerationStructure =
    device.makeAccelerationStructure(size: sizes.accelerationStructureSize)!

// Allocate scratch buffer
let scratchBuffer = device.makeBuffer(length: sizes.buildScratchBufferSize,
                                      options: .storageModePrivate)!

// Create command buffer/encoder
let commandBuffer = commandQueue.makeCommandBuffer()!
let commandEncoder = commandBuffer.makeAccelerationStructureCommandEncoder()!

// Encode acceleration structure build
commandEncoder.build(accelerationStructure: accelerationStructure,
                     descriptor: accelerationStructureDescriptor,
                     scratchBuffer: scratchBuffer,
                     scratchBufferOffset: 0)

// Commit command buffer
commandEncoder.endEncoding()
commandBuffer.commit()

[[kernel]]
void rtKernel(primitive_acceleration_structure accelerationStructure [[buffer(0)]],
              /* ... */)
{
    // generate ray, create intersector...

    intersection = intersector.intersect(r, accelerationStructure);

    // shading...
}

computeEncoder.setAccelerationStructure(accelerationStructure, bufferIndex: 0)

[[intersection(triangle, triangle_data)]]
bool alphaTestIntersectionFunction(uint primitiveIndex        [[primitive_id]],
                                   uint geometryIndex         [[geometry_id]],
                                   float2 barycentricCoords   [[barycentric_coord]],
                                   device Material *materials [[buffer(0)]])
{
    texture2d<float> alphaTexture = materials[geometryIndex].alphaTexture;

    float2 UV = interpolateUVs(materials[geometryIndex].UVs,
        primitiveIndex, barycentricCoords);

    float alpha = alphaTexture.sample(sampler, UV).x;

    return alpha > 0.5f;
}

// Create a primitive acceleration structure descriptor
let accelerationStructureDescriptor = MTLPrimitiveAccelerationStructureDescriptor()

// Create one or more bounding box geometry descriptors:
let geometryDescriptor = MTLAccelerationStructureBoundingBoxGeometryDescriptor()

geometryDescriptor.boundingBoxBuffer = boundingBoxBuffer
geometryDescriptor.boundingBoxCount = boundingBoxCount

accelerationStructureDescriptor.geometryDescriptors = [ geometryDescriptor ]

struct BoundingBoxResult {
    bool accept [[accept_intersection]];
    float distance [[distance]];
};

[[intersection(bounding_box)]]
BoundingBoxResult sphereIntersectionFunction(float3 origin            [[origin]],
                                             float3 direction         [[direction]],
                                             float minDistance        [[min_distance]],
                                             float maxDistance        [[max_distance]],
                                             uint primitiveIndex      [[primitive_id]],
                                             device Sphere *spheres   [[buffer(0)]])
{
    float distance;

    if (!intersectRaySphere(origin, direction, spheres[primitiveIndex], &distance))
        return { false, 0.0f };

    if (distance < minDistance || distance > maxDistance)
        return { false, 0.0f };

    return { true, distance };
}

[[intersection(bounding_box)]]
BoundingBoxResult sphereIntersectionFunction(/* ... */,
                                             ray_data float3 & normal [[payload]])
{
    // ...

    if (distance < minDistance || distance > maxDistance)
        return { false, 0.0f };

    float3 intersectionPoint = origin + direction * distance;
    normal = normalize(intersectionPoint - spheres[primitiveIndex].origin);

    return { true, distance };
}

[[kernel]]
void rtKernel(/* ... */)
{
    // generate ray, create intersector...

  float3 normal;

    intersection = intersector.intersect(r, accelerationStructure, functionTable, normal);

    // shading...
}

// Load functions from Metal library
let sphereIntersectionFunction = library.makeFunction(name: “sphereIntersectionFunction”)!
// other functions...

// Attach functions to ray tracing compute pipeline descriptor
let linkedFunctions = MTLLinkedFunctions()

linkedFunctions.functions = [ sphereIntersectionFunction, alphaTestFunction, ... ]

computePipelineDescriptor.linkedFunctions = linkedFunctions

// Compile and link ray tracing compute pipeline
let computePipeline = try device.makeComputePipeline(descriptor: computePipelineDescriptor,
                                                     options: [],
                                                     reflection: nil)

class MTLAccelerationStructureGeometryDescriptor : NSObject {

    var intersectionFunctionTableOffset: Int

// ...

}

struct MTLAccelerationStructureInstanceDescriptor {
    var intersectionFunctionTableOffset: UInt32
    // ...
};

// Allocate intersection function table
let descriptor = MTLIntersectionFunctionTableDescriptor()

descriptor.functionCount = intersectionFunctions.count

let functionTable = computePipeline.makeIntersectionFunctionTable(descriptor: descriptor)

for i in 0 ..< intersectionFunctions.count {
    // Get a handle to the linked intersection function in the pipeline state
    let functionHandle = computePipeline.functionHandle(function: intersectionFunctions[i])

    // Insert the function handle into the table
    functionTable.setFunction(functionHandle, index: i)
}

// Bind intersection function resources
functionTable.setBuffer(sphereBuffer, offset: 0, index: 0)

[[kernel]]
void rtKernel(primitive_acceleration_structure accelerationStructure   [[buffer(0)]],
              intersection_function_table<triangle_data> functionTable [[buffer(1)]],
              /* ... */)
{
    // generate ray, create intersector...

    intersection = intersector.intersect(r, accelerationStructure, functionTable);

    // shading...
}

encoder.setIntersectionFunctionTable(functionTable, bufferIndex: 1)