Can a library describe the memory hierarchy?

Aditya Kumar · 8 October 2026

TL;DR — A question we keep getting: why make memory placement part of the language, when a library can encode it? Part of it can. A Rust library can put the memory space in a tensor's type, and the Rust compiler will then reject a host read of GPU memory. What a library cannot do is the rest of the hierarchy: add up how much memory a call path uses, compute shapes from other shapes on stable Rust, or check the code that runs on the device. That is the part Vx puts in the compiler.

An earlier post, float * fails to distinguish host DRAM from GPU HBM, argued that the memory space belongs in the type. This one asks whether that type has to come from the language.

The question

The usual example is a Rust tensor library such as Burn. Its Tensor<B, 2> carries the backend and the rank in its type. The shape and the device are values, checked when the program runs: multiply a [4, 4] tensor by a [4, 3] one the wrong way round, or mix tensors from two devices, and you find out at run time.

That is a choice Burn made, not a limit of Rust. A library can move both checks to compile time.

The Rust version

One empty type per memory space, a trait for "the CPU can read this", and const generics for the shape:

use std::marker::PhantomData;

struct CpuDram;
struct GpuHbm;

trait HostVisible {}
impl HostVisible for CpuDram {}

struct Tensor<const R: usize, const C: usize, M> {
    data: Vec<f32>,
    space: PhantomData<M>,
}

fn sum_on_host<const R: usize, const C: usize, M: HostVisible>(t: &Tensor<R, C, M>) -> f32 {
    t.data.iter().sum()
}

fn matmul<const M: usize, const K: usize, const N: usize, S>(
    a: &Tensor<M, K, S>,
    b: &Tensor<K, N, S>,
) -> Tensor<M, N, S> { /* ... */ }

With rustc 1.98, a host read of a Tensor<4, 4, GpuHbm> fails with error[E0277]: the trait bound `GpuHbm: HostVisible` is not satisfied, and a matmul of a [4, 3] tensor by a [4, 4] one fails with error[E0308]: mismatched types. Tensors in two different spaces cannot be multiplied either, because S must be the same type for both.

So for "this tensor lives in GPU memory and the CPU cannot read it", the critic is right: a library does it. Vx does not need to be a language for that part alone.

Where the library stops

Shapes computed from other shapes. matmul works because its result only reuses its inputs' sizes. A concat whose result has A + B rows does not:

fn concat<const A: usize, const B: usize, M>(
    x: Tensor<A, 1, M>,
    y: Tensor<B, 1, M>,
) -> Tensor<{ A + B }, 1, M>

error: generic parameters may not be used in const operations

Arithmetic on const generics needs the generic_const_exprs feature, which is still unstable. Reshape, concat, split and padding all need it. Vx's compiler does the arithmetic itself: t.reshape([15]) on a [4, 4] tensor is refused with reshape arithmetic mismatch: source has 16 elements, target has 15.

How much memory a call path uses. A GPU's shared memory is 228 KB per multiprocessor on an H100. Whether a kernel's tiles fit depends on what each function holds while it calls the next one. Here each function fits in a 4 MiB space by itself, and together they do not:

Memory CPU_DRAM {}
Memory SMEM {
  within: Memory::CPU_DRAM, capacity: 4 MiB, managed: explicit
}
Topology Dev {
  arch: nvptx64,
  memory: Memory::SMEM,
  visible: [Memory::SMEM],
  transfer Memory::CPU_DRAM -> Memory::SMEM : 10
  transfer Memory::SMEM -> Memory::CPU_DRAM : 10
}

fn inner() -> i32 {
  let y = Tensor<f32, [1024, 768]>::uninit();
  let _tile = transfer(y, Memory::SMEM);
  return 1;
}

fn outer() -> i32 {
  let x = Tensor<f32, [1024, 768]>::uninit();
  let tile = transfer(x, Memory::SMEM);
  let r = inner();
  let _back = transfer(tile, Memory::CPU_DRAM);
  return r;
}

fn main() -> i32 {
  return outer();
}

Error[E6027] at 22:11: the working set along call path 'outer -> inner' in memory space
  'SMEM' peaks at 6291456 bytes, over its 4194304 byte capacity: 'outer' holds 3145728 bytes
  across its call, 'inner' itself peaks at 3145728 bytes

Rust's types check each function on its own. To get this error from a library, every function would have to state its peak memory in its signature, and every caller would have to add its own by hand. No one would write that. The Vx compiler sees every function and adds them up.

The machine. In the Rust version, the hardware is the list of impl HostVisible lines. They say who can read what, and nothing more. The first lines of the Vx example above are a machine description: what each memory holds, which processors see it, and what a transfer between two spaces costs. The compiler reads it from a machine file, so describing a new chip means writing a new file, not a new set of types.

The code that runs on the device. Rust compiles a crate for one target. GPU code lives in a separate crate built for the GPU (rust-gpu, or the nvptx64 target), or in a smaller language inside macros, which is what Burn's CubeCL does. The host's types do not reach inside it. In Vx a spawn on block is in the same function as the code around it, and the compiler checks it as code for the device it runs on:

fn main() -> i32 {
  let host = Tensor<f32, [4]>::new();
  let gpu = transfer(host, Memory::GPU_HBM);
  let _p = spawn on(Topology::NPU[0]) {
    let _x = gpu[0];
  };
  return 0;
}

Error[E6003] at 5:14: 'gpu' lives in GPU_HBM but NPU[0] sees only [NPU_HBM];
  insert an explicit transfer to NPU_HBM (cost 100 on the declared path)

Transfer to NPU_HBM instead, as the error says, and the program compiles and runs.

So which is it?

Putting a memory space in a type is a library feature, and Rust does it well. A memory hierarchy is more than that: sizes that add up across calls, shape arithmetic, costs read from a description of the machine, and device code checked against what the device can see. Each of those needs the compiler to see the whole program and the machine together. A library gets one function's types at a time.

The heterogeneous model covers placement and transfer, and E6027 and E6003 have their own pages. The Rust examples were checked with rustc 1.98.1 on stable.