Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
GPGPU with METAL in Swift
#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];
}
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.
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.