Skip to content

gpu.md

Latest commit

 

History

History
177 lines (139 loc) · 7.14 KB

gpu.md

File metadata and controls

177 lines (139 loc) · 7.14 KB
Raw

GPU Offload

libmtlc has two GPU code generators, NVIDIA PTX and SPIR-V (OpenCL), both emitted from the same IR with no nvcc, no cudart, and no LLVM. Through the reference frontend, kernels are written in Mettle, compiled to a .ptx module with --emit-ptx (or a .spv module with --emit-spirv), and, for the CUDA path, launched from a normal Mettle host program via the std/gpu bindings and the dispatch statement. (A frontend driving libmtlc directly reaches the same generators through mtlc_emit; see Writing a frontend for libmtlc.)

The model is two-stage and explicit: kernels live in their own file, the host manages device memory itself, and dispatch only performs the launch. This mirrors how real GPU code manages persistent VRAM.

Writing a kernel

A kernel file is compiled with mettle --emit-ptx. Use the kernel keyword for GPU entry points (it parses like fn and is emitted as a PTX .entry):

// kernels.mettle  ->  mettle --emit-ptx kernels.mettle -o kernels.ptx
kernel vadd(a: float32*, b: float32*, c: float32*, n: int32) {
  var i: int32 = block.x * block_dim.x + thread.x;
  if (i < n) {
    c[i] = a[i] + b[i];
  }
}

Index built-ins

Inside --emit-ptx compiles, the GPU thread/block indices are built-in member expressions that mirror CUDA:

Mettle CUDA PTX special register
thread.x threadIdx.x %tid.x
block.x blockIdx.x %ctaid.x
block_dim.x blockDim.x %ntid.x
grid_dim.x gridDim.x %nctaid.x

.x, .y, and .z are all available. The canonical global-thread index is:

var i: int32 = block.x * block_dim.x + thread.x;

These built-ins are only active under --emit-ptx, so member access on an ordinary struct named block in a CPU program is unaffected.

Supported kernel constructs

Kernels use the same syntax as CPU code: arithmetic, comparisons, if/while, pointer indexing, casts, and a set of GPU math intrinsics declared as extern: sqrtf, rsqrtf, fabsf, sinf, cosf, logf, expf (lowered to PTX sqrt.rn / rsqrt.approx / ex2.approx etc.), plus h2f / f2h for fp16 conversion. The PTX backend is validated by round-tripping emitted PTX through ptxas and by differential execution against a CPU reference on real hardware.

Launching from the host

The host is a normal Mettle program. Import std/gpu, set up device buffers explicitly, then launch with dispatch:

import "std/io";
import "std/mem";
import "std/gpu";

fn main() -> int32 {
  if (gpu_init() == 0) { println(cstr("GPU init failed")); return 1; }

  // load the emitted PTX and resolve the kernel
  var fp: cstring = fopen(cstr("kernels.ptx"), cstr("rb"));
  var ptx: uint8* = (uint8*)malloc(65536);
  var len: int64 = fread((cstring)ptx, 1, 65535, fp); fclose(fp); ptx[len] = 0;
  var mod: int64 = gpu_module(ptx);
  var vadd: int64 = gpu_func(mod, cstr("vadd"));

  var n: int32 = 1 << 20;
  var bytes: int64 = (int64)n * 4;
  var ha: float32* = (float32*)malloc(bytes);
  var hb: float32* = (float32*)malloc(bytes);
  var hc: float32* = (float32*)malloc(bytes);
  var i: int32 = 0;
  while (i < n) { ha[i] = (float32)i; hb[i] = (float32)(2 * i); i = i + 1; }

  // device buffers (you own VRAM)
  var da: int64 = gpu_malloc(bytes);
  var db: int64 = gpu_malloc(bytes);
  var dc: int64 = gpu_malloc(bytes);
  gpu_to_device(da, (uint8*)ha, bytes);
  gpu_to_device(db, (uint8*)hb, bytes);

  // launch: one line replaces param-packing + cuLaunchKernel + sync
  dispatch vadd[(n + 255) / 256, 256](da, db, dc, n);

  gpu_to_host((uint8*)hc, dc, bytes);
  gpu_free(da); gpu_free(db); gpu_free(dc);
  return 0;
}

The dispatch statement

dispatch KERNEL[grid, block](arg0, arg1, ...);
  • KERNEL is a handle (the int64 returned by gpu_func).
  • grid and block are integer expressions: the number of blocks and the number of threads per block (1-D).
  • The arguments are passed by value. Device pointers are int64 handles; scalars (int32, float32, ...) are forwarded with their natural width.

dispatch desugars to argument marshalling plus a call to gpu_launch, which issues cuLaunchKernel and then cuCtxSynchronize. It is launch-only: device allocation and host/device copies remain explicit (the gpu_malloc / gpu_to_device / gpu_to_host calls above).

Building

# 1. compile the kernels to a PTX module
mettle --emit-ptx kernels.mettle -o kernels.ptx

# 2. build the host, linking the CUDA driver import stub (build-time only)
mettle --build host.mettle -o host \
  --link-arg "<CUDA>/lib/x64/cuda.lib"        # Windows: cuda.lib; Linux: -lcuda

The host links nvcuda (the OS driver), exactly as a Mettle program links kernel32 or libc; there is no bundled CUDA DLL. At run time the driver JITs the PTX to SASS for the installed GPU.

SPIR-V (OpenCL) target

The same kernels compile to SPIR-V with --emit-spirv, targeting the OpenCL 1.2 execution environment (Physical64 addressing, the Kernel capability, the OpenCL memory model). This is the flavor that fits Mettle's kernel ABI unchanged: kernels take raw typed pointers and do pointer arithmetic plus loads/stores, which is the OpenCL/CUDA model, not the Vulkan descriptor-buffer model.

mettle --emit-spirv kernels.mettle -o kernels.spv

The output is a binary SPIR-V module (one OpEntryPoint ... Kernel per kernel) that an OpenCL runtime consumes with clCreateProgramWithIL. The same source constructs as the PTX path are supported: arithmetic, comparisons, if/while (including &&/|| and nesting), pointer indexing, casts, the gpu_* index built-ins (mapped to the OpenCL work-item built-ins, so thread reads LocalInvocationId, block reads WorkgroupId, block_dim reads WorkgroupSize, and grid_dim reads NumWorkgroups), gpu_barrier() (an OpControlBarrier), the f32 math intrinsics (an OpExtInst from OpenCL.std), h2f/f2h, and the unsigned atomics.

Control flow maps directly onto SPIR-V basic blocks (OpBranch / OpBranchConditional), exactly as the PTX path maps it onto bra. SPIR-V's structured-control-flow rules (OpSelectionMerge/OpLoopMerge) are mandated only by the Shader capability, so Kernel (OpenCL) modules may branch freely, which spirv-val --target-env opencl1.2 confirms.

Notes and limits

  • The PTX emitter targets .target sm_90, which is forward-compatible: the driver JITs it to newer architectures (e.g. sm_120 / Blackwell).
  • dispatch grids are 1-D for now (grid, block). Multi-dimensional launches go through cuLaunchKernel in std/gpu directly.
  • Kernels and host code live in separate files (the kernel file is compiled with --emit-ptx; the host with --build).

See examples/gpu_vadd/ for the complete, runnable version of the program above, and examples/llm/qwen3/gpu/ for a full set of LLM inference kernels.