A fine-grained JIT Interface: `func.new` bytecode

[titzer]

Problem

We consider the problem of dynamically generating Wasm code while running on a Wasm engine.

Harvard architecture, no JIT

Core WebAssembly disallows using runtime data as code or generating new code at runtime. This important property of Wasm allows efficient validation of all code in a module and establishes other important security properties. Further, this restriction allows static analysis against the state declared in a module but not exported, admitting sound closed-world optimizations, which is common in tooling today.

Guest runtimes want to JIT new code

However, guest runtimes running on top of Wasm, such as language runtimes like Python or Lua, and hardware emulators, such as QEMU, can benefit tremendously by generating specialized code at runtime. The performance benefit of dynamic code generation can be extreme; as much as 10 to 100 times faster than running in interpreted mode. Today, these runtimes can run only in interpreted mode and are prohibitively slow.

Inconsistent host support for new modules

To date, generating new Wasm code has been an optional host capability provided by the embedder. On the web, this is done through JavaScript APIs for creating new modules directly from bytes. The wasm-c-api allows embedding a Wasm engine in a C/C++ program and creating new modules dynamically. Other host platforms may offer the ability to create new modules from bytes. Generally, the granularity of generation is Wasm modules. Some experiments have been done with function-at-a-time mechanisms, but no standard mechanism has been proposed and nothing is portable across platforms.

Today, guest runtimes that generate new code are few, and they make a number of compromises, such as batching generated functions, in order to workaround limitations and cost of new modules on host platforms.

Proposing a lightweight core-Wasm mechanism

To address this problem, we propose to add a new mechanism to core Wasm. Doing so exposes the mechanism to security features/mitigations, makes behavior explicit and documented, allows sound toolchain analysis and transformations, and exposes it to engine optimizations.

A new instruction: func.new

The main component of this idea is a new core Wasm bytecode, func.new, which creates a new function at runtime from bytecode stored in Wasm memory.

  func.new $mt $ft $scope: [at at] -> [(ref $ft)]
  where:
  - $mt = memory code at limits
  - $ft = func [t1*] -> [t2*]
  - $scope = export*

This bytecode has immediates:

• a memory index, indicating the memory which contains the code for the new function
• a function type, indicating the signature of the new function
• a scope, indicating a list of functions, tables, globals, tags, and memories that new code may legally access

At runtime, this instruction takes operands:

• start, an integer indicating the start offset within the memory of the code
• length, an integer indicating the length of the code

The memory index and function type are straightforward. The $scope immediate is an explicit enumeration of the module contents the new code may reference. In particular, $scope produces a new index space for types, globals, functions, memories, etc. For compactness of the bytecode, and to reuse the same scope for multiple different func.new instructions, $scope will be factored out to its own section rather than inline as immediates.

To execute the instruction, the engine first copies the bytes of code from the Wasm memory, then validates the bytecode under a module context corresponding to $scope and a function context corresponding to the expected function signature $ft. Upon out of bounds memory access or validation error, the instruction traps. If validation of the code is successful, the engine creates an internal representation of a new function and pushes a non-null reference to it onto the operand stack. The function's store is derived from the store under which the instruction was executed; i.e. the new function's instance is (a subset of) the caller's instance.

A new section: JIT scope

A scope defines what declarations are accessible to new code passed to a func.new instruction. We introduce a new "scope" section to factor out the list of accessible declarations so that func.new may refer to a scope by index rather than listing declarations individually. A scope section allows declaring multiple scopes to allow different func.new instructions different accessibility. A scope consists of a list of declarations (similar to export declarations, but without names). The order of declarations in a scope creates new index spaces, i.e. it renumbers the declarations from the surrounding module, starting over from 0.

A new memory flag: code

Generating new code into a module is a powerful feature. It's important that modules have mechanisms to limit and control access to the feature. With the proposed scoping mechanism, a module can limit the interface that new code has access to by explicit enumeration of allowable declarations. This makes sense for modules with security requirements, e.g. where a threat model consider dynamically-generated code potentially compromised or even malicious. Another possible threat is that dynamically-generated code could be corrupted while in Wasm memory, either intentionally or unintentionally, before being supplied to a func.new invocation.

Supporting dynamically-generated code also requires significant runtime support, including a validator and at least one execution tier. Some host platforms may simply be unable to support dynamically-generated code because of memory constraints, lack of execution tiers, or security controls.

For many reasons, we might want to consider the feature "opt-in". We propose a new flag for memories, the code flag, which is part of a memory declaration. Like the shared flag, it explicitly marks a memory that has the capability to be used with func.new. For many of the same reasons as shared memory, this capability thus requires that declarer of the memory and the user of the memory both agree it can be used for code.

An example

Putting the pieces together, we can write an example that uses a code memory, a scope, and creates new functions at runtime.

(module
   (type $t1 (func))  ;; the type for new functions
   (func $f1 ...)
   (func $f2 ...)
   (func $f3 ...)
   (memory $m1 code 1 1)  ;; the memory used to temporarily store code for func.new
   (memory $m2 1 1)       ;; a memory accessible to new code

   (scope $s1             ;; the scope a new function may use
     (func $f1 $f2)       ;; expose $f1 and $f2 to new code
     (memory $m2))        ;; expose only $m2 to new code

   (func $gen
     (local $n (ref $t1)) ;; a variable to hold the new funcref
     ...
     (local.set $n 
       (func.new $m1 $t1 $s1                ;; code lives in $m1, result sig is $t1, scope is $s1
          (i32.const 1024) (i32.const 10))) ;; code is stored at address 1024 and is 10 bytes long
     ...
     (call_ref $t1 (local.get $n))     ;; call the new function!!
   )
)

Sound static analysis

Sound static analysis of Wasm code is vital for module-level optimizations. This proposal preserves reasoning about a module's internal behavior by making its interactions with runtime-generated code explicit. Since runtime generated code can only access the declarations explicitly provided by its scope, analysis can soundly treat declarations mentioned in scopes similarly to exported functions. However, unlike exported functions, scopes cannot be accessed outside a module. Thus a module's public interface is not polluted with its internal use of dynamically-generated code.

Toolchain transformations

Sound static analysis implies that toolchains making closed-world assumptions for optimization (e.g. dead-code elimination) can account for all potential future uses by dynamically-generated code by simply considering scopes. Reorganization of modules resulting from DCE is sound, as scopes can be rewritten for updated indices. Since dynamically-generated code must use scoped indices, it is not affected by the renumbering of the containing module. Thus DCE and other aggressive transformations can be made sound and not affect dynamically-generated code.

Engine optimizations

The possibility of dynamically-generated code implies that an engine has runtime capability to parse function bodies and perform code validation. However, the func.new bytecode doesn't require parsing and validating sections, so the runtime system doesn't need a fully-featured Wasm module parser.

New code implies the need for at least one execution tier, such as interpreter or compiler. This implies AOT scenarios need to either disable the feature, e.g. reject code memories, trap on func.new, or integrate a new tier. Note that since a scope defines a module context, the execution tier only needs to support runtime features that could be legally used under that context--basically, no memories implies no load/store instructions, no GC types implies no GC instructions.

In the future, we could consider additional restrictions, such as an explicit list of allowable bytecodes for new functions. That could allow a module to limit language features for new functions and greatly reduce the runtime system requirement. Since this is a restriction a module imposes on itself, it need not be standardized. As an example of the usefulness of this, a module could limit its func.new capabilities to only allow i32 arithmetic and control flow. An AOT implementation could then provide a runtime execution tier that only supported that restricted set, keeping it both simple and optimized for the use case.

[rossberg]

This looks good! I'm curious to see how far this would take us and what people could do with it.

I have one question and two naming nits:

• Where would the "scope" section go? I assume between data count and code section?
• The more correct name for what's called "scope" here would be "environment" (i.e., a set of bindings, while a scope really is a program region).
• What's written "export" in the syntax above is called an "externidx" in the spec (an index with a namespace prefix, but without the export keyword and name).

[titzer]

Yeah, agree that "environment" is a better name, since there isn't nesting or even embedding. I was also thinking the section would go best just before the code section.

[conrad-watt]

Can the concept of scope/environment be unified with the concept of "nested modules" that showed up in earlier drafts of interface types/the component model? Is an environment essentially a nested module made up of just imports (plus the JIT'd function)?

[kmiller68]

Some random thoughts in no particular order:

Does it really make sense to have a different memory flag? If a producer can't control the code enough to trust a particular memory index is isolated they presumably couldn't guard against malicious use of the flag where it wouldn't be safe. It also wouldn't protect the host wasm VM in any way.

Is there any reason it only takes a linear memory rather than a GC array of bytes? It seems like that would be just as efficient (if not more so, especially if immutable) in most GC environments and possibly more convenient.

Can the scope/environment contain tables/globals? If so then this seems likely to cause memory leaks without some cycle breaking. GC runtimes would get this for free but other environments might not. Maybe this is already an issue though? If not, is that overly limiting for producers?

[titzer]

@conrad-watt Yes, I think conceptually they are very limited nested modules. I didn't want to unduly incorporate the generality of nested modules in this discussion, but that would probably be an elegant treatment in specification language and formalism.

[titzer]

@kmiller68 The choice of linear memory was only for simplicity and discussion. An additional instruction that takes a GC byte array would also be fine with me.

As for the flag, my primary thinking was to exclude accidental use of a memory for creating new code, so that modules have to be new and know about dynamic code generation in order to use the feature. The declaration of the memory and the use of the memory must agree. For example, importing a memory to be used for code generation would require the flag, and the provider of that memory would have to also have the flag, like the shared attribute.

The environment can contain globals and tables. But I don't think this creates any cycles, as the instance for the new function is just a subset of the instance that invokes func.new.

[cfallin]

Would it make sense to fit the notion of asynchronous/background compilation into this design?

Engines today often have mechanisms to kick off a compilation, and allow the background thread to fill in a function pointer asynchronously. For a multitude of reasons, a JIT inside a Wasm guest may wish to preserve that asynchrony. Perhaps there could be a variant (func.new_async or similar) that places a funcref in a designated table index (of compatible nullable funcref type) when complete?

Separately, this would enable a deferred mechanism/workflow, similar to how we process function additions in weval today: one could recover a pseudo-AOT workflow by collecting func.news and batching the function additions into a post-processing step.

[kg]

We currently ship a WASM JIT in production and would love to adopt an opcode like this. It might be helpful to provide some context to inform your thinking here, so I'll try and summarize what we built:

• Our baseline is a combination of AOT-compiled WASM with an interpreter for parts of the application that are not AOTd. So the core of an application is one wasm module that contains our "runtime" and then the AOT'd application code that gets looked up and invoked dynamically as appropriate.
• For anything that isn't AOTd we have transition helpers to move in and out of native WASM code into an interpreter. These have a limited set of supported signatures known at compile time.
• The interpreter has performance instrumentation to identify hot code, and when hot code is identified we tier it up twice, first with an optimization pass for the interpreter IR, and then second for a transition into jitted WASM.

Now for the actual JIT. It has a few modes but the most important one is described below.

• The jitted WASM is generated in small chunks (not mandatory; this was forced on us by Chrome's old 4KB limit for synchronous compilation, which is now gone). We generate a unique module for each chunk. Each chunk corresponds to a block of interpreter code, so that instead of interpreting code we can enter a chunk and then exit the chunk to resume interpretation later.
• Each module has a dynamically generated list of imports, a mix of statically known C helpers (like 'copy GC object in linear memory with write barrier') and libc functions like fmod or printf. This seems a bit different than the scope system you have in mind, but I don't think it would be a dealbreaker for every jitted function to have to share the same imports - this dynamic list is mostly motivated by the 4KB limit, to keep the import table small.
• It would be nice if the import list could be dynamic, as that would allow JITted functions to directly refer to each other at runtime. But if this isn't possible, we can always emulate it with call_indirect.
• Each module has one export since it contains one chunk of code. Now that there's no 4KB limit, we could choose to group them into larger multi-chunk modules, but this would purely be a VM focused optimization - it wouldn't really make life easier for us. So func.new operating on a single function is great.
• Async compilation wasn't an option for us because we need the ability to synchronously 'tier up' a hot loop. During startup an application may want to parse 2MB of JSON or a whole tzdb in one synchronous operation, even if that operation is running on a thread. So if we had to compile asynchronously, we wouldn't be able to receive (through Promise.then) the tiered up WASM code until after the loop finished, making it worthless.
• If func.new_async had a mechanism to allow receiving a compiled fn before an event loop turn, though, it would be a fantastic improvement for us. We have the ability to safely do this kind of transition, because we pre-fill a big block of function pointers with a safe 'dummy' function with the appropriate signature that does nothing. (This is necessary for thread safety, since function pointer tables are per-thread.)
• We import the whole linear memory from the application module, by necessity. I don't see a way around this. I think it would be fine if we had to set a flag on our application module's memory, as long as that flag doesn't require SharedArrayBuffer to be available. The need to enable SAB is a huge blocker for many of our users due to what it makes them do to webserver config.
• If we need to use a separate memory (and multiple memories) to store our JITcode that's probably workable, as long as clang/emscripten expose a way to copy data into that memory. Right now our JIT is implemented in typescript, so the JITcode is conceptually in a separate memory anyway.
• We used to import the function pointer table from the application module, but had to stop because it caused a large memory leak per-module in V8. I believe this is fixed now but we haven't had a chance to re-evaluate it. This means any indirect function calls required the use of a C helper in the application module. It looks like func.new will allow sharing tables, so it would just be important to make sure table sharing isn't expensive.
• All our compiled chunks share the same signature, so we wouldn't need the ability to define new types for a function. I can imagine cases where you'd want to be able to do this though - a full blown JIT for arbitrary methods can't know every possible signature in advance, so it would be ideal if there were a way to define functions of constructed signatures, even if only other func.new'd functions are able to invoke them. This is totally optional though, I can just imagine it being a barrier to use of func.new in a production JIT, especially if you want to expose JITed functions to JS.
• Re the earlier mentions of threads, if func.new could be made thread safe somehow - perhaps some mechanism where it patches the table in all threads at once - that would be really big for making thread-safe JITs. The need to ensure that a function pointer is 'safe' in all threads before exposing it creates really difficult synchronization problems, and the workarounds for it are expensive (pre-filling the fn pointer table with safe dummies takes multiple milliseconds at startup.)

This comment got longer than I wanted so I'll end it here. Feel free to ask questions. Really happy to see this proposal.

[yuri91]

A few points from the perspective of another production Wasm JIT engine (https://webvm.io):

• Asynchronous compilation is very important. I like the idea of asking the embedder to put the func ref in a given table slot when ready, because it does not require a round trip to the event loop, but returning some kind of promise would also be ok.
• We would like to be able to add metadata to the new function, such as branch hints, compilation hints, a name (to help in debugging). This would be easy if the unit of compilation was a module (or a "nested module"?), but if it's just a function it's unclear where to put this info

[titzer]

Ok, summarizing some above points for further discussion:

1. Should the interface accept GC arrays of bytes?
2. How to handle asynchronous compilation? @cfallin @kg and @yuri91 all mentioned this as being important.
3. Could we support new functions referring to another new function in a direct call instruction (e.g. importing them)?
4. Is it possible to add compilation hints to new functions?

For 2) I am curious if the need for asynchrony is to avoid a large pause from the engine compiling a Wasm function to machine code, or if the latency of validation is already prohibitively long. I think in any case the engine should make a synchronous copy of the bytes to avoid racing (early on, we had a security bug in V8 caused by racing on a SAB during module compilation using the synchronous JS API). In my experience validation is fast (hundreds of megabytes per second), so that makes me wonder if engines could copy+validate the bytecode synchronously and provide a valid funcref as a return value, but then block if the function is invoked before machine code is available.

[fgmccabe]

Async compilation sounds like a classic scenario where JSPI is the right approach (on the web at least).

Barring that, a pure Wasm solution for async compilation should wait for stack switching to land IMO. However, there are additional issues in using stack switching for this: in particular, at the moment, there is no intra-WebAssembly event model.

This seems to lead to a model similar to that for JS Strings: realizing jit compilation via special imports rather than adding new instructions.

[titzer]

Also, for reference, Wizard implements a function-at-a-time compilation as a host API. The API produces a new function in the calling instance, similar to func.new. As a lower bound, for a simple function containing nop, with synchronous copy+validate, Wizard can generate 3.5 million new functions per second.

[cfallin]

so that makes me wonder if engines could copy+validate the bytecode synchronously and provide a valid funcref as a return value, but then block if the function is invoked before machine code is available.

One use-case that this doesn't cover, at least, is the desire in the guest to continue using some lower tier (e.g., an interpreter) until the background compilation finishes. The two paths I see to provide that are the asynchronous write to a table slot, as I suggested above, or another opcode (something like func.ready) with signature funcref -> i32. Either way the end-result is the moral equivalent of what lazily compiling native JITs do: dynamically check whether a function pointer is non-null and call it if so.

[titzer]

Interesting idea @cfallin ! func.ready could even be spec'd to have more than one return value, so that 0 = not ready, 1 = ready, 2 = validation error. Then validation could even be made asynchronous if the engine chooses to do so.

[cfallin]

Indeed, and to state it explicitly, before the function is ready, calling the funcref could still have the blocking semantics you describe. That behavior actually makes it possible to do away with the func.new/func.new_async distinction, which is nice as well.

Async validation is interesting: in that case what are the semantics of calling the funcref? A trap? (For the record I don't see async validation as being as important as async compilation, so if there is any pushback here we should separate them; but it's nice to provide the possibility in any case.) A value of 2 to indicate a validation error feels like a little bit of a hack (we should either expose validation programmatically, which would include a synchronous/blocking query, or maintain it as a compiles-to-a-trapping-function semantics) but at the very least it's non-zero so a branch-and-call would succeed and hit that trap.