cranelift/egraphs: allow simplifying trapping arithmetic

[jameysharp]

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]

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]

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]

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

• Call the simplify ISLE term for instructions which have idempotent side-effects, just like we do for pure instructions.
• If any of the equivalent values are produced by pure instructions, then we have some rewrite rule that proves that the side effects can't actually happen. Build the e-class using only the pure alternatives and don't add this instruction to the side-effecting skeleton.
• Otherwise, put all the side-effecting alternatives in an e-class and record that this e-class must be evaluated at this point in the side-effecting skeleton if it hasn't already been evaluated at a point which dominates this one.

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:

• For each basic block, keep a list of side-effects. Each side effect is either an Inst or a Value. The last side effect will be the block's terminator Inst.
• Stop removing instructions from the Layout during equality saturation. Instead of tracking whether an instruction has side-effects by whether it's still in the layout, use the above explicit list.
• During equality saturation, add a side-effect to the list incrementally, after processing a side-effecting instruction.
• Between equality saturation and elaboration, clear all instructions out of the Layout.
• During elaboration, use the side-effect lists as the depth-first traversal roots, but otherwise treat the Value case as just like any other demanded value.

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.

[fitzgen]

I'd love to get to this one day.

My Wasm GC work is producing two patterns that require implementing the optimization of trapping arithmetic to get good code:

• (uadd_overflow_trap (iconst a) (iconst b)) should turn into either an unconditional trap, or more likely, into (iconst a+b)
• (uadd_overflow_trap (uextend a) (uextend b)) should turn into (iadd a b)

FWIW, these patterns can be matched and optimized during lowering, where we don't distinguish between side effectful vs non-side effectful lowerings. For example the Pulley backend optimizes the lowering of (trap{z,nz} (iconst ...)):

wasmtime/cranelift/codegen/src/isa/pulley_shared/lower.isle

Lines 98 to 114 in 5c8cd22

	(rule 2 (lower (trapz (iconst (u64_from_imm64 (u64_nonzero _))) code))
	(output_none))
	(rule 2 (lower (trapnz (iconst (u64_from_imm64 0)) code))
	(output_none))
	;; TODO: These rules are disabled because they insert a block terminator into
	;; the middle of the current block, which leads to regalloc errors. We should
	;; ideally be able to lower conditional traps that will always trap into
	;; unconditional traps though. This isn't very high priority though because
	;; traps, pretty much by definition, are not hot paths.
	;;
	;; (rule 3 (lower (trapnz (iconst (u64_from_imm64 (u64_nonzero _))) code))
	;; (side_effect (pulley_trap code)))
	;;
	;; (rule 3 (lower (trapz (iconst (u64_from_imm64 0)) code))
	;; (side_effect (pulley_trap code)))

However, we still can't insert block terminators early into the block (which should presumably truncate the rest of the block's lowering) so we can't turn (trapz (const 0)) into (trap), only (trapz (iconst (non_zero ...))) into the empty instruction sequence.

Additionally, when performed at lowering time, these rules' results cannot feed into other rules to create cascading optimizations.