[RFC] MLIR Dialect for WebAssembly (original) (raw)

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.)

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

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

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

Driver

Status and plan

In its actual state, the project consists of :

The proposed implementation covers Wasm 2.0 specification

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. :slight_smile:

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:

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 <luc.forget@woven.toyota>
Ferdinand Lemaire <ferdinand.lemaire@woven.toyota>
Jessica Paquette <jessica.paquette@woven.toyota>