1319358 - Wasm baseline: Generate return code in-line at first return point

Status: RESOLVED WONTFIX

(Core :: JavaScript: WebAssembly, defect, P3)

Description

[Lars T Hansen [:lth]]

If it is true that many functions are small - and we expect this to be the case - then we should focus a little bit on trimming boilerplate code size where we can.

:dannas put a nicely annotated listing of a trivial wasm function here:
https://bug1316803.bmoattachments.org/attachment.cgi?id=8812902

In that listing, there's only one point where the function returns, but that return is implemented as a jump to a common return point.  It should be straightforward to bind the returnLabel at the point where we encounter the first return statement, if any, and to emit the return code there, so that we avoid that jump.  If there is only the one return then the function will flow naturally into the return.

(Of course, a good wasm compiler will omit the (return E) node and just emit E here, so it's unclear how much this matters, but it seems like good hygiene and should cost very little.)

Comment 1

[Lars T Hansen [:lth]]

Testing shows that many functions are small: bug 1320374 comment 8.  This optimization may be quite useful.

Comment 2

[(Away)]

Attachment: Naive patch (broken)

I tried a simple patch to replace the first instance of "jmp returnLabel_" with the in-line return code, but I ran into problems:

* The value of |localSize_| can change between doReturn() and endFunction(), e.g. in https://dxr.mozilla.org/mozilla-central/rev/741a720c98cdb92c229376be0badbf036f653bff/js/src/jit-test/tests/wasm/regress/misc-control-flow.js#64,78. This leads to a complaint at MOZ_ASSERT(masm.framePushed() == framePushed) and test failures. Maybe we could patch in the correct localSize_ later, like the proposal in bug 1319388?

* Some unreachable-code situations never call doReturn(), so I kept the old epilogue code in place, but it feels like I shouldn't have to.

Comment 3

[Lars T Hansen [:lth]]

(In reply to David Major [:dmajor] from comment #2)
> Created attachment 8815957 [details] [diff] [review]
> Naive patch (broken)
> 
> * The value of |localSize_| can change between doReturn() and endFunction(),
> e.g. in
> https://dxr.mozilla.org/mozilla-central/rev/
> 741a720c98cdb92c229376be0badbf036f653bff/js/src/jit-test/tests/wasm/regress/
> misc-control-flow.js#64,78. This leads to a complaint at
> MOZ_ASSERT(masm.framePushed() == framePushed) and test failures. Maybe we
> could patch in the correct localSize_ later, like the proposal in bug
> 1319388?

|localSize_| should only be changed by (the mis-named) pushLocal() which is only called from init() which should be finished at that point, so localSize_ should be constant.  Details as to how this failure mode works?

Anyway that assertion comes from GenerateFunctionEpilogue, which is probably reacting to masm.framePushed() being incorrect, which is probably a result of not properly popping the stack here.  That's mysterious since returnCleanup should have done the right amount of cleanup.

> * Some unreachable-code situations never call doReturn(), so I kept the old
> epilogue code in place, but it feels like I shouldn't have to.

Yes, we should be able to avoid that, though I'm not sure it's much of a hardship.  (Maybe checking deadCode_ will be useful.)

Comment 4

[(Away)]

> |localSize_| should only be changed by (the mis-named) pushLocal() which is
> only called from init() which should be finished at that point, so
> localSize_ should be constant.  Details as to how this failure mode works?

Er, I wrote the wrong thing. localSize_ stays constant, but masm.framePushed_ changes at the stack below, between doReturn() and endFunction(). The effect is the same: mismatch between localSize_ and framePushed_.

00 js!js::jit::MacroAssembler::freeStack
01 js!js::jit::MacroAssembler::adjustStack
02 js!js::wasm::BaseCompiler::popStackOnBlockExit
03 js!js::wasm::BaseCompiler::endBlock
04 js!js::wasm::BaseCompiler::emitEnd
05 js!js::wasm::BaseCompiler::emitBody
06 js!js::wasm::BaseCompiler::emitFunction

The function that causes this is pretty hairy; I was surprised the language permits this:

 (func (result i32)
  (call 0
   (i32.const 42)
   (i32.const 53)
   (call 0 (i32.const 100) (i32.const 13) (i32.const 37) (i32.const 128))
   (return (i32.const 42))
  )

Comment 5

[Benjamin Bouvier [:bbouvier] (inactive)]

(In reply to David Major [:dmajor] from comment #4)
> The function that causes this is pretty hairy; I was surprised the language
> permits this:
> 
>  (func (result i32)
>   (call 0
>    (i32.const 42)
>    (i32.const 53)
>    (call 0 (i32.const 100) (i32.const 13) (i32.const 37) (i32.const 128))
>    (return (i32.const 42))
>   )

For what it's worth, this is because wasm evolved from an AST to a stack machine. Before the text format for testing (that we're using here) was just a LISP embedding expressions, but now it is really a stack machine assembly. That being said, to keep backward compatibility, we still allowed LISP-like expressions which are converted to a stack machine representation. The TL;DR is that this code is equivalent to:

(func (result i32)
  i32.const 42
  i32.const 53
  i32.const 100
  i32.const 13
  i32.const 37
  i32.const 128
  call 0
  i32.const 42
  return
  call 0
)

Comment 6

[Lars T Hansen [:lth]]

Attachment: bug1319358-inline-return.patch

The trick is that the assembler's notion of the stack size must be adjusted before and after the return code is emitted: before the code, to bring the stack down to an "empty frame"; after the code, to return it to its original state while also taking into account the amount popped by the epilogue.

I don't want to ask for review on this yet because I want to spend a little more time checking that the code that's emitted is an improvement, but I think this is all that's needed.

Comment 7

[Lars T Hansen [:lth]]

This is looking fine, really.  We get rid of the jump-to-return to implement the return operation, which is probably small potatoes but good hygiene.

// Program returns the sum of its arguments

(func (param i32) (param i32) (result i32)
  (return (i32.add (get_local 0) (get_local 1))))

// Generated code

        subq       $24, %rsp
;; stack overflow check (patched with proper size later)
        movq       %rsp, %rax
        addq       $0x0000, %rax
        cmpq       %rax, 0x20(%r14)
        jae        L_stack_overflow_trap
;; save parameters
        movl       %edi, 0xc(%rsp)
        movl       %esi, 0x8(%rsp)
;; save tls
        movq       %r14, 0x0(%rsp)
;; body
        movl       0x8(%rsp), %eax
        movl       0xc(%rsp), %ecx
        addl       %eax, %ecx
;; return
        movl       %ecx, %eax
;; restore tls and discard frame
        movq       0x0(%rsp), %r14
        nop (2 byte)
        addq       $24, %rsp
        ret

Comment 8

[Lars T Hansen [:lth]]

Testing to be done:

- try
- AngryBots
- embenchen (ideally on x86, x64, ARM)

Comment 9

[Lars T Hansen [:lth]]

https://treeherder.mozilla.org/#/jobs?repo=try&revision=04138eb65fe0

Comment 10

[Lars T Hansen [:lth]]

Attachment: bug1319358-inline-return.patch

Passes trybot, AngryBots, and embenchen (on x64 only, so far), should be good to go.

Comment 11

[Luke Wagner [:luke]]

So in any wasm function:
  (function (result i32) ... (i32.const 1) return)
the final 'return' is strictly unnecessary since one achieves the same effect with:
  (function (result i32) ... (i32.const 1))
since the final contents of the stack are the return value (and validated to match (result)).  If we do see these redundant `return`s in binaryen output, let's file binaryen bugs; these are trivial size and decode-time wins.

So given that, I think if we want to optimize out the final `jmp`, we should:
 - remove the "if (!deadCode_) doReturn(...)" in emitFunction() and the breakpoint() at the start of endFunction()
 - masm.bind(&returnLabel_) *before* the masm.bind(&stackOverflowLabel_) and emit the epilogue here (this mirrors CodeGenerator::generateWasm())

Also, as a minor point since it's probably a negligible perf difference, it seems slightly better to have the return always show up at the end of the function under the assumption that "most" functions return from their last return "most" of the time.

Comment 12

[Luke Wagner [:luke]]

Comment on attachment 8822622 [details] [diff] [review]
bug1319358-inline-return.patch

(Clearing for now)

Comment 13

[Lars T Hansen [:lth]]

That's fine, it won't happen for FF53 given your pushback above and I'd like to wait for the dust to settle on the dead code discussion anyhow.

Comment 14

[Luke Wagner [:luke]]

That makes sense.

Comment 15

[BMO Automation]

Per policy at https://wiki.mozilla.org/Bug_Triage/Projects/Bug_Handling/Bug_Husbandry#Inactive_Bugs. If this bug is not an enhancement request or a bug not present in a supported release of Firefox, then it may be reopened.