metal-examples-tutorials/ch10/ch10new.playground/Contents.swift

104 lines
4.2 KiB
Swift

import MetalKit
import PlaygroundSupport
guard let device = MTLCreateSystemDefaultDevice(),
let commandQueue = device.makeCommandQueue() else {
fatalError("Metal is not supported on this device")
}
let textureDescriptor = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: .rgba8Unorm,
width: 600,
height: 600,
mipmapped: false)
textureDescriptor.usage = [.shaderWrite, .shaderRead]
guard let texture = device.makeTexture(descriptor: textureDescriptor) else {
fatalError("Failed to create texture")
}
let shaderSource = """
#include <metal_stdlib>
using namespace metal;
float dist(float2 point, float2 center, float radius)
{
return length(point - center) - radius;
}
kernel void compute(texture2d<float, access::write> output [[texture(0)]],
uint2 gid [[thread_position_in_grid]])
{
int width = output.get_width();
int height = output.get_height();
float2 uv = float2(gid) / float2(width, height);
uv = uv * 2.0 - 1.0;
float distToCircle = dist(uv, float2(0), 0.5);
float distToCircle2 = dist(uv, float2(-0.1, 0.1), 0.5);
bool inside = distToCircle2 < 0;
output.write(inside ? float4(0) : float4(1, 0.7, 0, 1) * (1 - distToCircle), gid);
}
"""
guard let library = try? device.makeLibrary(source: shaderSource, options: nil),
let kernelFunction = library.makeFunction(name: "compute"),
let computePipelineState = try? device.makeComputePipelineState(function: kernelFunction) else {
fatalError("Failed to create compute pipeline")
}
guard let commandBuffer = commandQueue.makeCommandBuffer(),
let computeEncoder = commandBuffer.makeComputeCommandEncoder() else {
fatalError("Failed to create command buffer or compute encoder")
}
let w = computePipelineState.threadExecutionWidth
let h = computePipelineState.maxTotalThreadsPerThreadgroup / w
let threadsPerThreadgroup = MTLSize(width: w, height: h, depth: 1)
let threadgroupsPerGrid = MTLSize(width: (texture.width + w - 1) / w,
height: (texture.height + h - 1) / h,
depth: 1)
computeEncoder.setComputePipelineState(computePipelineState)
computeEncoder.setTexture(texture, index: 0)
computeEncoder.dispatchThreadgroups(threadgroupsPerGrid, threadsPerThreadgroup: threadsPerThreadgroup)
computeEncoder.endEncoding()
commandBuffer.commit()
commandBuffer.waitUntilCompleted()
class MetalView: MTKView {
var texture: MTLTexture!
init(frame frameRect: CGRect, device: MTLDevice?, texture: MTLTexture?) {
self.texture = texture
super.init(frame: frameRect, device: device)
self.isPaused = true
self.enableSetNeedsDisplay = true
self.drawableSize = CGSize(width: 600, height: 600) // Set the drawable size here
}
required init(coder: NSCoder) { fatalError("init(coder:) has not been implemented") }
override func draw(_ rect: CGRect) {
guard let drawable = self.currentDrawable else { return }
if drawable.texture.width != texture.width || drawable.texture.height != texture.height {
print("Drawable texture size does not match Metal texture size.")
return
}
guard let commandBuffer = device?.makeCommandQueue()?.makeCommandBuffer(),
let blitEncoder = commandBuffer.makeBlitCommandEncoder() else { return }
blitEncoder.copy(from: texture,
sourceSlice: 0,
sourceLevel: 0,
sourceOrigin: MTLOrigin(x: 0, y: 0, z: 0),
sourceSize: MTLSize(width: texture.width, height: texture.height, depth: 1),
to: drawable.texture,
destinationSlice: 0,
destinationLevel: 0,
destinationOrigin: MTLOrigin(x: 0, y: 0, z: 0))
blitEncoder.endEncoding()
commandBuffer.present(drawable)
commandBuffer.commit()
}
}
let view = MetalView(frame: CGRect(x: 0, y: 0, width: 600, height: 600), device: device, texture: texture)
PlaygroundPage.current.liveView = view