Metal compute shader example

compute.metal

#include <metal_stdlib>
using namespace metal;

kernel void add(const device float2 *in [[ buffer(0) ]],
                device float  *out [[ buffer(1) ]],
                uint id [[ thread_position_in_grid ]]) {
  out[id] = in[id].x + in[id].y;
}

compute.swift

import MetalKit

let device = MTLCreateSystemDefaultDevice()!
let commandQueue = device.makeCommandQueue()!
let library = try device.makeLibrary(filepath: "compute.metallib")

let commandBuffer = commandQueue.makeCommandBuffer()!
let encoder = commandBuffer.makeComputeCommandEncoder()!
encoder.setComputePipelineState(try device.makeComputePipelineState(function: library.makeFunction(name: "add")!))

let input: [Float] = [1.0, 2.0]
encoder.setBuffer(device.makeBuffer(bytes: input as [Float], length: MemoryLayout<Float>.stride * input.count, options: []),
                  offset: 0, index: 0)
let outputBuffer = device.makeBuffer(length: MemoryLayout<Float>.stride, options: [])!
encoder.setBuffer(outputBuffer, offset: 0, index: 1)

let numThreadgroups = MTLSize(width: 1, height: 1, depth: 1)
let threadsPerThreadgroup = MTLSize(width: 1, height: 1, depth: 1)
encoder.dispatchThreadgroups(numThreadgroups, threadsPerThreadgroup: threadsPerThreadgroup)
encoder.endEncoding()

commandBuffer.commit()
commandBuffer.waitUntilCompleted()

let result = outputBuffer.contents().load(as: Float.self)
print(String(format: "%f + %f = %f", input[0], input[1], result))

Makefile

SDK = xcrun -sdk macosx

all: compute.metallib compute

compute.metallib: Compute.metal
	# Metal intermediate representation (.air)
	$(SDK) metal -c -Wall -Wextra -std=osx-metal2.0 -o /tmp/Compute.air $^
	# Metal library (.metallib)
	$(SDK) metallib -o $@ /tmp/Compute.air

compute: compute.swift
	$(SDK) swiftc -o $@ $^

clean:
	rm -f compute.metallib compute