How to pass [[Int]] into MTLBuffer in Swift and receive it in Metal?

Question

What I'm trying to do

I have an [[Int]] in Swift and I'm trying to pass it into Metal through a buffer. Basically, I'm trying to use Metal to add two matrices multithreaded and give them back to Swift. So far, this is more difficult than expected.

The problem

Metal tells me that my graph isn't a pointer and I suspect I'm passing it into Metal wrong.

import MetalKit

let graph: [[Int]] = [
    [0, 1, 2, 999, 999, 999],
    [1, 0, 999, 5, 1, 999],
    [2, 999, 0, 2, 3, 999],
    [999, 5, 2, 0, 2, 2],
    [999, 1, 3, 2, 0, 1],
    [999, 999, 999, 2, 1, 0]]

func fooFunc(gra: [[Int]]) {
    
    let count = gra.count

    let device = MTLCreateSystemDefaultDevice()

    let commandQueue = device?.makeCommandQueue()

    let gpuFunctionLibrary = device?.makeDefaultLibrary()

    let funcGPUFunction = gpuFunctionLibrary?.makeFunction(name: "MetalFunc")

    var funcPipelineState: MTLComputePipelineState!
    do {
        funcPipelineState = try device?.makeComputePipelineState(function: funcGPUFunction!)
    } catch {
      print(error)
    }

    let graphBuff = device?.makeBuffer(bytes: gra,
                                      length: MemoryLayout<Int>.size * count * count,
                                      options: .storageModeShared)
    let resultBuff = device?.makeBuffer(length: MemoryLayout<Int>.size * count,
                                        options: .storageModeShared)

    let commandBuffer = commandQueue?.makeCommandBuffer()

    let commandEncoder = commandBuffer?.makeComputeCommandEncoder()
    commandEncoder?.setComputePipelineState(additionComputePipelineState)

    commandEncoder?.setBuffer(graphBuff, offset: 0, index: 0)
    commandEncoder?.setBuffer(resultBuff, offset: 0, index: 1)

    let threadsPerGrid = MTLSize(width: count, height: 1, depth: 1)
    let maxThreadsPerThreadgroup = additionComputePipelineState.maxTotalThreadsPerThreadgroup // 1024
    let threadsPerThreadgroup = MTLSize(width: maxThreadsPerThreadgroup, height: 1, depth: 1)
    commandEncoder?.dispatchThreads(threadsPerGrid,
                                    threadsPerThreadgroup: threadsPerThreadgroup)

    commandEncoder?.endEncoding()
    commandBuffer?.commit()
    commandBuffer?.waitUntilCompleted()
    
    let resultBufferPointer = resultBuff?.contents().bindMemory(to: Int.self,
                                                                capacity: MemoryLayout<Int>.size * count)

    print("Result: \(Int(resultBufferPointer!.pointee) as Any)")
}


gpuDijkstra(gra: graph)

#include <metal_stdlib>
using namespace metal;

kernel void MetalFunc(constant int *graph        [[ buffer(0) ]],
                      constant int *result   [[ buffer(1) ]],
                      uint   index [[ thread_position_in_grid ]]
{


    const int size = sizeof(*graph);
    int result[size][size];

    for(int k = 0; k<size; k++){
        result[index][k]=graph[index][k]+graph[index][k]; //ERROR: Subscripted value is not an array, pointer, or vector
    }
}

Answer 1

I don't know about the Swift implementation but your kernel function is wrong.

If you want to write to and read from the same buffer you should use Device address space attribute.

The device address space name refers to buffer memory objects allocated from the device memory pool that are both readable and writeable.

device int *result /* write and read */
constant int *graph /* read only */

This is the correct implementation:

    #include <metal_stdlib>
    using namespace metal;
    
    kernel void MetalFunc(constant int *graph  [[ buffer(0) ]],
                          device int *result   [[ buffer(1) ]],
                          uint   index [[ thread_position_in_grid ]]
    {
    
         /* Point to an array (graph buffer) */
         constant int(*r)[6][6] = (constant int(*)[6][6])graph;
         /* Point to an array (result buffer) */
         device int(*w)[6][6] = (device int(*)[6][6])result;
         /* Read the first element */
         int i = (*r)[0][0];

         (*w)[0][0] = 777; /* Write to the first element */
         (*w)[5][5] = 999; /* Write to the last element */
    }