vx-lang / vx-lang/Vx

A fixed-size local array is a heap allocation: `memref.alloc` where `memref.alloca` belongs

Open
#641 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug codegen
Dominant language
Rust
Stars
14
Forks
2
Avg merge
12h 42m
Merged PRs (30d)
61

Description

Every array literal becomes a malloc, however small and however plainly local.

fn main() -> i32 {
  let i = comptime {
    let mut a = 0;
    let l = [0, 1, 2, 3];
    for j in 1..10 {
        a = a + j;
    }
    l
  };
  return i[0];
}

The comptime fold is fine: the loop and a are gone from the output, exactly as #563
asked. What survives is the array, and it survives as a heap allocation.

define i32 @main() {
  call void @vx_init_signals()
  %1 = call ptr @malloc(i64 16)
  %2 = getelementptr inbounds nuw i32, ptr %1, i64 0
  store i32 0, ptr %2, align 4
  ...
  %7 = load i32, ptr %6, align 4
  ret i32 %7
}

MLIR is doing what it was asked. finalize-memref-to-llvm lowers memref.alloc to
malloc and memref.alloca to llvm.alloca; we ask for the first one. lower_array
(src/hir/flatten.rs:3468) turns the literal into an Opcode::TensorAlloc, and
src/codegen/flat/emit/tensor.rs:90 lowers that:

} else {
    self.body += &format!("  {n} = memref.alloc({sizes}) : {memty}\n");

Four i32s, size known at compile time, never leaving the frame. That is a stack slot.

-O2 hides it in the example above, because there the buffer is dead. It is not hidden
when the result is used -- see #642 for what it costs at run time.

The emitter already knows the shape of the fix. The shared-memory branch ten lines above
the line in question does exactly the right thing:

self.emit_slot(&format!("  {n} = memref.alloca() {{alignment = 16 : i64}} : {smty}\n"));

emit_slot places the op in the entry block, so an array declared inside a loop takes one
slot rather than a fresh one per iteration.

Two things gate the switch from alloc to alloca.

The buffer can escape. A tensor leaves the frame by being returned:

fn make() -> Tensor<i32, [4]> {
  let l = [1, 2, 3, 4];
  return l;
}

lowers to return %alloc : memref<4xi32>. An alloca there is a dangling stack pointer.
So this needs a check over the flat-HIR stream: a TensorAlloc whose register never
reaches a Ret, and is never stored into anything that outlives the frame, is the one
that converts. #643 would remove this gate for the returned case specifically, by
giving the callee a caller-allocated slot to write into -- at which point almost every
local buffer qualifies.

Size. memref.alloca on a Tensor<f32, [4096, 4096]> is 64 MB of stack, which is how
#440 fails. So the rule wants both conditions: statically shaped and under a threshold
becomes alloca; anything else stays alloc and owes a free.

The unit test emits_verifiable_tensor_alloc_store_read (src/codegen/flat.rs:2376) pins
memref.alloc() : memref<4xi32> and will need updating with the behavior.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at lower_array in src/hir/flatten.rs and the TensorAlloc handling in src/codegen/flat/emit/tensor.rs, then inspect emit_slot and flat-HIR return or storage paths for escape checks. Update src/codegen/flat.rs test emits_verifiable_tensor_alloc_store_read and verify that small, statically shaped nonescaping arrays use alloca while escaping, dynamic, or oversized arrays retain alloc.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend, compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
72/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.