A tiny scale-free kernel language — Joa Ebert

TLDR: I’ve built Phobos, a tiny kernel language inspired by Triton. It lowers to PTX and runs on NVIDIA GPUs. It achieves acceptable performance at 76% of cuBLAS SGEMM GFLOP/s on a 2080 SUPER (or 74% of the theoretical GFLOP/s peak). Phobos maps naturally to a distributed tile-DAG design. I only validated the cluster prototype on a single machine; there are no multi-node benchmarks here. This was a personal research project to gain a better understanding of low-level GPU concepts. I only realized the potential for distributed computing when implementing local tile optimizations.Ever since I was a kid and started to learn about (artificial) neural networks, I was fascinated by the human brain and AI to the point where I built object recognition software in my bedroom, using shitty webcams and a 1.8 GHz CPU for training. When I finished school I enrolled in Cognitive Informatics at Bielefeld University. Though, after a mere four days, I dropped out because I got a job offer from a startup in Paris and wanted to see the world.In the following years, I worked a lot on optimization, compilers and other machine learning related topics. I still follow the research, but I am an outsider to frontier AI research and tech.Here I am, back again with a shitty 2080 SUPER. Eager to learn.I strongly believe, in order to fully understand something, you really need to do the work and get your hands dirty. And because I have experience with compilers, optimization and LLVM, I thought writing a tiny language and compiler to target GPUs is approachable for me.I also have enough experience with software engineering, and compilers, to know that I must tread carefully. I want to keep this a learning exercise. It could very well snowball into a project that would consume too much of my available time if I take a wrong turn. Compilers are particularly tricky. Without being able to compile a program from start to end, you won’t see anything. Compiling a program is not trivial per se, so I am going to take some shortcuts. There won’t be any sort of linking or a phase model.The goal: create a small language that compiles and runs on NVIDIA GPUs. In doing so, I want to pick a model well suited for AI computations.

It should make optimizations possible without having to reconstruct the semantics of typical AI kernels. I know that I want to support first class tensor types / tiles. In fact, the entire language is oriented around tiles.Here are two example programs I want to be able to compile. I’ll explain in more detail when we get to the language.Vector AdditionThis is \(\vec{c} = \vec{a} + \vec{b}\).@autotune(BLOCK in [1024]) kernel add(a: tensor<f32>[N], b: tensor<f32>[N], c: tensor<f32>[N]) { let base = program_id(0) * BLOCK c[base :+ BLOCK] = a[base :+ BLOCK] + b[base :+ BLOCK] } Matrix MultiplicationThis is \(\mathbf{C}_{M \times N} = \alpha \mathbf{A}_{M \times K} \mathbf{B}_{K \times N} + \beta \mathbf{C}_{M \times N}\).@autotune(TILE_M in [32, 128], TILE_N in [64, 256], TILE_K in [16, 64]) kernel sgemm(A: tensor<f32>[M, K], B: tensor<f32>[K, N], C: tensor<f32>[M, N], alpha: f32, beta: f32) { let pm = program_id(0) let pn = program_id(1) var acc: tile<f32>[TILE_M, TILE_N] = 0.0 for kt in range(0, K, TILE_K) { var a = A[pm * TILE_M :+ TILE_M, kt :+ TILE_K] var b = B[kt :+ TILE_K, pn * TILE_N :+ TILE_N]

acc += dot(a, b) } let c_old = C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N] = alpha * acc + beta * c_old } Side Note: If you are familiar with kernel development, then yes: it’s basically a mini-Triton.Unfamiliar TerritoryLet’s start to get to know CUDA and its target format. If we look at resources such as CUDA Refresher: Getting started with CUDA, we are greeted with samples that look like this:/** CUDA kernel device code - CUDA Sample Codes * Computes the vector addition of A and B into C. * The three vectors have the same number of elements as numElements. */ __global__ void vectorAdd( float *A, float *B, float *C, int numElements) { int i = blockDim.x * blockIdx.x + threadIdx.x; if (i < numElements) { C[i] = A[i] + B[i]; } } We are not interested in that. This is high-level code. We want the x64 of the GPU, not the C++. I first looked at the NVIDIA CUDA Compiler (NVCC) docs1. However this is very convoluted as nvcc is what links entire executables. The CUDA docs2 themselves have a section that’s more in our domain: PTX and low-level programming. Low-level programming is what we want.PTX provides a stable programming model and instruction set for general purpose parallel programming. It is designed to be efficient on NVIDIA GPUs supporting the computation features defined by the NVIDIA Tesla architecture. High level language compilers for languages such as CUDA and C/C++ generate PTX instructions, which are optimized for and translated to native target-architecture instructions. 3There’s an existing LLVM backend we can use, called NVPTX. We have an idea for our pipeline now.Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ LLVM IR ⟶ NVPTX ⟶ PTX ⟶ GPUPTX is not the machine code the GPU runs but it is as close as we can get. SASS is the actual machine code4, but undocumented.

But there’s more. There always is. The feature creep. I can feel it. MLIR is the way to go now5. Let’s use it instead of producing LLVM IR. In fact, I do not plan to use MLIR features, like a custom dialect. But it gives me the option to reach for CUTLASS-like performance for free…a man can dream. Seriously though, MLIR exposes many instructions that I’ll need.Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ PTX ⟶ GPUFor this project I am also going to use Rust because it has bindings for LLVM with inkwell, MLIR with the melior crate and it supports pattern matching, which makes writing compilers 100x more pleasant. I had hopes that using MLIR would be as “simple” as using inkwell. But alas, it took me many hours of yak-shaving, patching the melior, mlir-sys and tblgen-rs crates, and other horrors, until I was able to compile a simple program.When I started this little adventure, I generated some PTX straight from LLVM IR just to get to know everything. That was the obvious path forward for me. But I quickly learned about MLIR. Getting started with inkwell took maybe a couple of hours. Getting melior and MLIR to work took me two days.Right-to-LeftBecause many things can break here, I am going from right-to-left. I am going to start with the most isolated task: run PTX code on the GPU. It is inevitably something I will have to do later.Let’s write the vectorAdd from above in PTX and run it with the cust crate.First we need to understand the compute cap of our GPU. In my case, it’s a 7.5 for the Turing family. Yes, I have a pretty old GPU with only 8GB VRAM. Sadly, it also lacks float32 tensor cores — but I do have 16-bit precision available. We can work with that.$ nvidia-smi --query-gpu=name,compute_cap name, compute_cap NVIDIA GeForce RTX 2080 SUPER, 7.5 We can now start to write PTX. This should feel familiar if you’ve ever written LLVM IR by hand.

At least we don’t have to write all of this in SSA form. In any case, this was the only PTX code I planned to write manually. Running the assembler directly with ptxas add.ptx was a useful sanity check to isolate any issues from Rust or the CUDA crate I am using. We can also compile kernels to PTX via nvcc -ptx in.cu -o out.ptx which is very helpful. We can write a kernel, compile it with the official toolchain, and inspect the generated PTX output..version 9.0 // 9.3 should work, but does not? .target sm_75 // target arch (compute cap) .address_size 64 .visible .entry add( .param .u64 a, .param .u64 b, .param .u64 c, .param .u32 numElements ) { .reg .pred %p; .reg .f32 %f<4>; .reg .b32 %r<5>; .reg .b64 %rd<11>; ld.param.u64 %rd1, [a]; ld.param.u64 %rd2, [b]; ld.param.u64 %rd3, [c]; ld.param.u32 %r1, [numElements]; // i = blockIdx.x * blockDim.x + threadIdx.x mov.u32 %r2, %ctaid.x; mov.u32 %r3, %ntid.x; mov.u32 %r4, %tid.x; mad.lo.s32 %r2, %r2, %r3, %r4; // %r2 is i // if i >= numElements, return setp.ge.s32 %p, %r2, %r1; @%p bra DONE; // byte offset = i * 4 mul.wide.s32 %rd4, %r2, 4; // %f1 = a[i * 4]

cvta.to.global.u64 %rd5, %rd1; add.s64 %rd6, %rd5, %rd4; ld.global.f32 %f1, [%rd6]; // %f2 = b[i * 4] cvta.to.global.u64 %rd7, %rd2; add.s64 %rd8, %rd7, %rd4; ld.global.f32 %f2, [%rd8]; // %f3 = %f1 + %f2 add.f32 %f3, %f1, %f2; // c[i * 4] = %f3 cvta.to.global.u64 %rd9, %rd3; add.s64 %rd10, %rd9, %rd4; st.global.f32 [%rd10], %f3; DONE: ret; } $ cargo run Compiling phobos v0.1.0 (C:\Users\joaeb\code\phobos) Finished `dev` profile [unoptimized + debuginfo] target(s) in 0.63s Running `target\debug\phobos.exe` 1048576 elements correct I include the Rust code here as we will come back to it later for testing purposes.Executing PTX on the GPU with Rust & custuse cust::prelude::*; const PTX: &str = include_str!("../ptx/add.ptx"); fn main() -> anyhow::Result<()> { let _ctx = cust::quick_init()?; // init driver let module = Module::from_ptx(PTX, &[])?; // compile the code let stream = Stream::new(StreamFlags::NON_BLOCKING, None)?; let func = module.get_function("add")?; let n: i32 = 1 << 20; // create some data let a: Vec<f32> = (0..n).map(|i| i as f32).collect(); let b: Vec<f32> = (0..n).map(|i| (2 * i) as f32).collect(); let mut c = vec![0.0f32; n as usize]; let a_dev = a.as_slice().as_dbuf()?; //

write to gpu let b_dev = b.as_slice().as_dbuf()?; let c_dev = c.as_slice().as_dbuf()?; let block: u32 = 256; // run our code let grid: u32 = (n as u32 + block - 1) / block; unsafe { launch!(func<<<grid, block, 0, stream>>>( a_dev.as_device_ptr(), b_dev.as_device_ptr(), c_dev.as_device_ptr(), n ))?; } // wait for gpu to complete stream.synchronize()?; // read from gpu c_dev.copy_to(&mut c)?; // sanity check for i in 0..n as usize { assert_eq!(c[i], a[i] + b[i]); } println!("{n} elements correct"); Ok(()) } The source code we provide (PTX) is getting JIT compiled by the driver for the specific GPU to SASS. This makes sense as the same code should run on many different GPUs. If we introduce errors to the PTX, a mere “JIT error” message will be thrown at us. ptxas can help us identify what’s actually broken. Profiling our binary with ncu works as well. Nice.I have certainly a better understanding of the stack now. I can run low-level code on the GPU and verify it on the CPU. Let’s revisit the pipeline from earlier.Source Code ⟶ Lexer ⟶ Parser ⟶ Codegen ⟶ MLIR ⟶ PTX ⟶ GPUWe have started at the finish line. But that’s good. At least for me, this was the most uncertain part. Everything else is known territory from now on, mostly.The Devil You KnewI really like LLVM and I used to work a lot with it for a sea-of-nodes IR where I started with Java bytecode, parsed that to an IR of mine, did my transformations and then lowered to LLVM IR which I could then use to emit to x64 or WASM. MLIR did not exist back then. This was now more than 10 years ago. MLIR would have been a good fit to model Java’s high-level constructs. For example, property access throws a NullPointerException. The high-level IR I created contained nodes like null-check, to then emit many LLVM IR nodes out of these.