Hello Metal Revisited

Same drill, different uniform

Recall: Metal has no GPU printf. So the configurable-battalion version on Apple Silicon writes per-thread identity into a buffer, the CPU prints. Same mission as the CUDA version.

Two key differences from CUDA you'll feel:

The 'block size' lives in two places — the kernel's hard-coded constant and the host's threadsPerThreadgroup. They must agree, and Metal won't always tell you when they don't.
The host driver explicitly waits with commandBuffer.waitUntilCompleted(); that's Metal's cudaDeviceSynchronize.

Code

ant_battalion.metal — kernel writes ID, host prints·metal

#include <metal_stdlib>
using namespace metal;

// out[gid] = (block_id, thread_in_block)
kernel void ant_battalion(
    device uint2 *out               [[buffer(0)]],
    uint  tid                       [[thread_index_in_threadgroup]],
    uint  bid                       [[threadgroup_position_in_grid]],
    uint  threads_per_block         [[threads_per_threadgroup]])
{
    uint gid = bid * threads_per_block + tid;
    out[gid] = uint2(bid, tid);
}

ant_battalion.swift — variable launch + verify·swift

import Metal
import Foundation

func main() throws {
    let args = CommandLine.arguments
    guard args.count == 3,
          let blocks  = Int(args[1]),
          let threads = Int(args[2]) else {
        print("Usage: ant_battalion <blocks> <threads_per_block>")
        exit(1)
    }
    let device = MTLCreateSystemDefaultDevice()!
    let lib = try device.makeLibrary(URL: URL(fileURLWithPath: "ant_battalion.metallib"))
    let fn  = lib.makeFunction(name: "ant_battalion")!
    let pipe = try device.makeComputePipelineState(function: fn)

    let total = blocks * threads
    let buf = device.makeBuffer(
        length: total * MemoryLayout<SIMD2<UInt32>>.stride,
        options: .storageModeShared)!

    let q  = device.makeCommandQueue()!
    let cb = q.makeCommandBuffer()!
    let e  = cb.makeComputeCommandEncoder()!
    e.setComputePipelineState(pipe)
    e.setBuffer(buf, offset: 0, index: 0)
    e.dispatchThreadgroups(
        MTLSize(width: blocks, height: 1, depth: 1),
        threadsPerThreadgroup: MTLSize(width: threads, height: 1, depth: 1))
    e.endEncoding()
    cb.commit(); cb.waitUntilCompleted()

    let p = buf.contents().bindMemory(to: SIMD2<UInt32>.self, capacity: total)
    for i in 0..<total {
        print("Ant \(i): block \(p[i].x), slot \(p[i].y)")
    }
}

do { try main() } catch { print("Error: \(error)"); exit(1) }

Build + run·bash

mkdir -p build
xcrun -sdk macosx metal -c ant_battalion.metal -o build/ant.air
xcrun -sdk macosx metallib build/ant.air -o build/ant_battalion.metallib
xcrun -sdk macosx swiftc ant_battalion.swift \
      -framework Metal -framework Foundation \
      -o build/ant_battalion

cd build && ./ant_battalion 4 32   # 128 ants
         ./ant_battalion 32 64    # 2048 ants — output is in 'CPU read' order

Exercise

Run ./ant_battalion 4 32 and confirm 128 lines come out in index order. Then change the host to dispatch blocks=4, threads=64 WITHOUT changing the kernel — run again. The output buffer is still 128 entries (you allocated for the original count), but you'll see only some block IDs because the launch geometry is now 256 threads. This is the silent-mismatch failure mode the callout warns about; reading threads_per_threadgroup from the built-in fixes it.

Same drill, different uniform

Code

External links

Exercise

Progress

Comments 0