metal-examples-tutorials/raytracing/Renderer.swift

307 lines
15 KiB
Swift

//
// Renderer.swift
//
// Created by Marius Horga on 7/7/18.
//
import MetalKit
import MetalPerformanceShaders
class Renderer: NSObject {
let device: MTLDevice!
let queue: MTLCommandQueue!
var rayPipeline: MTLComputePipelineState!
var shadePipeline: MTLComputePipelineState!
var shadowPipeline: MTLComputePipelineState!
var accumulatePipeline: MTLComputePipelineState!
var copyPipeline: MTLRenderPipelineState!
let maxFramesInFlight = 3
let alignedUniformsSize = (MemoryLayout<Uniforms>.stride + 255) & ~255
let rayStride = 48 // shouldn't be hardcoded... if possible
let intersectionStride = MemoryLayout<MPSIntersectionDistancePrimitiveIndexCoordinates>.stride
var accelerationStructure: MPSTriangleAccelerationStructure!
var intersector: MPSRayIntersector!
var vertexPositionBuffer: MTLBuffer!
var vertexNormalBuffer: MTLBuffer!
var vertexColorBuffer: MTLBuffer!
var rayBuffer: MTLBuffer!
var shadowRayBuffer: MTLBuffer!
var intersectionBuffer: MTLBuffer!
var uniformBuffer: MTLBuffer!
var randomBuffer: MTLBuffer!
var triangleMaskBuffer: MTLBuffer!
var renderTarget: MTLTexture!
var accumulationTarget: MTLTexture!
var semaphore: DispatchSemaphore!
var size: CGSize!
var randomBufferOffset: Int!
var uniformBufferOffset: Int!
var uniformBufferIndex: Int = 0
var frameIndex: uint!
init?(metalKitView: MTKView) {
metalKitView.colorPixelFormat = .rgba16Float
metalKitView.sampleCount = 1
metalKitView.drawableSize = metalKitView.frame.size
guard let device = metalKitView.device else {
return nil
}
self.device = device
queue = self.device.makeCommandQueue()
super.init()
semaphore = DispatchSemaphore.init(value: maxFramesInFlight)
createPipelines(device: self.device, view: metalKitView)
createScene()
createBuffers()
createIntersector()
mtkView(metalKitView, drawableSizeWillChange: metalKitView.frame.size)
}
func createPipelines(device: MTLDevice, view: MTKView) {
guard let library = device.makeDefaultLibrary() else {
return
}
let computeDescriptor = MTLComputePipelineDescriptor()
computeDescriptor.threadGroupSizeIsMultipleOfThreadExecutionWidth = true
let vertexFunction = library.makeFunction(name: "vertexShader")
let fragmentFunction = library.makeFunction(name: "fragmentShader")
let pipelineDescriptor = MTLRenderPipelineDescriptor()
pipelineDescriptor.sampleCount = view.sampleCount
pipelineDescriptor.vertexFunction = vertexFunction
pipelineDescriptor.fragmentFunction = fragmentFunction
pipelineDescriptor.colorAttachments[0].pixelFormat = view.colorPixelFormat
do {
computeDescriptor.computeFunction = library.makeFunction(name: "rayKernel")
rayPipeline = try device.makeComputePipelineState(descriptor: computeDescriptor, options: [], reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(name: "shadeKernel")
shadePipeline = try device.makeComputePipelineState(descriptor: computeDescriptor, options: [], reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(name: "shadowKernel")
shadowPipeline = try device.makeComputePipelineState(descriptor: computeDescriptor, options: [], reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(name: "accumulateKernel")
accumulatePipeline = try device.makeComputePipelineState(descriptor: computeDescriptor, options: [], reflection: nil)
copyPipeline = try device.makeRenderPipelineState(descriptor: pipelineDescriptor)
} catch let error {
print(error.localizedDescription)
}
}
func createScene() {
// light source
var transform = translate(tx: 0, ty: 1, tz: 0.3) * scale(sx: 0.5, sy: 1.98, sz: 0.5)
createCube(faceMask: .positiveY, color: float3(1), transform: transform, inwardNormals: true, triangleMask: TRIANGLE_MASK_LIGHT)
// top wall
transform = translate(tx: 0, ty: 1, tz: 0) * scale(sx: 2, sy: 2, sz: 2)
createCube(faceMask: .positiveY, color: float3(0.02, 0.4, 0.02), transform: transform, inwardNormals: true, triangleMask: TRIANGLE_MASK_GEOMETRY)
// bottom and back walls
createCube(faceMask: [.negativeY, .negativeZ], color: float3(1.0), transform: transform, inwardNormals: true, triangleMask: TRIANGLE_MASK_GEOMETRY)
// left wall
createCube(faceMask: .negativeX, color: float3(1.0, 0.02, 0.02), transform: transform, inwardNormals: true, triangleMask: TRIANGLE_MASK_GEOMETRY)
// right wall
createCube(faceMask: [.positiveX], color: float3(0.02, 0.02, 0.2), transform: transform, inwardNormals: true, triangleMask: TRIANGLE_MASK_GEOMETRY)
// short box
transform = translate(tx: 0.35, ty: 0.3, tz: 0.3725) * rotate(radians: -0.3, axis: float3(0.0, 1.0, 0.0)) * scale(sx: 0.6, sy: 0.6, sz: 0.6)
createCube(faceMask: .all, color: float3(1.0, 1.0, 0.3), transform: transform, inwardNormals: false, triangleMask: TRIANGLE_MASK_GEOMETRY)
// tall box
transform = translate(tx: -0.4, ty: 0.6, tz: -0.29) * rotate(radians: 0.3, axis: float3(0.0, 1.0, 0.0)) * scale(sx: 0.6, sy: 1.2, sz: 0.6)
createCube(faceMask: .all, color: float3(1.0, 1.0, 0.3), transform: transform, inwardNormals: false, triangleMask: TRIANGLE_MASK_GEOMETRY)
}
func createBuffers() {
let uniformBufferSize = alignedUniformsSize * maxFramesInFlight
uniformBuffer = device.makeBuffer(length: uniformBufferSize, options: .storageModeManaged)
randomBuffer = device.makeBuffer(length: 256 * MemoryLayout<float2>.stride * maxFramesInFlight, options: .storageModeManaged)
vertexPositionBuffer = device.makeBuffer(bytes: &vertices, length: vertices.count * MemoryLayout<float3>.stride, options: .storageModeManaged)
vertexColorBuffer = device.makeBuffer(bytes: &colors, length: colors.count * MemoryLayout<float3>.stride, options: .storageModeManaged)
vertexNormalBuffer = device.makeBuffer(bytes: &normals, length: normals.count * MemoryLayout<float3>.stride, options: .storageModeManaged)
triangleMaskBuffer = device.makeBuffer(bytes: &masks, length: masks.count * MemoryLayout<uint>.stride, options: .storageModeManaged)
vertexPositionBuffer?.didModifyRange(0..<(vertexPositionBuffer?.length)!)
vertexColorBuffer?.didModifyRange(0..<(vertexColorBuffer?.length)!)
vertexNormalBuffer?.didModifyRange(0..<(vertexNormalBuffer?.length)!)
triangleMaskBuffer?.didModifyRange(0..<(triangleMaskBuffer?.length)!)
}
func createIntersector() {
intersector = MPSRayIntersector(device: device)
intersector?.rayDataType = .originMaskDirectionMaxDistance
intersector?.rayStride = rayStride
intersector?.rayMaskOptions = .primitive
accelerationStructure = MPSTriangleAccelerationStructure(device: device)
accelerationStructure?.vertexBuffer = vertexPositionBuffer
accelerationStructure?.maskBuffer = triangleMaskBuffer
accelerationStructure?.triangleCount = vertices.count / 3
accelerationStructure?.rebuild()
}
func updateUniforms() {
uniformBufferOffset = alignedUniformsSize * uniformBufferIndex
let pointer = uniformBuffer!.contents().advanced(by: uniformBufferOffset)
let uniforms = pointer.bindMemory(to: Uniforms.self, capacity: 1)
uniforms.pointee.camera.position = float3(0.0, 1.0, 3.38)
uniforms.pointee.camera.forward = float3(0.0, 0.0, -1.0)
uniforms.pointee.camera.right = float3(1.0, 0.0, 0.0)
uniforms.pointee.camera.up = float3(0.0, 1.0, 0.0)
uniforms.pointee.light.position = float3(0.0, 1.98, 0.0)
uniforms.pointee.light.forward = float3(0.0, -1.0, 0.0)
uniforms.pointee.light.right = float3(0.25, 0.0, 0.0)
uniforms.pointee.light.up = float3(0.0, 0.0, 0.25)
uniforms.pointee.light.color = float3(12.0)
let fieldOfView = 45.0 * (Float.pi / 180.0)
let aspectRatio = Float(size.width) / Float(size.height)
let imagePlaneHeight = tanf(fieldOfView / 2.0)
let imagePlaneWidth = aspectRatio * imagePlaneHeight
uniforms.pointee.camera.right *= imagePlaneWidth
uniforms.pointee.camera.up *= imagePlaneHeight
uniforms.pointee.width = uint(size.width)
uniforms.pointee.height = uint(size.height)
uniforms.pointee.blocksWide = (uniforms.pointee.width + 15) / 16
uniforms.pointee.frameIndex = frameIndex
frameIndex += 1
uniformBuffer?.didModifyRange(uniformBufferOffset..<(uniformBufferOffset + alignedUniformsSize))
randomBufferOffset = 256 * MemoryLayout<float2>.stride * uniformBufferIndex
let p = randomBuffer!.contents().advanced(by: randomBufferOffset)
var random = p.bindMemory(to: float2.self, capacity: 1)
for _ in 0..<256 {
random.pointee = float2(Float(drand48()), Float(drand48()) )
random = random.advanced(by: 1)
}
randomBuffer?.didModifyRange(randomBufferOffset..<(randomBufferOffset + 256 * MemoryLayout<float2>.stride))
uniformBufferIndex = (uniformBufferIndex + 1) % maxFramesInFlight
}
}
extension Renderer: MTKViewDelegate {
func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
self.size = size
let rayCount = Int(size.width * size.height)
rayBuffer = device.makeBuffer(length: rayStride * rayCount, options: .storageModePrivate)
shadowRayBuffer = device.makeBuffer(length: rayStride * rayCount, options: .storageModePrivate)
intersectionBuffer = device.makeBuffer(length: intersectionStride * rayCount, options: .storageModePrivate)
let renderTargetDescriptor = MTLTextureDescriptor()
renderTargetDescriptor.pixelFormat = .rgba32Float
renderTargetDescriptor.textureType = .type2D
renderTargetDescriptor.width = Int(size.width)
renderTargetDescriptor.height = Int(size.height)
renderTargetDescriptor.storageMode = .private
renderTargetDescriptor.usage = [.shaderRead, .shaderWrite]
renderTarget = device.makeTexture(descriptor: renderTargetDescriptor)
accumulationTarget = device.makeTexture(descriptor: renderTargetDescriptor)
frameIndex = 0
}
func draw(in view: MTKView) {
semaphore.wait()
guard let commandBuffer = queue.makeCommandBuffer() else {
return
}
commandBuffer.addCompletedHandler { cb in
self.semaphore.signal()
}
updateUniforms()
let width = Int(size.width)
let height = Int(size.height)
let threadsPerThreadgroup = MTLSizeMake(8, 8, 1)
let threadgroups = MTLSizeMake(
(width + threadsPerThreadgroup.width - 1) / threadsPerThreadgroup.width,
(height + threadsPerThreadgroup.height - 1) / threadsPerThreadgroup.height,
1)
// 1st compute pipeline - Primary Rays
var computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.setBuffer(uniformBuffer, offset: uniformBufferOffset, index: 0)
computeEncoder?.setBuffer(rayBuffer, offset: 0, index: 1)
computeEncoder?.setBuffer(randomBuffer, offset: randomBufferOffset, index: 2)
computeEncoder?.setTexture(renderTarget, index: 0)
computeEncoder?.setComputePipelineState(rayPipeline!)
computeEncoder?.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder?.endEncoding()
// 2nd and 3rd compute pipelines
for _ in 0..<3 {
// Shade pipeline
intersector?.intersectionDataType = .distancePrimitiveIndexCoordinates
intersector?.encodeIntersection(commandBuffer: commandBuffer,
intersectionType: .nearest,
rayBuffer: rayBuffer!,
rayBufferOffset: 0,
intersectionBuffer: intersectionBuffer!,
intersectionBufferOffset: 0,
rayCount: width * height,
accelerationStructure: accelerationStructure!)
computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.setBuffer(uniformBuffer, offset: uniformBufferOffset, index: 0)
computeEncoder?.setBuffer(rayBuffer, offset: 0, index: 1)
computeEncoder?.setBuffer(shadowRayBuffer, offset: 0, index: 2)
computeEncoder?.setBuffer(intersectionBuffer, offset: 0, index: 3)
computeEncoder?.setBuffer(vertexColorBuffer, offset: 0, index: 4)
computeEncoder?.setBuffer(vertexNormalBuffer, offset: 0, index: 5)
computeEncoder?.setBuffer(randomBuffer, offset: randomBufferOffset, index: 6)
computeEncoder?.setBuffer(triangleMaskBuffer, offset: 0, index: 7)
computeEncoder?.setTexture(renderTarget, index: 0)
computeEncoder?.setComputePipelineState(shadePipeline!)
computeEncoder?.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder?.endEncoding()
// Shadows pipeline
intersector?.intersectionDataType = .distance
intersector?.encodeIntersection(commandBuffer: commandBuffer,
intersectionType: .any,
rayBuffer: shadowRayBuffer!,
rayBufferOffset: 0,
intersectionBuffer: intersectionBuffer!,
intersectionBufferOffset: 0,
rayCount: width * height,
accelerationStructure: accelerationStructure!)
computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.setBuffer(uniformBuffer, offset: uniformBufferOffset, index: 0)
computeEncoder?.setBuffer(shadowRayBuffer, offset: 0, index: 1)
computeEncoder?.setBuffer(intersectionBuffer, offset: 0, index: 2)
computeEncoder?.setTexture(renderTarget, index: 0)
computeEncoder?.setComputePipelineState(shadowPipeline!)
computeEncoder?.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder?.endEncoding()
}
// 4th compute pipeline
computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.setBuffer(uniformBuffer, offset: uniformBufferOffset, index: 0)
computeEncoder?.setTexture(renderTarget, index: 0)
computeEncoder?.setTexture(accumulationTarget, index: 1)
computeEncoder?.setComputePipelineState(accumulatePipeline!)
computeEncoder?.dispatchThreadgroups(threadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder?.endEncoding()
guard let descriptor = view.currentRenderPassDescriptor,
let renderEncoder = commandBuffer.makeRenderCommandEncoder(
descriptor: descriptor) else {
return
}
renderEncoder.setRenderPipelineState(copyPipeline!)
renderEncoder.setFragmentTexture(accumulationTarget, index: 0)
renderEncoder.drawPrimitives(type: .triangle, vertexStart: 0, vertexCount: 6)
renderEncoder.endEncoding()
guard let drawable = view.currentDrawable else {
return
}
commandBuffer.present(drawable)
commandBuffer.commit()
}
}