GPGPU with METAL in Swift

Shaders.metal

#include <metal_stdlib>

using namespace metal;

struct FloatVector {
    float x = 0;
    float y = 0;
    float z = 0;
    float w = 1;
};

kernel void transform(const device float4 *vectorIn [[ buffer(0) ]],
                    const device float4x4 &matrix [[ buffer(1) ]],
                    device FloatVector *vectorOut [[ buffer(2) ]],
                    uint id [[ thread_position_in_grid ]]) {
    
    float4 outV = matrix * vectorIn[id];
    
    vectorOut[id].x = outV[0];
    vectorOut[id].y = outV[1];
    vectorOut[id].z = outV[2];
    vectorOut[id].w = outV[3];
}

ViewController.Swift

import UIKit
import Metal

class ViewController: UIViewController {
    private var metalDevice: MTLDevice! = nil
    private var metalCommandQueue: MTLCommandQueue! = nil
    private var metalDefaultLibrary: MTLLibrary! = nil
    private var metalFunction: MTLFunction! = nil
    
    struct FloatVector {
        var x: Float = 0
        var y: Float = 0
        var z: Float = 0
        var w: Float = 1
        
        var array: [Float] {
            get {
                return [x, y, z, w]
            }
        }
    }
    
    override func viewDidLoad() {
        super.viewDidLoad()
        
        self.initMetal()

        var myVectors = [FloatVector]()
        
        var myMatrix: [Float] = [
            0, 0, 0, 1,
            0, 0, 1, 0,
            0, 1, 0, 0,
            1, 0, 0, 0]
        
        for var i: Int = 0; i < 50; ++i {
            var vector = FloatVector()
            vector.x = Float(i)
            vector.y = Float(i + 1)
            vector.z = Float(i + 2)
            myVectors.append(vector)
        }

        println("input vectors")
        for vector in myVectors {
            println(vector.array)
        }
        
        self.getVectorsWithMetalFrom(myVectors, transformMatrix: myMatrix)
    }
    
    private func initMetal() {
        self.metalDevice = MTLCreateSystemDefaultDevice()
        self.metalCommandQueue = self.metalDevice.newCommandQueue()
        self.metalDefaultLibrary = self.metalDevice.newDefaultLibrary()
        self.metalFunction = self.metalDefaultLibrary.newFunctionWithName("transform")
    }
    
    private func getVectorsWithMetalFrom(input: [FloatVector], transformMatrix: [Float]) -> [FloatVector] {
        let commandBuffer = self.metalCommandQueue.commandBuffer()
        var errorsPointer = NSErrorPointer()
        var computePipelineFilter = self.metalDevice.newComputePipelineStateWithFunction(
            self.metalFunction,
            error: errorsPointer)
        let computeCE = commandBuffer.computeCommandEncoder()
        computeCE.setComputePipelineState(computePipelineFilter!)
        
        var inputVectors = input
        var vectorsByteLength = input.count * sizeofValue(input[0])
        var inputBuffer = self.metalDevice.newBufferWithBytes(
            &inputVectors,
            length: vectorsByteLength,
            options: nil)
        inputBuffer.label = "inVector"
        
        var matrix = transformMatrix
        var matrixBuffer = self.metalDevice.newBufferWithBytes(
            &matrix,
            length: transformMatrix.count * sizeofValue(transformMatrix[0]),
            options: nil)
        matrixBuffer.label = "transformMatrix"
        
        var outputVectors = [FloatVector](count:input.count, repeatedValue: FloatVector())
        var outputBuffer = self.metalDevice.newBufferWithBytes(
            &outputVectors,
            length: vectorsByteLength,
            options: nil)
        outputBuffer.label = "outVector"
        
        computeCE.setBuffer(inputBuffer, offset: 0, atIndex: 0)
        computeCE.setBuffer(matrixBuffer, offset: 0, atIndex: 1)
        computeCE.setBuffer(outputBuffer, offset: 0, atIndex: 2)
        
        var threadsPerGroup = MTLSize(width: 32, height: 1, depth: 1)
        var numThreadgroups = MTLSize(width: (input.count + 31) / 32, height: 1, depth: 1)
        computeCE.dispatchThreadgroups(numThreadgroups, threadsPerThreadgroup: threadsPerGroup)
        
        computeCE.endEncoding()
        commandBuffer.commit()
        commandBuffer.waitUntilCompleted()

        var resultData = NSData(bytesNoCopy: outputBuffer.contents(),
        length: vectorsByteLength,
        freeWhenDone: false)
        
        resultData.getBytes(&outputVectors, length: vectorsByteLength)
        
        println("output vectors")
        for vector in outputVectors {
        println(vector.array)
        }
        
        return outputVectors
    }
    
    override func didReceiveMemoryWarning() {
        super.didReceiveMemoryWarning()
        // Dispose of any resources that can be recreated.
    }
}