1868521 - Support direct call inlining in Ion

Status: RESOLVED FIXED

(Core :: JavaScript: WebAssembly, enhancement, P1)

Description

[Julian Seward [:jseward]]

Bug 1864209 tracks experimentation with a purely-MIR-level inliner for use
with wasm. That scheme seems difficult to make work reliably, so this bug
tracks experimentation of an inliner for wasm structured in the same way that
inlining is done in Warp; that is, each callee to be inlined is generated
directly from wasm bytecode every time it is needed.

Attached is an initial patch that does this, for direct, non-catchable calls
only. It is very much WIP and contains a lot of leftover junk from the
bug 1864209 attempt that can now be removed. The patch will inline
arbitrarily deep, stopping only when a callee has more than 24 bytes of
bytecode. Multi-value returns don't work yet.

Comment 1

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-1

Comment 2

[Julian Seward [:jseward]]

Attachment: badness124.js

A simple test case.

Comment 3

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-4

Greatly cleaned up, and now handles args that would have been passed
on the stack. Multi-value returns still don't work though.

Comment 4

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-6

Able to inline any call-direct, non-catchable call, providing the callee's
bytecode is available. This includes calls to functions requiring args to be
passed on the stack (eg x86_32) and calls to functions that return multiple
values.

Comment 5

[Julian Seward [:jseward]]

Attachment: inlining-test-cases-2023Dec19.tar

Test cases for the comment 4 implementation.

Comment 6

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-7

A very minor cleanup of the comment 4 patch. Can process all of the
comment 5 test cases correctly, at least on x86_{32,64}-linux, although you'll
need to run with IONFLAGS=dump-mir-expr to see that. Applies against m-c
684483:5d6699b34edc.

Comment 7

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-10

Now with significantly fewer bugs.
Also, IONFLAGS=wasm-inlining now produces a 1-line summary
of inlining activity for each (top-level) function processed.

Comment 8

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-11

Now with loop depths handled properly. This fixes assertion
failures to do with phi nodes mysteriously disappearing.

Comment 9

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-13

Now doesn't cause crashing in the presence of return calls, and has a space
leak fix. Pretty much all of the jit_tests wasm tests pass now, and large
files (clang.wasm, dart.test.wasm, etc) can be processed. Is usable for
playing around assessing the perf effects of inlining (direct) calls. The
only major missing functionality now is that calls inside try blocks are not
(yet) inlined.

Comment 10

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-15

Rebased to m-c 691458:33d47d05c368 (9 Jan 2024).

Comment 11

[Julian Seward [:jseward]]

Attachment: bug1868521-inliner-take2.diff-16

With per-function-stack state split out of FunctionCompiler into new class
FunctionStackState.

Comment 12

[Ryan Hunt [:rhunt]]

Once we have bug 1885363 and bug 1876155 landed, I believe we could update this could to be in a state that's landable. I'm renaming this bug to move the focus from experimentation to something we could land.

Comment 13

[Ryan Hunt [:rhunt]]

Attachment: Bug 1868521 - wasm: Cleanup CodeMetadata and ModuleMetadata. r?jseward

1. Move declareFuncExported to FuncDesc
2. Move addDefinedFunc/addImportedFunc to ModuleMetadata
3. Make sure we serialize everything
4. Move branch hint parsing information to BranchHintCollection
5. Re-order all of the fields on CodeMetada to be grouped logically and roughly follow initialization order
6. namePayload doesn't need to be mutable after we made CodeMetadata be owned by ModuleMetadata

Comment 14

[Ryan Hunt [:rhunt]]

Attachment: Bug 1868521 - wasm: Move bytecode to CodeMetadata. r?jseward

We need bytecode() for inlining, move it to CodeMetadata
instead of Code.

Depends on D216750

Comment 15

[Ryan Hunt [:rhunt]]

Attachment: Bug 1868521 - wasm: Support direct call inlining in Ion when lazy tiering. r?jseward

This commit adds direct call inlining to Ion when we're lazy tiering,
using simple heuristcs.

The overall design is similar to how JS does inlining with Warp. When
we see a call instruction we want to inline, we recursively create a
new FunctionCompiler and have that generate MIR for the inlined callee
function.

The following things needed to be changed:

1. FunctionCompiler gains a notion of a 'callerCompiler' which tracks
   the outer caller function when inlined.
2. The top-level function in a compilation now holds alive a set of
   jit::CompileInfo for all the inlined functions. This is because
   this data structure is pointed to by MIR and so it must live until
   the MIR is compiled.
3. FunctionCompiler::initInlined is added which sets up the params to
   come from a caller, instead of using the MWasmABIArg defs.
4. FunctionCompiler::returnValues is changed to accumulate the returns
   when inlining, instead of using MWasmReturn instructions.
5. FunctionCompiler::finishInlinedCallDirect is added to bind the pending
   returns to a join block.
6. A shouldInlineCallDirect heuristic is added, partially controlled by
   new prefs. This is naive and could use some tuning.
7. Tail calls are not supported in the inlined function, but we have no
   good way to rule them out, so we just crash for now. The next commit
   fixes this.

There is a bit of redundant state on the stack of FunctionCompiler's
which could be factored out into some sort of 'SharedFunctionCompiler'.
This commit does not do that, but a future commit could.

Depends on D216752

Comment 16

[Ryan Hunt [:rhunt]]

Attachment: Bug 1868521 - wasm: Use FeatureUsage to prevent inlining tail calls. r?jseward

Add a per-funcdef FeatureUsage tracker that we can use when lazy
tiering to prevent inlining into a tail call. This lets us remove
the crash and replace it with an assert.

Depends on D216753

Comment 17

[Pulsebot]

Pushed by [e-mail]:
https://hg.mozilla.org/integration/autoland/rev/7a6190776114
wasm: Cleanup CodeMetadata and ModuleMetadata. r=jseward
https://hg.mozilla.org/integration/autoland/rev/e6767af87508
wasm: Move bytecode to CodeMetadata. r=jseward
https://hg.mozilla.org/integration/autoland/rev/a94fa14b655e
wasm: Support direct call inlining in Ion when lazy tiering. r=jseward
https://hg.mozilla.org/integration/autoland/rev/48e881f47b12
wasm: Use FeatureUsage to prevent inlining tail calls. r=jseward
https://hg.mozilla.org/integration/autoland/rev/bde59718573f
apply code formatting via Lando

Comment 18

[Iulian Moraru]

https://hg.mozilla.org/mozilla-central/rev/7a6190776114
https://hg.mozilla.org/mozilla-central/rev/e6767af87508
https://hg.mozilla.org/mozilla-central/rev/a94fa14b655e
https://hg.mozilla.org/mozilla-central/rev/48e881f47b12
https://hg.mozilla.org/mozilla-central/rev/bde59718573f