bytecodealliance / bytecodealliance/wasmtime

cranelift: Sink loads after lowering

Open
#6,097 2 comments 1 reaction 0 assignees View on GitHub
cranelift cranelift:area:machinst cranelift:area:x64
Dominant language
Rust
Stars
18.6k
Forks
1.8k
Avg merge
1d 18h
Merged PRs (30d)
126

Description

#### Feature

On x86-64, unlike RISC-style instruction sets, many instructions have the ability to do a memory load. We can "sink" a load instruction into one of these other instructions, under the right conditions. One of those conditions is that the result of the load must be used exactly once. If it's used multiple times, then sinking the load means that the load is duplicated.

Our current mechanism for checking whether load-sinking is safe has been quite error-prone. We check for unique uses in the input CLIF, but during lowering, a single use of a CLIF `Value` may turn into multiple uses of the corresponding register. To deal with that, we try to ensure that lowering rules only convert a `Value` into a register once if they need to use the register multiple times. However this has been difficult to ensure because there's an implicit conversion from `Value` to types like `GprMem` and `XmmMem`.

This would all be much easier if we could count the number of times that the value is used _after_ lowering, since that's what we actually want to know.

#### Benefit

The primary benefit is to eliminate this class of bugs from the x86-64 backend. (In principle other architectures could do load-sinking too, but we don't currently support any that do.)

Also, I think we can delete a bunch of code in `cranelift/codegen/src/machinst/lower.rs`, and possibly one linear-time pass over the input CLIF, by specializing this.

#### Implementation

Besides the single-use rule, the other rule is that a load must not be moved past another side-effecting instruction. The machinst::lower module currently uses "instruction coloring" to enforce this. As far as I can tell, there are no other ways for instructions to be re-ordered during lowering, so this is the only use of these colors.

It's important to note that machinst lowers basic blocks in post-order, so all blocks which could have used a value are lowered before the block which defines that value. In addition, it lowers the instructions in each block in reverse order, so again the uses are seen before their definition.

So here's the approach I have in mind:

1. In enums like `GprMem`, have the `Mem` variant keep not only the address-mode, but also the `Value` that we got the input from.
2. For each `Value` which is the result of a load, track whether it is used more than once during lowering. Currently machinst does this in `value_lowered_uses`, but I think it might not need that after this.
3. When converting a `Value` produced by a load in the current block into e.g. `GprMem`, assume that we're going to sink the load, unless we've already seen other uses of this value. (We won't try to sink loads across block boundaries.) Add the `Value` to a queue of tentatively-sunk loads.
4. When we're about to lower a side-effecting instruction, check its opcode:
- If it's a load, then pop `Value`s from the queue until we find the result of _this_ load. If it is still only used once, then we can forget about it: we've successfully sunk the load. Afterward we can forget about how many times this load is used, too: now that we've found its definition, we won't see any more uses.
- If it isn't a load, then pop every `Value` from the queue. We can't sink any loads past this instruction.
5. For every `Value` popped from the queue without proving that sinking that load was safe, add it to a set of failed attempts so we can clean it up later.
6. During register allocation and again while emitting machine code, any time a `Mem` variant comes up, check it against the set of failed attempts. If it's there, treat it as the equivalent `Reg` variant instead. Alternatively, rewrite the current block when we're done lowering it, so we can discard the failed-attempt set early.

The queue is necessary because we may have tentatively sunk multiple loads by the time we see the actual load instructions, in which case we need to make sure that they're all in the right order. If we wanted to allow reordering loads with respect to each other but not with any other side effects, this could be a set instead of a queue. We could even keep a queue for atomic loads and a set for other loads if we want to implement a hybrid policy…

Note that the queue must be empty by the time we finish lowering a block, because we only add loads to it which are in the current block.

I don't think any of the above needs any machinery in `machinst` beyond the guarantee that we always lower uses before definitions. So other backends can avoid paying the cost of tracking colors and unique uses. If/when we add another backend that supports sinking loads, we can factor it out to common code then.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.