Stack coloring enabled by lifetime markers

TL;DR — A function needs one stack slot for each local array. Two arrays can share a slot only when the compiler knows that they are not in use at the same time. LLVM merges two slots when lifetime markers show that their lifetimes do not overlap. Clang and rustc end each lifetime at the end of a scope, so two arrays in one scope never share a slot. Clang also emits no markers at -O0, or for a variable when a jump enters its scope after its declaration. Vx frees a tensor after its last use. Since #2507, the lifetime of each tensor ends there too. For the same function, C++ and Rust use 8 KB of stack, and Vx uses 4 KB. On an H100, the stack of each GPU thread is in HBM. Since #2587, a Vx kernel gives two tensors one slot there too.

A task in Vx runs on a stack of fixed size, so each frame on that stack must be small. vxc puts a small tensor on the stack. Until #2507, it gave each tensor its own slot for the whole call. This post shows how Vx now gives two tensors one slot, in a function where C++ and Rust use two slots.

What a lifetime marker tells LLVM

A front end such as clang, rustc or vxc puts each local array in an alloca: a stack slot of the function. Two LLVM intrinsics mark when a slot is in use:

call void @llvm.lifetime.start.p0(ptr %first)   ; the slot is in use from here
call void @llvm.lifetime.end.p0(ptr %first)     ; the slot is free from here

A pass named StackColoring reads these markers. Its own description says that it “attempts to merge disjoint stack slots and reduce the used stack space”. When the marked lifetimes of two slots do not overlap, StackColoring gives them one address.

StackColoring runs on machine instructions, after instruction selection. At that point there are no variables and no scopes, only stack slots. So the pass learns from the markers when a slot becomes free. How much stack the pass saves depends on what the front end knows about the values.

This pass is not shrink wrapping. LLVM's ShrinkWrap pass “looks for safe point where the prologue and epilogue can be inserted”. It moves the code that creates the frame and saves registers. It does not merge stack slots.

One program in three languages

Each program below fills an array of 500 integers, reads it backwards, and then does the same with a second array. The backward read keeps the arrays in memory. In a version that read forward, LLVM at -O2 kept no tensor in memory, and the function had no frame to compare. Here is the Vx program:

import core::convert;

fn one_after_the_other(n : i64) -> i64 {
  let mut first : Tensor<i64, [500]> = Tensor<i64, [500]>::fill(0i64);
  for i in 0..500 {
    first[i] = n * i64::from(i) % 7;
  }
  let mut sum : i64 = 0;
  for i in 0..500 {
    sum = sum * 3 + first[499 - i];
  }
  let mut second : Tensor<i64, [500]> = Tensor<i64, [500]>::fill(0i64);
  for i in 0..500 {
    second[i] = sum * i64::from(i) % 11;
  }
  for i in 0..500 {
    sum = sum * 5 + second[499 - i];
  }
  return sum;
}

fn main() -> i32 {
  println!("sum {}", one_after_the_other(1));
  return 0;
}

The C++ program does the same work. A signed overflow is undefined behavior in C++, so this program wraps through uint64_t, as the Vx and Rust programs wrap. It needs C++20, which defines the conversion back to int64_t.

#include <cstdint>
#include <cstdio>

// Wrapping arithmetic, as in Vx and Rust. Signed overflow is undefined in C++.
int64_t mul(int64_t a, int64_t b) { return int64_t(uint64_t(a) * uint64_t(b)); }
int64_t add(int64_t a, int64_t b) { return int64_t(uint64_t(a) + uint64_t(b)); }

int64_t one_after_the_other(int64_t n) {
  int64_t first[500];
  for (int64_t i = 0; i < 500; i++)
    first[i] = mul(n, i) % 7;
  int64_t sum = 0;
  for (int64_t i = 0; i < 500; i++)
    sum = add(mul(sum, 3), first[499 - i]);
  int64_t second[500];
  for (int64_t i = 0; i < 500; i++)
    second[i] = mul(sum, i) % 11;
  for (int64_t i = 0; i < 500; i++)
    sum = add(mul(sum, 5), second[499 - i]);
  return sum;
}

int main() {
  std::printf("sum %lld\n", (long long)one_after_the_other(1));
  return 0;
}

The Rust program uses wrapping_mul and wrapping_add for the same reason:

#[no_mangle]
pub fn one_after_the_other(n: i64) -> i64 {
    let mut first = [0i64; 500];
    for i in 0..500 {
        first[i] = n.wrapping_mul(i as i64) % 7;
    }
    let mut sum: i64 = 0;
    for i in 0..500 {
        sum = sum.wrapping_mul(3).wrapping_add(first[499 - i]);
    }
    let mut second = [0i64; 500];
    for i in 0..500 {
        second[i] = sum.wrapping_mul(i as i64) % 11;
    }
    for i in 0..500 {
        sum = sum.wrapping_mul(5).wrapping_add(second[499 - i]);
    }
    sum
}

fn main() {
    println!("sum {}", one_after_the_other(1));
}

The three programs are equivalent. Each one prints sum -3181038691461447954.

The table gives the frame size of one_after_the_other in each program: the bytes that the sub sp, sp, ... instructions at its start subtract. In the rows with an inner scope, braces enclose each array and its two loops.

ProgramCompilerLowest level-O2
C++: both arrays in one scopeclang 22.1.48064 bytes8016 bytes
C++: each array in its own scopeclang 22.1.48064 bytes4016 bytes
Rust: both arrays in one scoperustc 1.95.08208 bytes8000 bytes
Rust: each array in its own scoperustc 1.95.08224 bytes4000 bytes
Rust: drop(first) before secondrustc 1.95.012224 bytes8000 bytes
Vx: before #2507vxc at e66345d38016 bytes8000 bytes
Vx: with #2507vxc at e299ee964016 bytes4000 bytes

The lowest level is -O0 for clang and vxc, and -C opt-level=0 for rustc. I measured on an arm64 Mac with macOS 27. For Rust, I measured the function in a library object, because the optimized executable keeps no separate copy of it. Each array or tensor holds 4000 bytes. So one slot needs about 4 KB, and two slots need about 8 KB.

C++: a lifetime ends at the end of a scope

Clang starts the lifetime of a local variable at its declaration, and ends it at the end of its scope. Here are the markers that clang emits at -O2 for the C++ function:

call void @llvm.lifetime.start.p0(ptr nonnull %2)   ; first
call void @llvm.lifetime.start.p0(ptr nonnull %3)   ; second
call void @llvm.lifetime.end.p0(ptr nonnull %3)
call void @llvm.lifetime.end.p0(ptr nonnull %2)

The function stops using first before it declares second. But both arrays are in the scope of the function body. So both lifetimes end at the end of the function, the two lifetimes overlap, and StackColoring keeps two slots.

A C++ programmer can fix this by hand: put each array in its own inner scope, with braces.

int64_t one_after_the_other(int64_t n) {
  int64_t sum = 0;
  {
    int64_t first[500];
    for (int64_t i = 0; i < 500; i++)
      first[i] = mul(n, i) % 7;
    for (int64_t i = 0; i < 500; i++)
      sum = add(mul(sum, 3), first[499 - i]);
  }
  {
    int64_t second[500];
    for (int64_t i = 0; i < 500; i++)
      second[i] = mul(sum, i) % 11;
    for (int64_t i = 0; i < 500; i++)
      sum = add(mul(sum, 5), second[499 - i]);
  }
  return sum;
}

Then the frame size is 4016 bytes at -O2. But the fix depends on the programmer, and it does not help at -O0. At that level, clang emits no lifetime markers, and LLVM does not run StackColoring. So a debug build always keeps both slots.

Clang also emits no markers for a variable when a jump enters its scope after its declaration. In this switch, the jump to case 1 enters the scope of a after the declaration of a:

int64_t by_kind(int kind, int64_t n) {
  switch (kind) {
  case 0:
    int64_t a[500];
    fill(a, n);
    return total(a);
  case 1:
    int64_t b[500];
    fill(b, n);
    return total(b);
  }
  return 0;
}

Here fill and total are in another file. At -O2, clang marks b but not a, and the frame size is 8016 bytes. With braces around each case, clang marks both arrays, and the frame size is 4016 bytes. A comment in clang's source gives the reason. With a jump into its scope, the lifetime of the variable “gets broken up into several regions in IR”, and that “requires more work to handle correctly”. So clang omits the markers.

Rust: the same rule

rustc also ends the lifetime of a local variable at the end of its scope. In its markers for the Rust program, both starts come before both ends, as in the markers of clang. The frame size is 8000 bytes at -C opt-level=2, and the second array is at sp + 4000, in its own slot.

The fix is the same as in C++. With an inner scope around each array, the frame size is 4000 bytes at -C opt-level=2. It stays at 8224 bytes at -C opt-level=0.

A call to drop(first) before second does not help. An array of i64 is Copy, so drop gets a copy, and rustc warns: “calls to std::mem::drop with a value that implements Copy does nothing”. At -C opt-level=0, the frame size grows to 12224 bytes, probably for the copy that the call gets.

Vx: a lifetime ends after the last use

Vx frees a tensor after its last use, as the post How Vx frees memory describes. The type checker writes these drops. The compiler prints them when you set VX_PRINT_DROPS=1:

drop in one_after_the_other: second, after the statement on line 16
drop in one_after_the_other: first, after the statement on line 9

Line 9 is the loop that reads first for the last time. Line 16 is the loop that reads second for the last time. Since #2507, each drop is also the end of a lifetime:

llvm.intr.lifetime.start %13 : !llvm.ptr   // first
llvm.intr.lifetime.end %13 : !llvm.ptr     // first, after the loop on line 9
llvm.intr.lifetime.start %11 : !llvm.ptr   // second
llvm.intr.lifetime.end %11 : !llvm.ptr     // second, after the loop on line 16

The lifetime of first ends before the lifetime of second starts. So StackColoring gives the two tensors one slot, and the frame size is 4000 bytes at -O2. The program needs no inner scope, and vxc gives the two tensors one slot at -O0 too. With the compiler from before #2507, the program printed the same sum, and the frame size was 8000 bytes at -O2.

How #2507 does it

The change is about 80 lines of C++ in two MLIR passes:

  1. The markers. The pass vx-promote-buffers-to-stack moves a small tensor from the heap to the stack. It adds a lifetime start for each tensor where the heap allocation was. Each drop, which freed the heap buffer before, becomes a lifetime end.
  2. The alloca. At that point, the stack slot is a memref and not an llvm.alloca yet. So the marker names the pointer of the buffer. A second pass, vx-lifetime-markers-on-allocas, runs after the lowering to the LLVM dialect. It traces the pointer through the memref descriptor to the alloca, because LLVM accepts a marker only on an alloca.
  3. The safe default. When the second pass cannot trace a pointer, it erases the marker. The tensor then keeps its slot for the whole call, as before.

When a loop body creates a tensor, the lifetime of that tensor starts and ends in each iteration. The test two_tensors_used_one_after_the_other_share_a_stack_slot.vx checks the markers in the IR, the frame size in the machine code, and the result.

On an H100, the stack is in HBM

A GPU thread has a stack too. NVIDIA calls it local memory, and its guide says that “local memory is off-chip”. So the stack is in the device memory of the GPU, with the tensors of the program. On an H100, device memory is HBM (high-bandwidth memory). The driver makes the stack of each thread larger for a launch of a kernel that needs more. It does not make the stack smaller after the launch.

Before #2587, vxc gave a stack slot only to a tensor that a kernel makes before its first loop or if. That slot lived until the kernel returned. vxc took each other tensor from device malloc: a malloc that a kernel calls on the GPU. #2587 gives the tensors of a kernel stack slots and lifetime markers, as #2507 does for a function.

Here is the program as a kernel. Each block has 256 lanes, and each lane makes its own two tensors. A lane in Vx is one thread of a block, which CUDA calls a thread. vxc does not accept a call such as i64::from in a lane loop, so the loop index cannot become an i64. So each tensor holds 1000 i32 values: 4000 bytes, as in the CPU program.

pub fn in_each_lane(out : &mut Tensor<i32, [?], Memory::HBM>, blocks : i32) -> void {
  spawn on(Topology::Device) {
    for b in 0..blocks across blocks {
      for l in 0..256 across lanes {
        let t = b * 256 + l;
        let mut first : Tensor<i32, [1000]> = Tensor<i32, [1000]>::fill(0);
        for i in 0..1000 {
          first[i] = t * i % 7;
        }
        let mut sum : i32 = 0;
        for i in 0..1000 {
          sum = sum * 3 + first[999 - i];
        }
        let mut second : Tensor<i32, [1000]> = Tensor<i32, [1000]>::fill(0);
        for i in 0..1000 {
          second[i] = sum * i % 11;
        }
        for i in 0..1000 {
          sum = sum * 5 + second[999 - i];
        }
        out[t] = sum;
      }
    }
  }
  return;
}

The CUDA kernel does the same work:

#include <cstdint>

// Wrapping arithmetic, as in Vx. Signed overflow is undefined in C++.
__device__ int32_t mul(int32_t a, int32_t b) { return int32_t(uint32_t(a) * uint32_t(b)); }
__device__ int32_t add(int32_t a, int32_t b) { return int32_t(uint32_t(a) + uint32_t(b)); }

// The host launches it with blocks of 256 threads: in_each_lane<<<blocks, 256>>>(out).
__global__ void in_each_lane(int32_t *out) {
  int32_t t = blockIdx.x * 256 + threadIdx.x;
  int32_t first[1000] = {};
  for (int32_t i = 0; i < 1000; i++)
    first[i] = mul(t, i) % 7;
  int32_t sum = 0;
  for (int32_t i = 0; i < 1000; i++)
    sum = add(mul(sum, 3), first[999 - i]);
  int32_t second[1000] = {};
  for (int32_t i = 0; i < 1000; i++)
    second[i] = mul(sum, i) % 11;
  for (int32_t i = 0; i < 1000; i++)
    sum = add(mul(sum, 5), second[999 - i]);
  out[t] = sum;
}

The table gives the local memory of each thread. That is the size of the .local array that holds the stack of a thread in PTX, the assembly language of NVIDIA GPUs. The PTX is for an H100 (sm_90). I compiled each kernel to PTX but did not run it, because I have no GPU.

KernelCompiler-O0-O2
CUDA: both arrays in one scopeclang 22.1.48032 bytes8000 bytes
CUDA: each array in its own scopeclang 22.1.48032 bytes4000 bytes
Vx: before #2587vxc at 4ce0ffc4none: malloc and free for each tensornone: malloc and free for each tensor
Vx: with #2587vxc at 05bb8e7d4000 bytes4000 bytes

Clang ends each lifetime at the end of a scope in CUDA too. So the two arrays share a slot only with an inner scope for each, and only at -O2. vxc compiles the device code at LLVM's -O2 for each vxc level, so both of its columns are the same. HIP is the version of CUDA for AMD GPUs. For an AMD MI300 (gfx942), clang gives the HIP versions of the two CUDA kernels 8004 and 4004 bytes of scratch memory at -O2. Scratch memory is the name that AMD uses for local memory.

Before #2587, each lane freed first before it allocated second, so a lane held 4000 bytes of the heap at a time. NVIDIA's guide says that the heap for device malloc holds 8 MB by default. When the heap has too little memory, malloc returns null. The Vx runtime does not set a larger heap. So at most about 2000 lanes can hold a tensor at the same time. An H100 runs up to 2048 threads on each of its 132 streaming multiprocessors (SMs).

With #2587, the two tensors of a lane use the same 4000 bytes of local memory, one after the other. The lane calls no malloc. The test two_tensors_in_a_gpu_kernel_share_local_memory.vx checks this in the PTX.

What it does not do yet

The pass is the same

Clang, rustc and vxc all use the same StackColoring pass. The pass keeps a slot reserved until the end of its marked lifetime. Clang and rustc end the lifetime of a local variable at the end of its scope. vxc ends the lifetime of a tensor at its last use, where its type checker already puts the drop. So a function needs fewer stack slots, and the stack of each task stays smaller. On an NVIDIA GPU, the same markers make the local memory of each thread smaller. On an H100, that memory is HBM.

Sources

Vx is Apache 2.0 with the LLVM exception — install it and try the examples.