This is a proposal to support WebAssembly module representation in MLIR.
We have a work-in-progress implementation available here.
TL;DR: We have been working on an MLIR representation for WebAssembly programs, and tools to import such programs in MLIR and convert them up to the LLVM IR dialect, in order to be able to compile Wasm binaries to native code ahead-of-time.
The longer term goal is to provide support for Wasm usage on embedded systems.
The dialect covers a good percentage of the Wasm constructs, but is still in early stages of development. Open questions remain on what the LLVM lowering should look like. Thus, we think it’s a good time to go to the community, receive feedback on our current approach, and shift ongoing work into the community for cross-industry collaboration.
Context and Motivation
WebAssembly (abbreviated Wasm) is a well specified binary instruction format for an abstract stack machine.
An embedder is used to execute Wasm code in a given environment, sandboxing the execution.
LLVM already supports Wasm as a code generation target. This proposal is about doing the opposite: taking Wasm binary code as an input and producing native machine code for it.
Outside of web browser execution, there is also an interest in using Wasm as a binary format. There are several reasons for this.
Firstly, Wasm was designed from the ground up for safety and security. Out of the box, Wasm offers sandboxing, control flow integrity enforcement, memory safety guarantees, and so on. For developers interested in security, these guarantees are compelling.
Secondly, Wasm enables “compile once, run anywhere”; the concept that you can run a Wasm binary anywhere where its embedder is available. The embedder can be thought of as a Wasm interpreter. A close-to-home example of this is the 2024 RFC to build LLVM itself for Wasm.
However, there is a tradeoff to using Wasm as a portable binary format: this code isn’t native, and so it comes with a performance hit. For example, the YoWASP project uses Wasm to provide FPGA tools to many different platforms, but the cost of roughly a 50% performance hit.
We could keep most of that benefit if we added a lightweight AoT Wasm compiler that takes pre-compiled Wasm binaries, and compiles them directly to native code. In addition, AoT compiling allows the reduction of the amount of extra software that needs to be provided on the target platforms. The full interpreter / JIT compiler can be replaced by only a Wasm runtime that implements the sandboxing etc.This is especially relevant for safety critical applications for which the embedded software needs to undergo a qualification process (e.g. in aerospace, automotive, and other safety-critical contexts.)
Related Work
AOT compilers for Wasm already exist, such as WAMR or Wasmtime, but they are tightly coupled with their respective embedders. The proposed dialect is embedder-agnostic, and should allow for multiple lowering paths if required.
There are existing LLVM-based Wasm compilers (e.g.Wasker). However, existing solutions use LLVM as an external component. Our approach brings Wasm->native support directly into the LLVM monorepo.
High-level Approach
This is a high-level overview of how we get from a binary Wasm file to LLVM IR (and subsequently, native code.)
Suppose you have the following Wasm text format (WAT) code.
(module
(func (export "or_i32") (result i32)
i32.const 10
i32.const 3
i32.or)
)
Let’s say you have placed this WAT in the Wasm binary format as foo.wasm, using some external tool such as wat2wasm.
Via mlir-translate --import-wasm foo.wasm, we can lower this to code in the Wasm MLIR dialect like so:
wasm.func nested @or_i32() -> i32 {
%0 = wasm.const 10 : i32
%1 = wasm.const 3 : i32
%2 = wasm.or %0 %1 : i32
wasm.return %2 : i32
}
We’ve designed the Wasm dialect to mirror the opcodes, data structures, and so on that exist in Wasm. So, for example, each Wasm opcode has an associated opcode in the Wasm dialect.
Readers familiar with Wasm might be wondering at this point, “what about the Wasm stack?” The Wasm stack is tracked at parsing time to correctly map operation results to uses. So, while parsing a Wasm binary, we can verify that code like this is invalid:
(module
(func (export "or_i32") (result i32)
i32.const 10
i32.or)
)
In the above case, the importer will recognize that there are insufficient values on the stack for the i32.or, emit an error, and produce no Wasm dialect code.
Once we’ve finished translation, we can pass the Wasm dialect code to mlir-opt --convert-wasm-to-standard. This will convert the Wasm dialect code to the appropriate standard dialect(s) like so:
func.func @or_i32() -> i32 {
%c10_i32 = arith.constant 10 : i32
%c3_i32 = arith.constant 3 : i32
%0 = arith.ori %c10_i32, %c3_i32 : i32
return %0 : i32
}
This provides a path to LLVM IR, and subsequently, native code on the user’s platform of choice.
For each implemented Wasm construct, we have included
- Translation unit tests in mlir/test/Target/WebAssembly
- Conversion unit tests in mlir/test/Conversion/WasmToStandard
To automate this process, we also have implemented a simple driver, called Wasabi that performs the translation and conversion phases.
Via Wasabi, the user can automatically convert from a Wasm binary to the MLIR dialect, standard dialects, or LLVM IR.
# Translate foo.wasm to Wasm MLIR and output the result.
$ wasabi /path/to/foo.wasm --dump --output-type wasm-mlir
# Translate foo.wasm to Wasm MLIR, convert to standard dialects, and output the result.
$ wasabi /path/to/foo.wasm --dump --output-type std-mlir
# Translate foo.wasm to Wasm MIR, convert to standard dialects, lower to LLVM IR, and output the result.
$ wasabi /path/to/foo.wasm --dump
We’ve included LLVM lit tests which serve as additional examples for all components.
What Will Be Upstreamed?
We propose that we upstream the
- Wasm dialect
- Translation
- Conversions
- Wasabi driver
Into llvm-project as quickly as possible. We have been following LLVM and MLIR coding standards and conventions with the plan of upstreaming from the beginning.
We are happy to make any changes to the implementation that are necessary to bring this work in-tree.
Why Now?
We believe that multiple parties can and will benefit from this work in the embedded space and beyond. With the first Asian LLVM conference happening this week, we think that this is an excellent time to begin pushing this work into the open.
The work is not 100% complete, but is in a “mostly LLVM OSS programming style state”. We think it’s at a state where it’s easy for contributors to join the effort and push it forward, and we think there is enough work there that it can definitely be pushed to completion in OSS.
Git Integration Approach
We’ve put the current code up for community consideration and review on this fork branch.
We don’t have strong opinions on how this should be upstreamed into LLVM, so we would like community input on the best path forward.
Code Layout
The current implementation is structured like so:
Dialect
Translation
Conversion
- mlir/lib/Conversion/WasmToStandard
- mlir/include/mlir/Conversion/WasmToStandard
- mlir/test/Conversion/WasmToStandard
Driver
Status and plan
In its actual state, the project consists of :
- A new WebAssembly dialect for MLIR which represents Wasm programs
- A translation target which allows to import Wasm binary to this dialect with mlir-translate
- A (set of) conversions to lower this dialect to standard MLIR dialects (memref, arith, math, cf, func, …) that can themselves be lowered to the LLVM dialect and then compiled to native code. At the moment this is a monolithic conversion, targeting multiple standard dialects at the same time.
The proposed implementation covers Wasm 2.0 specification
- Structural constructs of Wasm that are represented in the dialect and can be imported:
- Functions (type definitions and partial support of Wasm instructions for function code, see below for more detail)
- Globals
- Memory (no initializers at the moment)
- Tables (no initializer at the moment)
- Imports
- Exports
- Start point → not yet supported
- Instruction support:
- Numeric instructions on scalar types: Almost 100% coverage
- Vector instructions: no support at the moment
- Global and local variable instructions: 100% support
- Control flow instructions: 100% support
- Trap: embedder dependent, not yet supported
- Memory and Table instructions: not implemented yet, as it can also be embedder dependent.
Technical Debt
There are some validations in the Wasm Binary importer for which no negative tests have been implemented. Existing tools to generate Wasm binaries from WAT verify that the input WAT is valid. Therefore, it is difficult to generate broken Wasm binaries today.
Plan for this is to first reach some level of agreement on what should be tested in the importer vs delegated to dialect verifier, specify this somewhere and write the test cases accordingly afterwards.
It might be interesting to actually create a fuzzing tool for this, but a first implementation can be as simple as manually editing valid Wasm binary files to produce invalid ones.
Currently we’re not consistent on where we call it “WebAssembly” and where we call it “Wasm”. Since naming things is one of the hardest problems in computer science, we would like to ask the community for assistance in choosing the most appropriate name. 
Design Considerations
Representing the Wasm Stack
While Wasm targets an abstract stack machine, the proposed WebAssembly dialect abstracts away the stack and uses an SSA representation of the program.
It is of the importer’s responsibility to track the stack state during the import in order to map correctly results to consumers.
As a result, the WebAssembly Dialect can represent incorrect Wasm programs.
For instance, this code snippet
%0 = wasm.const 1: i32
%1 = wasm.add %0 %0 : i32
Would represent an invalid program, as the stack machine semantics forbids reusing the same value twice (an operation consumes its operands), so here a second distinct wasm.const 1 would be necessary to have a program with valid stack state.
There is however no conceptual difficulty to implement additional verification to enforce the validity of a represented module. Part of the validation would be to ensure that operands of Wasm operations that are taken on the stack have exactly one consumer.
An alternative consisting in having operations producing and consuming an opaque-typed stack value has been considered, but this doesn’t solve the issue (a !wasm.stack value can be reused multiple times if no verification forbids it) and just adds complexity for tracking which result is reused where, makes implementing the conversion to standard dialects challenging (a shared state to track which is the actual type of the value(s) added / consumed on the stack would be required, which isn’t easily done within the Conversion Rewriter framework) and make the type checking more complex.
Embedder-Dependent Constructs
In addition, open questions remain on how to lower some constructs that are dependent on the embedder. Indeed, while arithmetic operations poses no great difficulty as they are “naturally” mapped to operations from Math or Arith dialects, operations on memory or function table for instance (indirect call of a function stored in a function table) is trickier as it is actually dependent on the embedder API.
For instance:
- the memory management could be done internally to the module, with a memref global storing a pointer to a memref storing a vector of bytes, but in the case the embedder manages the memory it should be lowered to call to the embedder management API instead
- The way to register an initialization function isn’t clear either. Registration of the entry point function is also something dependent on the embedder mechanism.
Detailed Wasm Dialect Example
This example shows how Wasm concepts such as locals, function arguments, control flow operations, and arithmetic operations are represented in SSA form via the Wasm dialect.
Given the following WAT code:
(module
(func (param $i i32) (result i32)
(loop $my_loop (result i32)
local.get $i
i32.const 1
i32.add
local.set $i
local.get $i
i32.const 10
i32.lt_s
br_if $my_loop
local.get $i
)
)
)
We can produce the following Wasm dialect code:
module {
wasm.func nested @func_0(%arg0: i32) -> i32 {
%0 = wasm.local_from_arg %arg0 : i32
wasm.loop : {
%2 = wasm.local_get %0 : memref<i32>
%3 = wasm.const 1 : i32
%4 = wasm.add %2 %3 : i32
wasm.local_set %0 : memref<i32> to %4 : i32
%5 = wasm.local_get %0 : memref<i32>
%6 = wasm.const 10 : i32
%7 = wasm.lt_si %5 %6 : i32 -> i32
wasm.branch_if %7 to level 0 else ^bb1
^bb1: // pred: ^bb0
%8 = wasm.local_get %0 : memref<i32>
wasm.block_return %8 : i32
}> ^bb1
^bb1(%1: i32): // pred: ^bb0
wasm.return %1 : i32
}
};
Local Variables
Since Wasm “local” values are variables that can be overwritten with “get” and “set” on the stack, they can’t be represented as normal SSA values. We have chosen to represent them as memref in order to benefit from the pointer-like semantic.
%2 = wasm.local_get %0 : memref<i32>
Function Arguments
Function arguments are mapped to Wasm locals using the local_from_arg operation. A function argument is effectively a value that “comes in” on the Wasm stack, passed along by the caller to the callee function.
%0 = wasm.local_from_arg %arg0 : i32
Control Flow Operations
Blocks, If, and Loops, are terminator operations that take some operands based on their function type. They define regions containing the operations that are subject to the control flow (e.g. conditionally executed for If)..
Control flow operations have a successor block which is the usual block to execute when the loop / block / if has finished, and which will be the destination of the jump of the block_return operation.
Control flow operations also implement the WasmLabelLevelInterface which defines a label target block, which is the one that should be reached when branch instructions bubble up to this operation.
Branch instructions define an integer attribute that specifies how many nesting levels they should exit. They are terminator instructions, but due to the multi-region level exiting capabilities, storing the destination block as successor was not an option.
wasm.loop : {
// ...
wasm.branch_if %7 to level 0 else ^bb1
^bb1: // pred: ^bb0
%8 = wasm.local_get %0 : memref<i32>
wasm.block_return %8 : i32
}> ^bb1
^bb1(%1: i32): // pred: ^bb0
wasm.return %1 : i32
}
Arithmetic Operations and Constants
For arithmetic operations, we tried to stay as close as possible to existing standard dialects. This makes it trivial to convert to other dialects in a majority of cases.
%3 = wasm.const 1 : i32
// ...
%4 = wasm.add %2 %3 : i32
For operations with no representation available in the standard dialects (for example, rotates), we translate into bit-twiddling operations within the Wasm dialect that do have representations within the standard dialects. The simpler Wasm dialect operations are then subsequently lowered to operations within the target dialect.
The testcase mlir/test/Conversion/WasmToStandard/wasm-rotr-to-arith.mlir serves as an example of this. The RotateOpConversion struct in WasmToStandard.cpp shows a concrete implementation.
Future and Ongoing Work
Future work includes of course the implementation of Wasm operations that are not supported in the current state of the project (vector instructions, table operations, …).
In addition, conversion to standard dialects for the embedder-specific operations should be rethought, taking into consideration how the integration should be done.
As mentioned in the technical debt section, some negative tests for validation need to be implemented. Currently, the project does not contain googletest-style tests. As the project matures, we would like to add broader testing.
Finally, another missing feature is the ability to import Wasm text format, as currently only import from binary Wasm files is supported. Adding a WAT parser would make it much easier to validate the project via LLVM lit style tests, and also increase the breadth of the project.
Thank you for reading,
Luc Forget <[email protected]>
Ferdinand Lemaire <[email protected]>
Jessica Paquette <[email protected]>