Add an initial wasi-nn implementation for Wasmtime

[abrown]

This change adds a crate, wasmtime-wasi-nn, that uses wiggle to expose the current state of the wasi-nn API and openvino to implement the exposed functions. It includes an end-to-end test demonstrating how to do classification using wasi-nn:

• crates/wasi-nn/examples/classification-example contains Rust code that is compiled to the wasm32-wasi target and run with a Wasmtime embedding that exposes the wasi-nn calls
• the example uses Rust bindings for wasi-nn contained in crates/wasi-nn/tests/wasi-nn-rust-bindings; this crate contains code generated by witx-bindgen and eventually should be its own standalone crate

This change does not wire up wasi-nn to Wasmtime in any way that an external user would be able to observe.

Things to discuss:

• the end-to-end classification test contains a 95MB model weight file which adds considerable weight to a git clone--what to do? Removing the test would mean that the CI could not verify that wasmtime-wasi-nn works
• the openvino crate is not yet published but should be when Add a Rust bridge openvinotoolkit/openvino#2342 is merged; this means that builds of this PR will fail until then [edit: I ended up publishing them prior to merging that PR; we are still trying to figure out where that code will end up]
• the wasi-nn-rust-bindings crate should probably be its own repository but I cannot move it there until I wade through Intel's external release process (for reference, it took me several weeks for the openvino crate); it could be moved to its own repository in a separate PR
• the last time we talked about this, the consensus was to hide the wasi-nn functionality behind a Cargo feature; I will check this item off once I figure out how to do that correctly but I wanted to get feedback on the current approach

[leonwanghui]

Hi @abrown, I think it's a great idea to add neural network module in WASI and Wasmtime, and I'm also interested in expanding other inference frameworks into wasi-nn backend plugin. Could u give some guidelines to develop a new backend plugin? Thanks!

[abrown]

Could u give some guidelines to develop a new backend plugin? Thanks!

I think the process currently looks like this (at a very high level):

1. create a specification (e.g. wasi-nn) and present it to the the WASI subgroup (see their meetings); you might even want to discuss the idea in the subgroup before creating the specification in case others have had similar ideas
2. use the specification to generate a Rust trait using wiggle as in witx.rs, implement the generated trait as in impl.rs, and then create a way to link in the module lib.rs. You will need a way to contain the module state as well: ctx.rs.
3. The next part is still TBD since wasi-nn is the first attempt to add a WASI module that is not wasi-common: you probably need a way to conditionally compile the module if it is something like wasi-nn that users will want to opt in to and, once this is figured out, the Wasmtime linker should be modified in several places to include the new functionality (see an example here).

Hopefully that covers the basics but as you might guess there are details that could bear more documentation.

[abrown]

Tagging a few people as reviewers to solicit some early feedback on this approach to adding optional WASI modules to Wasmtime.

[alexcrichton]

This all looks pretty reasonable to me! It seems like a relatively straightforward application of what we already have with wasi-nn using wiggle to bridge OpenVINO and Wasmtime. The wasi-nn bindings crate looks pretty simple as well being generated by witx-bindgen. Overall looks great!

In terms of integration I don't think that the .github folder in a subdirectory will work, I think it needs to be up near the top. This is a lot of code though about stuff I don't really know about (the wasi-nn proposal), and I suspect most others don't as well. Do you feel that there's benefit from being in this repository itself, vs external? Externally would make CI easier, but the cost would be that you'd need to update it whenever Wasmtime itself updates.

[sunfishcode]

I'm also not familiar with openvino specifically, but at a high level the integration with wiggle looks reasonable.

Thinking about this more, it occurs to me that that using a cargo feature to guard wasi-nn has a potential downside. If people not working on wasi-nn change Wasmtime's APIs, they'll still have to build wasi-nn to avoid CI failures. As such, I started thinking about what this would look like in a separate repo. I'm not convinced it's the right thing yet, but here's the idea:

• Move Wasmtime's src/commands into a proper library crate, say wasmtime-commands
• Add a way to add additional modules to load in RunCommand
• Make wasi-nn a separate repo which uses this wasmtime-commands crate to build a CLI program just like wasmtime but which includes wasi-nn.

That way, wasi-nn would be in control of when it sees Wasmtime API changes. That would also avoid the 95 MB file in the Wasmtime repo. Thoughts?

[alexcrichton]

That way, wasi-nn would be in control of when it sees Wasmtime API changes. That would also avoid the 95 MB file in the Wasmtime repo. Thoughts?

Belatedly, but this sounds good to me!

[github-actions]

Subscribe to Label Action

cc @kubkon

This issue or pull request has been labeled: "cranelift", "wasi"

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

• kubkon: wasi

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

Learn more.

[abrown]

@pchickey, @sunfishcode, @alexcrichton: I think this is ready for an official review (pending any crazy CI failures... I think I've seen the majority of them already). What you will see with the current state of this PR is two commits:

• one to add the wasi-nn module hidden behind a Cargo feature; when enabled, the Linker will be informed of the new function names--this is the approach we discussed in the last Wasmtime meeting
• another to add a script, ci/run-wasi-nn-example.sh, that will build and run an example using wasi-nn without any of the large test fixtures (these have been moved to a Gist and are retrieved by the script)

[alexcrichton]

Seems reasonable to me, but are you thinking we should land all the support in this repository? I may have misunderstood but it sounded like we were shooting instead for landing this as a separate repository with integration hooks provided here.

[alexcrichton]

@sunfishcode and @abrown what do y'all think of making a wasmtime-nn (or similar) repository in the bytecodealliance org to house this integration?

[alexcrichton]

Ok everything looks good to me, thanks again for the patience here!