Implement AOT compilation for components #5160

alexcrichton · 2022-10-31T14:38:29Z

This series of commits is an implementation of AOT-compiling components, namely creating a byte-based artifact from a Component which can, at a later date, be fed back into Wasmtime. The intention of this support is to match the AOT support for Module, only with components instead.

At a high level the serialized artifact is the same for modules. The artifact is an ELF object file with standard .text/data sections as well as Wasmtime-specific sections with bincode-encoded data or otherwise binary-encoded data. The difference with normal modules is that a component squashes together all of its component modules, so the .text section contains all executable code for all modules contained within. Additionally the .rodata section contains the concatenation of all module data segments. This enables having one mmap for the entire component, including all modules contained within.

Lots of refactoring internally was necessary for this, namely around the end-part of compilation where an object is emitted and loaded into modules/components. The general shape of the Module compilation pipeline is now quite similar to the Component compilation pipeline. There were two caveats from this implementation, however:

Modules within a component sharing the same mmap means that the module registration system within stores must now account for modules with overlapping text sections. The .text() accessor returns the same slice for multiple modules, meaning there's now a second-level of indirection for maps where we might want frame information from them.
Components can export modules, and in the embedding API this returns an actual Module. This module, however, shares the exact same serialized artifact as the component itself, meaning that if the Module itself were serialized it couldn't be deserialized back into a module. For that reason I've at least for now disabled Module::serialize dynamically for modules acquired from a component. This should be fixable but it didn't seem worth it at this time.

Otherwise this is quite a large change so I've tried to break apart the commits as best I could. Each individual one should pass at least cargo test and represents an incremental step towards AOT-compiling components. If desired I can further split out commits into separate PRs as well.

As a final note, the route I've chosen here of squashing everything into one object file is notably distinct from previously intended strategies. For example with module linking modules were serialized by concatenating all of their objects together. This was not easily introspectable so my other idea was to use an ar archive for storing a concatenation of objects. In the end I settled on the "one object file" approach for a few reasons:

Using an ar format with the alignment requirements for our object files from mmap'ing the original file to an executable mmap was not going to be easy and I never bottomed out this idea.
Each object would have its own alignment requirements which was going to add up to a lot of padding. I was afraid of excessive padding in a component with multiple small modules internally.
Conceptually to me at least a component encompassing modules felt better represented as one large mmap with one shared allocation between all the types. This commit for example cuts down on the individul number of Arcs required and additionally fixes longstanding issues such as only creating one SignatureCollection for a component.
Overall it felt easier at the time to implement shared systems for loading all of this and felt like the more reasonable course of action.

Of course the serialization format is just an internal detail of Wasmtime and if necessitated we can change it in the future as well.

Closes #5119

This commit is the first in what will likely be a number towards preparing for serializing a compiled component to bytes, a precompiled artifact. To that end my rough plan is to merge all of the compiled artifacts for a component into one large object file instead of having lots of separate object files and lots of separate mmaps to manage. To that end I plan on eventually using `ModuleTextBuilder` to build one large text section for all core wasm modules and trampolines, meaning that `ModuleTextBuilder` is no longer specific to one module. I've extracted out functionality such as function name calculation as well as relocation resolving (now a closure passed in) in preparation for this. For now this just keeps tests passing, and the trajectory for this should become more clear over the following commits.

This commit removes the `ComponentCompiler::emit_obj` function in favor of `Compiler::emit_obj`, now renamed `append_code`. This involved significantly refactoring code emission to take a flat list of functions into `append_code` and the caller is responsible for weaving together various "families" of functions and un-weaving them afterwards.

This commit moves the ELF file parsing and section iteration from `CompiledModule` into `CodeMemory` so one location keeps track of section ranges and such. This is in preparation for sharing much of this code with components which needs all the same sections to get tracked but won't be using `CompiledModule`. A small side benefit from this is that the section parsing done in `CodeMemory` and `CompiledModule` is no longer duplicated.

Previously components would generate an "always trapping" function and the metadata around which pc was allowed to trap was handled manually for components. With recent refactorings the Wasmtime-standard trap section in object files is now being generated for components as well which means that can be reused instead of custom-tracking this metadata. This commit removes the manual tracking for the `always_trap` functions and plumbs the necessary bits around to make components look more like modules.

Not expected to have any measurable impact on performance, but complexity-wise this should make it a bit easier to understand the internals since there's no longer any need to store this somewhere else than its owner's location.

This commit is a large refactoring of the component compilation process to produce a single artifact instead of multiple binary artifacts. The core wasm compilation process is refactored as well to share as much code as necessary with the component compilation process. This method of representing a compiled component necessitated a few medium-sized changes internally within Wasmtime: * A new data structure was created, `CodeObject`, which represents metadata about a single compiled artifact. This is then stored as an `Arc` within a component and a module. For `Module` this is always uniquely owned and represents a shuffling around of data from one owner to another. For a `Component`, however, this is shared amongst all loaded modules and the top-level component. * The "module registry" which is used for symbolicating backtraces and for trap information has been updated to account for a single region of loaded code holding possibly multiple modules. This involved adding a second-level `BTreeMap` for now. This will likely slow down instantiation slightly but if it poses an issue in the future this should be able to be represented with a more clever data structure. This commit additionally solves a number of longstanding issues with components such as compiling only one host-to-wasm trampoline per signature instead of possibly once-per-module. Additionally the `SignatureCollection` registration now happens once-per-component instead of once-per-module-within-a-component.

This commit adds support for AOT-compiled components in the same manner as `Module`, specifically adding: * `Engine::precompile_component` * `Component::serialize` * `Component::deserialize` * `Component::deserialize_file` Internally the support for components looks quite similar to `Module`. All the prior commits to this made adding the support here (unsurprisingly) easy. Components are represented as a single object file as are modules, and the functions for each module are all piled into the same object file next to each other (as are areas such as data sections). Support was also added here to quickly differentiate compiled components vs compiled modules via the `e_flags` field in the ELF header.

The current representation of a module within a component means that the implementation of `Module::serialize` will not work if the module is exported from a component. The reason for this is that `serialize` doesn't actually do anything and simply returns the underlying mmap as a list of bytes. The mmap, however, has `.wasmtime.info` describing component metadata as opposed to this module's metadata. While rewriting this section could be implemented it's not so easy to do so and is otherwise seen as not super important of a feature right now anyway.

github-actions · 2022-10-31T14:47:47Z

Subscribe to Label Action

cc @peterhuene

This issue or pull request has been labeled: "cranelift", "cranelift:area:aarch64", "cranelift:area:machinst", "cranelift:area:x64", "wasmtime:api"

Thus the following users have been cc'd because of the following labels:

peterhuene: wasmtime:api

To subscribe or unsubscribe from this label, edit the .github/subscribe-to-label.json configuration file.

Learn more.

fitzgen · 2022-11-01T16:59:45Z

(Looking at this today, but might take a little while cause this is not small)

fitzgen

LGTM!

fitzgen · 2022-11-01T19:53:22Z

cranelift/codegen/src/isa/mod.rs

@@ -285,7 +285,7 @@ pub trait TargetIsa: fmt::Display + Send + Sync {
    /// The `num_labeled_funcs` argument here is the number of functions which
    /// will be "labeled" or might have calls between them, typically the number
    /// of defined functions in the object file.
-    fn text_section_builder(&self, num_labeled_funcs: u32) -> Box<dyn TextSectionBuilder>;
+    fn text_section_builder(&self, num_labeled_funcs: usize) -> Box<dyn TextSectionBuilder>;


Not totally clear to me what the motivation for changing this was. Not saying its necessarily wrong or anything but I'm just curious.

Yeah there ended up being a lot of casts around this and having everything be usize removed all the casts, so I went ahead and updated this to usize.

fitzgen · 2022-11-01T20:07:17Z

crates/cranelift/src/compiler.rs

@@ -367,16 +368,23 @@ impl wasmtime_environ::Compiler for Compiler {
            .map(|f| *f.downcast().unwrap())
            .collect();

-        let mut builder = ModuleTextBuilder::new(obj, &translation.module, &*self.isa);
+        let num_funcs = funcs.len() + compiled_trampolines.len();
+        let mut builder = ModuleTextBuilder::new(obj, &*self.isa, num_funcs);


I guess its just because we are passing in .len()-ish things as input? And isn't necessarily tied to a single u32-sized index space anymore?

crates/environ/src/compilation.rs

Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

alexcrichton added 9 commits October 29, 2022 10:21

Remove a now-unnecessary Arc in Module

7ef1613

Not expected to have any measurable impact on performance, but complexity-wise this should make it a bit easier to understand the internals since there's no longer any need to store this somewhere else than its owner's location.

Fix compile errors from prior commits

153bba1

alexcrichton added 2 commits October 31, 2022 08:15

Fix windows build

abcedc7

Fix an unused function warning

923c47a

alexcrichton requested a review from fitzgen October 31, 2022 16:57

fitzgen approved these changes Nov 1, 2022

View reviewed changes

Update crates/environ/src/compilation.rs

f14666c

Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

alexcrichton enabled auto-merge (squash) November 2, 2022 14:30

alexcrichton merged commit cd53bed into bytecodealliance:main Nov 2, 2022

alexcrichton deleted the refactor-compilation branch November 2, 2022 15:26

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement AOT compilation for components #5160

Implement AOT compilation for components #5160

alexcrichton commented Oct 31, 2022 •

edited

Loading

github-actions bot commented Oct 31, 2022

fitzgen commented Nov 1, 2022

fitzgen left a comment

fitzgen Nov 1, 2022

alexcrichton Nov 2, 2022

fitzgen Nov 1, 2022

Implement AOT compilation for components #5160

Implement AOT compilation for components #5160

Conversation

alexcrichton commented Oct 31, 2022 • edited Loading

github-actions bot commented Oct 31, 2022

Subscribe to Label Action

fitzgen commented Nov 1, 2022

fitzgen left a comment

Choose a reason for hiding this comment

fitzgen Nov 1, 2022

Choose a reason for hiding this comment

alexcrichton Nov 2, 2022

Choose a reason for hiding this comment

fitzgen Nov 1, 2022

Choose a reason for hiding this comment

alexcrichton commented Oct 31, 2022 •

edited

Loading