bytecodealliance / wasmtime

A fast and secure runtime for WebAssembly
https://wasmtime.dev/
Apache License 2.0
15.21k stars 1.28k forks source link

cranelift/egraphs: Allow rewriting to unconditional trap #6080

Open jameysharp opened 1 year ago

jameysharp commented 1 year ago

Feature

When a mid-end optimization rule in ISLE matches an instruction which has a result value, it should be possible to replace that instruction with an unconditional trap. For example, (udiv _ _ (iconst _ 0)) should rewrite to a trap with code int_divz.

This is only possible once we resolve #5908; until then, ISLE rules never fire for instructions which could trap.

Benefit

This particular rewrite doesn't fit in our current framework, which only supports replacing a value with another value.

Traps, specifically, are special. They make the rest of the current block unreachable, as well as any block dominated by the current block. It's important to drop all dominated blocks because those are exactly the blocks which may have used the result of the original instruction. It's also useful to drop all dominated instructions because then we can avoid running all the egraph machinery on any of them.

Like the branch optimizations that we aren't doing yet, discarding dominated branches may move other blocks down the dominator tree and make some block parameters known. We don't have to update the dominator tree when that happens, but it's useful to do because it makes more information available to the affected blocks, which can lead to better optimizations.

Implementation

The ISLE simplify term can only produce instructions which have a result Value, so it can't directly produce a trap instruction. I think we should change its return type to a new ValueOrTrap enum, and define an implicit conversion from Value to ValueOrTrap. When a rule returns the Trap variant with a trap code, the caller needs to remove the remaining instructions in the current block and insert the appropriate trap instruction into the data-flow graph.

If any rewrite rule says that an instruction is equivalent to a trap, then we can ignore all the other rewrites and just take the trap. This is sort of like how subsume works.

Things would get a little weird if a trap were generated somewhere other than the top-level right-hand side of a rule. All instructions which use the result would be unreachable, and those instructions' other operands would be unused, so we'd end up deleting all the instructions created by the rule except for the trap. Making the return type of simplify be the only place where a trap can appear means we can statically prevent writing rules which do this extra work.

It's also weird if this happens while we're doing recursive simplification on instructions which were newly-created by other rewrites. However, I think that case isn't actually useful and we should just prohibit it, since that means a rule created an instruction which would trap, without being equivalent to any possibly-trapping instruction in the original program. We could panic if we see a rewrite to a trap during recursive simplification.

Alternatives

I haven't thought of any besides keeping the status quo, but I think we should do something like this.

jameysharp commented 1 year ago

Also, maybe the block where the trap occurs should be automatically marked cold. It obviously is not going to execute more than once per invocation of the program; blocks don't get much more cold than that.