bytecodealliance / wasmtime

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

cranelift/egraphs: allow simplifying trapping arithmetic #5908

Open jameysharp opened 1 year ago

jameysharp commented 1 year ago

Feature

We currently have simplification rules in opts/algebraic.isle to rewrite x/1 to x, and in opts/cprop.isle to constant-fold division when both operands are constant. But these rules can't fire, because simplify is only called on instructions which never have side effects, and division may trap if the divisor is 0.

I'd like to be able to use simplification rules like these on trapping arithmetic.

Benefit

We could optimize away more instructions. There are comments in algebraic.isle suggesting more desired optimizations that we can't implement today either:

;; TODO: strength reduction: div to shifts
;; TODO: div/rem by constants -> magic multiplications

In addition to replacing expressions with simpler equivalents, we could sometimes remove them entirely if they're unused. We currently can't do that because if the instruction would trap then we need to ensure that trap occurs. But if we have a valid rewrite from a possibly-trapping instruction to an expression which can't trap, then it's actually safe to treat it as a pure instruction.

I suspect this might be particularly valuable for uadd_overflow_trap, which is used in bounds checks for dynamic heaps in cranelift-wasm.

Implementation

One idea is in #5796, but that approach was not selected in #5800 because it "is a little more complex than I would like". I still like the idea of a force instruction to pull otherwise-pure instructions into the side-effecting skeleton, but it's possible some other approach is better.

Alternatives

We could maybe call simplify from the is_mergeable_for_egraph branch of optimize_skeleton_inst, or introduce a new ISLE term for simplifying idempotent instructions.

cfallin commented 1 year ago

Strong +1 to this -- this is the right way to represent idempotent side-effects IMHO. (Doing it the simpler way was a matter of expediency on my part as I wanted to get the optimization we disabled back in, but I fully support reworking it.)

Making side-effecting-but-idempotent ops part of the egraph has potential beyond arithmetic too, IMHO: the simplest case I can think of is actually just trapnz vN where we've rewritten vN to a constant 0. If a trapnz is really a (force (trapnz vN)) and (trapnz 0) is rewritten to 0, or (nop), or something appropriately-typed, then we can elide traps lowered/legalized from other logic as well if it's safe to do so.

jameysharp commented 1 year ago

I suspect we want to optimize trapnz the same way we do for branches, whatever that turns out to be. I suspect that will be different than this force mechanism because I think force needs to take a Value operand, and traps/branches have no results.

But definitely, I think it'll be important that we have some story for optimizing conditional traps too!

jameysharp commented 1 year ago

After discussion with @cfallin and @fitzgen today, I think this is our current plan:

In earlier discussion we talked about introducing a force CLIF instruction as the way to connect the side-effecting skeleton to an e-class. But here's an alternative that I think we've agreed is reasonable:

Elaboration already duplicates a pure instruction if its existing placement doesn't dominate the new place where it's needed, which is exactly what we need for instructions with idempotent side effects too. Given that, GVN is safe for these instructions as well.

Open question: what should happen if a simplify rule creates a side-effecting instruction somewhere other than at top-level of the right-hand side? I think that needs to be prohibited because we wouldn't know there were still side effects somewhere in the expression tree, so we'd treat it as pure. We can use the type system in the same way as #6080 to express this constraint statically.