[totaland]

Fixes #1393. Fixes #1394.

Why

On the GPU, a ! job dies with bend: memory fault (machine stack overflow?) once a lane's stack fills. The same binary with --gpu off finishes. Two shapes hit it:

• Metal: non-tail recursion under a bang faults at 1023 list elements; CPU completes through 100000 #1393: a bang's root that recurses non-tail (List.append over a list) faults at 1023 elements.
• Metal: one ! over a two-leaf fork tree (a 365-fork spine beside a 3-fork spine) faults; each leaf alone, a balanced 365+365, and --gpu off all pass #1394: a fork tree with one long spine beside a short one (365 + 3) faults, while each spine alone passes.

The limit is the same on CUDA (see the comment by @nuxyel): it hits at 3 + w × depth > 2048 words, where w is the number of words a level keeps.

Both programs from the issues, built from upstream 60fa05dc and from this branch, on an M1 Pro:

run	upstream, --gpu off	upstream, Metal	this PR, --gpu off	this PR, Metal
deep 1022	521738	521738	521738	521738
deep 1023	522760	fault	522760	522760
deep 2048	2096135	fault	2096135	2096135
deep 100000	704982711	fault	704982711	704982711
deep 1000000	1783293671	fault	1783293671	1783293671
lsum 2048	2096128	fault	2096128	2096128
lsum 1000000	1783293664	fault	1783293664	1783293664
forks 64 800 365	4280936846	4280936846	4280936846	4280936846
forks 64 800 365 365	4266906396	4266906396	4266906396	4266906396
forks 64 800 365 3	2063974529	fault	2063974529	2063974529
forks 64 800 174 3	3985679902	fault	3985679902	3985679902
forks 64 418 365 3	875539697	fault	875539697	875539697
forks 64 800 365 3 3 3 3	4003022170	fault	4003022170	4003022170

lsum is deep with a hand-written non-tail sum (h + LSum(t)) in place of List.append. Its source is at the end.

How it failed

A bang's root and a fork kid run on the lane's own STAK_LEN-word (2048) stack on the device, as tests/compile/fork_leaf_budget.bend says. A non-tail call pushes its continuation frame through WL_ROOM. When the next push would not fit, WL_ROOM posted ERR_DEEP and the lane returned. Nothing moved the frames anywhere else, so the depth limit was about 1022 frames of List.append. In #1394's unbalanced tree, a leaf job runs above a deeper fork spine on the same lane, so it gets less of the stack.

Change

WL_ROOM(N, D) now takes D, the number of words that the calling def holds on the stack (emit_fork passes depth - ws.length). When a push would overflow on DEVICE, it calls lane_spill instead of failing:

1. If the def's D words, the frame and a 3-word exit frame would not fit even on an empty stack (d + n + 3 > STAK_LEN), it returns 0 and WL_ROOM posts ERR_DEEP as before. This check comes before any heap block is taken.
2. Otherwise it copies the lane's stack to a heap block of cls_fit(STAK_LEN) words.
3. It rebuilds the bottom of the stack as one exit frame: the task word term_tsk(FID_EXIT, b) points at the block, the next word holds k, the number of words below the def's own, and the third word is FID_EXIT.
4. It copies the def's D words back above that frame with WL_ARGS and returns the new sp. The push then goes ahead on an almost empty stack.

When a return reaches that frame, the FID_EXIT case in work_loop sees a continuation whose function id is FID_EXIT. It copies the k words back with WL_ARGS, frees the block and returns into the restored top frame. Spills nest, because each spilled block keeps the exit frame below it, so the depth is bounded by the heap. When the heap runs out, the run stops with the usual out-of-memory message (--gpu 512MB, 80,000,000 levels). The host lane is unchanged, because WL_ROOM only checks on DEVICE.

lane_spill is FAR, so Metal emits one copy of it instead of inlining it at every WL_ROOM site. An earlier version inlined it: symreg's .gpu grew from 220 KB to 283 KB, and the COMPILER column of the perf gate failed (symreg 1.24x, merkle 1.17x).

tests/compile/bang_lane_spill.bend covers:

• deep!(5000n): one bang's root, 5001 List.append frames (Metal: non-tail recursion under a bang faults at 1023 list elements; CPU completes through 100000 #1393).
• two!(365n): a 365-fork spine beside a 3-fork spine (Metal: one ! over a two-leaf fork tree (a 365-fork spine beside a 3-fork spine) faults; each leaf alone, a balanced 365+365, and --gpu off all pass #1394).
• pair!(600n): the first spill happens at a fork call where the def still holds two words on the stack (D > 0), so a build that drops those words fails it.

On main, its compiled run faults on the GPU.

Verification

Merged onto main cece1cbf, Apple M4 Pro and the mini cluster:

• repo PASS: 54 / 54. bend2/comp.ts is 65,723 ttok, under the 66,000 cap (main is 65,424, so +299).
• test PASS: 1693 / 1693.
• ping PASS: 63 / 63.
• safe: 1098 agree, 0 false rejects, 0 timeouts.
• perf PASS: 124 / 124, the same as main. No bench fills a lane, so no spill runs, and every runtime cell is within noise of main. In the COMPILER column, symreg is 1.14x on this branch against 1.09x on main, and merkle 1.15x against 1.11x.
• deep!(1000000n) (deep(p) + 1) prints 1000000 on Metal, the same as --gpu off.

CUDA: @nuxyel built 6764c33 on an RTX 4050 (CUDA 12.9). It runs every depth past the old limit with the CPU's results, up to deep!(100000n), and the two loop kernels they time run at the same speed.

Earlier, by @totaland, on an M1 Pro against upstream 60fa05dc: every test in tests/ gave byte-identical output on both trees (interpreter, compiled and JS), except the new test's compiled lane, which faults upstream. The Metal benches (mandelbrot, tree-matmul, tree-bitonic, nbody) were within noise.

Limits

• Spill cost: a spill copies up to STAK_LEN words to the heap and back. A loop of calls at exactly the depth where the stack fills spills and restores the whole stack on every call: 100,000 calls there take 24.7 s, against 0.27 s one frame shallower. Main faults at that depth, so this is not a regression. Avoiding it would need a partial restore.
• Not tested: Apple GPUs other than the M1 Pro and M4 Pro.

lsum.bend

import Base

def Gen(n: Nat, +k: U32, acc: List<&2, U32>) -> List<&2, U32>:
  match n:
    case 0n:
      acc
    case 1n+p:
      Gen(p, (k + 1 : U32), k <> acc)

# Non-tail sum: one frame per element.
def LSum(xs: List<&2, U32>) -> U32:
  match xs:
    case Nil{}:
      0
    case Con{h, t}:
      (h + LSum(t) : U32)

def Run(m: Maybe<&2, U32>) -> IO(Unit):
  match m:
    case Some{n}:
      IO.print(U32.show(LSum!(Gen(U32.to_nat(n), 0, []))))
    case None{}:
      IO.die(Unit, 2, "usage: lsum <n>")

def Args(args: List<String>) -> IO(Unit):
  match args:
    case Con{_, Con{n, _}}:
      Run(U32.read(n))
    case _:
      IO.die(Unit, 2, "usage: lsum <n>")

def main() -> IO(Unit):
  do IO<Unit>:
    args : List<String> <- IO.args()
    Args(args)

[Lorenzobattistela]

Thanks @totaland. I reproduced both issues on an M4 Pro at the same boundaries as the reports: List.append passes at 1022 and faults at 1023, and the fork tree passes at 173+3 and faults at 174+3 (d 417 passes, 418 faults). With this PR merged on main 8578b7c, every case prints the same result as --gpu off, up to 1,000,000 frames. The new test fails on main as it should, and ping (48/48) and test (1598/1598) pass.

What blocks the merge:

1. bend2/comp.ts is 64,275 ttok, over its 64,000 cap, so the repo gate fails (53/54). The caps are permanent, so the change has to fit under 64,000.

2. On Metal, lane_spill gets inlined at every WL_ROOM site. symreg's .gpu grows from 220 KB to 283 KB, its build takes 13% longer, and the perf gate fails on the COMPILER column (symreg 1.24x, merkle 1.17x).

3. The test never spills at a call where the def still holds words on top of the stack (D > 0). A build that skips copying those D words back still passes it.

Smaller things: lane_spill takes its heap block before the ERR_DEEP check. Also, a loop of calls at exactly the depth where the stack fills spills and restores the whole stack on every call: 100,000 calls there take 24.7 s, against 0.27 s one frame shallower. Main faults at that depth, so this isn't a regression, but it's worth knowing about. We have no CUDA machine, so the CUDA lane is untested here.

[totaland]

I have CUDA I will test again and let you know. Also will have a look again at the comp.ts

[Lorenzobattistela]

If comp.ts is not possible to shrink that's not a no-go though, no need to push for simplifications that hurt our code style. Mostly I'm looking for code simplifications. If it's not possible, just let me know, we're discussing the possibility to raise the cap of comp.ts specifically.

[totaland]

That will be great. Also when I put the kernel code out of comp.ts and into packages it reduce decoding time and reduce tok/s about 1.6%. Can we have a discussion in team about that too? I am building inference engine with bend.

[Lorenzobattistela]

Thanks @totaland, and good to hear you're building an inference engine on Bend.

On moving the kernels: Base and the compiler are only for really basic things, the primitives every Bend program needs. Inference kernels (rmsnorm, quantized matmul, attention and so on) are too specific for that.

The right place for them is your own package. If you need hand-written C there, a foreign effect can import it (def e(x: A) -> IO(B): with import "./e.c", see the GUIDE), and that C can drive Metal or CUDA itself. It runs from the event loop on the host, not inside !, but for an inference engine that is usually how kernels get launched anyway.

About the 1.6%: if you share a repro (the program, the machine and before/after numbers over several runs), we're happy to look at why the package version is slower. If it comes from something general in the compiler, we'd rather fix that for everyone than move model-specific code into it.

[nuxyel]

The same limit holds on CUDA (Bend 2.0.37, and main 881cb2e today; RTX 4050 Laptop, driver 610.57.04, CUDA 12.9): a non-tail recursion inside a ! faults with bend: memory fault (machine stack overflow?) once it needs more than the lane's 2,048 stack words, and the same binary with --gpu off finishes. The smallest program I found:

import Base

def deep(n: Nat) -> U32:
  match n:
    case 0n:
      0
    case 1n+p:
      (deep(p) + 1 : U32)

def main() -> IO(Unit):
  IO.print(U32.show(deep!(5000n)))

GPU: the fault, exit 1. --gpu off: 5000. Without the !: 5000. The thresholds match STAK_LEN = 1 << 11, at 3 + w x depth > 2048 words, w the words a level keeps:

recursion	words a level	last depth that runs	first that faults
deep(p) + 1	1	2045	2046
List.append	2	1022	1023
deep3(p, a, b) + a + b	3	681	682

Every run was deterministic, and --gpu 4GB does not change it (the stack is a fixed block outside the heap). I hit it with a matrix product that appended row lists at each join of its fork tree (fine on the CPU, a fault on CUDA); replacing List.append with two tail loops fixed it. Two small things that would have saved me time: the message says "machine stack", but it is the device lane's stack, and neither bend guide nor SHADERS.md gives the 2,048-word figure.

I built this PR (6764c33) and its base (3f9c40e) from source and ran the programs above on CUDA. The base faults at every depth past the limit; the PR runs all of them with the CPU's results, my matrix product included (9031405), and a depth of 100,000 (deep!(100000n)). It did not change the speed of two loop kernels I time (a flat loop over 16,384 leaves: 213-219 ms against 213-259; a tiled matmul: 172-175 ms both), measured with the CPU busy.

[Lorenzobattistela]

Gates green merged onto main cece1cb; the spill-at-fill-depth cost is tracked in #1790.