metalTexture.swift

//
//  main.swift
//  ComputeTexWrite
//
//  Created by Alexander Tarasikov on 23/12/2016.
//  Copyright © 2016 test. All rights reserved.
//

import Foundation
import Metal

func byteSize(array : [Int32] ) -> Int {
    return array.count * MemoryLayout.size(ofValue: array[0])
}

func runMetal() {
    print("Hello, World!")
    let dataLength = 1 << 8
    let useTexture = true
    
    var input = [Int32](repeating:0, count: dataLength)
    for i in (0...(dataLength - 1)) {
        input[i] = Int32(arc4random_uniform(1024))
    }
    let dataByteLength = dataLength * MemoryLayout.size(ofValue: input[0])
    
    /**
     * The buffer to store the "stage" and "pass" indices (arguments for the kernel).
     */
    var paramData = [Int32](repeating:0, count: 2)
    
    let device = MTLCreateSystemDefaultDevice()
    print("Metal Device:", device as Any)
    
    let defaultLibrary = device?.newDefaultLibrary()
    let sigmoidFunction = defaultLibrary?.makeFunction(name: "stripe_Kernel")
    
    let cmdQueue = device?.makeCommandQueue()
    
    let width = 1024
    let height = 768
    
    let threadsPerGroup = MTLSize(width: 8, height: 8, depth: 1)
    let numThreadGroups = MTLSize(width: width / threadsPerGroup.width, height: height / threadsPerGroup.height, depth: 1)
    
    let inBuffer = device?.makeBuffer(bytes: &input, length:dataByteLength, options: MTLResourceOptions())
    let paramBuffer = device?.makeBuffer(bytes: &paramData, length:byteSize(array:paramData), options: MTLResourceOptions.storageModeManaged)
    let texBuffer = device?.makeBuffer(length:4 * width * height, options: MTLResourceOptions.storageModeShared)
    
    var pipeline : MTLComputePipelineState? = nil
    do {
        pipeline = try device?.makeComputePipelineState(function: sigmoidFunction!)
    }
    catch {
        print("Failed to make compute pipeline")
        return
    }
    
    var textureDesc = MTLTextureDescriptor.texture2DDescriptor(pixelFormat: MTLPixelFormat.r32Float, width: width, height: height, mipmapped: false)
    textureDesc.usage = MTLTextureUsage.shaderWrite
    let metalTexture = device?.makeTexture(descriptor: textureDesc)
    
    paramData[0] = Int32(0)
    paramData[1] = Int32(0)
    memcpy(paramBuffer?.contents(), &paramData, byteSize(array: paramData))
    paramBuffer?.didModifyRange(NSMakeRange(0, byteSize(array: paramData)))
    
    let cmdBuffer = cmdQueue?.makeCommandBuffer()
    let cmdEncoder = cmdBuffer?.makeComputeCommandEncoder()
    cmdEncoder?.setComputePipelineState(pipeline!)
    cmdEncoder?.setTexture(metalTexture, at: 0)
    cmdEncoder?.setBuffer(inBuffer, offset: 0, at: 0)
    cmdEncoder?.setBuffer(paramBuffer, offset: 0, at: 1)
    cmdEncoder?.setBuffer(paramBuffer, offset: 0, at: 2)
    cmdEncoder?.setBuffer(texBuffer, offset:0, at:3)
    cmdEncoder?.dispatchThreadgroups(numThreadGroups, threadsPerThreadgroup: threadsPerGroup)
    cmdEncoder?.endEncoding()
    
    let blitEncoder = cmdBuffer?.makeBlitCommandEncoder()
    blitEncoder?.synchronize(texture: metalTexture!, slice: 0, level: 0)
    blitEncoder?.endEncoding()
    
    cmdBuffer?.commit()
    cmdBuffer?.waitUntilCompleted()
    
    let textureBufferSize = 4 * width * height
    var result = [Float](repeating:0, count: textureBufferSize)
    
    if useTexture {
        let region = MTLRegionMake2D(0, 0, width, height)
        metalTexture?.getBytes(&result, bytesPerRow: width * 4, from: region, mipmapLevel: 0)
    }
    else {
        let data = NSData(bytesNoCopy: texBuffer!.contents(), length: textureBufferSize, freeWhenDone: false)
        data.getBytes(&result, length: textureBufferSize)
    }
    
    print("Done")
    for item in result {
        if (item != 0.0) {
            print(String(format:"%4.4f", item))
        }
    }
}

runMetal()

TestShader.metal

#include <metal_stdlib>
using namespace metal;

kernel void stripe_Kernel(texture2d<float, access::write> outTexture [[ texture(0) ]],
                          device const float *inBuffer [[ buffer(0) ]],
                          device const ushort * imageWidth [[ buffer(1) ]],
                          device const ushort * imageHeight [[ buffer(2) ]],
                          device float *outBuffer [[ buffer(3) ]],
                          uint2 gid [[ thread_position_in_grid ]])
{
    const ushort imageW = *imageWidth;
    const ushort imageH = *imageHeight;
    
    const uint32_t textureW = outTexture.get_width();  // eg. 2048
    
    uint32_t posX = gid.x;  // eg. 0...2047
    uint32_t posY = gid.y;  // eg. 0...895
    
    uint32_t sourceX = ((int)(posY/imageH)*textureW + posX) % imageW;
    uint32_t sourceY = (int)(posY% imageH);
    
    const uint32_t ptr = (sourceX + sourceY* imageW);
    float pixel = inBuffer[ptr];
    
    outTexture.write(42.0f, gid);
    outTexture.write(44.0f, uint2(0, 0));
    
    outBuffer[0] = 33.0f;
    outBuffer[1] = 11.0f;
}