Skip to content

Instantly share code, notes, and snippets.

@tanishiking
Last active September 9, 2024 12:49
Show Gist options
  • Save tanishiking/5b0141a7effa65f071a30ee09f327eaf to your computer and use it in GitHub Desktop.
Save tanishiking/5b0141a7effa65f071a30ee09f327eaf to your computer and use it in GitHub Desktop.

LoopInvariantCodeMotion - wasm-opt

Flattening

TL;DR: We would need running --flatten pass before --licm to make LICM work.

;; test.wat
(module
  (global $x (mut i32) (i32.const 0))
  (func (export "test")
    (result i32)
    (loop $1
      (global.set $x
        (i32.add (i32.const 0) (i32.const 1))
      )
    )
    (global.get $x)
  )
)

When applying loop-invariant code motion to the above code, I expected the pass to move (i32.add (i32.const 0) (i32.const 1)) out of the loop, but it doesn't.

;; wasm-opt -all --licm test.wat -S -o -
(module
 (type $0 (func (result i32)))
 (global $x (mut i32) (i32.const 0))
 (export "test" (func $0))
 (func $0 (type $0) (result i32)
  (loop $1
   (global.set $x
    (i32.add
     (i32.const 0)
     (i32.const 1)
    )
   )
  )
  (global.get $x)
 )
)

This is because the licm pass is designed to move entire expressions, not sub-expressions. In this case, the entire global.set expression is considered as a unit. https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/passes/LoopInvariantCodeMotion.cpp#L116

And, global.set won't move because it's unsafe to move:

bool unsafeToMove = effects.writesGlobalState() ||
                    effectsSoFar.invalidates(effects) ||
                    (effects.readsMutableGlobalState() &&
                     loopEffects.writesGlobalState());

https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/passes/LoopInvariantCodeMotion.cpp#L125-L128

To move the i32.add sub-expression out of the loop, we need to flatten it by using local.set to store the result of i32.add and then local.get to retrieve the result later. We can do this automatically by flatten pass as documented:

// Flattening is not necessary here, but may help (as separating // out expressions may allow moving at least part of a larger whole). https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/passes/LoopInvariantCodeMotion.cpp#L21-L22

// too, as with more nesting moving code is harder - so something // like -O --flatten --licm -O may be best). https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/passes/LoopInvariantCodeMotion.cpp#L212-L213

❯ wasm-opt -all --flatten --licm loop.wat -S -o -
(module
 (type $0 (func (result i32)))
 (global $x (mut i32) (i32.const 0))
 (export "test" (func $0))
 (func $0 (type $0) (result i32)
  (local $0 i32)
  (local $1 i32)
  (local $2 i32)
  (local $3 i32)
  (block
   (block
    (local.set $0
     (i32.add
      (i32.const 0)
      (i32.const 1)
     )
    )
    (loop $1
     (nop)
     (global.set $x
      (local.get $0)
     )
    )
   )
   (local.set $1
    (global.get $x)
   )
   (local.set $2
    (local.get $1)
   )
  )
  (local.set $3
   (local.get $2)
  )
  (return
   (local.get $3)
  )
 )
)

SideEffects

Next, let's see the loop that contains a function call:

;; cat test.wat
(module
  (func $loop
    (local $x i32)
    (local $y i32)
    (loop $loop
      (local.set $x (i32.eqz (local.get $y)))
      (call $foo (local.get $x))
    )
  )
  (func $foo
    (param $x i32)
    (nop)
  )
)

Since the loop body is already flattened, I expect:

  • (local.set $x (i32.eqz (local.get $y))) to be moved out, as it is side-effect free.
  • (call $foo (local.get $x)) to remain inside the loop, as function calls can cause side effects.

Unfortunately, local.set didn't move:

;; wasm-opt -all --licm test.wat -S -o -
(module
 (type $0 (func))
 (type $1 (func (param i32)))
 (func $loop (type $0)
  (local $x i32)
  (local $y i32)
  (loop $loop
   (local.set $x
    (i32.eqz
     (local.get $y)
    )
   )
   (call $foo
    (local.get $x)
   )
  )
 )
 (func $foo (type $1) (param $x i32)
  (nop)
 )
)

Why? The code that determines whether the given expression can be moved or not is located around here:

      if (interestingToMove(curr)) {
        // Let's see if we can move this out.
        // Global state changes would prevent this - we might end up
        // executing them just once.
        // And we must also move across anything not moved out already,
        // so check for issues there too.
        // The rest of the loop's effects matter too, we must also
        // take into account global state like interacting loads and
        // stores.
        bool unsafeToMove = effects.writesGlobalState() ||
                            effectsSoFar.invalidates(effects) ||
                            (effects.readsMutableGlobalState() &&
                             loopEffects.writesGlobalState());
        // TODO: look into optimizing this with exceptions. for now, disallow
        if (effects.throws() || loopEffects.throws()) {
          unsafeToMove = true;
        }
        if (!unsafeToMove) {
          // So far so good. Check if our local dependencies are all
          // outside of the loop, in which case everything is good -
          // either they are before the loop and constant for us, or
          // they are after and don't matter.
          if (effects.localsRead.empty() ||
              !hasGetDependingOnLoopSet(curr, loopSets)) {

https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/passes/LoopInvariantCodeMotion.cpp#L116-L139

  • effect is the analyzed effects of the current Expression
  • effectsSoFar is the combined effects from the beginning of loop to (before) the current Expression
  • loopEffects is the all effects caused from the loop

TL;DR: loopEffects.throws() will be true when a loop contains a function call and Wasm EH is enabled. wasm-opt conservetively mark function call can throw

      // When EH is enabled, any call can throw. Skip this for return calls
      // because the throw is already more precisely captured by the combination
      // of `hasReturnCallThrow` and `branchesOut`.
      if (parent.features.hasExceptionHandling() && parent.tryDepth == 0 &&
          !curr->isReturn) {
        parent.throws_ = true;
      }

https://github.com/WebAssembly/binaryen/blob/7e117284029bfbfc2b638caa53335e1b2c53490e/src/ir/effects.h#L564-L570

Can we (Scala.js) move out table dispatch?

  • First and foremost, as long as we have call_indirect or call withing a loop, any expressions in the loop is unsafe to loop
    • We have to somehow fix the TODO in wasm-opt
    • // TODO: look into optimizing this with exceptions. for now, disallow

  • Do we want to make itables immutable?
    • It's better to do. If it's immutable, retrieving itables makes effects.readsMutableGlobalState() false. However, it doesn't matter if loopEffects.writesGlobalState() is false
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment