Recurse Center Week One

I started my first week at Recurse Center. The overall goal for the retreat is to understand the Large Language Model Serving Stack. The first week was not as productive as I would have liked but there was no shortage of creativity from the batch. The week opened with a set of pair programming sessions where we paired on Conway’s Game of Life. The range of implementations included: a shader implementation, a Haskell implementation, a red versus blue implementation and more.

For the second half I spent time looking through a GPU Mode lecture on Metal programming and the tinygrad notes. I also tried to obtain an overview of Metal, the GPU programming framework for Apple. During the week I stumbled upon a set of puzzles.

The starting task was to compile a Metal Kernel, which required installing Xcode and generating code which is reminiscent of C++.

#include <metal_stdlib>
using namespace metal;

kernel void add_arrays(device const float* a   [[buffer(0)]],
                       device const float* b   [[buffer(1)]],
                       device float* out       [[buffer(2)]],
                       uint id [[thread_position_in_grid]]) {
    out[id] = a[id] + b[id];
}

The accompanying starter code was more intimidating at first glance.

import Metal
import Foundation

let n = 1_000_000
let device = MTLCreateSystemDefaultDevice()!
let queue = device.makeCommandQueue()!

// Load the kernel. Compiling from source at runtime avoids library-bundling issues in a command line tool.
let path = "<path>"   // change this
let source = try String(contentsOfFile: path, encoding: .utf8)
let library = try device.makeLibrary(source: source, options: nil)
let function = library.makeFunction(name: "add_arrays")!
let pipeline = try device.makeComputePipelineState(function: function)

// Data
let a = [Float](repeating: 1.0, count: n)
let b = [Float](repeating: 2.0, count: n)
let size = n * MemoryLayout<Float>.stride
let bufA = device.makeBuffer(bytes: a, length: size, options: .storageModeShared)!
let bufB = device.makeBuffer(bytes: b, length: size, options: .storageModeShared)!
let bufOut = device.makeBuffer(length: size, options: .storageModeShared)!

// Encode and run
let cmd = queue.makeCommandBuffer()!
let enc = cmd.makeComputeCommandEncoder()!
enc.setComputePipelineState(pipeline)
enc.setBuffer(bufA, offset: 0, index: 0)
enc.setBuffer(bufB, offset: 0, index: 1)
enc.setBuffer(bufOut, offset: 0, index: 2)

let threadsPerGroup = MTLSize(width: min(pipeline.maxTotalThreadsPerThreadgroup, 256), height: 1, depth: 1)
enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1), threadsPerThreadgroup: threadsPerGroup)
enc.endEncoding()
cmd.commit()
cmd.waitUntilCompleted()

// Check the result
let result = bufOut.contents().bindMemory(to: Float.self, capacity: n)
print(result[0], result[n - 1])   // expect 3.0 3.0

After a second reading, it’s actually not too far from CUDA and digestible with time. I spent time working on the Metal puzzles (I’m using a Mac).

The next step is to complete the puzzles and figure out the rough structure of the Metal Intermediate representation and how it interfaces with a compiler like tinygrad. Hopefully, I can return the following week with a post that is heavier on content.