MetalByTutorials/21-metal-performance-shaders/projects/final/Raytracing/Raytracing Shared/Renderer.swift

427 lines
16 KiB
Swift

//
/**
* Copyright (c) 2018 Razeware LLC
*
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
*
* Notwithstanding the foregoing, you may not use, copy, modify, merge, publish,
* distribute, sublicense, create a derivative work, and/or sell copies of the
* Software in any work that is designed, intended, or marketed for pedagogical or
* instructional purposes related to programming, coding, application development,
* or information technology. Permission for such use, copying, modification,
* merger, publication, distribution, sublicensing, creation of derivative works,
* or sale is expressly withheld.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
import MetalKit
import MetalPerformanceShaders
typealias float2 = SIMD2<Float>
typealias float3 = SIMD3<Float>
class Renderer: NSObject {
var device: MTLDevice!
var commandQueue: MTLCommandQueue!
var library: MTLLibrary!
var renderPipeline: MTLRenderPipelineState!
var rayPipeline: MTLComputePipelineState!
var rayBuffer: MTLBuffer!
var shadowRayBuffer: MTLBuffer!
var shadePipelineState: MTLComputePipelineState!
var accumulatePipeline: MTLComputePipelineState!
var accumulationTarget: MTLTexture!
var accelerationStructure: MPSTriangleAccelerationStructure!
var shadowPipeline: MTLComputePipelineState!
var vertexPositionBuffer: MTLBuffer!
var vertexNormalBuffer: MTLBuffer!
var vertexColorBuffer: MTLBuffer!
var indexBuffer: MTLBuffer!
var uniformBuffer: MTLBuffer!
var randomBuffer: MTLBuffer!
var intersectionBuffer: MTLBuffer!
let intersectionStride =
MemoryLayout<MPSIntersectionDistancePrimitiveIndexCoordinates>.stride
var intersector: MPSRayIntersector!
let rayStride =
MemoryLayout<MPSRayOriginMinDistanceDirectionMaxDistance>.stride
+ MemoryLayout<float3>.stride
let maxFramesInFlight = 3
let alignedUniformsSize = (MemoryLayout<Uniforms>.size + 255) & ~255
var semaphore: DispatchSemaphore!
var size = CGSize.zero
var randomBufferOffset = 0
var uniformBufferOffset = 0
var uniformBufferIndex = 0
var frameIndex: uint = 0
var renderTarget: MTLTexture!
lazy var vertexDescriptor: MDLVertexDescriptor = {
let vertexDescriptor = MDLVertexDescriptor()
vertexDescriptor.attributes[0] =
MDLVertexAttribute(name: MDLVertexAttributePosition,
format: .float3,
offset: 0, bufferIndex: 0)
vertexDescriptor.attributes[1] =
MDLVertexAttribute(name: MDLVertexAttributeNormal,
format: .float2,
offset: 0, bufferIndex: 1)
vertexDescriptor.layouts[0] = MDLVertexBufferLayout(stride: MemoryLayout<float3>.stride)
vertexDescriptor.layouts[1] = MDLVertexBufferLayout(stride: MemoryLayout<float3>.stride)
return vertexDescriptor
}()
var vertices: [float3] = []
var normals: [float3] = []
var colors: [float3] = []
init(metalView: MTKView) {
guard let device = MTLCreateSystemDefaultDevice() else {
fatalError("GPU not available")
}
metalView.device = device
metalView.colorPixelFormat = .rgba16Float
metalView.sampleCount = 1
metalView.drawableSize = metalView.frame.size
self.device = device
commandQueue = device.makeCommandQueue()!
library = device.makeDefaultLibrary()
super.init()
metalView.delegate = self
mtkView(metalView, drawableSizeWillChange: metalView.bounds.size)
semaphore = DispatchSemaphore.init(value: maxFramesInFlight)
buildPipelines(view: metalView)
createScene()
createBuffers()
buildIntersector()
buildAccelerationStructure()
}
func buildAccelerationStructure() {
accelerationStructure =
MPSTriangleAccelerationStructure(device: device)
accelerationStructure?.vertexBuffer = vertexPositionBuffer
accelerationStructure?.triangleCount = vertices.count / 3
accelerationStructure?.rebuild()
}
func buildIntersector() {
intersector = MPSRayIntersector(device: device)
intersector?.rayDataType = .originMinDistanceDirectionMaxDistance
intersector?.rayStride = rayStride
}
func buildPipelines(view: MTKView) {
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
let computeDescriptor = MTLComputePipelineDescriptor()
computeDescriptor.threadGroupSizeIsMultipleOfThreadExecutionWidth = true
do {
computeDescriptor.computeFunction = library.makeFunction(
name: "shadowKernel")
shadowPipeline = try device.makeComputePipelineState(
descriptor: computeDescriptor,
options: [],
reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(
name: "shadeKernel")
shadePipelineState = try device.makeComputePipelineState(
descriptor: computeDescriptor,
options: [],
reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(
name: "accumulateKernel")
accumulatePipeline = try device.makeComputePipelineState(
descriptor: computeDescriptor,
options: [],
reflection: nil)
computeDescriptor.computeFunction = library.makeFunction(
name: "primaryRays")
rayPipeline = try device.makeComputePipelineState(
descriptor: computeDescriptor,
options: [],
reflection: nil)
renderPipeline = try device.makeRenderPipelineState(descriptor: pipelineDescriptor)
} catch {
print(error.localizedDescription)
}
}
func createScene() {
loadAsset(name: "train", position: [-0.3, 0, 0.4], scale: 0.5)
loadAsset(name: "treefir", position: [0.5, 0, -0.2], scale: 0.7)
loadAsset(name: "plane", position: [0, 0, 0], scale: 10)
loadAsset(name: "sphere", position: [-1.9, 0.0, 0.3], scale: 1)
loadAsset(name: "sphere", position: [2.9, 0.0, -0.5], scale: 2)
loadAsset(name: "plane-back", position: [0, 0, -1.5], scale: 10)
}
func createBuffers() {
let uniformBufferSize = alignedUniformsSize * maxFramesInFlight
let options: MTLResourceOptions = {
#if os(iOS)
return .storageModeShared
#else
return .storageModeManaged
#endif
} ()
uniformBuffer = device.makeBuffer(length: uniformBufferSize, options: options)
randomBuffer = device.makeBuffer(length: 256 * MemoryLayout<float2>.stride * maxFramesInFlight, options: options)
vertexPositionBuffer = device.makeBuffer(bytes: &vertices, length: vertices.count * MemoryLayout<float3>.stride, options: options)
vertexColorBuffer = device.makeBuffer(bytes: &colors, length: colors.count * MemoryLayout<float3>.stride, options: options)
vertexNormalBuffer = device.makeBuffer(bytes: &normals, length: normals.count * MemoryLayout<float3>.stride, options: options)
}
func update() {
updateUniforms()
updateRandomBuffer()
uniformBufferIndex = (uniformBufferIndex + 1) % maxFramesInFlight
}
func updateUniforms() {
uniformBufferOffset = alignedUniformsSize * uniformBufferIndex
let pointer = uniformBuffer!.contents().advanced(by: uniformBufferOffset)
let uniforms = pointer.bindMemory(to: Uniforms.self, capacity: 1)
var camera = Camera()
camera.position = float3(0.0, 1.0, 3.38)
camera.forward = float3(0.0, 0.0, -1.0)
camera.right = float3(1.0, 0.0, 0.0)
camera.up = float3(0.0, 1.0, 0.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
camera.right *= imagePlaneWidth
camera.up *= imagePlaneHeight
var light = AreaLight()
light.position = float3(0.0, 1.98, 0.0)
light.forward = float3(0.0, -1.0, 0.0)
light.right = float3(0.25, 0.0, 0.0)
light.up = float3(0.0, 0.0, 0.25)
light.color = float3(4.0, 4.0, 4.0)
uniforms.pointee.camera = camera
uniforms.pointee.light = light
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
#if os(OSX)
uniformBuffer?.didModifyRange(uniformBufferOffset..<(uniformBufferOffset + alignedUniformsSize))
#endif
}
func updateRandomBuffer() {
randomBufferOffset = 256 * MemoryLayout<float2>.stride * uniformBufferIndex
let pointer = randomBuffer!.contents().advanced(by: randomBufferOffset)
var random = pointer.bindMemory(to: float2.self, capacity: 256)
for _ in 0..<256 {
random.pointee = float2(Float(drand48()), Float(drand48()) )
random = random.advanced(by: 1)
}
#if os(OSX)
randomBuffer?.didModifyRange(randomBufferOffset..<(randomBufferOffset + 256 * MemoryLayout<float2>.stride))
#endif
}
}
extension Renderer: MTKViewDelegate {
func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {
self.size = size
frameIndex = 0
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)
let rayCount = Int(size.width * size.height)
rayBuffer = device.makeBuffer(length: rayStride * rayCount,
options: .storageModePrivate)
shadowRayBuffer = device.makeBuffer(length: rayStride * rayCount,
options: .storageModePrivate)
accumulationTarget = device.makeTexture(
descriptor: renderTargetDescriptor)
intersectionBuffer = device.makeBuffer(
length: intersectionStride * rayCount,
options: .storageModePrivate)
}
func draw(in view: MTKView) {
semaphore.wait()
guard let commandBuffer = commandQueue.makeCommandBuffer() else {
return
}
commandBuffer.addCompletedHandler { cb in
self.semaphore.signal()
}
update()
// MARK: generate rays
let width = Int(size.width)
let height = Int(size.height)
let threadsPerGroup = MTLSizeMake(8, 8, 1)
let threadGroups = MTLSizeMake((width + threadsPerGroup.width - 1)
/ threadsPerGroup.width,
(height + threadsPerGroup.height - 1)
/ threadsPerGroup.height,
1)
var computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.label = "Generate Rays"
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: threadsPerGroup)
computeEncoder?.endEncoding()
for _ in 0..<3 {
// MARK: generate intersections between rays and model triangles
intersector?.intersectionDataType = .distancePrimitiveIndexCoordinates
intersector?.encodeIntersection(
commandBuffer: commandBuffer,
intersectionType: .nearest,
rayBuffer: rayBuffer,
rayBufferOffset: 0,
intersectionBuffer: intersectionBuffer,
intersectionBufferOffset: 0,
rayCount: width * height,
accelerationStructure: accelerationStructure)
// MARK: shading
computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.label = "Shading"
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?.setTexture(renderTarget, index: 0)
computeEncoder?.setComputePipelineState(shadePipelineState!)
computeEncoder?.dispatchThreadgroups(
threadGroups,
threadsPerThreadgroup: threadsPerGroup)
computeEncoder?.endEncoding()
// MARK: shadows
intersector?.label = "Shadows Intersector"
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?.label = "Shadows"
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: threadsPerGroup)
computeEncoder?.endEncoding()
}
// MARK: accumulation
computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.label = "Accumulation"
computeEncoder?.setBuffer(uniformBuffer, offset: uniformBufferOffset,
index: 0)
computeEncoder?.setTexture(renderTarget, index: 0)
computeEncoder?.setTexture(accumulationTarget, index: 1)
computeEncoder?.setComputePipelineState(accumulatePipeline)
computeEncoder?.dispatchThreadgroups(threadGroups,
threadsPerThreadgroup: threadsPerGroup)
computeEncoder?.endEncoding()
guard let descriptor = view.currentRenderPassDescriptor,
let renderEncoder = commandBuffer.makeRenderCommandEncoder(
descriptor: descriptor) else {
return
}
renderEncoder.setRenderPipelineState(renderPipeline!)
// MARK: draw call
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()
}
}