From e93187a2585d65dd6a5da8f2af22758c7008ed71 Mon Sep 17 00:00:00 2001 From: Andrew Brown Date: Mon, 16 Nov 2020 10:54:00 -0800 Subject: [PATCH] Add an initial wasi-nn implementation for Wasmtime (#2208) * Add an initial wasi-nn implementation for Wasmtime 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/tests/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 * Test wasi-nn as a CI step This change adds: - a GitHub action for installing OpenVINO - a script, `ci/run-wasi-nn-example.sh`, to run the classification example --- .github/actions/install-openvino/README.md | 8 + .github/actions/install-openvino/action.yml | 14 ++ .github/actions/install-openvino/install.sh | 16 ++ .github/workflows/main.yml | 20 ++ .gitmodules | 3 + Cargo.lock | 127 ++++++++++ Cargo.toml | 2 + ci/run-experimental-x64-ci.sh | 1 + ci/run-wasi-nn-example.sh | 35 +++ crates/wasi-nn/Cargo.toml | 29 +++ crates/wasi-nn/LICENSE | 220 ++++++++++++++++++ crates/wasi-nn/README.md | 38 +++ crates/wasi-nn/build.rs | 10 + crates/wasi-nn/examples/README.md | 1 + .../classification-example/Cargo.lock | 13 ++ .../classification-example/Cargo.toml | 15 ++ .../examples/classification-example/README.md | 2 + .../classification-example/src/main.rs | 54 +++++ .../examples/wasi-nn-rust-bindings/.gitignore | 1 + .../examples/wasi-nn-rust-bindings/Cargo.lock | 5 + .../examples/wasi-nn-rust-bindings/Cargo.toml | 14 ++ .../examples/wasi-nn-rust-bindings/LICENSE | 201 ++++++++++++++++ .../examples/wasi-nn-rust-bindings/README.md | 65 ++++++ .../wasi-nn-rust-bindings/src/error.rs | 76 ++++++ .../wasi-nn-rust-bindings/src/generated.rs | 199 ++++++++++++++++ .../examples/wasi-nn-rust-bindings/src/lib.rs | 3 + crates/wasi-nn/spec | 1 + crates/wasi-nn/src/ctx.rs | 125 ++++++++++ crates/wasi-nn/src/impl.rs | 176 ++++++++++++++ crates/wasi-nn/src/lib.rs | 26 +++ crates/wasi-nn/src/witx.rs | 40 ++++ scripts/publish.rs | 6 +- src/commands/run.rs | 9 + 33 files changed, 1554 insertions(+), 1 deletion(-) create mode 100644 .github/actions/install-openvino/README.md create mode 100644 .github/actions/install-openvino/action.yml create mode 100755 .github/actions/install-openvino/install.sh create mode 100755 ci/run-wasi-nn-example.sh create mode 100644 crates/wasi-nn/Cargo.toml create mode 100644 crates/wasi-nn/LICENSE create mode 100644 crates/wasi-nn/README.md create mode 100644 crates/wasi-nn/build.rs create mode 100644 crates/wasi-nn/examples/README.md create mode 100644 crates/wasi-nn/examples/classification-example/Cargo.lock create mode 100644 crates/wasi-nn/examples/classification-example/Cargo.toml create mode 100644 crates/wasi-nn/examples/classification-example/README.md create mode 100644 crates/wasi-nn/examples/classification-example/src/main.rs create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs create mode 100644 crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs create mode 160000 crates/wasi-nn/spec create mode 100644 crates/wasi-nn/src/ctx.rs create mode 100644 crates/wasi-nn/src/impl.rs create mode 100644 crates/wasi-nn/src/lib.rs create mode 100644 crates/wasi-nn/src/witx.rs diff --git a/.github/actions/install-openvino/README.md b/.github/actions/install-openvino/README.md new file mode 100644 index 000000000000..a494f57fdfbc --- /dev/null +++ b/.github/actions/install-openvino/README.md @@ -0,0 +1,8 @@ +# install-openvino + +A GitHub action to install OpenVINO from a package repository. This is only necessary for `wasi-nn` support but there +are enough steps here to package the functionality separately and avoid cluttering the CI. + +Future improvements: + - make this installer work for different OS/distributions (e.g. https://docs.openvinotoolkit.org/latest/openvino_docs_install_guides_installing_openvino_windows.html) + - it would be nice to output the install directory (i.e. `/opt/intel/openvino`) diff --git a/.github/actions/install-openvino/action.yml b/.github/actions/install-openvino/action.yml new file mode 100644 index 000000000000..e4858c2d298c --- /dev/null +++ b/.github/actions/install-openvino/action.yml @@ -0,0 +1,14 @@ +name: 'Install OpenVINO' +description: 'Install OpenVINO binaries from a package repository; this is significantly faster than building from source' + +inputs: + version: + description: 'The release version of OpenVINO to install' + required: false + default: '2020.4.287' + +runs: + using: composite + steps: + - run: ${{ github.action_path }}/install.sh ${{ inputs.version }} + shell: bash diff --git a/.github/actions/install-openvino/install.sh b/.github/actions/install-openvino/install.sh new file mode 100755 index 000000000000..16e523385411 --- /dev/null +++ b/.github/actions/install-openvino/install.sh @@ -0,0 +1,16 @@ +#!/bin/bash + +set -e + +# Retrieve OpenVINO checksum. +wget https://apt.repos.intel.com/openvino/2020/GPG-PUB-KEY-INTEL-OPENVINO-2020 +echo '5f5cff8a2d26ba7de91942bd0540fa4d GPG-PUB-KEY-INTEL-OPENVINO-2020' > CHECKSUM +md5sum --check CHECKSUM + +# Add OpenVINO repository (deb). +sudo apt-key add GPG-PUB-KEY-INTEL-OPENVINO-2020 +echo "deb https://apt.repos.intel.com/openvino/2020 all main" | sudo tee /etc/apt/sources.list.d/intel-openvino-2020.list +sudo apt update + +# Install OpenVINO package. +sudo apt install -y intel-openvino-runtime-ubuntu18-$1 diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 910a032c7ccf..3bd17dbdca6d 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -51,6 +51,7 @@ jobs: runs-on: ubuntu-latest env: RUSTDOCFLAGS: -Dbroken_intra_doc_links + OPENVINO_SKIP_LINKING: 1 steps: - uses: actions/checkout@v2 with: @@ -252,6 +253,7 @@ jobs: --all \ --exclude lightbeam \ --exclude wasmtime-lightbeam \ + --exclude wasmtime-wasi-nn \ --exclude peepmatic \ --exclude peepmatic-automata \ --exclude peepmatic-fuzzing \ @@ -304,6 +306,23 @@ jobs: CARGO_VERSION: "+nightly" RUST_BACKTRACE: 1 + # Build and test the wasi-nn module. + test_wasi_nn: + name: Test wasi-nn module + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + with: + submodules: true + - uses: ./.github/actions/install-rust + with: + toolchain: nightly-2020-08-25 + - run: rustup target add wasm32-wasi + - uses: ./.github/actions/install-openvino + - run: ./ci/run-wasi-nn-example.sh + env: + RUST_BACKTRACE: 1 + # Verify that cranelift's code generation is deterministic meta_determinist_check: name: Meta deterministic check @@ -411,6 +430,7 @@ jobs: --all \ --exclude lightbeam \ --exclude wasmtime-lightbeam \ + --exclude wasmtime-wasi-nn \ --exclude peepmatic \ --exclude peepmatic-automata \ --exclude peepmatic-fuzzing \ diff --git a/.gitmodules b/.gitmodules index 1586fcc47e06..0eb1df4e0f86 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,6 @@ [submodule "WASI"] path = crates/wasi-common/WASI url = https://github.com/WebAssembly/WASI +[submodule "crates/wasi-nn/spec"] + path = crates/wasi-nn/spec + url = https://github.com/WebAssembly/wasi-nn diff --git a/Cargo.lock b/Cargo.lock index 156fd62c97df..0dc6b5835e2c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -128,6 +128,30 @@ dependencies = [ "serde", ] +[[package]] +name = "bindgen" +version = "0.55.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75b13ce559e6433d360c26305643803cb52cfbabbc2b9c47ce04a58493dfb443" +dependencies = [ + "bitflags", + "cexpr", + "cfg-if 0.1.10", + "clang-sys", + "clap", + "env_logger 0.7.1", + "lazy_static", + "lazycell", + "log", + "peeking_take_while", + "proc-macro2", + "quote", + "regex", + "rustc-hash", + "shlex", + "which", +] + [[package]] name = "bit-set" version = "0.5.2" @@ -210,6 +234,15 @@ dependencies = [ "jobserver", ] +[[package]] +name = "cexpr" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4aedb84272dbe89af497cf81375129abda4fc0a9e7c5d317498c15cc30c0d27" +dependencies = [ + "nom", +] + [[package]] name = "cfg-if" version = "0.1.10" @@ -235,6 +268,17 @@ dependencies = [ "winapi", ] +[[package]] +name = "clang-sys" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9da1484c6a890e374ca5086062d4847e0a2c1e5eba9afa5d48c09e8eb39b2519" +dependencies = [ + "glob", + "libc", + "libloading", +] + [[package]] name = "clap" version = "2.33.3" @@ -1048,6 +1092,12 @@ version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" +[[package]] +name = "lazycell" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b294d6fa9ee409a054354afc4352b0b9ef7ca222c69b8812cbea9e7d2bf3783f" + [[package]] name = "leb128" version = "0.2.4" @@ -1070,6 +1120,16 @@ dependencies = [ "cc", ] +[[package]] +name = "libloading" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2443d8f0478b16759158b2f66d525991a05491138bc05814ef52a250148ef4f9" +dependencies = [ + "cfg-if 0.1.10", + "winapi", +] + [[package]] name = "lightbeam" version = "0.21.0" @@ -1184,6 +1244,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0debeb9fcf88823ea64d64e4a815ab1643f33127d995978e099942ce38f25238" +[[package]] +name = "nom" +version = "5.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffb4262d26ed83a1c0a33a38fe2bb15797329c85770da05e6b828ddb782627af" +dependencies = [ + "memchr", + "version_check", +] + [[package]] name = "num-integer" version = "0.1.44" @@ -1242,6 +1312,26 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5" +[[package]] +name = "openvino" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87a74f90f07f134153e3ad2ffa724a3ebda92cdc6e099f7fe7d9185cf960f028" +dependencies = [ + "openvino-sys", + "thiserror", +] + +[[package]] +name = "openvino-sys" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e72a2e5bd353bd3cf39b2663767e0ae0325a7588c47fd496cbf9a09237ef7ca8" +dependencies = [ + "bindgen", + "cmake", +] + [[package]] name = "os_pipe" version = "0.9.2" @@ -1252,6 +1342,12 @@ dependencies = [ "winapi", ] +[[package]] +name = "peeking_take_while" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b17cddbe7ec3f8bc800887bab5e717348c95ea2ca0b1bf0837fb964dc67099" + [[package]] name = "peepmatic" version = "0.68.0" @@ -1828,6 +1924,12 @@ dependencies = [ "dirs", ] +[[package]] +name = "shlex" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fdf1b9db47230893d76faad238fd6097fd6d6a9245cd7a4d90dbd639536bbd2" + [[package]] name = "smallvec" version = "1.4.2" @@ -2344,6 +2446,7 @@ dependencies = [ "wasmtime-obj", "wasmtime-runtime", "wasmtime-wasi", + "wasmtime-wasi-nn", "wasmtime-wast", "wat", ] @@ -2547,6 +2650,21 @@ dependencies = [ "wiggle", ] +[[package]] +name = "wasmtime-wasi-nn" +version = "0.21.0" +dependencies = [ + "anyhow", + "log", + "openvino", + "thiserror", + "wasmtime", + "wasmtime-runtime", + "wasmtime-wasi", + "wasmtime-wiggle", + "wiggle", +] + [[package]] name = "wasmtime-wast" version = "0.21.0" @@ -2604,6 +2722,15 @@ dependencies = [ "wast 27.0.0", ] +[[package]] +name = "which" +version = "3.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d011071ae14a2f6671d0b74080ae0cd8ebf3a6f8c9589a2cd45f23126fe29724" +dependencies = [ + "libc", +] + [[package]] name = "wig" version = "0.21.0" diff --git a/Cargo.toml b/Cargo.toml index f4adcf55e27e..f1a538d8c703 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ wasmtime-jit = { path = "crates/jit", version = "0.21.0" } wasmtime-obj = { path = "crates/obj", version = "0.21.0" } wasmtime-wast = { path = "crates/wast", version = "0.21.0" } wasmtime-wasi = { path = "crates/wasi", version = "0.21.0" } +wasmtime-wasi-nn = { path = "crates/wasi-nn", version = "0.21.0", optional = true } wasi-common = { path = "crates/wasi-common", version = "0.21.0" } structopt = { version = "0.3.5", features = ["color", "suggestions"] } object = { version = "0.22.0", default-features = false, features = ["write"] } @@ -80,6 +81,7 @@ default = ["jitdump", "wasmtime/wat", "wasmtime/parallel-compilation"] lightbeam = ["wasmtime/lightbeam"] jitdump = ["wasmtime/jitdump"] vtune = ["wasmtime/vtune"] +wasi-nn = ["wasmtime-wasi-nn"] # Try the experimental, work-in-progress new x86_64 backend. This is not stable # as of June 2020. diff --git a/ci/run-experimental-x64-ci.sh b/ci/run-experimental-x64-ci.sh index a7e650b13023..8b41831b15f2 100755 --- a/ci/run-experimental-x64-ci.sh +++ b/ci/run-experimental-x64-ci.sh @@ -13,6 +13,7 @@ cargo $CARGO_VERSION \ --features experimental_x64 \ --all \ --exclude wasmtime-lightbeam \ + --exclude wasmtime-wasi-nn \ --exclude peepmatic \ --exclude peepmatic-automata \ --exclude peepmatic-fuzzing \ diff --git a/ci/run-wasi-nn-example.sh b/ci/run-wasi-nn-example.sh new file mode 100755 index 000000000000..4d82e8494383 --- /dev/null +++ b/ci/run-wasi-nn-example.sh @@ -0,0 +1,35 @@ +#!/bin/bash + +# The following script demonstrates how to execute a machine learning inference using the wasi-nn module optionally +# compiled into Wasmtime. Calling it will download the necessary model and tensor files stored separately in $FIXTURE +# into $TMP_DIR (optionally pass a directory with existing files as the first argument to re-try the script). Then, +# it will compile the example code in crates/wasi-nn/tests/example into a Wasm file that is subsequently +# executed with the Wasmtime CLI. +set -e +WASMTIME_DIR=$(dirname "$0" | xargs dirname) +FIXTURE=https://gist.github.com/abrown/c7847bf3701f9efbb2070da1878542c1/raw/07a9f163994b0ff8f0d7c5a5c9645ec3d8b24024 + +# Inform the environment of OpenVINO library locations. Then we use OPENVINO_INSTALL_DIR below to avoid building all of +# OpenVINO from source (quite slow). +source /opt/intel/openvino/bin/setupvars.sh + +# Build Wasmtime with wasi-nn enabled; we attempt this first to avoid extra work if the build fails. +OPENVINO_INSTALL_DIR=/opt/intel/openvino cargo build -p wasmtime-cli --features wasi-nn + +# Download all necessary test fixtures to the temporary directory. +TMP_DIR=${1:-$(mktemp -d -t ci-XXXXXXXXXX)} +wget --no-clobber --directory-prefix=$TMP_DIR $FIXTURE/frozen_inference_graph.bin +wget --no-clobber --directory-prefix=$TMP_DIR $FIXTURE/frozen_inference_graph.xml +wget --no-clobber --directory-prefix=$TMP_DIR $FIXTURE/tensor-1x3x300x300-f32.bgr + +# Now build an example that uses the wasi-nn API. +pushd $WASMTIME_DIR/crates/wasi-nn/examples/classification-example +cargo build --release --target=wasm32-wasi +cp target/wasm32-wasi/release/wasi-nn-example.wasm $TMP_DIR +popd + +# Run the example in Wasmtime (note that the example uses `fixture` as the expected location of the model/tensor files). +OPENVINO_INSTALL_DIR=/opt/intel/openvino cargo run --features wasi-nn -- run --mapdir fixture::$TMP_DIR $TMP_DIR/wasi-nn-example.wasm + +# Clean up. +rm -rf $TMP_DIR \ No newline at end of file diff --git a/crates/wasi-nn/Cargo.toml b/crates/wasi-nn/Cargo.toml new file mode 100644 index 000000000000..c05fd1d0522f --- /dev/null +++ b/crates/wasi-nn/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "wasmtime-wasi-nn" +version = "0.21.0" +authors = ["The Wasmtime Project Developers"] +description = "Wasmtime implementation of the wasi-nn API" +documentation = "https://docs.rs/wasmtime-wasi-nn" +license = "Apache-2.0 WITH LLVM-exception" +categories = ["wasm", "computer-vision"] +keywords = ["webassembly", "wasm", "neural network"] +repository = "https://github.com/bytecodealliance/wasmtime" +readme = "README.md" +edition = "2018" + +[dependencies] +# These dependencies are necessary for the witx-generation macros to work: +anyhow = "1.0" +log = { version = "0.4", default-features = false } +wasmtime = { path = "../wasmtime", version = "0.21.0", default-features = false } +wasmtime-runtime = { path = "../runtime", version = "0.21.0" } +wasmtime-wiggle = { path = "../wiggle/wasmtime", version = "0.21.0" } +wasmtime-wasi = { path = "../wasi", version = "0.21.0" } +wiggle = { path = "../wiggle", version = "0.21.0" } + +# These dependencies are necessary for the wasi-nn implementation: +openvino = "0.1.5" +thiserror = "1.0" + +[badges] +maintenance = { status = "experimental" } diff --git a/crates/wasi-nn/LICENSE b/crates/wasi-nn/LICENSE new file mode 100644 index 000000000000..f9d81955f4bc --- /dev/null +++ b/crates/wasi-nn/LICENSE @@ -0,0 +1,220 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +--- LLVM Exceptions to the Apache 2.0 License ---- + +As an exception, if, as a result of your compiling your source code, portions +of this Software are embedded into an Object form of such source code, you +may redistribute such embedded portions in such Object form without complying +with the conditions of Sections 4(a), 4(b) and 4(d) of the License. + +In addition, if you combine or link compiled forms of this Software with +software that is licensed under the GPLv2 ("Combined Software") and if a +court of competent jurisdiction determines that the patent provision (Section +3), the indemnity provision (Section 9) or other Section of the License +conflicts with the conditions of the GPLv2, you may retroactively and +prospectively choose to deem waived or otherwise exclude such Section(s) of +the License, but only in their entirety and only with respect to the Combined +Software. + diff --git a/crates/wasi-nn/README.md b/crates/wasi-nn/README.md new file mode 100644 index 000000000000..165b5063d002 --- /dev/null +++ b/crates/wasi-nn/README.md @@ -0,0 +1,38 @@ +# wasmtime-wasi-nn + +This crate enables support for the [wasi-nn] API in Wasmtime. Currently it contains an implementation of [wasi-nn] using +OpenVINO™ but in the future it could support multiple machine learning backends. Since the [wasi-nn] API is expected +to be an optional feature of WASI, this crate is currently separate from the [wasi-common] crate. This crate is +experimental and its API, functionality, and location could quickly change. + +[examples]: examples +[openvino]: https://crates.io/crates/openvino +[wasi-nn]: https://github.com/WebAssembly/wasi-nn +[wasi-common]: ../wasi-common + +### Use + +Use the Wasmtime APIs to instantiate a Wasm module and link in the `WasiNn` implementation as follows: + +``` +let wasi_nn = WasiNn::new(&store, WasiNnCtx::new()?); +wasi_nn.add_to_linker(&mut linker)?; +``` + +### Build + +This crate should build as usual (i.e. `cargo build`) but note that using an existing installation of OpenVINO™, rather +than building from source, will drastically improve the build times. See the [openvino] crate for more information + +### Example + +An end-to-end example demonstrating ML classification is included in [examples]: + - `tests/wasi-nn-rust-bindings` contains ergonomic bindings for writing Rust code against the [wasi-nn] APIs + - `tests/classification-example` contains a standalone Rust project that uses the [wasi-nn] APIs and is compiled to the + `wasm32-wasi` target using the `wasi-nn-rust-bindings` + +Run the example from the Wasmtime project directory: + +``` +ci/run-wasi-nn-example.sh +``` diff --git a/crates/wasi-nn/build.rs b/crates/wasi-nn/build.rs new file mode 100644 index 000000000000..aeced29a0fda --- /dev/null +++ b/crates/wasi-nn/build.rs @@ -0,0 +1,10 @@ +//! This build script: +//! - has the configuration necessary for the wiggle and witx macros. + +use std::path::PathBuf; + +fn main() { + // This is necessary for Wiggle/Witx macros. + let wasi_root = PathBuf::from("./spec").canonicalize().unwrap(); + println!("cargo:rustc-env=WASI_ROOT={}", wasi_root.display()); +} diff --git a/crates/wasi-nn/examples/README.md b/crates/wasi-nn/examples/README.md new file mode 100644 index 000000000000..94711fe4f456 --- /dev/null +++ b/crates/wasi-nn/examples/README.md @@ -0,0 +1 @@ +See `ci/run-wasi-nn-example.sh` for how the classification example is tested during CI. \ No newline at end of file diff --git a/crates/wasi-nn/examples/classification-example/Cargo.lock b/crates/wasi-nn/examples/classification-example/Cargo.lock new file mode 100644 index 000000000000..f98a1e19a803 --- /dev/null +++ b/crates/wasi-nn/examples/classification-example/Cargo.lock @@ -0,0 +1,13 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +[[package]] +name = "wasi-nn" +version = "0.1.0" + +[[package]] +name = "wasi-nn-example" +version = "0.19.0" +dependencies = [ + "wasi-nn 0.1.0", +] + diff --git a/crates/wasi-nn/examples/classification-example/Cargo.toml b/crates/wasi-nn/examples/classification-example/Cargo.toml new file mode 100644 index 000000000000..52bf2dac321b --- /dev/null +++ b/crates/wasi-nn/examples/classification-example/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "wasi-nn-example" +version = "0.19.0" +authors = ["The Wasmtime Project Developers"] +readme = "README.md" +edition = "2018" +publish = false + +[dependencies] +wasi-nn = { path = "../wasi-nn-rust-bindings", version = "0.1.0" } + +# This crate is built with the wasm32-wasi target, so it's separate +# from the main Wasmtime build, so use this directive to exclude it +# from the parent directory's workspace. +[workspace] diff --git a/crates/wasi-nn/examples/classification-example/README.md b/crates/wasi-nn/examples/classification-example/README.md new file mode 100644 index 000000000000..aa56ad0cbaf7 --- /dev/null +++ b/crates/wasi-nn/examples/classification-example/README.md @@ -0,0 +1,2 @@ +This example project demonstrates using the `wasi-nn` API to perform ML inference. It consists of Rust code that is +built using the `wasm32-wasi` target. See `ci/run-wasi-nn-example.sh` for how this is used. diff --git a/crates/wasi-nn/examples/classification-example/src/main.rs b/crates/wasi-nn/examples/classification-example/src/main.rs new file mode 100644 index 000000000000..898a4bfff370 --- /dev/null +++ b/crates/wasi-nn/examples/classification-example/src/main.rs @@ -0,0 +1,54 @@ +use std::convert::TryInto; +use std::fs; +use wasi_nn; + +pub fn main() { + let xml = fs::read_to_string("fixture/frozen_inference_graph.xml").unwrap(); + println!("First 50 characters of graph: {}", &xml[..50]); + + let weights = fs::read("fixture/frozen_inference_graph.bin").unwrap(); + println!("Size of weights: {}", weights.len()); + + let graph = unsafe { + wasi_nn::load( + &[&xml.into_bytes(), &weights], + wasi_nn::GRAPH_ENCODING_OPENVINO, + wasi_nn::EXECUTION_TARGET_CPU, + ) + .unwrap() + }; + println!("Graph handle ID: {}", graph); + + let context = unsafe { wasi_nn::init_execution_context(graph).unwrap() }; + println!("Execution context ID: {}", context); + + // Load a tensor that precisely matches the graph input tensor (see + // `fixture/frozen_inference_graph.xml`). + let tensor_data = fs::read("fixture/tensor-1x3x300x300-f32.bgr").unwrap(); + println!("Tensor bytes: {}", tensor_data.len()); + let tensor = wasi_nn::Tensor { + dimensions: &[1, 3, 300, 300], + r#type: wasi_nn::TENSOR_TYPE_F32, + data: &tensor_data, + }; + unsafe { + wasi_nn::set_input(context, 0, tensor).unwrap(); + } + + // Execute the inference. + unsafe { + wasi_nn::compute(context).unwrap(); + } + + // Retrieve the output (TODO output looks incorrect). + let mut output_buffer = vec![0f32; 1 << 20]; + unsafe { + wasi_nn::get_output( + context, + 0, + &mut output_buffer[..] as *mut [f32] as *mut u8, + (output_buffer.len() * 4).try_into().unwrap(), + ); + } + println!("output tensor: {:?}", &output_buffer[..1000]) +} diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore b/crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore new file mode 100644 index 000000000000..ea8c4bf7f35f --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore @@ -0,0 +1 @@ +/target diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock b/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock new file mode 100644 index 000000000000..b8e596de2fc4 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock @@ -0,0 +1,5 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +[[package]] +name = "wasi-nn" +version = "0.1.0" diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml b/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml new file mode 100644 index 000000000000..42256e8204a9 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "wasi-nn" +version = "0.1.0" +authors = ["The Wasmtime Project Developers"] +readme = "README.md" +edition = "2018" +publish = false + +[dependencies] + +# This crate is only used when building the example, so it's separate +# from the main Wasmtime build, so use this directive to exclude it +# from the parent directory's workspace. +[workspace] \ No newline at end of file diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE b/crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE new file mode 100644 index 000000000000..16fe87b06e80 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright [yyyy] [name of copyright owner] + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md b/crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md new file mode 100644 index 000000000000..b11ef16773ab --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md @@ -0,0 +1,65 @@ +wasi-nn Rust Bindings +===================== + +This crate contains API bindings for [wasi-nn] system calls in Rust. It is similar in purpose to the [wasi bindings] but +this crate provides access to the optional neural network functionality from WebAssembly. + +[wasi-nn]: https://github.com/WebAssembly/wasi-nn +[wasi bindings]: https://github.com/bytecodealliance/wasi + +> __NOTE__: These bindings are experimental (use at your own risk) and subject to upstream changes in the wasi-nn +> specification. + +> __NOTE__: In the future this crate may be (should be) moved to its own repository, like the [wasi bindings]. + +### Use + +Depend on this crate in your `Cargo.toml`: + +```toml +[dependencies] +wasi-nn = "0.1.0" +``` + +Use the wasi-nn APIs in your application: + +```rust +use wasi_nn; + +unsafe { + wasi_nn::load( + &[&xml.into_bytes(), &weights], + wasi_nn::GRAPH_ENCODING_OPENVINO, + wasi_nn::EXECUTION_TARGET_CPU, + ) + .unwrap() +} +``` + +Compile the application to WebAssembly: + +```shell script +cargo build --target=wasm32-wasi +``` + +Run the generated Wasm in a runtime supporting wasi-nn. Currently Wasmtime has experimental support using the Wasmtime +APIs; see [main.rs](../main.rs) for an example of how this is accomplished. + +### Generation + +This crate contains code ([`src/generated.rs`](src/generated.rs)) generated by +[`witx-bindgen`](https://github.com/bytecodealliance/wasi/tree/main/crates/witx-bindgen). +To regenerate this code, run `witx-bindgen` against the [`wasi-nn` WITX file](https://github.com/WebAssembly/wasi-nn/blob/master/phases/ephemeral/witx/wasi_ephemeral_nn.witx): + +```shell script +.../crates/witx-bindgen$ cargo run .../wasi-nn/phases/ephemeral/witx/wasi_ephemeral_nn.witx +``` + +### License + +This project is licensed under the Apache 2.0 license. See [LICENSE](LICENSE) for more details. + +### Contribution + +Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in this project by you, +as defined in the Apache-2.0 license, shall be licensed as above, without any additional terms or conditions. diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs new file mode 100644 index 000000000000..7ae1a50edb74 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs @@ -0,0 +1,76 @@ +use super::NnErrno; +use core::fmt; +use core::num::NonZeroU16; + +/// A raw error returned by wasi-nn APIs, internally containing a 16-bit error +/// code. +#[derive(Copy, Clone, PartialEq, Eq, Ord, PartialOrd)] +pub struct Error { + code: NonZeroU16, +} + +impl Error { + /// Constructs a new error from a raw error code, returning `None` if the + /// error code is zero (which means success). + pub fn from_raw_error(error: NnErrno) -> Option { + Some(Error { + code: NonZeroU16::new(error)?, + }) + } + + /// Returns the raw error code that this error represents. + pub fn raw_error(&self) -> u16 { + self.code.get() + } +} + +impl fmt::Display for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{} (error {})", strerror(self.code.get()), self.code)?; + Ok(()) + } +} + +impl fmt::Debug for Error { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Error") + .field("code", &self.code) + .field("message", &strerror(self.code.get())) + .finish() + } +} + +/// This should be generated automatically by witx-bindgen but is not yet for enums other than +/// `Errno` (this API uses `NnErrno` to avoid naming conflicts). TODO: https://github.com/bytecodealliance/wasi/issues/52. +fn strerror(code: u16) -> &'static str { + match code { + super::NN_ERRNO_SUCCESS => "No error occurred.", + super::NN_ERRNO_INVALID_ARGUMENT => "Caller module passed an invalid argument.", + super::NN_ERRNO_MISSING_MEMORY => "Caller module is missing a memory export.", + super::NN_ERRNO_BUSY => "Device or resource busy.", + _ => "Unknown error.", + } +} + +#[cfg(feature = "std")] +extern crate std; +#[cfg(feature = "std")] +impl std::error::Error for Error {} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn error_from_success_code() { + assert_eq!(None, Error::from_raw_error(0)); + } + + #[test] + fn error_from_invalid_argument_code() { + assert_eq!( + "Caller module passed an invalid argument. (error 1)", + Error::from_raw_error(1).unwrap().to_string() + ); + } +} diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs new file mode 100644 index 000000000000..d03fd1ed0cd6 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs @@ -0,0 +1,199 @@ +// This file is automatically generated, DO NOT EDIT +// +// To regenerate this file run the `crates/witx-bindgen` command + +use core::mem::MaybeUninit; + +pub use crate::error::Error; +pub type Result = core::result::Result; +pub type BufferSize = u32; +pub type NnErrno = u16; +/// No error occurred. +pub const NN_ERRNO_SUCCESS: NnErrno = 0; +/// Caller module passed an invalid argument. +pub const NN_ERRNO_INVALID_ARGUMENT: NnErrno = 1; +/// Caller module is missing a memory export. +pub const NN_ERRNO_MISSING_MEMORY: NnErrno = 2; +/// Device or resource busy. +pub const NN_ERRNO_BUSY: NnErrno = 3; +pub type TensorDimensions<'a> = &'a [u32]; +pub type TensorType = u8; +pub const TENSOR_TYPE_F16: TensorType = 0; +pub const TENSOR_TYPE_F32: TensorType = 1; +pub const TENSOR_TYPE_U8: TensorType = 2; +pub const TENSOR_TYPE_I32: TensorType = 3; +pub type TensorData<'a> = &'a [u8]; +#[repr(C)] +#[derive(Copy, Clone, Debug)] +pub struct Tensor<'a> { + /// Describe the size of the tensor (e.g. 2x2x2x2 -> [2, 2, 2, 2]). To represent a tensor containing a single value, + /// use `[1]` for the tensor dimensions. + pub dimensions: TensorDimensions<'a>, + pub r#type: TensorType, + /// Contains the tensor data. + pub data: TensorData<'a>, +} +pub type GraphBuilder<'a> = &'a [u8]; +pub type GraphBuilderArray<'a> = &'a [GraphBuilder<'a>]; +pub type Graph = u32; +pub type GraphEncoding = u8; +/// TODO document buffer order +pub const GRAPH_ENCODING_OPENVINO: GraphEncoding = 0; +pub type ExecutionTarget = u8; +pub const EXECUTION_TARGET_CPU: ExecutionTarget = 0; +pub const EXECUTION_TARGET_GPU: ExecutionTarget = 1; +pub const EXECUTION_TARGET_TPU: ExecutionTarget = 2; +pub type GraphExecutionContext = u32; +/// Load an opaque sequence of bytes to use for inference. +/// +/// This allows runtime implementations to support multiple graph encoding formats. For unsupported graph encodings, +/// return `errno::inval`. +/// +/// ## Parameters +/// +/// * `builder` - The bytes necessary to build the graph. +/// * `encoding` - The encoding of the graph. +/// * `target` - Where to execute the graph. +pub unsafe fn load( + builder: GraphBuilderArray, + encoding: GraphEncoding, + target: ExecutionTarget, +) -> Result { + let mut graph = MaybeUninit::uninit(); + let rc = wasi_ephemeral_nn::load( + builder.as_ptr(), + builder.len(), + encoding, + target, + graph.as_mut_ptr(), + ); + if let Some(err) = Error::from_raw_error(rc) { + Err(err) + } else { + Ok(graph.assume_init()) + } +} + +/// TODO Functions like `describe_graph_inputs` and `describe_graph_outputs` (returning +/// an array of `$tensor_description`s) might be useful for introspecting the graph but are not yet included here. +/// Create an execution instance of a loaded graph. +/// TODO this may need to accept flags that might affect the compilation or execution of the graph. +pub unsafe fn init_execution_context(graph: Graph) -> Result { + let mut context = MaybeUninit::uninit(); + let rc = wasi_ephemeral_nn::init_execution_context(graph, context.as_mut_ptr()); + if let Some(err) = Error::from_raw_error(rc) { + Err(err) + } else { + Ok(context.assume_init()) + } +} + +/// Define the inputs to use for inference. +/// +/// This should return an $nn_errno (TODO define) if the input tensor does not match the expected dimensions and type. +/// +/// ## Parameters +/// +/// * `index` - The index of the input to change. +/// * `tensor` - The tensor to set as the input. +pub unsafe fn set_input(context: GraphExecutionContext, index: u32, tensor: Tensor) -> Result<()> { + let rc = wasi_ephemeral_nn::set_input(context, index, &tensor as *const _ as *mut _); + if let Some(err) = Error::from_raw_error(rc) { + Err(err) + } else { + Ok(()) + } +} + +/// Extract the outputs after inference. +/// +/// This should return an $nn_errno (TODO define) if the inference has not yet run. +/// +/// ## Parameters +/// +/// * `index` - The index of the output to retrieve. +/// * `out_buffer` - An out parameter to which to copy the tensor data. The caller is responsible for allocating enough memory for +/// the tensor data or an error will be returned. Currently there is no dynamic way to extract the additional +/// tensor metadata (i.e. dimension, element type) but this should be added at some point. +/// +/// ## Return +/// +/// * `bytes_written` - The number of bytes of tensor data written to the `$out_buffer`. +pub unsafe fn get_output( + context: GraphExecutionContext, + index: u32, + out_buffer: *mut u8, + out_buffer_max_size: BufferSize, +) -> Result { + let mut bytes_written = MaybeUninit::uninit(); + let rc = wasi_ephemeral_nn::get_output( + context, + index, + out_buffer, + out_buffer_max_size, + bytes_written.as_mut_ptr(), + ); + if let Some(err) = Error::from_raw_error(rc) { + Err(err) + } else { + Ok(bytes_written.assume_init()) + } +} + +/// Compute the inference on the given inputs (see `set_input`). +/// +/// This should return an $nn_errno (TODO define) if the inputs are not all defined. +pub unsafe fn compute(context: GraphExecutionContext) -> Result<()> { + let rc = wasi_ephemeral_nn::compute(context); + if let Some(err) = Error::from_raw_error(rc) { + Err(err) + } else { + Ok(()) + } +} + +pub mod wasi_ephemeral_nn { + use super::*; + #[link(wasm_import_module = "wasi_ephemeral_nn")] + extern "C" { + /// Load an opaque sequence of bytes to use for inference. + /// + /// This allows runtime implementations to support multiple graph encoding formats. For unsupported graph encodings, + /// return `errno::inval`. + pub fn load( + builder_ptr: *const GraphBuilder, + builder_len: usize, + encoding: GraphEncoding, + target: ExecutionTarget, + graph: *mut Graph, + ) -> NnErrno; + /// TODO Functions like `describe_graph_inputs` and `describe_graph_outputs` (returning + /// an array of `$tensor_description`s) might be useful for introspecting the graph but are not yet included here. + /// Create an execution instance of a loaded graph. + /// TODO this may need to accept flags that might affect the compilation or execution of the graph. + pub fn init_execution_context(graph: Graph, context: *mut GraphExecutionContext) + -> NnErrno; + /// Define the inputs to use for inference. + /// + /// This should return an $nn_errno (TODO define) if the input tensor does not match the expected dimensions and type. + pub fn set_input( + context: GraphExecutionContext, + index: u32, + tensor: *mut Tensor, + ) -> NnErrno; + /// Extract the outputs after inference. + /// + /// This should return an $nn_errno (TODO define) if the inference has not yet run. + pub fn get_output( + context: GraphExecutionContext, + index: u32, + out_buffer: *mut u8, + out_buffer_max_size: BufferSize, + bytes_written: *mut BufferSize, + ) -> NnErrno; + /// Compute the inference on the given inputs (see `set_input`). + /// + /// This should return an $nn_errno (TODO define) if the inputs are not all defined. + pub fn compute(context: GraphExecutionContext) -> NnErrno; + } +} diff --git a/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs new file mode 100644 index 000000000000..86fdcd009145 --- /dev/null +++ b/crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs @@ -0,0 +1,3 @@ +mod error; +mod generated; +pub use generated::*; diff --git a/crates/wasi-nn/spec b/crates/wasi-nn/spec new file mode 160000 index 000000000000..68e73cf612eb --- /dev/null +++ b/crates/wasi-nn/spec @@ -0,0 +1 @@ +Subproject commit 68e73cf612eb8671e0c141b107bf1593491bdc16 diff --git a/crates/wasi-nn/src/ctx.rs b/crates/wasi-nn/src/ctx.rs new file mode 100644 index 000000000000..2857f2cf99f9 --- /dev/null +++ b/crates/wasi-nn/src/ctx.rs @@ -0,0 +1,125 @@ +//! Implements the base structure (i.e. [WasiNnCtx]) that will provide the implementation of the +//! wasi-nn API. +use crate::r#impl::UsageError; +use crate::witx::types::{Graph, GraphExecutionContext}; +use openvino::InferenceError; +use std::cell::RefCell; +use std::collections::HashMap; +use std::hash::Hash; +use thiserror::Error; +use wiggle::GuestError; + +/// Possible errors for interacting with [WasiNnCtx]. +#[derive(Debug, Error)] +pub enum WasiNnError { + #[error("guest error")] + GuestError(#[from] GuestError), + #[error("openvino error")] + OpenvinoError(#[from] InferenceError), + #[error("usage error")] + UsageError(#[from] UsageError), +} + +pub(crate) type WasiNnResult = std::result::Result; + +pub struct Table { + entries: HashMap, + next_key: u32, +} + +impl Default for Table { + fn default() -> Self { + Self { + entries: HashMap::new(), + next_key: 0, + } + } +} + +impl Table +where + K: Eq + Hash + From + Copy, +{ + pub fn insert(&mut self, value: V) -> K { + let key = self.use_next_key(); + self.entries.insert(key, value); + key + } + + pub fn remove(&mut self, key: K) -> Option { + self.entries.remove(&key) + } + + pub fn get(&self, key: K) -> Option<&V> { + self.entries.get(&key) + } + + pub fn get_mut(&mut self, key: K) -> Option<&mut V> { + self.entries.get_mut(&key) + } + + pub fn len(&self) -> usize { + self.entries.len() + } + + fn use_next_key(&mut self) -> K { + let current = self.next_key; + self.next_key += 1; + K::from(current) + } +} + +pub struct ExecutionContext { + pub(crate) graph: Graph, + pub(crate) request: openvino::InferRequest, +} + +impl ExecutionContext { + pub(crate) fn new(graph: Graph, request: openvino::InferRequest) -> Self { + Self { graph, request } + } +} + +/// Capture the state necessary for calling into `openvino`. +pub struct Ctx { + pub(crate) core: openvino::Core, + pub(crate) graphs: Table, + pub(crate) executions: Table, +} + +impl Ctx { + /// Make a new `WasiNnCtx` with the default settings. + pub fn new() -> WasiNnResult { + Ok(Self { + core: openvino::Core::new(None)?, + graphs: Table::default(), + executions: Table::default(), + }) + } +} + +/// This structure provides the Rust-side context necessary for implementing the wasi-nn API. At the +/// moment, it is specialized for a single inference implementation (i.e. OpenVINO) but conceivably +/// this could support more than one backing implementation. +pub struct WasiNnCtx { + pub(crate) ctx: RefCell, +} + +impl WasiNnCtx { + /// Make a new `WasiNnCtx` with the default settings. + pub fn new() -> WasiNnResult { + Ok(Self { + ctx: RefCell::new(Ctx::new()?), + }) + } +} + +#[cfg(test)] +mod test { + use super::*; + + #[test] + fn instantiate() { + WasiNnCtx::new().unwrap(); + } +} diff --git a/crates/wasi-nn/src/impl.rs b/crates/wasi-nn/src/impl.rs new file mode 100644 index 000000000000..a8828c1fe9d7 --- /dev/null +++ b/crates/wasi-nn/src/impl.rs @@ -0,0 +1,176 @@ +//! Implements the wasi-nn API. +use crate::ctx::{ExecutionContext, WasiNnResult as Result}; +use crate::witx::types::{ + ExecutionTarget, Graph, GraphBuilderArray, GraphEncoding, GraphExecutionContext, Tensor, + TensorType, +}; +use crate::witx::wasi_ephemeral_nn::WasiEphemeralNn; +use crate::WasiNnCtx; +use openvino::{Layout, Precision, TensorDesc}; +use thiserror::Error; +use wiggle::GuestPtr; + +#[derive(Debug, Error)] +pub enum UsageError { + #[error("Only OpenVINO's IR is currently supported, passed encoding: {0}")] + InvalidEncoding(GraphEncoding), + #[error("OpenVINO expects only two buffers (i.e. [ir, weights]), passed: {0}")] + InvalidNumberOfBuilders(u32), + #[error("Invalid graph handle; has it been loaded?")] + InvalidGraphHandle, + #[error("Invalid execution context handle; has it been initialized?")] + InvalidExecutionContextHandle, + #[error("Not enough memory to copy tensor data of size: {0}")] + NotEnoughMemory(u32), +} + +impl<'a> WasiEphemeralNn for WasiNnCtx { + fn load<'b>( + &self, + builders: &GraphBuilderArray<'_>, + encoding: GraphEncoding, + target: ExecutionTarget, + ) -> Result { + if encoding != GraphEncoding::Openvino { + return Err(UsageError::InvalidEncoding(encoding).into()); + } + if builders.len() != 2 { + return Err(UsageError::InvalidNumberOfBuilders(builders.len()).into()); + } + let builders = builders.as_ptr(); + let xml = builders.read()?.as_slice()?; + let weights = builders.add(1)?.read()?.as_slice()?; + let graph = self + .ctx + .borrow_mut() + .core + .read_network_from_buffer(&xml, &weights)?; + let executable_graph = self + .ctx + .borrow_mut() + .core + .load_network(&graph, map_execution_target_to_string(target))?; + let id = self + .ctx + .borrow_mut() + .graphs + .insert((graph, executable_graph)); + Ok(id) + } + + fn init_execution_context(&self, graph: Graph) -> Result { + let request = + if let Some((_, executable_graph)) = self.ctx.borrow_mut().graphs.get_mut(graph) { + executable_graph.create_infer_request()? + } else { + return Err(UsageError::InvalidGraphHandle.into()); + }; + + let execution_context = ExecutionContext::new(graph, request); + let handle = self.ctx.borrow_mut().executions.insert(execution_context); + Ok(handle) + } + + fn set_input<'b>( + &self, + context: GraphExecutionContext, + index: u32, + tensor: &Tensor<'b>, + ) -> Result<()> { + let graph = if let Some(execution) = self.ctx.borrow_mut().executions.get_mut(context) { + execution.graph + } else { + return Err(UsageError::InvalidExecutionContextHandle.into()); + }; + + let input_name = if let Some((graph, _)) = self.ctx.borrow().graphs.get(graph) { + graph.get_input_name(index as usize)? + } else { + unreachable!("It should be impossible to attempt to access an execution's graph and for that graph not to exist--this is a bug.") + }; + + // Construct the blob structure. + let dimensions = tensor + .dimensions + .as_slice()? + .iter() + .map(|d| *d as u64) + .collect::>(); + let precision = match tensor.type_ { + TensorType::F16 => Precision::FP16, + TensorType::F32 => Precision::FP32, + TensorType::U8 => Precision::U8, + TensorType::I32 => Precision::I32, + }; + // TODO There must be some good way to discover the layout here; this should not have to default to NHWC. + let desc = TensorDesc::new(Layout::NHWC, &dimensions, precision); + let data = tensor.data.as_slice()?; + let blob = openvino::Blob::new(desc, &data)?; + + // Actually assign the blob to the request (TODO avoid duplication with the borrow above). + if let Some(execution) = self.ctx.borrow_mut().executions.get_mut(context) { + execution.request.set_blob(&input_name, blob)?; + } else { + return Err(UsageError::InvalidExecutionContextHandle.into()); + } + + Ok(()) + } + + fn compute(&self, context: GraphExecutionContext) -> Result<()> { + if let Some(execution) = self.ctx.borrow_mut().executions.get_mut(context) { + Ok(execution.request.infer()?) + } else { + return Err(UsageError::InvalidExecutionContextHandle.into()); + } + } + + fn get_output<'b>( + &self, + context: GraphExecutionContext, + index: u32, + out_buffer: &GuestPtr<'_, u8>, + out_buffer_max_size: u32, + ) -> Result { + let graph = if let Some(execution) = self.ctx.borrow_mut().executions.get_mut(context) { + execution.graph + } else { + return Err(UsageError::InvalidExecutionContextHandle.into()); + }; + + let output_name = if let Some((graph, _)) = self.ctx.borrow().graphs.get(graph) { + graph.get_output_name(index as usize)? + } else { + unreachable!("It should be impossible to attempt to access an execution's graph and for that graph not to exist--this is a bug.") + }; + + // Retrieve the tensor data. + let (mut blob, blob_size) = + if let Some(execution) = self.ctx.borrow_mut().executions.get_mut(context) { + let mut blob = execution.request.get_blob(&output_name)?; // TODO shouldn't need to be mut + let blob_size = blob.byte_len()? as u32; + if blob_size > out_buffer_max_size { + return Err(UsageError::NotEnoughMemory(blob_size).into()); + } + (blob, blob_size) + } else { + return Err(UsageError::InvalidExecutionContextHandle.into()); + }; + + // Copy the tensor data over to the `out_buffer`. + let mut out_slice = out_buffer.as_array(out_buffer_max_size).as_slice()?; + (&mut out_slice[..blob_size as usize]).copy_from_slice(blob.buffer()?); + + Ok(blob_size) + } +} + +/// Return the execution target string expected by OpenVINO from the `ExecutionTarget` enum provided +/// by wasi-nn. +fn map_execution_target_to_string(target: ExecutionTarget) -> &'static str { + match target { + ExecutionTarget::Cpu => "CPU", + ExecutionTarget::Gpu => "GPU", + ExecutionTarget::Tpu => unimplemented!("OpenVINO does not support TPU execution targets"), + } +} diff --git a/crates/wasi-nn/src/lib.rs b/crates/wasi-nn/src/lib.rs new file mode 100644 index 000000000000..e604ec768a5f --- /dev/null +++ b/crates/wasi-nn/src/lib.rs @@ -0,0 +1,26 @@ +mod ctx; +mod r#impl; +mod witx; + +pub use ctx::WasiNnCtx; + +// Defines a `struct WasiNn` with member fields and appropriate APIs for dealing with all the +// various WASI exports. +wasmtime_wiggle::wasmtime_integration!({ + // The wiggle code to integrate with lives here: + target: witx, + // This must be the same witx document as used above: + witx: ["$WASI_ROOT/phases/ephemeral/witx/wasi_ephemeral_nn.witx"], + // This must be the same ctx type as used for the target: + ctx: WasiNnCtx, + // This macro will emit a struct to represent the instance, with this name and docs: + modules: { + wasi_ephemeral_nn => { + name: WasiNn, + docs: "An instantiated instance of the wasi-nn exports.", + function_override: {} + } + }, + // Error to return when caller module is missing memory export: + missing_memory: { witx::types::Errno::MissingMemory }, +}); diff --git a/crates/wasi-nn/src/witx.rs b/crates/wasi-nn/src/witx.rs new file mode 100644 index 000000000000..6ec48800c538 --- /dev/null +++ b/crates/wasi-nn/src/witx.rs @@ -0,0 +1,40 @@ +//! Contains the macro-generated implementation of wasi-nn from the its witx definition file. +use crate::ctx::WasiNnCtx; +use crate::ctx::WasiNnError; + +// Generate the traits and types of wasi-nn in several Rust modules (e.g. `types`). +wiggle::from_witx!({ + witx: ["$WASI_ROOT/phases/ephemeral/witx/wasi_ephemeral_nn.witx"], + ctx: WasiNnCtx, + errors: { errno => WasiNnError } +}); + +use types::Errno; + +/// Wiggle generates code that performs some input validation on the arguments passed in by users of +/// wasi-nn. Here we convert the validation error into one (or more, eventually) of the error +/// variants defined in the witx. +impl types::GuestErrorConversion for WasiNnCtx { + fn into_errno(&self, e: wiggle::GuestError) -> Errno { + eprintln!("Guest error: {:?}", e); + Errno::InvalidArgument + } +} + +impl<'a> types::UserErrorConversion for WasiNnCtx { + fn errno_from_wasi_nn_error(&self, e: WasiNnError) -> Errno { + eprintln!("Host error: {:?}", e); + match e { + WasiNnError::OpenvinoError(_) => unimplemented!(), + WasiNnError::GuestError(_) => unimplemented!(), + WasiNnError::UsageError(_) => unimplemented!(), + } + } +} + +/// Additionally, we must let Wiggle know which of our error codes represents a successful operation. +impl wiggle::GuestErrorType for Errno { + fn success() -> Self { + Self::Success + } +} diff --git a/scripts/publish.rs b/scripts/publish.rs index af12d22dbade..5713b9f05a4d 100644 --- a/scripts/publish.rs +++ b/scripts/publish.rs @@ -64,6 +64,7 @@ const CRATES_TO_PUBLISH: &[&str] = &[ "wasmtime", "wasmtime-wiggle", "wasmtime-wasi", + "wasmtime-wasi-nn", "wasmtime-rust-macro", "wasmtime-rust", "wasmtime-wast", @@ -308,7 +309,10 @@ fn verify(crates: &[Crate]) { .arg("--manifest-path") .arg(&krate.manifest) .env("CARGO_TARGET_DIR", "./target"); - if krate.name.contains("lightbeam") || krate.name == "witx" { + if krate.name.contains("lightbeam") + || krate.name == "witx" + || krate.name.contains("wasi-nn") + { cmd.arg("--no-verify"); } let status = cmd.status().unwrap(); diff --git a/src/commands/run.rs b/src/commands/run.rs index 359686031e2f..056e55ba039a 100644 --- a/src/commands/run.rs +++ b/src/commands/run.rs @@ -15,6 +15,9 @@ use wasi_common::{preopen_dir, WasiCtxBuilder}; use wasmtime::{Engine, Func, Linker, Module, Store, Trap, Val, ValType}; use wasmtime_wasi::Wasi; +#[cfg(feature = "wasi-nn")] +use wasmtime_wasi_nn::{WasiNn, WasiNnCtx}; + fn parse_module(s: &OsStr) -> Result { // Do not accept wasmtime subcommand names as the module name match s.to_str() { @@ -353,6 +356,12 @@ fn populate_with_wasi( let wasi = Wasi::new(linker.store(), cx); wasi.add_to_linker(linker)?; + #[cfg(feature = "wasi-nn")] + { + let wasi_nn = WasiNn::new(linker.store(), WasiNnCtx::new()?); + wasi_nn.add_to_linker(linker)?; + } + // Repeat the above, but this time for snapshot 0. let mut cx = wasi_common::old::snapshot_0::WasiCtxBuilder::new(); cx.inherit_stdio().args(argv).envs(vars);