wasm-opt doesn't support loops with parameters

[rockwalrus]

Assemble the following with wabt's wat2wasm:

(module
  (func (param i32 i32)
    (i32.const 0)
    (loop $loop (param i32)
      (if
        (then
          (local.set 0
            (i32.add
              (local.get 0)
              (i32.const 2)))))
      (block
        (local.set 1
          (i32.add
            (local.get 1)
            (i32.const -1)))
        (br $loop
          (i32.const 1))))))

An optimization-free pass through wasm-opt produces:

(module
 (type $i32_i32_=>_none (func (param i32 i32)))
 (func $0 (param $0 i32) (param $1 i32)
  (if
   (i32.const 0)
   (local.set $0
    (i32.add
     (local.get $0)
     (i32.const 2)
    )
   )
  )
  (loop $label$1
   (block $label$3
    (local.set $1
     (i32.add
      (local.get $1)
      (i32.const -1)
     )
    )
    (drop
     (i32.const 1)
    )
    (br $label$1)
   )
  )
 )
)

which isn't the same thing at all.

I'd love for it to be smart enough to rerelooper this code as:

(module
  (func (param i32 i32)
    (loop $loop
      (block
        (local.set 1
          (i32.add
            (local.get 1)
            (i32.const -1)))
        (local.set 0
          (i32.add
            (local.get 0)
            (i32.const 2)))
        (br $loop)))))

but I'd be quite satisfied with it just understanding loop parameters and generating correct code for them.

[rockwalrus]

Come to think of it, a local variable->loop parameter pass is probably all you need to rerelooper that if statement to the right spot. A local variable->loop parameter transformation could also save bytes when applicable.

[MaxGraey]

An optimization-free pass through wasm-opt produces

I don't think it's optimization-free. It definitely -O1 or something like this which includes LICM (loop invariant code motion).

[rockwalrus]

Here's the output with -O0:

$ wasm-opt test2.wasm -O0 --print          
(module
 (type $i32_i32_=>_none (func (param i32 i32)))        
 (func $0 (param $0 i32) (param $1 i32)
  (if                                                      
   (i32.const 0)
   (local.set $0                                            
    (i32.add
     (local.get $0)                                          
     (i32.const 2)
    )                                                      
   )
  )                                                       
  (loop $label$1
   (block $label$3                                          
    (local.set $1
     (i32.add
      (local.get $1)
      (i32.const -1)
     )
    )
    (drop
     (i32.const 1)
    )
    (br $label$1)
   )
  )
 )
)
warning: no output file specified, not emitting output

[tlively]

The problem here is that Binaryen does not even support parsing loops (or blocks) with input parameters. The fix would not be in the relooper code, but in the parsing code, similar to how the binary parser inserts extra locals to parse stacky Wasm into Binaryen IR. In the long term, it would be even better if we supported block and loop parameters natively in Binaryen IR.

[LunaAmora]

Anything is planned at the moment to fix this problem? This issue breaks my generated code as wasm-opt removes the params from the loops.

[tlively]

No, nothing concretely planned at the moment, unfortunately. Can you share more about your use case?

[LunaAmora]

No, nothing concretely planned at the moment, unfortunately. Can you share more about your use case?

I'm writing a compiler targeting Wasm to learn more how it works, and the way my language works translates really cleanly to using params on blocks on the generated code. I was using wasm-opt before implementing the codegen of loops, so it's a little of a bother that a supported feature of the lang on the runtimes i tested wasn't working here.

[tlively]

It would be nice to support parsing block input types by introducing new locals to hold the values as they enter the block. It would also be nice to optimize to emit block input types at the end of the optimization pipeline, for example using Stack IR or Poppy IR. Natively representing block input types in Binaryen IR, however, would be much more challenging and probably is not worth doing at all.

[rockwalrus]

@tlively That sounds like the right approach to me.

[LunaAmora]

@tlively Yeah, i think at least supporting parsing to not break codes that relies on the feature would be the more worthwhile approach.

Or at least give an error when it encounters a block input, now you just get a error about what happens to the code without the parameter at all: (so i thought the problem was ifs with params, had to search for a while to find this issue explaning the problem)

[LunaAmora]

What is strange to me is that I tested with other code now that uses block params and it manages to extract it away correctly. Maybe the error with the other code is a wasm-opt bug?

The following code pass wasm-opt without issues (running through wat2wasm first):

(func $dup (param i32 ) (result i32 i32) local.get 0 local.get 0)

(func $checkNumber (param i32) (result i32)
  local.get 0
  block $case (param i32) (result i32)
    block $default (param i32)
      block $case0 (param i32)
        block $case1 (param i32)
          call $dup 
          i32.const 0 i32.lt_s br_if $case0
          call $dup
          i32.const 9 i32.lt_s br_if $case1
          br $default
        end
        i32.const 1
        br $case
      end
      i32.const 2
      br $case
    end
    i32.const 3
    br $case
  end
)

While these two versions of a While loop doesn't:

(func $countToZero (param i32) ;; This gives [wasm-validator error] when running wasm-opt
  local.get 0
  loop $while (param i32) (result i32)
    call $dup
    i32.const 0
    i32.gt_s
    if (param i32) (result i32)
      i32.const 1
      i32.sub
      br $while
    end
  end drop
)

(func $countToZero2 (param i32) ;; This makes the runtime executing the optmized code hang
  local.get 0
  block $out (param i32) (result i32)
    loop $while (param i32) (result i32)
      call $dup
      i32.const 0
      i32.gt_s i32.eqz 
      br_if $out
      i32.const 1
      i32.sub
      br $while
    end
  end drop
)

[tlively]

The loop example passes validation after Binaryen's binary parser gets through with it, but I'm pretty sure the parser completely changes the program's semantics. I double checked and the binary parser code does indeed discard the block input type with no special handling.