Your guide to Metal ray tracing

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

geometryDescriptor.vertexBuffer = vertexBuffer
geometryDescriptor.indexBuffer = indexBuffer
geometryDescriptor.triangleCount = triangleCount

// Create geometry descriptor:
let geometryDescriptor = MTLAccelerationStructureBoundingBoxGeometryDescriptor()

geometryDescriptor.boundingBoxBuffer = boundingBoxBuffer
geometryDescriptor.boundingBoxCount = boundingBoxCount

let geometryDescriptor = MTLAccelerationStructureCurveGeometryDescriptor()
  
geometryDescriptor.controlPointBuffer = controlPointBuffer
geometryDescriptor.radiusBuffer = radiusBuffer
geometryDescriptor.indexBuffer = indexBuffer

geometryDescriptor.controlPointCount = controlPointCount
geometryDescriptor.segmentCount = segmentCount
geometryDescriptor.curveType = .round
geometryDescriptor.curveBasis = .bezier
geometryDescriptor.segmentControlPointCount = 4

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

// Add geometry descriptor to acceleration structure descriptor
accelerationStructureDescriptor.geometryDescriptors = [ geometryDescriptor ]

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

// Query for size and alignment requirement in a heap
let heapSize: MTLSizeAndAlign
heapSize = device.heapAccelerationStructureSizeAndAlign(size: sizes.accelerationStructureSize)

// Allocate acceleration structure from heap
var accelerationStructure: MTLAccelerationStructure!
accelerationStructure = heap.makeAccelerationStructure(size: heapSize.size)

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

let commandEncoder = commandBuffer.makeAccelerationStructureCommandEncoder()!

commandEncoder.build(accelerationStructure: accelerationStructure,
                     descriptor: accelerationStructureDescriptor,
                     scratchBuffer: scratchBuffer,
                     scratchBufferOffset: 0)

commandEncoder.endEncoding()

var instanceASDesc = MTLInstanceAccelerationStructureDescriptor()

instanceASDesc.instanceCount = ...
instanceASDesc.instancedAccelerationStructures = [ mountainAS, treeAS, ... ]
instanceASDesc.instanceDescriptorType = .userID

let size = MemoryLayout<MTLAccelerationStructureUserIDInstanceDescriptor>.stride
let instanceDescriptorBufferSize = size * instanceASDesc.instanceCount

let instanceDescriptorBuffer = device.makeBuffer(length: instanceDescriptorBufferSize,
                                                 options: .storageModeShared)!
    
instanceASDesc.instanceDescriptorBuffer = instanceDescriptorBuffer

var instanceDesc = MTLAccelerationStructureUserIDInstanceDescriptor()

instanceDesc.accelerationStructureIndex = 0    // index into instancedAccelerationStructures
instanceDesc.transformationMatrix = ...
instanceDesc.mask = 0xFFFFFFFF

var instanceASDesc = MTLIndirectInstanceAccelerationStructureDescriptor()

instanceASDesc.instanceDescriptorType = .indirect
instanceASDesc.maxInstanceCount = ...
instanceASDesc.instanceCountBuffer = ...
instanceASDesc.instanceDescriptorBuffer = ...

device MTLIndirectAccelerationStructureInstanceDescriptor *instance_buffer = ...;
// ...
acceleration_structure<> as = ...;
instance_buffer[i].accelerationStructureID = as;
instance_buffer[i].transformationMatrix[0] = ...;
instance_buffer[i].transformationMatrix[1] = ...;
instance_buffer[i].transformationMatrix[2] = ...;
instance_buffer[i].transformationMatrix[3] = ...;
instance_buffer[i].mask = 0xFFFFFFFF;

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

// Create command buffer/encoder ...

// Refit acceleration structure
commandEncoder.refit(sourceAccelerationStructure: accelerationStructure,
                     descriptor: asDescriptor,
                     destinationAccelerationStructure: accelerationStructure,
                     scratchBuffer: scratchBuffer,
                     scratchBufferOffset: 0)

// Use compaction to reclaim memory

// Create command buffer/encoder ...

sizeCommandEncoder.writeCompactedSize(accelerationStructure: accelerationStructure,
                                      buffer: sizeBuffer,
                                      offset: 0,
                                      sizeDataType: .ulong)

// endEncoding(), commit command buffer and wait until completed ...

// Allocate new acceleration structure using UInt64 from sizeBuffer ...

compactCommandEncoder.copyAndCompact(sourceAccelerationStructure: accelerationStructure,
                             destinationAccelerationStructure: compactedAccelerationStructure)

encoder.setAccelerationStructure(primitiveAccelerationStructure, bufferIndex:0)

// Intersect rays with a primitive acceleration structure

[[kernel]]
void trace_rays(acceleration_structure<> as, /* ... */) {
  intersector<> i;

  ray r(origin, direction);

  intersection_result<> result = i.intersect(r, as);

  if (result.type == intersection_type::triangle) {
    float distance = result.distance;


    // shade triangle...
  }
}

// Intersect rays with a primitive acceleration structure

[[kernel]]
void trace_rays(acceleration_structure<> as, /* ... */) {
  intersector<triangle_data> i;

  ray r(origin, direction);

  intersection_result<triangle_data> result = i.intersect(r, as);

  if (result.type == intersection_type::triangle) {
    float distance = result.distance;
    float2 coords = result.triangle_barycentric_coord;

    // shade triangle...
  }
}

encoder.setAccelerationStructure(instanceAccelerationStructure, bufferIndex:0)
encoder.useHeap(accelerationStructureHeap);

// Intersect rays with an instance acceleration structure

[[kernel]]
void trace_rays(acceleration_structure<instancing> as, /* ... */) {
  intersector<instancing, max_levels<3>> i;

  ray r(origin, direction);

  intersection_result<instancing, max_levels<3>> result = i.intersect(r, as);

  if (result.type == intersection_type::triangle) {
    float distance = result.distance;

    // shade triangle...
  }
}

// Intersect rays with an instance acceleration structure

[[kernel]]
void trace_rays(acceleration_structure<instancing> as, /* ... */) {
  intersector<instancing, max_levels<3>> i;

  ray r(origin, direction);

  intersection_result<instancing, max_levels<3>> result = i.intersect(r, as);

  if (result.type == intersection_type::triangle) {
    float distance = result.distance;
    for (uint i = 0; i < result.instance_count; ++i) {
      uint id = result.instance_id[i];
      // ...
    }
    // shade triangle...
  }
}

// Intersect rays with curve primitives

[[kernel]]
void trace_rays(acceleration_structure<> as, /* ... */) {
  intersector<> i;

  i.assume_geometry_type(geometry_type::curve | geometry_type::triangle);

  ray r(origin, direction);

  intersection_result<> result = i.intersect(r, as);

  if (result.type == intersection_type::curve) {
    float distance = result.distance;
    // shade curve...
  }
}

// Intersect rays with curve primitives

[[kernel]]
void trace_rays(acceleration_structure<> as, /* ... */) {
  intersector<curve_data> i;

  i.assume_geometry_type(geometry_type::curve | geometry_type::triangle);

  ray r(origin, direction);

  intersection_result<curve_data> result = i.intersect(r, as);

  if (result.type == intersection_type::curve) {
    float distance = result.distance;
    float param = result.curve_parameter;
    // shade curve...
  }
}

// Intersect rays with curve primitives

[[kernel]]
void trace_rays(acceleration_structure<> as, /* ... */) {
  intersector<curve_data> i;

  i.assume_geometry_type(geometry_type::curve | geometry_type::triangle);
  i.assume_curve_type(curve_type::round);
  i.assume_curve_basis(curve_basis::bezier);
  i.assume_curve_control_point_count(3);

  ray r(origin, direction);

  intersection_result<curve_data> result = i.intersect(r, as);

  if (result.type == intersection_type::curve) {
    float distance = result.distance;
    float param = result.curve_parameter;
    // shade curve...
  }
}