Vx

Tiles, without a tile type

Aditya Kumar · 21 September 2026

TL;DR — Tile-based GPU programming structures kernels around sub-arrays of a tensor. Vx has no built-in tile type, yet nested loops over a Tensor<T, [..], Memory::SMEM> compile to PTX: 8 bytes of .shared storage, bar.sync instructions, and grid-strided execution over %ctaid.x and %tid.x. Tracking memory space in types turns shared-memory overcommit into a compile-time error. What Vx still lacks is a dedicated tile abstraction: there are no tile arithmetic operations, no tensor-core instructions, only one grid dimension, and loop bounds must be static literals.

A developer recently asked if Vx could support tile-based GPU programming like NVIDIA's cuTile. To find out, I wrote a shared-memory staging kernel in Vx and examined the compiler's PTX output.

Tiles versus threads

In standard CUDA programming, you write code from the perspective of an individual thread. The programmer coordinates the entire cooperative thread group: computing array offsets with threadIdx and blockIdx, staging data into __shared__ memory, synchronizing threads with __syncthreads(), and loading elements back out.

Tile-based programming frameworks like Triton or cuTile raise this level of abstraction. You express operations directly on multi-dimensional chunks of data called tiles. The compiler then handles the thread layout and schedules transfers between global DRAM and fast on-chip SRAM.

Working at the tile level eliminates common concurrency bugs. In CUDA, an off-by-one indexing error during staging silently corrupts numbers without crashing. A missing __syncthreads() creates a subtle data race that might only appear once in ten runs.

Vx takes a different angle by placing memory hierarchies directly into the type system. Every tensor type tracks which physical memory space holds it and which processor can access it. Can this same placement system compile a working tile-style kernel?

The staging kernel

To test this question, I wrote a minimal staging kernel. It copies slices of an array from GPU global memory (HBM) into shared memory (SMEM), synchronizes, and writes the results back to global memory. The arithmetic is intentionally simple so the memory patterns and loop structures remain easy to follow:

Memory CPU_DRAM {}
Memory GPU_HBM {
  within: Memory::CPU_DRAM, capacity: 40 GiB, bandwidth: 3 TB/s, managed: cached
}
Memory SMEM {
  within: Memory::GPU_HBM, capacity: 228 KiB, bandwidth: 128 B/cyc,
  granule: 1 KiB, managed: explicit, scope: sm
}
Topology Dev {
  arch: nvptx64,
  memory: Memory::GPU_HBM,
  visible: [Memory::GPU_HBM, Memory::SMEM],
  transfer Memory::CPU_DRAM -> Memory::GPU_HBM : 63 GB/s,
  transfer Memory::GPU_HBM -> Memory::SMEM
}

fn main() -> i32 {
  let mut q_h = Tensor<f32, [8]>::uninit();
  let mut o_h = Tensor<f32, [8]>::uninit();
  for i in 0..8 {
    q_h[i] = (i as f32);
    o_h[i] = 0.0;
  }
  let q = transfer(q_h, Memory::GPU_HBM);
  let mut o = transfer(o_h, Memory::GPU_HBM);

  spawn on(Topology::Dev) {
    // The tile. Its memory space is part of its type.
    let mut qs = Tensor<f32, [2], Memory::SMEM>::uninit();
    for bq in 0..4 {              // one iteration per tile
      for qi in 0..2 {            // cooperative fill
        qs[qi] = q[bq * 2 + qi] * 2.0;
      }
      barrier();
      for qi in 0..2 {            // cooperative consume
        o[bq * 2 + qi] = qs[qi] + 1.0;
      }
      barrier();
    }
  }
  print(o);
  return 0;
}

Inside spawn on(Topology::Dev), the source omits traditional CUDA boilerplate. It contains no blockIdx, no threadIdx, no launch configuration parameters, no __shared__ annotations, and no __syncthreads() calls. Instead, the kernel defines four elements:

How the compiler maps loops to PTX

To compile this kernel, the compiler must determine which loop runs over thread blocks and which loop runs over individual threads. It does this through a check called parallel_two_level. To guarantee that threads never race or leave unwritten gaps, the analysis verifies every index expression that writes to a captured array.

A captured write index must follow one of two patterns:

  1. The thread loop variable itself (qi).
  2. The affine expression bq * K + t, where K equals the exact trip count of that thread loop.

The exact multiplier is critical. If t spans 0..4, an expression like bq * 8 + t leaves unwritten gaps in the output tensor. If t spans 0..16, different iterations overlap and overwrite the same indices. The multiplier bq * 2 + qi is exact, ensuring every (block, thread) coordinate owns one distinct element.

The staging kernel satisfies this proof, emitting the following MLIR attributes:

{arch = "nvptx64", vx_parallel_threads = 2 : i64, vx_parallel_trip = 4 : i64, vx_parallel_two_level}

The compiler derives the launch geometry directly from the loop bounds: 4 blocks of 2 threads. It passes launch=4,2 to the runtime dispatch payload and lowers the loop body to PTX:

.visible .entry vx_npu_kernel_0(
        .shared .align 16 .b8 vx_npu_kernel_0_smem_0[8];
        mov.u32         %r11, %ctaid.x;
        mov.u32         %r1, %tid.x;
        mov.u32         %r2, %ntid.x;
        mov.u32         %r3, %nctaid.x;
        ...
        ld.global.b32   %r6, [%rd12];      // read one element out of HBM
        st.shared.b32   [%rd13], %r7;      // write it into the tile
        bar.sync        0;
        ld.shared.b32   %r8, [%rd14];      // read it back out of the tile
        st.global.b32   [%rd15], %r9;      // write the result to HBM
        bar.sync        0;

(An excerpt, with the address arithmetic between these lines cut. The comments are mine.)

The generated PTX has the ingredients of a hand-written staging kernel:

The compiled kernel also employs grid-strided loops: the outer block loop steps by %nctaid.x, and the inner thread loop steps by %ntid.x. Grid-striding allows the kernel to execute correctly across any hardware launch configuration. The launch=4,2 payload is a sizing hint: it is the shape that leaves no thread idle and no iteration doubled, and any other shape still computes the same answer.

At a 1×1 launch (one block of one thread), the parallel loops reduce directly to the original serial loop nest. That degeneracy is why the same source also runs on a CPU host with no GPU in it, producing the verified result:

$ vxc tile_kernel.vx
[1,  3,  5,  7,  9,  11,  13,  15]

That output came from the host, not from a GPU. The degeneracy is the reason it is the right number: one launch shape is the serial program, so the host path and the device path stay the same program instead of two implementations that have to be kept in agreement.

There is one case where this equivalence breaks down: placing a barrier() inside a thread loop. No serial execution order can satisfy a barrier within a single thread's loop iterations. In that scenario, the compiler marks the payload with coop=1, instructing the runtime to reject host fallback and avoid returning corrupted data.

Two compile-time checks CUDA lacks

Writing this kernel triggered two compiler errors that have no direct equivalents in traditional GPU programming.

The first error occurred when I initially attempted to copy a row of the HBM tensor directly into the shared-memory tile using a standard assignment (=). That draft used two-dimensional tensors, which is why the rows below are [4] wide and the kernel above uses a [2] tile:

Error[E3004]: Type mismatch in assignment: cannot assign Tensor<f32, [4], GPU, GpuHbm>
              to Tensor<f32, [4], Custom("SMEM"), Custom("SMEM")>

In CUDA C++, pointers do not track memory spaces at the type level. An assignment that crosses address spaces compiles without complaint, and the mistake surfaces when the program runs. In Vx, memory space is an explicit type parameter. An HBM tensor and an SMEM tile are incompatible types. Data movement across spaces requires an explicit transfer() call, which succeeds only along routes permitted by the machine topology. (The compiler renders internal space names like Custom("SMEM") in this message, which is a known rough edge.)

The second error prevents shared-memory overflow. Shared memory is a strict per-SM budget, and exceeding it is ordinarily something the device toolchain or the launch tells you, on a machine that has the part in it. In Vx the topology declares the capacity of Memory::SMEM and the tiles are statically sized, so the compiler adds up the tiles live at once before it emits any code:

Error[E6010]: the working set placed in memory space 'SMEM' (3 tiles) sums to 147456 bytes,
              over its 131072 byte capacity; place fewer/smaller tiles or declare it `overcommit`

This message comes from tests/frontend/fail/three_dtcm_tiles_do_not_fit_a_cortex_m7.vx. That test targets the tightly-coupled memory of an ARM Cortex-M7 microcontroller, not a GPU. The check keys on the declared capacity of a space rather than on what kind of part holds it, which is the point: it runs on a laptop, with no device toolchain involved.

What is still missing

The kernel above demonstrates that Vx can compile the physical mechanics of a shared-memory staging pipeline. However, Vx does not yet provide a complete tile programming model. Five major features remain unimplemented:

Where this leaves us

Vx can already lower a shared-memory staging pattern to a functional GPU device image. It does so by using static analysis to prove properties of nested loops. The developer still writes explicit element loops, and hardware tensor instructions remain unavailable.

Tile frameworks and Vx address complementary concerns:

These two concepts fit together cleanly. Building tile operations on top of typed memory spaces would combine their strengths: developers could write concise matrix operations while the type system guarantees memory safety across hardware hierarchies. The existing loop lowering and memory tracking provide a solid foundation for that abstraction.

The two-level block/thread mapping is Vx#379 and general device-kernel emission is Vx#251. The kernel above is the shape of tests/backend/pass/custom_topology_device_image.vx; the shared-memory budget check is tests/frontend/fail/three_dtcm_tiles_do_not_fit_a_cortex_m7.vx. Vx is Apache 2.0 with the LLVM exception — install it, or read What a pointer forgets for why the memory space is in the type in the first place.