nablatensor
← learnChapter 12 · Under the hood7 min read

Anatomy of a GPU kernel source file

A different GPU-facing subsystem from Chapter 12.1's AAD codegen: nablatensor-tensor's own catalog of hand-written CUDA-C kernels, shared byte-for-byte between the cuda and rocm backends. What one real kernel looks like end to end, how GpuKernel parses a name and a launch geometry out of the text so they can never drift apart, and a second, generated kind of kernel — fusedSource — whose common-subexpression cache turns out to be keyed by object identity, not by what an expression actually computes.

Chapter 12.1 watched a recorded tape turn into JVM bytecode. That's the AAD engine's own code generator. This page is a different subsystem entirely: nablatensor-tensor's bulk tensor operations — matmul, reductions, convolutions — run on cuda/rocm from a catalog of hand-written CUDA-C kernels, not a per-tape generator. What does that kind of GPU kernel source actually look like, and how is 20 kernels' worth of it kept from drifting out of sync with itself?

The whole story

Two panels contrasting before and after a refactor. Before: one roughly 360-line CUDA-C text block holding 20 concatenated kernels, plus a hand-maintained array of kernel names that has to stay in sync by hand — a typo only caught at module-load time on a machine with a GPU. After, what actually landed: five family files, each a list of GpuKernel records, each record parsing its own entry-point name and parameter types out of its own source text, so the kernel-name array is derived rather than typed twice. Below, one real, complete kernel — transpose2d, nine lines, one thread per element. Beside it, a second way kernel source gets made: fusedSource turns an elementwise expression tree into one SSA-style kernel, but its common-subexpression cache is keyed by object identity rather than by what the expression actually computes — reusing the same node object twice collapses to four lines, while two separate objects computing the identical thing produce seven, with the shared computation duplicated under two different names. A banner reports numbers measured this session in pure Java, no device needed: 20 kernels, 14,214 characters, 359 lines, matmul_tiled parsed to arity 7 and a 16 by 16 block. A second banner explains why every MAX/MIN case wraps fmaxf/fminf in explicit isnan checks — the raw C intrinsics silently discard a NaN operand instead of propagating it.

Did you know?

CUDA and HIP's fmaxf/fminf follow the C standard's fmax/fmin contract: given one NaN operand, they silently return the other one. That's the opposite of what a numerics library usually wants — a NaN in a tensor should stay a NaN, not vanish the moment it's maxed against something else. Every MAX/MIN case in ew_binary, ew_scalar and unaryExpr's RELU case wraps the intrinsic in exactly that check: isnan(x) ? x : (isnan(y) ? y : fmaxf(x, y)). The safety isn't in the hardware or the compiler; it's three extra tokens somebody had to remember to write, in every single case label that touches fmaxf/fminf.

Before: one string, one array, two sources of truth

GpuKernels.TENSOR_SOURCE used to be a single concatenated CUDA-C block — the cookbook's own description, docs/cookbook/gpu-kernel-source-layout.md — with a hand-maintained String[] of kernel names next to it. The two had to agree by convention, and nothing checked that they did until a backend tried to load a kernel that wasn't there.

What actually landed: five family files — ElementwiseKernels/MatmulKernels/ReductionKernels/ConvKernels/ RandomKernels — each a List<GpuKernel>, plus a DevicePrelude of __device__ helpers emitted first. GpuKernels itself is now a façade:

public static final List<GpuKernel> TENSOR_KERNELS = Stream.of(
        ElementwiseKernels.KERNELS, MatmulKernels.KERNELS, RandomKernels.KERNELS,
        ReductionKernels.KERNELS, ConvKernels.KERNELS)
    .flatMap(List::stream)
    .toList();

public static final String[] TENSOR_KERNEL_NAMES = BY_NAME.keySet().toArray(String[]::new);

public static final String TENSOR_SOURCE = DevicePrelude.SOURCE
    + TENSOR_KERNELS.stream().map(GpuKernel::source).collect(Collectors.joining());

The name array is a keySet(), not a literal. It cannot list a kernel that isn't in TENSOR_KERNELS, and building that index throws on the spot if two kernels ever claim the same name.

One kernel, in full

GpuKernel.of(source) reads the entry-point name and every parameter's type straight out of the C text with a regex — extern "C" __global__ void (\w+)\s*\(([^)]*)\) — so the name a backend loads can never be a different string than the name the source actually defines. Here is the plainest of the twenty, real source, complete:

extern "C" __global__ void transpose2d(float* out, const float* in, int rows, int cols) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= rows * cols) return;
  int r = i / cols, c = i % cols;
  out[c * rows + r] = in[i];
}

One thread per output element, a bounds guard against a grid rounded up past the tensor's actual size, no __shared__ memory, no cross-thread communication at all. Real numbers, measured this session in pure Java — generating and parsing this catalog needs no device, no driver, nothing but the JVM:

kernels: 20
TENSOR_SOURCE length: 14214 chars, 359 lines

Twenty kernels, the exact names the cookbook's own pinned test expects, in 359 lines — matching its own "~360 lines" estimate almost exactly.

A launch geometry that travels with the source

matmul_tiled is the one kernel worth reading past the "plainest" example above, because its own signature carries a real constraint. Parsed, not typed by hand:

matmul_tiled arity=7 paramTypes=[float*, const float*, const float*, int, int, int, unsigned long long] block=16x16

That seventh parameter — unsigned long long tile_offset — exists because the two tiled GEMM kernels use a 16×16 thread block matched to a __shared__ tile of the same shape, and MatmulKernels's own class doc gives the exact reason for tile_offset: "a grid wider than INT_MAX blocks can be launched in chunks." The block size (16, 16) lives on the same GpuKernel record as the source and the parsed parameter list — not as a magic 16 at the CudaBackend/RocmBackend call site the cookbook says it used to be.

💡 A second way kernel source gets made — and its blind spot

Not every kernel here is hand-written text. GpuKernels.fusedSource(expr, numInputs) walks an elementwise Expr tree and emits a single fused kernel: one SSA-style named local per node, post-order, so a whole chain of ops runs in one launch instead of one launch per op. Built the same (in0 + in1) * (in0 + in1) shape two different ways this session, and the generator's real output differs:

// same Expr object used for the (in0+in1) subtree in both places:
extern "C" __global__ void fused_kernel(float* out, const float* in0, const float* in1, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  float v0 = in0[i];
  float v1 = in1[i];
  float v2 = (v0 + v1);
  float v3 = (v2 * v2);
  out[i] = v3;
}

// two separate Expr objects, computing the identical (in0+in1):
extern "C" __global__ void fused_kernel(float* out, const float* in0, const float* in1, int n) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;
  if (i >= n) return;
  float v0 = in0[i];
  float v1 = in1[i];
  float v2 = (v0 + v1);
  float v3 = in0[i];
  float v4 = in1[i];
  float v5 = (v3 + v4);
  float v6 = (v2 * v5);
  out[i] = v6;
}

emitFused's cache is an IdentityHashMap<Expr, String> — it recognizes "I have already emitted code for this exact node object," not "I have already emitted code for a node that computes the same thing." Build the same subtree once and reuse the reference, and it collapses to one local. Build it twice, even with identical operands and operator, and the fused kernel recomputes it under a second name — seven lines instead of four, for a tree whose two halves are indistinguishable to a compiler doing value-based common-subexpression elimination, but not to this one.

Try it yourself

Pick any two-input elementwise expression and build it exactly the two ways above: construct the shared subtree once and pass the same Expr reference to both sides of the outer op, then rebuild it with new Expr.Binary(...) a second time for the other side. Print GpuKernels.fusedSource(tree, 2) both times. The first version's local count is always fewer than the second's — by exactly the size of whichever subtree you duplicated.

▶️ Run it

No test or example needs a GPU to see any of this — GpuKernel and GpuKernels are plain Java that parses and concatenates strings. Everything on this page came from a small standalone program run directly against the built classes:

System.out.println(GpuKernels.TENSOR_KERNEL_NAMES.length);       // 20
System.out.println(GpuKernels.TENSOR_SOURCE.lines().count());    // 359

GpuKernel matmul = GpuKernels.kernel("matmul_tiled");
System.out.println(matmul.arity() + " " + matmul.paramTypes());  // 7 [...]

Expr shared = new Expr.Binary(Op.ADD, new Expr.Input(0), new Expr.Input(1));
System.out.println(GpuKernels.fusedSource(
    new Expr.Binary(Op.MUL, shared, shared), 2));

The project's own GpuKernelsTest (nablatensor-tensor/src/test/java/.../GpuKernelsTest.java) checks the same catalog with mvn -o test, no GPU required — it pins the 20 names, checks every one is actually defined in TENSOR_SOURCE, and checks that DevicePrelude's helpers appear before their first caller.

⚠️ What this doesn't do

This page is entirely about nablatensor-tensor's own kernel catalog — matmul, reductions, convolutions, elementwise ops for the tensor library. It is not the AAD engine's tape-to-GPU path Chapter 4.3 covered (CudaAadCodegen/VulkanAadCodegen, which compile a recorded pricing tape, not a tensor expression), and the two share no code. It also doesn't run any of these kernels on an actual device: parsing and generating the source needs nothing but the JVM, but compiling it with NVRTC/HIPRTC and dispatching it is CudaBackend/RocmBackend's job, on a machine with the matching driver — out of scope here the same way every GPU-dispatch step has been since Chapter 4.1. And it only shows the two approaches this codebase actually chose (family files plus parsed records); the cookbook itself lists five more it considered and didn't take — an annotation-driven registry, an annotation processor, a typed CUDA-C builder DSL — each with its own effort/risk trade-off, for anyone who hits this catalog's limits later.

What's next

→ Deeper: docs/cookbook/gpu-kernel-source-layout.md has all seven approaches this project considered, with effort/risk ratings for each, not just the one that landed. → Chapter 12 (Under the hood) is complete for now. Next: Chapter 13 — Bonus: run it without writing Java.

Questions or corrections? open an issue