Metal Compute Shader Example in Swift and Objective C

add.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;
}

main.m

@import MetalKit;
#import <Foundation/Foundation.h>

int main(int argc, const char * argv[]) {
    @autoreleasepool {
        
        //----------------------------------------------------------------------
        // Setup
        id<MTLDevice> device = MTLCreateSystemDefaultDevice();
        id<MTLCommandQueue> commandQueue = [device newCommandQueue];
        id<MTLLibrary> library = [device newDefaultLibrary];
        id<MTLFunction> kernelFunction = [library newFunctionWithName:@"add"];
        //----------------------------------------------------------------------
        // pipeline
        NSError *error = NULL;
        [commandQueue commandBuffer];
        id<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];
        id<MTLComputeCommandEncoder> encoder = [commandBuffer computeCommandEncoder];
        [encoder setComputePipelineState:[device newComputePipelineStateWithFunction:kernelFunction error:&error]];
        //----------------------------------------------------------------------
        // Set Data
        float input[] = {1,2};
        NSInteger dataSize = sizeof(input);
        
        [encoder setBuffer:[device newBufferWithBytes:input length:dataSize options:0]
                    offset:0
                   atIndex:0];
        
        id<MTLBuffer> outputBuffer = [device newBufferWithLength:sizeof(float) options:0];
        [encoder setBuffer:outputBuffer offset:0 atIndex:1];
        //----------------------------------------------------------------------
        // Run Kernel
        MTLSize numThreadgroups = {1,1,1};
        MTLSize numgroups = {1,1,1};
        [encoder dispatchThreadgroups:numThreadgroups threadsPerThreadgroup:numgroups];
        [encoder endEncoding];
        [commandBuffer commit];
        [commandBuffer waitUntilCompleted];
        //----------------------------------------------------------------------
        // Results
        float *output = [outputBuffer contents];
        printf("result = %f\n", output[0]);
    }
    return 0;
}

main.swift

import MetalKit
//----------------------------------------------------------------------
// Setup
let device = MTLCreateSystemDefaultDevice()!
let commandQueue = device.makeCommandQueue()!
let library = try device.makeLibrary(filepath: "compute.metallib")
//----------------------------------------------------------------------
// pipeline
let commandBuffer = commandQueue.makeCommandBuffer()!
let encoder = commandBuffer.makeComputeCommandEncoder()!
encoder.setComputePipelineState(try device.makeComputePipelineState(function: library.makeFunction(name: "add")!))
//----------------------------------------------------------------------
// Set Data
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)
//----------------------------------------------------------------------
// Run Kernel
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()
//----------------------------------------------------------------------
// Results
let result = outputBuffer.contents().load(as: Float.self)
print(String(format: "%f + %f = %f", input[0], input[1], result))

Notes.md

Notes

Create a command line target in xcode either in Swift or Obj-C. add a metal file to the project and copy and paste the respective code into each file.