RFC: pass a stack pointer through handlers to avoid loads and round trips

[matthargett]

I made this an RFC because it looks like the acc branch is already poking around in this area. I wanted to get some exhaustive numbers up to help drive those decisions.

Problem

the flushes themselves are only 1.8% of pipeline slots in the model and 2.4% in tinywasm. Most of the cost is loads waiting on older stores. At 64 copies, back-end stalls rise from 5.3% to 25.3% of slots, and execution latency is 17× higher. In tinywasm they rise from 13.2% to 19.7% with 3× the execution latency. That matches iPhone 12, where audio DSP gained 89M back-end cycles per call with our changes.

Proposal

passing a stack pointer through the handlers, the way #75 passes the instruction slice. this should make every stack address computable without a load, as well as leaving nothing for the predictor to guess. A top-of-stack register should also remove the top value's round trip. I can submit a PR for this, but I didn't want to stomp on the acc branch activity.

Evidence

To prove/disprove, I made a reduction model of tinywasm's dispatch: K copies of one handler that updates the stack top in place, each copy running the same 25 instructions and dispatching with become through a table, in a loop.

copies	1–8	16	20	24	28	32	64
flushes / 1k dispatches	0.00	0.10	0.22	0.54	3.66	4.26	3.88
cycles / dispatch	5.40	5.45	5.72	5.84	6.53	6.64	7.14

with no aliasing between handlers, 64 copies run at 5.28 cycles per dispatch and never flush. In this testing, I figured out that live pairs count, not total pairs! the same 64 copies run in blocks of 64 (few pairs active at once) flush 0.12 times per 1,000 dispatches.

Inside tinywasm is a loop whose statements cycle through D distinct stack binops, so 5 to 45 distinct handler pairs, with instructions per dispatch unchanged. Flushes and cycles are flat up to 11 pairs, step up between 11 and 19, then plateau. Our faster frameless handlers from #77 pay more: +9.4% cycles across the range, against +5.2% on #75. The real loops in my benchmark suite are well past that: Audio DSP needs 40 distinct handler pairs for 90% of its dispatches, xmrsplayer 136 and graphql 226.

so this refines the original guess: the flushes themselves are only 1.8% of pipeline slots in the reduction and 2.4% in tinywasm. Most of the cost is loads waiting on older stores. At 64 copies, back-end stalls rise from 5.3% to 25.3% of slots, and execution latency is 17× higher. In tinywasm they rise from 13.2% to 19.7% with 3× the execution latency. That M4 e-core result matches the iPhone 12, where audio DSP gained 89M back-end cycles per call with our changes.

[explodingcamera]

Yeah, this is interesting. I’d like to try passing the stack pointer through the handlers, though I’m a little worried about adding another argument when I’d also like to keep room in registers for accumulators later.

No need to work around the acc branch, that was (outside of improvements to the parser that i pulled from there) a failed experiment 😄. My guess is that tinywasm’s fused opcodes already skip a lot of stack accesses, and my changes got in the way of that. I've tried caching the top of the stack in the past too, but the extra branch in the common path wiped out the benefit. I think I’ll need to rethink the parser and IR more thoroughly to make accumulators work, so it’d be useful to measure the stack-pointer idea separately.

[matthargett]

awesome. I measured both concerns you ran into during your experiments: the registers and the branch.

Register budget

A probe with 24 integer and 24 float arguments per function, compiled for each target, counts how many arrive in registers:

target	Rust ABI (int / float)	extern "rust-preserve-none"
arm64: iOS, tvOS, watchOS (incl. arm64_32), Android, Windows	8 / 8	23–24 / 8
x86-64 macOS and Linux	6 / 8	12 / 8
x86-64 Windows	4 / 4	12 / 8

The Unbudgeted handlers use five integer arguments today: the executor, the instruction slice's pointer and length, the instruction pointer, and the next instruction. Bounded uses four.

• arm64 has room for three more.
• x64 Mac/Linux has room for one.
• Windows x64 is already one over :( Every handler writes the next Instruction to the stack before its jmp, and i32.add pushes and pops rsi/rdi on every dispatch.

rust-preserve-none is nightly, like become. It gives 12 argument registers on every x86-64 OS, and since no register is callee-saved, no handler ever pushes. On arm64 it measured the same as the Rust ABI.

Top of stack

I built a reduction example in Rust with eight handler kinds, up to 64 copies each, dispatched like tinywasm. Only where the stack's height and top live varies. I ran it on an M4's efficiency cores and on the efficiency cores of an iPhone 12, SE and XS Max. Cycles per dispatch at 171 handlers (the real loops I profile use 136–226 handler pairs):

	M4	iPhone 12	iPhone SE	iPhone XS Max
today (Vec)	8.65	10.56	9.60	21.18
height in an arg	−34%	−32%	−16%	−6%
height + top, no reload when a pop empties the stack	−46%	−43%	−22%	−6%
top behind a flag	−43%	−37%	−15%	−11%
top only, length still in the Vec	+5%	0%	+1%	−7%
height passed on the stack (over budget)	+7%	+4%	+24%	+8%

• A top register gains nothing while the length stays in the Vec I think that's what you ran in your experiments? The length's store-to-load round trip is the critical path, so it has to move first.
• The height alone gives −34%, but flushes quadruple With the length chain gone, the value slots race. This is what I was running into when I posted the RFC.
• Height and top together give −46% flat as the handler count grows.
• No branch is needed, AFAICS Validation guarantees a pop only happens on a non-empty stack, so the top register is always valid when a handler reads it
  • A push writes the value through to its slot, so calls, returns, traps and host calls see the whole stack as they do today.
  • A pop reloads the top from the slot below. Where the lowering knows the pop empties the stack, a variant skips the reload. That reload is wasted, and it was most of the remaining flushes.
• That said, you were right about the branch cost
  • It does predict well across most of the microarchitectures I'm measuring on. On the A14 (iPhone 12) efficiency cores the flag version has 0.3–0.5 conditional mispredicts per 1k dispatches, the same as without it.
  • It's still 1–15% slower there at every handler count, and 7–16% on the A12 (iPhone XS) and A13 (iPhone SE 2) efficiency cores.
  • The counters put the cost in the front end: taken branches break up instruction fetch on every dispatch. The M-series silicon's microarchitecture absorbs it, even on its efficiency cores.
  • I'm betting that this is still end-to-end faster, even on A10X and A12. it's funny how moving things around can still have user-visible uplift.
• The flag version is about as fast here (the branch predicts well on M4), but it costs another register and still flushes 22 times per 1k.
• Passing the stack's slice as well (two more registers, so rust-preserve-none) gains another 8–11%.
• On the XS Max everything converges to about 20 cycles per dispatch past about 50 handlers. From 17 to 53 handlers the iPhone 12 mispredicts 13–18% of its dispatches, whatever the stack does. On the smallest cores, fewer dispatches will matter more than the stack.
• Passing the stack's slice too gives −51% and no flushes at all (0.05 per 1k). With two more registers, no slot address waits for a load of the base pointer. That's nine arguments, so it needs rust-preserve-none, or dropping the passed Instruction to fit arm64's eight.

In the hot loops I profile, nearly every dispatch is a 32-bit lane op or control flow. So the two registers to add are that lane's height and top: 7 of 8 on arm64. On x64 they fit with rust-preserve-none, or by dropping the passed Instruction and reloading it. That reload is 10% slower than passing it, but with the other things improved it's still 41% faster than today.

In other words, the stack's height and top get home registers, the way wasm3's ops all receive _sp and _r0 (which is also similar to other runtimes like Luau).

The model's handlers are simpler than tinywasm's, and the fused ops already skip the stack a lot, so the real-world gain in integrated will be well below these numbers. I can prototype the height first, so it's measured on its own as you suggested, and then add the top. Just lmk how you'd like me to proceed based on this data!