Phobos: A tiny scale-free kernel language — Joa EbertTLDR: 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.<br>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.<br>Here I am, back again with a shitty 2080 SUPER. Eager to learn.<br>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.<br>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&rsquo;t see anything. Compiling a program is not trivial per se, so I am going to take some shortcuts. There won&rsquo;t be any sort of linking or a phase model.<br>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.<br>Here are two example programs I want to be able to compile. I&rsquo;ll explain in more detail when we get to the language.<br>Vector Addition<br>This is \(\vec{c} = \vec{a} + \vec{b}\).<br>@autotune(BLOCK in [1024])<br>kernel add(a: tensor[N], b: tensor[N], c: tensor[N]) {<br>let base = program_id(0) * BLOCK<br>c[base :+ BLOCK] = a[base :+ BLOCK] + b[base :+ BLOCK]

Matrix Multiplication<br>This is \(\mathbf{C}_{M \times N} = \alpha \mathbf{A}_{M \times K} \mathbf{B}_{K \times N} + \beta \mathbf{C}_{M \times N}\).<br>@autotune(TILE_M in [32, 128], TILE_N in [64, 256], TILE_K in [16, 64])<br>kernel sgemm(A: tensor[M, K],<br>B: tensor[K, N],<br>C: tensor[M, N],<br>alpha: f32,<br>beta: f32) {<br>let pm = program_id(0)<br>let pn = program_id(1)<br>var acc: tile[TILE_M, TILE_N] = 0.0<br>for kt in range(0, K, TILE_K) {<br>var a = A[pm * TILE_M :+ TILE_M, kt :+ TILE_K]<br>var b = B[kt :+ TILE_K, pn * TILE_N :+ TILE_N]<br>acc += dot(a, b)<br>let c_old = C[pm * TILE_M :+ TILE_M, pn * TILE_N :+ TILE_N]<br>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&rsquo;s basically a mini-Triton.

Unfamiliar Territory<br>Let&rsquo;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:<br>/** CUDA kernel device code - CUDA Sample Codes<br>* Computes the vector addition of A and B into C.<br>* The three vectors have the same number of elements as numElements.<br>*/<br>__global__ void vectorAdd( float *A, float *B, float *C, int numElements) {<br>int i = blockDim.x * blockIdx.x + threadIdx.x;<br>if (i numElements) {<br>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&rsquo;s more in our domain: PTX and low-level programming. Low-level programming is what we want.<br>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. 3

There&rsquo;s an existing...