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
This commit is contained in:
29
crates/wasi-nn/Cargo.toml
Normal file
29
crates/wasi-nn/Cargo.toml
Normal file
@@ -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" }
|
||||
220
crates/wasi-nn/LICENSE
Normal file
220
crates/wasi-nn/LICENSE
Normal file
@@ -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.
|
||||
|
||||
38
crates/wasi-nn/README.md
Normal file
38
crates/wasi-nn/README.md
Normal file
@@ -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
|
||||
```
|
||||
10
crates/wasi-nn/build.rs
Normal file
10
crates/wasi-nn/build.rs
Normal file
@@ -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());
|
||||
}
|
||||
1
crates/wasi-nn/examples/README.md
Normal file
1
crates/wasi-nn/examples/README.md
Normal file
@@ -0,0 +1 @@
|
||||
See `ci/run-wasi-nn-example.sh` for how the classification example is tested during CI.
|
||||
13
crates/wasi-nn/examples/classification-example/Cargo.lock
generated
Normal file
13
crates/wasi-nn/examples/classification-example/Cargo.lock
generated
Normal file
@@ -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",
|
||||
]
|
||||
|
||||
15
crates/wasi-nn/examples/classification-example/Cargo.toml
Normal file
15
crates/wasi-nn/examples/classification-example/Cargo.toml
Normal file
@@ -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]
|
||||
2
crates/wasi-nn/examples/classification-example/README.md
Normal file
2
crates/wasi-nn/examples/classification-example/README.md
Normal file
@@ -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.
|
||||
54
crates/wasi-nn/examples/classification-example/src/main.rs
Normal file
54
crates/wasi-nn/examples/classification-example/src/main.rs
Normal file
@@ -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])
|
||||
}
|
||||
1
crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore
vendored
Normal file
1
crates/wasi-nn/examples/wasi-nn-rust-bindings/.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
/target
|
||||
5
crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock
generated
Normal file
5
crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.lock
generated
Normal file
@@ -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"
|
||||
14
crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml
Normal file
14
crates/wasi-nn/examples/wasi-nn-rust-bindings/Cargo.toml
Normal file
@@ -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]
|
||||
201
crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE
Normal file
201
crates/wasi-nn/examples/wasi-nn-rust-bindings/LICENSE
Normal file
@@ -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.
|
||||
65
crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md
Normal file
65
crates/wasi-nn/examples/wasi-nn-rust-bindings/README.md
Normal file
@@ -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.
|
||||
76
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs
Normal file
76
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/error.rs
Normal file
@@ -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<Error> {
|
||||
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()
|
||||
);
|
||||
}
|
||||
}
|
||||
199
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs
Normal file
199
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/generated.rs
Normal file
@@ -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<T, E = Error> = core::result::Result<T, E>;
|
||||
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<Graph> {
|
||||
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<GraphExecutionContext> {
|
||||
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<BufferSize> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
3
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs
Normal file
3
crates/wasi-nn/examples/wasi-nn-rust-bindings/src/lib.rs
Normal file
@@ -0,0 +1,3 @@
|
||||
mod error;
|
||||
mod generated;
|
||||
pub use generated::*;
|
||||
1
crates/wasi-nn/spec
Submodule
1
crates/wasi-nn/spec
Submodule
Submodule crates/wasi-nn/spec added at 68e73cf612
125
crates/wasi-nn/src/ctx.rs
Normal file
125
crates/wasi-nn/src/ctx.rs
Normal file
@@ -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<T> = std::result::Result<T, WasiNnError>;
|
||||
|
||||
pub struct Table<K, V> {
|
||||
entries: HashMap<K, V>,
|
||||
next_key: u32,
|
||||
}
|
||||
|
||||
impl<K, V> Default for Table<K, V> {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
entries: HashMap::new(),
|
||||
next_key: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<K, V> Table<K, V>
|
||||
where
|
||||
K: Eq + Hash + From<u32> + 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<V> {
|
||||
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<Graph, (openvino::CNNNetwork, openvino::ExecutableNetwork)>,
|
||||
pub(crate) executions: Table<GraphExecutionContext, ExecutionContext>,
|
||||
}
|
||||
|
||||
impl Ctx {
|
||||
/// Make a new `WasiNnCtx` with the default settings.
|
||||
pub fn new() -> WasiNnResult<Self> {
|
||||
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<Ctx>,
|
||||
}
|
||||
|
||||
impl WasiNnCtx {
|
||||
/// Make a new `WasiNnCtx` with the default settings.
|
||||
pub fn new() -> WasiNnResult<Self> {
|
||||
Ok(Self {
|
||||
ctx: RefCell::new(Ctx::new()?),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod test {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn instantiate() {
|
||||
WasiNnCtx::new().unwrap();
|
||||
}
|
||||
}
|
||||
176
crates/wasi-nn/src/impl.rs
Normal file
176
crates/wasi-nn/src/impl.rs
Normal file
@@ -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<Graph> {
|
||||
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<GraphExecutionContext> {
|
||||
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::<Vec<_>>();
|
||||
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<u32> {
|
||||
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"),
|
||||
}
|
||||
}
|
||||
26
crates/wasi-nn/src/lib.rs
Normal file
26
crates/wasi-nn/src/lib.rs
Normal file
@@ -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 },
|
||||
});
|
||||
40
crates/wasi-nn/src/witx.rs
Normal file
40
crates/wasi-nn/src/witx.rs
Normal file
@@ -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
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user