Vector Addition: CUDA & Metal

#include <cstdio>
#include <vector>
#include <cuda_runtime.h>

__global__ void vec_add(const float* __restrict__ A,
                        const float* __restrict__ B,
                        float*       __restrict__ C,
                        size_t N) {
    size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) C[i] = A[i] + B[i];
}

int main() {
    const size_t N = 1 << 24;
    std::vector<float> hA(N, 1.5f), hB(N, 2.5f), hC(N, 0.f);

    float *dA, *dB, *dC;
    cudaMalloc(&dA, N*sizeof(float));
    cudaMalloc(&dB, N*sizeof(float));
    cudaMalloc(&dC, N*sizeof(float));

    cudaMemcpy(dA, hA.data(), N*sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(dB, hB.data(), N*sizeof(float), cudaMemcpyHostToDevice);

    int threads = 256;
    int blocks  = (N + threads - 1) / threads;
    vec_add<<<blocks, threads>>>(dA, dB, dC, N);
    cudaDeviceSynchronize();

    cudaMemcpy(hC.data(), dC, N*sizeof(float), cudaMemcpyDeviceToHost);
    printf("hC[0] = %.2f, hC[N-1] = %.2f\n", hC[0], hC[N-1]);

    cudaFree(dA); cudaFree(dB); cudaFree(dC);
    return 0;
}

#include <metal_stdlib>
using namespace metal;

kernel void vec_add(const device float *A [[buffer(0)]],
                    const device float *B [[buffer(1)]],
                    device       float *C [[buffer(2)]],
                    constant     uint  &N [[buffer(3)]],
                    uint gid              [[thread_position_in_grid]]) {
    if (gid < N) C[gid] = A[gid] + B[gid];
}

import Metal
import Foundation

let device = MTLCreateSystemDefaultDevice()!
let n: UInt32 = 1 << 24
let bytes = Int(n) * MemoryLayout<Float>.stride

let A = device.makeBuffer(length: bytes, options: .storageModeShared)!
let B = device.makeBuffer(length: bytes, options: .storageModeShared)!
let C = device.makeBuffer(length: bytes, options: .storageModeShared)!
let pA = A.contents().bindMemory(to: Float.self, capacity: Int(n))
let pB = B.contents().bindMemory(to: Float.self, capacity: Int(n))
let pC = C.contents().bindMemory(to: Float.self, capacity: Int(n))
for i in 0..<Int(n) { pA[i] = 1.5; pB[i] = 2.5 }

let lib = try device.makeLibrary(URL: URL(fileURLWithPath: "vec_add.metallib"))
let pipe = try device.makeComputePipelineState(function: lib.makeFunction(name: "vec_add")!)

var nVar = n
let q = device.makeCommandQueue()!
let cb = q.makeCommandBuffer()!
let e  = cb.makeComputeCommandEncoder()!
e.setComputePipelineState(pipe)
e.setBuffer(A, offset: 0, index: 0)
e.setBuffer(B, offset: 0, index: 1)
e.setBuffer(C, offset: 0, index: 2)
e.setBytes(&nVar, length: MemoryLayout<UInt32>.size, index: 3)
let tg = MTLSize(width: 256, height: 1, depth: 1)
let grid = MTLSize(width: (Int(n) + 255) / 256, height: 1, depth: 1)
e.dispatchThreadgroups(grid, threadsPerThreadgroup: tg)
e.endEncoding()
cb.commit(); cb.waitUntilCompleted()

print("pC[0] = \(pC[0]), pC[n-1] = \(pC[Int(n)-1])")