MetalSmith is a helper library for OSX and iOS apps that use Metal. Quickly prototype and iterate on both render and compute kernels, in the context of your own, real, application. MetalSmith brings the Metal edit-compile-execute cycle down to less than half a second -- all without leaving XCode. Additionally, it provides a handful of SwiftUI components to easily invoke, debug, and visualize your shaders, buffers, and textures. Finally, it provides a code generation tool, similar to Sourcery, that reflects on your metal shader code and automatically generate Swift: to invoke kernels in a type-safe-ish way, to generate MTLVertexDescriptors, and so forth.

Because the library is very new, I recommend cloning the repo locally so that you can make enhancements.

Code generation

You do NOT need to use the code-gen tool to use the "instant" compiles, but this is a good place to start to understand how this all works. Consider a simple compute kernel like the following:

kernel void add_arrays(device const float* inA,
                       device const float* inB,
                       device float* result,
                       uint gid [[thread_position_in_grid]])
{
    result[gid] = inA[gid] + inB[gid];
}

The code-gen tool can automatically generate Swift bindings for this function. An example template is provided, but you can make your own. The easiest way to invoke the code-gen tool is to add a new "Run Script" "Build Phase" to your XCode target. This is quite similar to the process for Sourcery, see this link for detailed instructions. NOTE: Make sure your project has added MetalSmith swift package as a dependency, including the "mtlsmith" target. (There is no need to install a separate command line tool thanks to XCode 11). Finally, add the following as your "Run Script":

"$TARGET_BUILD_DIR/../$CONFIGURATION/mtlsmith" --sources $SRCROOT --templates $SRCROOT/XXX/Templates --output $SRCROOT/XXX/AutoGenerated

Now, simply building your project will generate a class called AddArraysEncoder, which you can use like this:

// The usual Metal setup code. (This can be a automated away too, see further below)
let device = MTLCreateSystemDefaultDevice()
let library = device?.makeDefaultLibrary()
let commandQueue = device?.makeCommandQueue()

// Create some test data
let inA_: [Float] = [1,2,3]
let inB_: [Float] = [4,5,6]
let count = inA_.count
let inA    = device?.makeBuffer(bytes: inA_, length: MemoryLayout<Float>.stride * count)
let inB    = device?.makeBuffer(bytes: inB_, length: MemoryLayout<Float>.stride * count)
let result = device?.makeBuffer(length: MemoryLayout<Float>.stride * count)

// AND HERE is the good stuff:

let addArrays = AddArraysEncoder(library: library)
let commandBuffer = commandQueue?.makeCommandBuffer()

addArrays.commandBuffer(commandBuffer)
    .inA(inA)
    .inB(inB)
    .result(result)
    .dispatch(width: count)

Again, you should take my code generation template just as a starting point since it's mostly a proof of concept. But it does "prove" that statically guaranteed named parameters are possible, eliminating the need for tedious argument index constants, and it does provide some type-safety as well. It eliminates the boiler-plate of setting up the command pipeline and so forth. Whether this is all worth the complexity of building a code-generation tool is a question that has yet to be answered. I think for compute kernels, it will be worthwhile when this code is mature. But for render kernels, I'm skeptical.

Visualizing your computation

Whether your kernel works with textures or buffers (as add_array), MetalSmith provides SwiftUI wrappers for visualizing your computation. This is most-useful for setting up a "playground" or "REPL" for iterating on your kernel. The code below sets up a preview canvas that visualizes three buffers (inA, inB, result) before the computation. It then invokes add_arrays twice with different arguments. Finally it displays the results.

struct AddArrays_Previews: PreviewProvider {
    static var previews: some View {
        let count       = 10
        let environment = MetalEnvironment()

        let inA: [Float] = Array(0..<count).map { Float($0) }
        let inB: [Float] = inA.map { $0 * 4 }

        let result = environment.device!.makeBuffer(length: MemoryLayout<Float>.stride * count, options: .storageModeShared)
        result?.label = "result"
        
        return Group {
            HStack {
                Buffer(inA)
                Buffer(inB)
                Buffer(result, of: Float.self, count: count).copy() // since result is mutated, we make a COPY before
            }
            .previewDisplayName("Buffers before function invocation")
            CommandBuffer() {
                AddArrays()
                    .inA(inA)
                    .inB(inB)
                    .result(result)
                    .dispatch(width: count)
                AddArrays()
                    .inA(result)
                    .inB(inB)
                    .result(result)
                    .dispatch(width: count)
            }
            Buffer(result, of: Float.self, count: count)
                .previewDisplayName("Result buffer after invocation")
        }
        .previewLayout(.sizeThatFits)
        .environmentObject(environment)
    }
}

In the above example `AddArrays` is a thin-wrapper class around the `AddArraysEncoder` class we saw earlier. It's also made with just code-generation. It conforms to SwiftUI's `View` protocol. Note, additionally, that we're passing in vanilla `[Float]` arrays for `inA` and `inB`. These are just convenience functions for testing your code in the SwiftUI preview canvas.

The Buffer View class visualizes MTLBuffers and UnsafeBufferPointers. It can visualize buffers of scalars like Float and (using some Swift-reflection-voodoo) buffers of structs, displaying their fields. Similarly there is a Texture view class that visualizes MTLTextures.

"Instant" compilation aka live-coding

Since the SwiftUI preview canvas ALREADY drastically reduces compilation and app-start times, you can already get very fast feedback using what just was presented above. Hitting ⌘S in shader code will pause the preview canvas, but this can be restarted by pressing ⌘⌥P -- the whole process takes probably 3 seconds. However, to get sub-second compilation, a more complicated solution is needed. We're going to use another shader in this example, just for the sake of variety.

Here is the shader:

kernel void colors(
                   const texture2d<float, access::write> image,
                   const constant float &time,
                   const uint2 gid [[thread_position_in_grid]])
{
    float2 uv = -1. + 2. * float2(gid) / float2(image.get_width(), image.get_height());
    float4 col = float4(
                        abs(sin(cos(time+3.*uv.y)*2.*uv.x+time)),
                        abs(cos(sin(time+2.*uv.x)*3.*uv.y+time)),
                        0.1,
                        1.0);
    image.write(col, gid);
}

We can set up a view to animate this shader by incrementing a clock (driven by an MTKViewDelegate.draw draw callback under the hood):

struct DynamicLibraryLoading: View {
    @EnvironmentObject var environment: MetalEnvironment

    var body: some View {
        let image = environment.device?.makeTexture(width: 480, height: 640, usage: [.shaderRead, .shaderWrite])
        var clock: Float = 1.0

        return Group {
            Texture(image)
                .onDraw { _ in
                    clock += 0.01
                    return CommandBuffer {
                        Colors().time(clock).image(image).dispatch(width: 480, height: 640)
                    }
            }
        }
    }
}

Next, we need to set up the MetalEnvironment to watch our source code for changes:

struct DynamicLibraryLoading_Previews: PreviewProvider {
    static var previews: some View {
        var environment = MetalEnvironment()
        if let srcroot = ProcessInfo.processInfo.environment["SRCROOT"] {
            environment = environment.watch(srcroot)
        }

        return DynamicLibraryLoading()
            .environmentObject(environment)
    }
}

Finally, ensure SRCROOT is available to your app, as described here.

That's basically it, but there are caveats. Firstly, your application needs to be an OSX application -- not iOS or Catalyst, because it shells out to the metal compiler. Secondly, while the above example will work in the SwiftUI preview canvas -- the animating texture will update as soon as you hit ⌘S in a .metal file -- if you want other views to update (such as the visualizations of buffers in the AddArray example above), you have to actually RUN the application. With the application running, everytime you save a metal file, the entire view that depends on the MetalEnvironment will be refreshed.

Given the these limitations, I think the most practical approach is to create a new OSX target in your project. Then you can use all this regardless of whether your main app is iOS or Catalyst or whatever. In any case, make sure to turn off the App Sandbox, or the application won't be allowed to read from your documents directory.

Acknowledgements

Much of the code is based on Sourcery!