metalkit / metal Goto Github PK

let shader =    
    "#include <metal_stdlib>\n" +
        "using namespace metal;" +
        "kernel void k(texture2d<float,access::write> o[[texture(0)]]," +
        "              uint2 gid[[thread_position_in_grid]]) {" +
        "   int width = o.get_width();" +
        "   int height = o.get_height();" +
        "   float2 uv = float2(gid) / float2(width, height);" +
        "   float3 color = mix(float3(1.0, 0.6, 0.1), float3(0.5, 0.8, 1.0), sqrt(1 - uv.y));" +
        "   float2 q = uv - float2(0.67, 0.25);" +
        "   float r = 0.2 + 0.1 * cos(atan2(q.x, q.y) * 9.0 + 20.0 * q.x);" +
        "   color *= smoothstep(r, r + 0.01, length(q));" +
        "   r = 0.03 + 0.002 * cos(120.0 * q.y) + exp(-50.0 * (1.0 - uv.y));" +
        "   color *= 1.0 - (1.0 - smoothstep(r, r + 0.002, abs(q.x - 0.25 * sin(2.0 * q.y)))) * smoothstep(0.0, 0.1, q.y);" +
        "   o.write(float4(color, 1.0), gid);" +
"}"

import MetalKit
import PlaygroundSupport

public class MetalView: NSObject, MTKViewDelegate {
    weak var view: MTKView!
    let commandQueue: MTLCommandQueue!
    let device: MTLDevice!
    let cps: MTLComputePipelineState!
    
    public init?(mtkView: MTKView, shader: String) {
        view = mtkView
        view.clearColor = MTLClearColorMake(0.5, 0.5, 0.5, 1)
        view.colorPixelFormat = .bgra8Unorm
        device = MTLCreateSystemDefaultDevice()!
        commandQueue = device.makeCommandQueue()
        let library = try! device.makeLibrary(source: shader, options: nil)
        let function = library.makeFunction(name:"k")!
        cps = try! device.makeComputePipelineState(function: function)
        
        super.init()
        view.delegate = self
        view.device = device
    }
    
    public func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {}
    
    public func draw(in view: MTKView) {
        if let drawable = view.currentDrawable,
            let commandBuffer = commandQueue.makeCommandBuffer(),
            let commandEncoder = commandBuffer.makeComputeCommandEncoder() {
            commandEncoder.setComputePipelineState(cps)
            commandEncoder.setTexture(drawable.texture, index: 0)
            let groups = MTLSize(width: Int(view.frame.width)/4, height: Int(view.frame.height)/4, depth: 1)
            let threads = MTLSize(width: 8, height: 8,depth: 1)
            commandEncoder.dispatchThreadgroups(groups,threadsPerThreadgroup: threads)
            commandEncoder.endEncoding()
            commandBuffer.present(drawable)
            commandBuffer.commit()
        }
    }
}

let frame = CGRect(x: 0, y: 0, width: 300, height: 300)
let view = MTKView(frame: frame)
let delegate = MetalView(mtkView: view, shader: shader)
view.delegate = delegate
PlaygroundPage.current.liveView = view