Tuples inhibiting optimizations

[kripken]

It seems many opts work better without tuples, because tuples hide their components from the optimizations.

Testcase

(module
  (global $g (mut i32) (i32.const 0))

  (export "g" (global $g))

  (func $func (result i32 i32)
    ;; Return an unknown value and 42.
    (return
      (tuple.make
        (global.get $g)
        (i32.const 42)
      )
    )
  )

  (func "caller" (result i32)
    (local $tuple (i32 i32))
    (local $sum i32)

    ;; Get a tuple from a call
    (local.set $tuple
      (call $func)
    )

    ;; Return the sum of the tuple elements.
    (local.set $sum
      (tuple.extract 0
        (local.get $tuple)
      )
    )
    (local.set $sum
      (i32.add
        (local.get $sum)
        (tuple.extract 1
          (local.get $tuple)
        )
      )
    )
    (local.get $sum)
  )
)

After --inlining

(module
 (type $0 (func (result i32 i32)))
 (type $1 (func (result i32)))
 (global $g (mut i32) (i32.const 0))
 (export "g" (global $g))
 (export "caller" (func $1))
 (func $1 (type $1) (result i32)
  (local $tuple (i32 i32))
  (local $sum i32)
  (local.set $tuple
   (block (result i32 i32)
    (block $__inlined_func$func (result i32 i32)
     (br $__inlined_func$func
      (tuple.make
       (global.get $g)
       (i32.const 42)
      )
     )
    )
   )
  )
  (local.set $sum
   (tuple.extract 0
    (local.get $tuple)
   )
  )
  (local.set $sum
   (i32.add
    (local.get $sum)
    (tuple.extract 1
     (local.get $tuple)
    )
   )
  )
  (local.get $sum)
 )
)

Now we have a block with multivalue and a br to it.

After -O3

(module
 (type $0 (func (result i32)))
 (global $g (mut i32) (i32.const 0))
 (export "g" (global $g))
 (export "caller" (func $1))
 (func $1 (type $0) (; has Stack IR ;) (result i32)
  (local $0 (i32 i32))
  (i32.add
   (tuple.extract 1
    (local.tee $0
     (tuple.make
      (global.get $g)
      (i32.const 42)
     )
    )
   )
   (tuple.extract 0
    (local.get $0)
   )
  )
 )
)

The code is simplified somewhat, but still has tuple operations.

After --roundtrip -O3

(module
 (type $0 (func (result i32)))
 (global $g (mut i32) (i32.const 0))
 (export "g" (global $g))
 (export "caller" (func $1))
 (func $1 (type $0) (; has Stack IR ;) (result i32)
  (i32.add
   (global.get $g)
   (i32.const 42)
  )
 )
)

This is the result we wanted before roundtripping.

Roundtripping lowers the tuple locals into normal ones, which optimizations can operate on, so we can do a lot more if we roundtrip atm.

One option here might be to add a pass to lower tuples into normal locals? cc @tlively

@askeksa is this the issue you were seeing?

[kripken]

Adding

    (if
      (i32.eqz
        (global.get $g)
      )
      (return
        (tuple.make
          (i32.const 1337)
          (i32.const 99999)
        )
      )
    )

At the top of $func, so there isn't a single return anymore, causes the end result even after roundtripping to still use tuples. But it seems like there as well, we might benefit from lowering them down (e.g. if one part of a tuple is not even used, we still keep the whole tuple atm). In general "non-escaping" tuples (that don't get sent as a return value) should perhaps always be lowered?

[tlively]

A pass to lower tuples away where possible might be useful, but it would probably only be useful after inlining, since tuple locals assigned to after a multivalue function call cannot be lowered away. Overall the approach of round-tripping to try to get rid of tuples seems reasonable.

I'm surprised round-tripping already works so well on these examples, though. The new IRBuilder creates fewer tuple scratch locals and tuple.make than the current binary parser, so finishing IRBuilder and using it in the binary parser will be one way to improve the results of round tripping. We could also be smarter in the binary emitter about combining local.get of tuple locals followed by tuple.extract into just a single local.get of the extracted component to further improve the results of round tripping.

[askeksa]

Yes, this is the issue we are running into. If we optimize mc.wasm.gz using our usual Binaryen pipeline:

wasm-opt -g -all --closed-world -tnh -O3 --type-ssa --gufa -O3 --type-merging -O1

then the _WasmDefaultMap. function becomes a mostrosity of gets and sets with more than 600 locals. Inserting a --roundtrip before the final -O1 reduces it to something a lot more reasonable:

 (func $_WasmDefaultMap. (type $ref|$_Type|_ref|$_Type|_=>_ref|$__WasmDefaultMap&_HashFieldBase&MapMixin|) (param $0 (ref $_Type)) (param $1 (ref $_Type)) (result (ref $__WasmDefaultMap&_HashFieldBase&MapMixin))
  (local $2 i32)
  (local $3 (ref $_WasmI32ArrayBase))
  (local $4 (ref $Object))
  (local.set $2
   (i32.const 0)
  )
  (local.set $3
   (block $label$1 (result (ref $_WasmI32ArrayBase))
    (br_on_non_null $label$1
     (global.get $global$119)
    )
    (call $_uninitializedIndex)
   )
  )
  (local.set $4
   (block $label$2 (result (ref $_ListBase))
    (br_on_non_null $label$2
     (global.get $global$140)
    )
    (call $_uninitializedData)
   )
  )
  (struct.new $__WasmDefaultMap&_HashFieldBase&MapMixin
   (i32.const 254)
   (local.get $2)
   (local.get $3)
   (i64.const 0)
   (local.get $4)
   (i64.const 0)
   (local.get $0)
   (local.get $1)
  )
 )

It can still be improved (inlining the temporary values into the struct.new) by using higher optimization levels for the final pass. We can get most of the way by just adding --simplify-locals --vacuum:

 (func $_WasmDefaultMap. (type $ref|$_Type|_ref|$_Type|_=>_ref|$__WasmDefaultMap&_HashFieldBase&MapMixin|) (param $0 (ref $_Type)) (param $1 (ref $_Type)) (result (ref $__WasmDefaultMap&_HashFieldBase&MapMixin))
  (local $2 i32)
  (local $3 (ref $_WasmI32ArrayBase))
  (local $4 (ref $Object))
  (struct.new $__WasmDefaultMap&_HashFieldBase&MapMixin
   (i32.const 254)
   (i32.const 0)
   (block $label$1 (result (ref $_WasmI32ArrayBase))
    (br_on_non_null $label$1
     (global.get $global$119)
    )
    (call $_uninitializedIndex)
   )
   (i64.const 0)
   (block $label$2 (result (ref $_ListBase))
    (br_on_non_null $label$2
     (global.get $global$140)
    )
    (call $_uninitializedData)
   )
   (i64.const 0)
   (local.get $0)
   (local.get $1)
  )
 )

Is there a pass to just remove unused locals?

Is it expected that -O1 leaves such potentially inlineable temporary values?

But yeah, being able to avoid the --roundtrip would be nice.

[askeksa]

Actually, in our particular case, the multi-value blocks are not targeted by any branches. If we don't emit the blocks, we get the good result without roundtripping, but omitting those blocks involves some unwieldy special cases in our code generator. So removing those redundant blocks early in the Binaryen pipeline could also be a way to fix the issue.

[kripken]

Is there a pass to just remove unused locals?

--reorder-locals reorders them by use and also removes unused ones (all those above a certain index, trivial after the sort).

Is it expected that -O1 leaves such potentially inlineable temporary values?

Sorry, I'm not sure what values you mean? (In your last code fragment, it looks fully optimized?)

So removing those redundant blocks early in the Binaryen pipeline could also be a way to fix the issue.

That sounds interesting, but I don't follow. I compiled mc.wasm with the flags you gave but I'm not sure I see such blocks. Can you show me a concrete example?

[askeksa]

--reorder-locals reorders them by use and also removes unused ones (all those above a certain index, trivial after the sort).

That does it. Thanks.

Is it expected that -O1 leaves such potentially inlineable temporary values?

Sorry, I'm not sure what values you mean? (In your last code fragment, it looks fully optimized?)

That fragment included --simplify-locals --vacuum. I'm talking about the first fragment, which is the output of -O1.

So removing those redundant blocks early in the Binaryen pipeline could also be a way to fix the issue.

That sounds interesting, but I don't follow. I compiled mc.wasm with the flags you gave but I'm not sure I see such blocks. Can you show me a concrete example?

From a no-op run, I get code like this inside _WasmDefaultMap.:

  (local.set $110
   (block $label$1 (result (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref null $Object) i64 (ref $Object) i64 i64 (ref $_Type) (ref $_Type))
    (local.set $4
     (block (result (ref $_Type))
      (local.set $69
       (local.get $2)
      )
      (local.set $5
       (local.get $3)
      )
      (local.get $69)
     )
    )
    (local.set $102
     (block $label$2 (result (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref $_Type) (ref null $Object) i64 (ref $Object) i64 i64 (ref $_Type) (ref $_Type))

and so on, with 6 multi-result blocks nested inside each other. If we omit those blocks, Binaryen optimizes the function to the optimal result without the roundtrip.

[kripken]

I'm talking about the first fragment, which is the output of -O1.

I see now, thanks. Yes, that fragment has some local values that can be applied at their uses. -O1 does do that (using --simplify-locals) but I think in this case another run of -O1 may be required. (Almost none of the passes converge until they find no more work, not unless you pass in --converge).

If we omit those blocks, Binaryen optimizes the function to the optimal result without the roundtrip.

How can those blocks be omitted? The output of the first block is written to local $110 (and there are uses of it later).

[askeksa]

The locals receiving the results of the blocks are synthetic tuple locals inserted by Binaryen. In the input, the blocks are redundant because they are not targeted by any branches. They just wrap pieces of code leaving a bunch of values on the stack. Some of those values are stored to locals, while others are left on the stack until the end of the surrounding block (where more values have been pushed onto the stack in the meantime).

If the blocks are omitted, Binaryen just sees some nested gets and sets, whcih it easily untangles.

The easiest way to remove the blocks might be to detect them during the Binaryen IR construction and construct the IR as if the blocks were not there. If this is not feasible (or a more general solution is desired), a pass to remove them could work in this particular (typical) scenario, where:

• The block is not targeted by any branches.
• The block result is written directly to a local $t
• $t is not written anywhere else
• All uses of $t are in tuple.extract

In this scenario, the transformation could replace $t with one new local per tuple element, insert local.sets where each element is computed (inside the block) and replace the tuple.extracts with corresponding local.gets, then remove the block (required up front to maintain validity of any non-defaultable element values).

[kripken]

I see, thanks @askeksa

I think that can work, but I'm worried about targeting blocks in an early phase, as later opts would prevent the opt from working. Instead I opened #5937 which lowers away trivial tuples in a more general way. More work but I think it's worthwhile.

When I run -O3 --gufa -O3 --closed-world -tnh -all on dart_mc.wasm I get this, which looks optimal and has no more tuples:

 (func $_WasmDefaultMap. (type $153) (; has Stack IR ;) (param $0 (ref $_Type)) (param $1 (ref $_Type)) (result (ref $__WasmDefaultMap&_HashFieldBase&MapMixin))
  (struct.new $__WasmDefaultMap&_HashFieldBase&MapMixin
   (i32.const 254)
   (i32.const 0)
   (block $label$7 (result (ref $_WasmI32ArrayBase))
    (br_on_non_null $label$7
     (global.get $global$38)
    )
    (call $_uninitializedIndex)
   )
   (i64.const 0)
   (block $label$8 (result (ref $_ListBase))
    (br_on_non_null $label$8
     (global.get $global$39)
    )
    (call $_uninitializedData)
   )
   (i64.const 0)
   (local.get $0)
   (local.get $1)
  )
 )

[askeksa]

That looks perfect, thanks!

[tlively]

I implemented the binary writer improvements I mentioned in #5923 (comment) in #5941 as well.