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:
- PyTorch. User writes Python; framework dispatches to opaque cuBLAS/cuDNN/CUTLASS kernels. Custom kernels require dropping into a separate language (CUDA C++ or Triton) with its own type system, build process, and integration story. The compiler does not reason about kernel composition.
- JAX. User writes NumPy-flavored code; XLA lowers to HLO; HLO compiles to a fixed set of kernels. Custom kernels via
pallasor by linking external code. Better than PyTorch on the algebraic side; equally bolted-on at the kernel-author tier. - TensorFlow. Similar to JAX in structure; same fundamental separation between user code and kernel code.
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:
- Divisibility constraints (
M % 128 == 0,K % 16 == 0) are checked at runtime, if at all. - The block-tile load is omitted but is significant code (with bounds-check, swizzling for bank conflicts, etc.).
- Tensor core constraints (FP16 inputs, FP32 accumulator, specific shapes) are enforced by template instantiation; mistakes produce template errors.
- Async memory not used (this is pre-Ampere style); a modern version would use
cp.async.
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:
- Pointers everywhere. The whole stride-and-offset arithmetic is at the user's fingertips.
- No type-level shape checking.
A_ptris just a pointer; whether the dimensions matchB_ptr's expected dimensions is the user's responsibility. - Tensor core selection is implicit —
tl.dotpicks the right tensor core variant based on input dtypes and Triton's autotuner. The user can't declare "this kernel must use Hopper WGMMA." - Hardware capabilities are not in the type. A Triton kernel that requires sm_80 features fails at compile time with a backend-specific error, not a type error.
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:
- Memory hierarchy in the type.
on=cuda:global,on=cuda:shared,on=cuda:register. Each has its own type constructor (TensorvsTile) reflecting the lifecycle. - Capability requirements are explicit.
requires capability(sm_80, mma_m16n16k16_f16_f16_f32)declares what hardware features the kernel needs. Targeting sm_70 fails to compile. - Block and warp structure are language constructs.
block_gridandwarp_gridare typed scopes;block_idandwarp_idare typed indices in those scopes. cp_asyncis a typed primitive. Its type ensures the source iscuda:globaland the destination iscuda:shared. Thecp_async_commitandcp_async_wait_alloperations are typed sync points.wmma_mmacarries its capability requirement. The compiler verifies the kernel satisfies it (it does — declared at the top).- No raw pointers. Tile addressing is done through bounded slicing (
a[block_m * 128 ..+ 128, k * 16 ..+ 16]), with the type system tracking shape compatibility.
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:
- CUDA Toolkit 12.x or later. Provides the NVCC compiler, libraries (cuBLAS for reference comparison), and the runtime API for the Rust runtime.
- MLIR's NVVM dialect. The lowering target for tile IR. Already part of LLVM/MLIR; just enable it.
- PTX assembler (
ptxas). Comes with the toolkit; invoked as a subprocess from the OCaml compiler (consistent with Section 5 §3). - Rust runtime extensions.
cudarcorcustcrate for CUDA Driver API bindings; device memory allocation, kernel launch, stream management.
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:
- Phase 0: type-level placement axis, movement primitives, NVVM dialect dependency.
- Phase 1: CUDA-target lowering for the algebraic operations (no tile authoring yet; just dispatch to existing kernels). Calls into cuBLAS/cuDNN for reference performance.
- Phase 5: full tile IR with kernel authoring; capability type system; custom tensor-core kernels.
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:
- Mojo's surface is Python-flavored; ours is OCaml-flavored. Surface preference; not technical.
- Mojo aims for full Python compatibility; we don't.
- Mojo's type system is being built; ours is committed in Phase 0 with refinement types and SMT discharge. We're more academically oriented; Mojo more engineering-oriented.
- Mojo's tile-level story is in active development; ours is forward-declared in Phase 0.
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
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 isTensor[..., 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.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.
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.
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.
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.