Specification

Phase 0 Companion: CUDA as a Native Concern

Companion to the Phase 0 specification. CUDA is a first-class design concern from day one, not a backend that gets bolted on later. This document establishes what that commitment means.


1. Framing

In every existing ML stack, CUDA is bolted on:

The pattern: two languages with a hard boundary. The "user language" is type-checked, abstract, portable. The "kernel language" is performance-focused, hardware-specific, and lives outside the type system. Crossing the boundary requires manual work; the compiler can't see across it.

This memo's commitment: CUDA is in the type system, in the IR, and in the surface language from day one. Kernels are first-class, type-checked, and composable. Device placement, memory hierarchy, hardware capabilities, async memory operations, and tensor cores are all typed primitives — not things that disappear into opaque library calls.

This is the position memo 6 (hardware abstraction) argued for in general; this document specializes it to CUDA and commits to it concretely for Phase 0.


2. What "native" means concretely

Five specific commitments distinguish "native" from "bolted on":

2.1 Device placement is in the tensor type

A tensor isn't just Tensor[F32, [B, S, D], R]. It is Tensor[F32, [B, S, D], R, on=cuda:global] or Tensor[F32, [B, S, D], R, on=cuda:shared] or Tensor[F32, [B, S, D], R, on=host]. The placement is part of the type, not a runtime attribute.

Consequence: operations between tensors on different placements either type-check (because the operation is defined across placements, like a host-to-device copy) or fail with a type error. There is no silent cross-device dispatch.

This is the fifth refinement axis on the tensor type, alongside the four established earlier: shape, effects, structure, sharding. Adding it to Phase 0 is a small change to Section 1's grammar and Section 2's operation typings.

2.2 The tile IR is a first-class layer

The compiler stack has at least three levels: algebraic IR (the tensor algebra of Section 2), tile IR (the kernel-author level), and target IR (PTX, AMDGCN). Each level is typed.

The tile IR is the language kernel authors write — the analog of Triton, but with the type-system properties of the algebraic level. Memory hierarchy, thread/block coordination, tile shapes, and hardware capabilities are all in scope at this level.

2.3 Hardware capabilities are typed

Operations that depend on specific hardware features carry capability requirements. A tensor-core matrix-multiply-accumulate operation has type roughly:

wmma_mma : forall (M, N, K) requires capability(sm_70+, mma_m16n16k16, F16, F16, F32).
           TileMem[F16, [M, K], on=cuda:shared] →
           TileMem[F16, [K, N], on=cuda:shared] →
           TileMem[F32, [M, N], on=cuda:register] →
           TileMem[F32, [M, N], on=cuda:register]

A kernel that calls wmma_mma carries the capability requirement transitively. If you target Volta (sm_70), it works. If you target Pascal (sm_60), it fails to compile with a clear error: "kernel requires sm_70+ tensor core support; target is sm_60."

This is the capability typing from memo 6, applied to CUDA specifically. Modern features (Hopper WGMMA, FP8 tensor cores, TMA) get their own capability declarations.

2.4 Async memory operations are typed primitives

cp.async, cp.async.bulk (TMA on Hopper), and the various memory-fence operations are typed primitives in the tile IR. They are not function calls into a runtime library; they are language constructs.

This means the type system can reason about completion ordering, fence requirements, and the in-flight queue. A kernel that loads but doesn't fence before consuming will get a type error.

2.5 Kernels are composable, not opaque

A kernel function can call another kernel function. The type system tracks the composition. If kernel A is launched on a tile of shape [128, 128] and kernel A calls kernel B with a sub-tile of shape [64, 64], the type system knows about the relationship and can reason about it.

This is unlike Triton, where one @triton.jit kernel cannot directly call another (you have to use tl.inline_function or refactor as device functions). It's also unlike CUDA C++, where the compiler sees the composition but the user's type system doesn't track tile shapes meaningfully.


3. The example: tile-level tensor-core GEMM

The canonical CUDA kernel: matrix multiplication using tensor cores. Show it in CUDA C++, in Triton, and in our language to make the difference visible.

3.1 CUDA C++ (simplified)

#include <mma.h>
using namespace nvcuda;

__global__ void gemm_tc(const half* A, const half* B, float* C,
                        int M, int N, int K) {
    constexpr int BM = 128, BN = 128, BK = 16;
    constexpr int WM = 64,  WN = 64;        // per-warp tile
    constexpr int TM = 16,  TN = 16, TK = 16; // tensor-core tile

    __shared__ half A_smem[BM][BK];
    __shared__ half B_smem[BK][BN];

    int block_m = blockIdx.x;
    int block_n = blockIdx.y;
    int warp_id = threadIdx.x / 32;
    int warp_m  = warp_id / 2;
    int warp_n  = warp_id % 2;

    wmma::fragment<wmma::accumulator, TM, TN, TK, float> c_frag;
    wmma::fill_fragment(c_frag, 0.0f);

    for (int k = 0; k < K; k += BK) {
        // Cooperative load A and B tiles into shared memory
        // [thread block load logic, omitted for brevity]
        __syncthreads();

        // Tensor core MMA on per-warp tiles
        wmma::fragment<wmma::matrix_a, TM, TN, TK, half, wmma::row_major> a_frag;
        wmma::fragment<wmma::matrix_b, TM, TN, TK, half, wmma::row_major> b_frag;
        wmma::load_matrix_sync(a_frag, &A_smem[warp_m * WM][0], BK);
        wmma::load_matrix_sync(b_frag, &B_smem[0][warp_n * WN], BN);
        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);

        __syncthreads();
    }

    // Store result
    int c_row = block_m * BM + warp_m * WM;
    int c_col = block_n * BN + warp_n * WN;
    wmma::store_matrix_sync(&C[c_row * N + c_col], c_frag, N, wmma::mem_row_major);
}

void launch_gemm(half* A, half* B, float* C, int M, int N, int K) {
    dim3 grid(M/128, N/128);
    dim3 block(128);  // 4 warps
    gemm_tc<<<grid, block>>>(A, B, C, M, N, K);
}

Notable things hidden or mismanaged here:

3.2 Triton (idiomatic)

@triton.jit
def gemm_tc(A_ptr, B_ptr, C_ptr,
            M, N, K,
            stride_am, stride_ak,
            stride_bk, stride_bn,
            stride_cm, stride_cn,
            BLOCK_M: tl.constexpr,
            BLOCK_N: tl.constexpr,
            BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_am = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_bn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k  = tl.arange(0, BLOCK_K)

    a_ptrs = A_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
    b_ptrs = B_ptr + offs_k[:, None]  * stride_bk + offs_bn[None, :] * stride_bn

    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k)
        accumulator += tl.dot(a, b)
        a_ptrs += BLOCK_K * stride_ak
        b_ptrs += BLOCK_K * stride_bk

    c_ptrs = C_ptr + offs_am[:, None] * stride_cm + offs_bn[None, :] * stride_cn
    tl.store(c_ptrs, accumulator)

Better than CUDA C++. The tile abstraction handles thread-level coordination implicitly. But still:

3.3 The same kernel in our language (Phase 5 surface, illustrative)

kernel gemm_tc
    [M, N, K]
    requires (M mod 128 = 0)
    requires (N mod 128 = 0)
    requires (K mod 16 = 0)
    requires capability(sm_80, mma_m16n16k16_f16_f16_f32)
    (a : Tensor[F16, [M, K], on=cuda:global])
    (b : Tensor[F16, [K, N], on=cuda:global])
    : Tensor[F32, [M, N], on=cuda:global] {

  // ---- Block-level tile structure ----
  block_grid [M/128, N/128] {
    let block_m = block_id 0 in
    let block_n = block_id 1 in

    // ---- Shared memory tiles ----
    shared a_smem : Tile[F16, [128, 16], on=cuda:shared]
    shared b_smem : Tile[F16, [16, 128], on=cuda:shared]

    // ---- Per-warp register accumulator ----
    let warps_per_block = 4 in
    warp_grid [2, 2] {
      let warp_m = warp_id 0 in
      let warp_n = warp_id 1 in

      register c_reg : Tile[F32, [64, 64], on=cuda:register]
      c_reg <- fill 0.0

      // ---- Tile loop ----
      for k in (0 .. K/16) {
        // Async cooperative load: global -> shared
        cp_async a_smem <- a[block_m * 128 ..+ 128, k * 16 ..+ 16]
        cp_async b_smem <- b[k * 16 ..+ 16, block_n * 128 ..+ 128]
        cp_async_commit
        cp_async_wait_all

        block_sync

        // Tensor core MMA over the warp's sub-tile
        c_reg <- wmma_mma c_reg
                          a_smem[warp_m * 64 ..+ 64, 0 ..+ 16]
                          b_smem[0 ..+ 16, warp_n * 64 ..+ 64]

        block_sync
      }

      // Store register tile back to global
      a[block_m * 128 + warp_m * 64 ..+ 64,
        block_n * 128 + warp_n * 64 ..+ 64] <- c_reg
    }
  }
}

This is speculative syntax — Phase 5 commits a final design — but the structural commitments are real. Notable features:

The kernel is roughly the same length as the CUDA C++ version, longer than the Triton version. The win is type-level guarantees — divisibility, capability, memory hierarchy, sync ordering. The error mode is "type error at compile time," not "wrong result at runtime."


4. What this changes about Phase 0

Several specific updates to the Phase 0 spec are required to make CUDA a native concern from day one. None of them are large; all of them are necessary.

4.1 Add the device refinement axis (Section 1)

Section 1's grammar grows by one axis:

Type ::= Tensor[Dtype, Shape, Refinements, Placement]
       | ...

Placement ::= host
            | cuda:global(device : Int)
            | cuda:shared
            | cuda:register
            | TileMem[Placement]    -- tile-IR-only types

Tensors at the algebraic IR level have placements host or cuda:global(d) for some device index d. Tile IR types have the additional placements cuda:shared and cuda:register.

Operations in the algebraic IR are placement-polymorphic by default — matmul works on tensors regardless of placement, as long as both operands are on the same placement. Operations between different placements require explicit movement primitives (copy_to_device, copy_to_host).

4.2 Add device-movement primitives to the operation set (Section 2)

Two new primitive operations:

copy_to_device : forall (dt, s, R, d). 
                 Tensor[dt, s, R, host] → Tensor[dt, s, R, cuda:global(d)]

copy_to_host   : forall (dt, s, R, d). 
                 Tensor[dt, s, R, cuda:global(d)] → Tensor[dt, s, R, host]

These are explicit in the IR. They lower to cudaMemcpy (synchronous) or cudaMemcpyAsync (asynchronous) depending on context. The runtime handles the actual transfers.

4.3 Forward-declare the tile IR (Section 5)

Section 5's implementation infrastructure now describes a multi-tier compilation pipeline:

Algebraic IR (typed tensor algebra)
    ↓ lowering pass A: kernel selection / fusion
Tile IR (typed tile-level kernels)
    ↓ lowering pass B: hardware-specific code generation
Target IR (PTX via MLIR NVVM dialect, eventually also AMDGCN, SPIR-V)
    ↓ assembler
Executable

Phase 0 does not implement passes A or B. It commits to the structure and to the type-system properties (placement axis, capability typing). Phase 5 implements the tile IR design and the lowering passes.

4.4 Implementation infrastructure additions (Section 5)

The build adds CUDA toolchain dependencies:

CI implications: at least some testing requires a CUDA-capable runner. GitHub Actions has limited GPU support; alternatives are self-hosted runners or cloud GPU instances. Phase 0 doesn't need GPU CI — algebraic-level correctness is testable without hardware. Phase 5 will need it.

4.5 Updated Phase plans

The roadmap's Phase 5 (Hardware Abstraction) was the original home for tile IR work. The CUDA-native commitment moves some of that earlier:

The intermediate Phase 1 step is new. It exists because the CUDA-native commitment requires something working at Phase 1 — a forward pass should be able to run on a GPU, even if the kernels it dispatches to are external libraries. The full kernel-authoring story arrives at Phase 5.


5. Notes on the design

5.1 Compared to Triton

Triton is the closest analog. Differences:

Triton This language
Memory hierarchy Implicit (compiler decides) Explicit in types (on=cuda:shared, etc.)
Capability declarations Implicit Explicit (requires capability(sm_80, ...))
Cross-target portability NVIDIA-first; AMD via separate backend Capability-typed; portable across targets that supply the capability
Kernel composition Limited (tl.inline_function) First-class via tile-IR functions
Type checking Runtime (Python decorators) Compile-time
Error messages on hardware mismatch Backend-specific compilation error Type error with clear cause

Triton is the right model for what a tile-level kernel looks like. Our extension adds the type system properties Triton chose not to take on.

5.2 Compared to CUDA C++

CUDA C++ is the maximum-control reference. Differences:

CUDA C++ This language
Thread-level coordination Explicit (threadIdx.x etc.) Tile-level (block_id, warp_id); thread-level hidden
Memory hierarchy Storage-class qualifiers (__shared__) Type-level placement
Tensor cores nvcuda::wmma template metaprogramming Typed primitive operations
Hardware capabilities Compute capability flags + template specialization Types and requires clauses
Type checking C++ template system Refinement-typed

CUDA C++ is more flexible (you can write anything you want, including things that shouldn't be written). The type system trades flexibility for compile-time guarantees. The win is in productive iteration: a kernel that doesn't compile for a clear reason is more useful than one that compiles silently and produces wrong results.

5.3 Compared to Mojo

Mojo's ambition is similar — Python-superset with hardware-aware compilation, MLIR-based lowering, ownership for memory safety. Differences:

The two designs would converge on similar tile-level abstractions. The differences are emphasis: where Mojo is willing to leave correctness to runtime checks for ergonomic reasons, we commit to compile-time guarantees.

5.4 The non-goal: hiding CUDA

This language does not try to abstract away CUDA. It tries to type-check CUDA. A kernel author writing tile-level code is writing CUDA-shaped code — block grids, warp tiles, shared memory, tensor cores. The language gives them a better type system to write it in, not a different programming model.

This is a deliberate choice. Higher-level abstractions (frameworks like Keras, PyTorch's nn.Module) have their place; they're built on top of tile-level kernels, not instead of them. The language supports both layers; the high-level layer hides the low-level layer for users who don't need it; but the low-level layer is in the same language, with the same type system, accessible when needed.

The bolted-on alternative is to have one language for the high level and a completely different language for the low level. This is the status quo. The native commitment rejects it.


6. Open issues introduced by the CUDA commitment

  1. Multi-device meshes interact with placement. A tensor sharded across multiple cuda:global(d) devices needs the placement axis to integrate with the sharding axis (memo 4). Both are refinements on the tensor type; they should compose. Concretely, a sharded tensor is Tensor[..., Sharded(..., mesh=M), on=cuda:global] where the placement applies per-device-in-the-mesh. The compositional rule needs to be worked out; not a Phase 0 deliverable.

  2. AMD and TPU as later targets. The capability typing is meant to be portable, but Phase 0's commitment is CUDA-first. Adding AMD support means defining the AMD capability set; adding TPU means accepting that TPUs don't fit the SIMT model and need a different tile-IR shape. Memo 6 §9 flagged this as a hard subproblem.

  3. JIT compilation for AD. Reverse-mode AD over kernels (Phase 2) generates new kernels at AD-application time. This means JIT compilation of tile-IR programs to PTX. The textual-emission-via-subprocess model from Section 5 §3 is too slow for AD's fine-grained kernel synthesis. A direct API path may be needed for Phase 2+; flagged.

  4. Cooperative groups / cluster-level features (Hopper, Blackwell). Modern CUDA has features beyond block-level scope (thread block clusters, distributed shared memory). The tile IR design needs to accommodate them. Specifying which features are in scope vs. deferred is a Phase 5 design decision.

  5. PTX vs. SASS visibility. The compiler emits PTX (via MLIR NVVM); the GPU driver compiles PTX to SASS at load time. For maximum performance, some users want SASS-level control (via inline asm). The language should support an inline-PTX escape hatch; details are Phase 5.


7. Summary

The commitment: CUDA is in the type system, the IR, and the surface language from day one. Memory hierarchy is typed. Hardware capabilities are typed. Async memory operations are typed primitives. Kernels are first-class, type-checked, composable. Tile-level kernel authoring lives in the same language as algebraic-level model code, with the same type system applied throughout.

What changes about Phase 0: a fifth refinement axis (placement) on the tensor type; two new primitive operations (copy_to_device, copy_to_host); a forward-declared tile IR layer; CUDA toolchain dependencies in the implementation infrastructure; an intermediate Phase 1 milestone for "GPU forward pass via library kernels" before the full Phase 5 tile-IR work.

Why this is the right call: every existing ML stack pays a permanent ergonomic and reliability tax for the bolted-on architecture. Custom kernels live outside the type system; integration is brittle; new hardware features require library updates that the compiler doesn't see. A type system that knows about CUDA is a type system that can give compile-time guarantees about kernel correctness, capability requirements, and memory-hierarchy correctness — guarantees Triton and CUDA C++ do not provide.

The trade: the design now carries CUDA-shaped commitments from day one. AMD, TPU, dataflow, and other architectures are second-class in early phases. Memo 6's vision of fully-portable capability typing remains the long-term goal; Phase 0 commits to CUDA first because it's where the user demand is, where the hardware features are most aggressive, and where the existing tooling pain is greatest.

This is the same trade Mojo, Triton (informally), and IREE made — and the right one. Pretending to be hardware-agnostic from day one would mean deferring exactly the design work that matters most.


End of CUDA-native walkthrough.