Stream: git-wasmtime

Topic: wasmtime / issue #7727 Cranelift: let regalloc manage cal...


view this post on Zulip Wasmtime GitHub notifications bot (Dec 25 2023 at 22:03):

jameysharp added the cranelift label to Issue #7727.

view this post on Zulip Wasmtime GitHub notifications bot (Dec 25 2023 at 22:03):

jameysharp added the cranelift:goal:optimize-speed label to Issue #7727.

view this post on Zulip Wasmtime GitHub notifications bot (Dec 25 2023 at 22:03):

jameysharp opened issue #7727:

Feature

Instead of writing target/ABI-specific code in each backend to spill callee-saved registers to the stack in the function prologue and restore them in the epilogue, we've been discussing that we should let the register allocator decide whether, when, and where to move those registers.

Benefit

The register allocator has enough information about register pressure to be able to decide whether it needs to make use of the callee-saved registers at all, or to delay spilling them in case there's a simple early-return path with low register pressure.

In addition, when the callee-saved registers must be spilled, the register allocator already knows how to do that, so we can delete a bunch of code from the backends which handles the special case of spilling callee-saved registers.

Implementation

Currently, every backend has fake "instructions" called args and rets. These pseudo-instructions don't generate any real instructions when they're emitted to machine code. They represent the ABI requirements for how arguments and return values are passed between functions, and they only exist to record information for register allocation. They work by pairing up a virtual register representing an argument with a physical register that is where that argument must live on entry to the function. The virtual register is then used everywhere else in the function so that the register allocator is allowed to change which physical register the value lives in, or even move it to the stack if necessary.

We can use those pseudo-instructions to describe how callee-saved registers work too. Callee-saved registers act like extra arguments to a function, which are never used inside the function but must be returned in the same registers at the end. By tying them to virtual registers, we can let the register allocator decide whether to leave those values in-place or move them out of the way to make those registers available for other uses.

I don't think we can correctly use the rets instruction before return_call instructions used for tail-calls. So a first version of this probably shouldn't try to change our tail-call ABI, which currently has no callee-saved registers. However, once the basic idea is working, it should be straightforward to add the functionality of the rets instruction into the return_call instruction, and then we can fix #6759.

Some backends (aarch64 at least) are currently able to make use of special "store-multiple" and "load-multiple" instructions for pushing and popping multiple callee-saved registers at once. A good enhancement here would be to find a way to let RA2 make use of those instructions any time it needs to spill or reload multiple registers at once, even if they're not the callee-saved register spills. However that shouldn't be necessary for a first version of this idea.

Alternatives

The current implementation works fine, it just seems more complex than necessary. In particular, @fitzgen and I found callee-saved registers were difficult to deal with when implementing tail-call support in Cranelift, and that led to the issue in #6759.

Credits

This started from a suggestion by @sunfishcode and I've discussed it further with @elliottt and others.

view this post on Zulip Wasmtime GitHub notifications bot (Dec 26 2023 at 02:03):

cfallin commented on issue #7727:

I really like this idea! It would definitely simplify things overall. A few thoughts to add:

I don't think either of those are necessarily fundamental (or, well, maybe the Windows part is), just two potential issues we'd have to work through!

view this post on Zulip Wasmtime GitHub notifications bot (Jan 08 2024 at 17:10):

ghostway0 commented on issue #7727:

I think I've got a solution after talking with Jamey for a bit. Where does Cranelift push the clobbered callee-saved regs?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 08 2024 at 17:11):

ghostway0 edited a comment on issue #7727:

I think I've got a solution after talking with Jamey for a bit. Where does Cranelift push the clobbered callee-saved regs? I couldn't find it, tried for the past ~day.

view this post on Zulip Wasmtime GitHub notifications bot (Jan 08 2024 at 17:11):

ghostway0 edited a comment on issue #7727:

I think I've got a solution after talking with Jamey for a bit. Where does Cranelift push the clobbered callee-saved regs?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 08 2024 at 21:57):

cfallin commented on issue #7727:

This is implemented in the machinst::abi module and the corresponding trait impls for each ISA; grep for gen_clobber_save. Could you say more about your idea?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 11 2024 at 17:35):

ghostway0 commented on issue #7727:

Thanks!
Sure, I plan to add a method for the ABI trait to get the callee-saved registers, save the ArgPairs and add it to args and rets. How does that sound? I think I've implemented it already - as in, correctly - just didn't test it yet.
Sorry for the delayed response

view this post on Zulip Wasmtime GitHub notifications bot (Jan 11 2024 at 18:47):

cfallin commented on issue #7727:

Ah, that could work. My main concerns would be:

Thoughts?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 11 2024 at 19:12):

ghostway0 commented on issue #7727:

(a) Isn't the argument that RA2 is smart enough to handle those?
(b) That was the plan, I have a TODO there. How should I handle that? Windows is annoying :)

view this post on Zulip Wasmtime GitHub notifications bot (Jan 12 2024 at 07:26):

cfallin commented on issue #7727:

(a) Isn't the argument that RA2 is smart enough to handle those?

It certainly can handle function-long liveranges, sure. My main concern described above is performance (of the compiler, and of its output). The appeal of the more general solution is its simplicity but it may present a tradeoff because it uses more general code (the main RA2 solver) to replace a very specific and hardcoded mechanism. Fortunately it's an empirical question and one can benchmark it! :-)

Another aspect of generated-code perf that just occurred to me: on aarch64, we can do more efficient "bulk saves" with STP/LDP (store/load pair instructions). RA2 isn't able to coalesce adjacent spills into combined store instructions. Not impossible to do, but then we're just moving complexity rather than eliminating it.

(b) That was the plan, I have a TODO there. How should I handle that? Windows is annoying :)

Unfortunately I haven't been able to come up with a good idea here. The fastcall ABI is pretty specific about what the prologue should look like -- save frame pointer, save clobbered callee-saves, with nothing else in the way (because its unwinder does abstract interpretation of the prologue code, rather than interpreting something like DWARF metadata, which is pretty wild).

We could get that form by putting specific code generation for it in the args pseudoinst, but then we've just moved the explicit codegen for saving clobbers to a different place, rather than eliminated it (while the original motivation was to eliminate it).

So I guess what I'm coming around to is: Windows fastcall support may force the issue; and if that weren't the case, there are some performance concerns I hypothesize, but those can be answered empirically. Maybe all this does motivate why we have specific code for clobber-saves: it's just different enough in requirements that sharing with general spills doesn't reduce complexity; and also it's just different enough that we can address it with a different algorithm, possibly with better performance. I'm curious if anyone else has ideas on this though, especially the fastcall bit!

view this post on Zulip Wasmtime GitHub notifications bot (Jan 12 2024 at 14:32):

ghostway0 commented on issue #7727:

Another aspect of generated-code perf that just occurred to me: on aarch64, we can do more efficient "bulk saves" with STP/LDP (store/load pair instructions). RA2 isn't able to coalesce adjacent spills into combined store instructions. Not impossible to do, but then we're just moving complexity rather than eliminating it.
Doesn't RA2 use parallel moves everywhere for these? If those bulk saves would be in the same bulk mov...
Unfortunately I haven't been able to come up with a good idea here. The fastcall ABI is pretty specific about what the prologue should look like -- save frame pointer, save clobbered callee-saves, with nothing else in the way (because its unwinder does abstract interpretation of the prologue code, rather than interpreting something like DWARF metadata, which is pretty wild).
Oof. The easiest solution would probably be to check the abi wherever machinst does clobber saves, no?

In a backtracking allocator you could set live ranges' priority, right? Why not have callee-saved registers' bundles at a lower priority? I think there's a constraint forcing live ranges to live on the stack, too. What do you think?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 12 2024 at 14:32):

ghostway0 edited a comment on issue #7727:

Another aspect of generated-code perf that just occurred to me: on aarch64, we can do more efficient "bulk saves" with STP/LDP (store/load pair instructions). RA2 isn't able to coalesce adjacent spills into combined store instructions. Not impossible to do, but then we're just moving complexity rather than eliminating it.

Doesn't RA2 use parallel moves everywhere for these? If those bulk saves would be in the same bulk mov...

Unfortunately I haven't been able to come up with a good idea here. The fastcall ABI is pretty specific about what the prologue should look like -- save frame pointer, save clobbered callee-saves, with nothing else in the way (because its unwinder does abstract interpretation of the prologue code, rather than interpreting something like DWARF metadata, which is pretty wild).

Oof. The easiest solution would probably be to check the abi wherever machinst does clobber saves, no?

In a backtracking allocator you could set live ranges' priority, right? Why not have callee-saved registers' bundles at a lower priority? I think there's a constraint forcing live ranges to live on the stack, too. What do you think?

view this post on Zulip Wasmtime GitHub notifications bot (Jan 13 2024 at 07:43):

cfallin commented on issue #7727:

Another aspect of generated-code perf that just occurred to me: on aarch64, we can do more efficient "bulk saves" with STP/LDP (store/load pair instructions). RA2 isn't able to coalesce adjacent spills into combined store instructions. Not impossible to do, but then we're just moving complexity rather than eliminating it.

Doesn't RA2 use parallel moves everywhere for these? If those bulk saves would be in the same bulk mov...

Yes, a parallel move is used for all spills at the same program point (at the same priority); the resolver could try to group moves together to make use of store/load-pair, at the cost of more complexity.

Unfortunately I haven't been able to come up with a good idea here. The fastcall ABI is pretty specific about what the prologue should look like -- save frame pointer, save clobbered callee-saves, with nothing else in the way (because its unwinder does abstract interpretation of the prologue code, rather than interpreting something like DWARF metadata, which is pretty wild).

Oof. The easiest solution would probably be to check the abi wherever machinst does clobber saves, no?

In a backtracking allocator you could set live ranges' priority, right? Why not have callee-saved registers' bundles at a lower priority? I think there's a constraint forcing live ranges to live on the stack, too. What do you think?

Yes, one could probably construct constraints to get things mostly into the right shape, but my main concern is that this is a fragile "just-so" construction: the necessary prologue shape would then be an emergent behavior of constraints attached to the pseudoinsts, handling of those constraints, and RA2's particular implementation choices for spills. Changes anywhere in that stack would break things, and in a way that could be somewhat hard to detect and debug.

Both of these actually get to the heart of my concern that has developed after working through these details: we're coming up with complex ways to adjust the general solver and RA backend (e.g. parallel move resolver), and feed it the right constraints, to emit code in just the right shape; to replace a small handwritten prologue generation function that emits code in just the right shape.

The original idea had a lot of appeal to me because it truly seemed simpler: remove the open-coded sequence, add some RA2 operands instead. But we've quickly gone the other way here to recover all the necessary properties: careful sorting and coalescing of adjacent spillslot pairs in parallel moves (for aarch64 ldp/stp), including a notion of pair stores/loads in the RA2 API where we only had generic moves before; careful use of subpriorities to push spills to a certain location and make the prologue into the right shape for fastcall; perhaps some other as-yet-unknown optimizations to make this match the original performance (e.g., special-casing these LRs in RA2's main loop and trying not to split them -- actually I think we need to ensure we don't split them, given that we describe their location only once, and that creates additional special cases in the solver core); and then reliance on all of these pieces continuing to interoperate in order to satisfy the ABI. In contrast, today we have ~100 LoC of code per backend to generate clobber-saves. From a "least global complexity" standpoint the latter seems better to me, at least.

view this post on Zulip Wasmtime GitHub notifications bot (Jan 13 2024 at 08:32):

jameysharp commented on issue #7727:

It sounds to me like at minimum we shouldn't try to make RA2 handle callee-saved register spills on Windows. In fact this makes me think that tail-calls need very careful thought on Windows, since it seems we need callee-saved registers for general function-call performance (per #6759) but I bet unwinding isn't prepared to deal with tail-calls; but that's a separate concern. Also I wonder whether s390x would cause trouble for this plan.

I still think it's worth testing this approach on the other calling conventions, primarily because giving regalloc room to consider alternate plans should allow generating better code, and it would be nice to characterize how much better and what it costs in compile time. Even if we can't have that win on Windows, we do have a lot of people using Cranelift on non-Windows.

Also, if we could make RA2 use stp/ldp for parallel moves, that could be a win for code outside the prologue/epilogue too. If so it would be a good experiment independent of this one.

I certainly was hoping that this idea would let us delete all the special handling of callee-saved registers and I'm disappointed that Windows sinks that plan, but I'm not yet fully convinced that means we should throw out this idea entirely.

view this post on Zulip Wasmtime GitHub notifications bot (Jan 15 2024 at 06:04):

cfallin commented on issue #7727:

For sure, I'd be interested in the possible performance upside of this too -- @ghostway0, would you be able to do some compile+runtime tests with your prototype? We're happy to provide details on how to do this if you need.

Also, if we could make RA2 use stp/ldp for parallel moves, that could be a win for code outside the prologue/epilogue too. If so it would be a good experiment independent of this one.

In principle that'd be nice, yep -- but since the paired memory accesses need adjacent spillslots, RA2's effectively random assignment of slots (depending on priorities and processing order and what happens to split or backtrack and ...) would eliminate most opportunistic uses of these, I suspect. To make this work well one would need to influence both the assignment of spillslots and the order of emitted moves, so that stores/loads to adjacent pairs of spillslots happen next to each other. That seems like a painfully complex change to RA2... actually I think possibly exponential complexity (we have the n! permutations of spillslot assignments, and for each permutation, certain spill/reload pairs will or won't be ldp/stp depending on whether adjacent addresses. Perhaps a greedy algorithm would do OK, I'm not sure). It also has weird phasing implications: spillslots are assigned before move insertion, but we need to know which moves are adjacent if we want this pair-aware assignment. All that to say, I'm not sure how to do this; I'd love to be proven wrong though.


Last updated: Dec 23 2024 at 12:05 UTC