README.md (#3)

* Tensor::from_array

* README.md

README.md

@@ -0,0 +1,126 @@
+# ExecuTorch-rs
+
+`executorch` is a Rust library for executing PyTorch models in Rust.
+It is a Rust wrapper around the [ExecuTorch C++ API](https://pytorch.org/executorch).
+It depends on version `0.2.1` of the CPP API, but will advance as the API does.
+The underlying C++ library is still in alpha, and its API is subject to change together with the Rust API.
+
+## Usage
+Create a model in `Python` and export it:
+```python
+import torch
+from executorch.exir import to_edge
+from torch.export import export
+
+class Add(torch.nn.Module):
+ def __init__(self):
+ super(Add, self).__init__()
+
+ def forward(self, x: torch.Tensor, y: torch.Tensor):
+ return x + y
+
+
+aten_dialect = export(Add(), (torch.ones(1), torch.ones(1)))
+edge_program = to_edge(aten_dialect)
+executorch_program = edge_program.to_executorch()
+with open("model.pte", "wb") as file:
+ file.write(executorch_program.buffer)
+```
+Execute the model in Rust:
+```rust
+use executorch::{EValue, Module, Tag, Tensor, TensorImpl};
+use ndarray::array;
+
+let mut module = Module::new("model.pte");
+
+let data1 = array![1.0_f32];
+let input_tensor1 = TensorImpl::from_array(data1.view());
+let input_evalue1 = EValue::from_tensor(Tensor::new(input_tensor1.as_ref()));
+
+let data2 = array![1.0_f32];
+let input_tensor2 = TensorImpl::from_array(data2.view());
+let input_evalue2 = EValue::from_tensor(Tensor::new(input_tensor2.as_ref()));
+
+let outputs = module.forward(&[input_evalue1, input_evalue2]).unwrap();
+assert_eq!(outputs.len(), 1);
+let output = outputs.into_iter().next().unwrap();
+assert_eq!(output.tag(), Some(Tag::Tensor));
+let output = output.as_tensor().as_array::<f32>();
+
+println!("Output tensor computed: {:?}", output);
+assert_eq!(output, array![2.0].into_dyn());
+```
+See `example/hello_world_add` and `example/hello_world_add_module` for the complete examples.
+
+## Build
+To build the library, you need to build the C++ library first.
+The C++ library allow for great flexibility with many flags, customizing which modules, kernels, and extensions are built.
+Multiple static libraries are built, and the Rust library links to them.
+In the following example we build the C++ library with the necessary flags to run example `hello_world_add`:
+```bash
+# Clone the C++ library
+cd ${TEMP_DIR}
+git clone --depth 1 --branch v0.2.1 https://github.com/pytorch/executorch.git
+cd executorch
+git submodule sync --recursive
+git submodule update --init --recursive
+
+# Install requirements
+./install_requirements.sh
+
+# Build C++ library
+mkdir cmake-out && cd cmake-out
+cmake \
+ -DDEXECUTORCH_SELECT_OPS_LIST=aten::add.out \
+ -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \
+ -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=OFF \
+ -DBUILD_EXECUTORCH_PORTABLE_OPS=ON \
+ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \
+ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \
+ -DEXECUTORCH_ENABLE_PROGRAM_VERIFICATION=ON \
+ -DEXECUTORCH_ENABLE_LOGGING=ON \
+ ..
+make -j
+
+# Static libraries are in cmake-out/
+# core:
+# cmake-out/libexecutorch.a
+# cmake-out/libexecutorch_no_prim_ops.a
+# kernels implementations:
+# cmake-out/kernels/portable/libportable_ops_lib.a
+# cmake-out/kernels/portable/libportable_kernels.a
+# extension data loader, enabled with EXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON:
+# cmake-out/extension/data_loader/libextension_data_loader.a
+# extension module, enabled with EXECUTORCH_BUILD_EXTENSION_MODULE=ON:
+# cmake-out/extension/module/libextension_module_static.a
+
+# Run example
+# We set EXECUTORCH_RS_EXECUTORCH_LIB_DIR to the path of the C++ build output
+cd ${EXECUTORCH_RS_DIR}/examples/hello_world_add
+python export_model.py
+EXECUTORCH_RS_EXECUTORCH_LIB_DIR=${TEMP_DIR}/executorch/cmake-out cargo run
+```
+
+The `executorch` crate will always look for the following static libraries:
+- `libexecutorch.a`
+- `libexecutorch_no_prim_ops.a`
+
+Additional libs are required if feature flags are enabled (see next section):
+- `libextension_data_loader.a`
+- `libextension_module_static.a`
+
+The static libraries of the kernels implementations are required only if your model uses them, and they should be **linked manually** by the binary that uses the `executorch` crate.
+For example, the `hello_world_add` example uses a model with a single addition operation, so it compile the C++ library with `DEXECUTORCH_SELECT_OPS_LIST=aten::add.out` and contain the following lines in its `build.rs`:
+```rust
+println!("cargo::rustc-link-lib=static=portable_kernels");
+println!("cargo::rustc-link-lib=static:+whole-archive=portable_ops_lib");
+
+let libs_dir = std::env::var("EXECUTORCH_RS_EXECUTORCH_LIB_DIR").unwrap();
+println!("cargo::rustc-link-search={}/kernels/portable/", libs_dir);
+```
+Note that the `portable_ops_lib` is linked with `+whole-archive` to ensure that all symbols are included in the binary.
+
+## Cargo Features
+- `default`: disables all features.
+- `extension-data-loader`: include the `FileDataLoader` strut. The `libextension_data_loader.a` static library is required, compile C++ `executorch` with `EXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON`.
+- `extension-module`: include the `Module` strut. The `libextension_module_static.a` static library is required, compile C++ `executorch` with `EXECUTORCH_BUILD_EXTENSION_MODULE=ON`.

examples/hello_world_add/src/main.rs

@@ -1,3 +1,5 @@
+#![deny(warnings)]
+
 use executorch::{
  EValue, FileDataLoader, HierarchicalAllocator, MallocMemoryAllocator, MemoryManager, Program,
  ProgramVerification, Span, Tag, Tensor, TensorImpl,
@@ -38,19 +40,13 @@ fn main() {
 
  let mut method = program.load_method("forward", &mut memory_manager).unwrap();
 
- let mut data1 = vec![1.0_f32; 1];
- let sizes1 = [1];
- let data_order1 = [0];
- let strides1 = [1];
- let mut input_tensor1 = TensorImpl::new(&sizes1, &mut data1, &data_order1, &strides1);
- let input_evalue1 = EValue::from_tensor(Tensor::new(&mut input_tensor1));
-
- let mut data2 = vec![1.0_f32; 1];
- let sizes2 = [1];
- let data_order2 = [0];
- let strides2 = [1];
- let mut input_tensor2 = TensorImpl::new(&sizes2, &mut data2, &data_order2, &strides2);
- let input_evalue2 = EValue::from_tensor(Tensor::new(&mut input_tensor2));
+ let data1 = array![1.0_f32];
+ let input_tensor1 = TensorImpl::from_array(data1.view());
+ let input_evalue1 = EValue::from_tensor(Tensor::new(input_tensor1.as_ref()));
+
+ let data2 = array![1.0_f32];
+ let input_tensor2 = TensorImpl::from_array(data2.view());
+ let input_evalue2 = EValue::from_tensor(Tensor::new(input_tensor2.as_ref()));
 
  let mut method_exe = method.start_execution();
 

examples/hello_world_add_module/src/main.rs

@@ -1,6 +1,7 @@
+#![deny(warnings)]
+
 use executorch::{EValue, Module, Tag, Tensor, TensorImpl};
 use ndarray::array;
-use std::vec;
 
 fn main() {
  env_logger::Builder::new()
@@ -11,19 +12,13 @@ fn main() {
 
  let mut module = Module::new("model.pte");
 
- let mut data1 = vec![1.0_f32; 1];
- let sizes1 = [1];
- let data_order1 = [0];
- let strides1 = [1];
- let mut input_tensor1 = TensorImpl::new(&sizes1, &mut data1, &data_order1, &strides1);
- let input_evalue1 = EValue::from_tensor(Tensor::new(&mut input_tensor1));
-
- let mut data2 = vec![1.0_f32; 1];
- let sizes2 = [1];
- let data_order2 = [0];
- let strides2 = [1];
- let mut input_tensor2 = TensorImpl::new(&sizes2, &mut data2, &data_order2, &strides2);
- let input_evalue2 = EValue::from_tensor(Tensor::new(&mut input_tensor2));
+ let data1 = array![1.0_f32];
+ let input_tensor1 = TensorImpl::from_array(data1.view());
+ let input_evalue1 = EValue::from_tensor(Tensor::new(input_tensor1.as_ref()));
+
+ let data2 = array![1.0_f32];
+ let input_tensor2 = TensorImpl::from_array(data2.view());
+ let input_evalue2 = EValue::from_tensor(Tensor::new(input_tensor2.as_ref()));
 
  let outputs = module.forward(&[input_evalue1, input_evalue2]).unwrap();
  assert_eq!(outputs.len(), 1);

src/tensor.rs

@@ -1,6 +1,6 @@
-use std::marker::PhantomData;
+use std::{marker::PhantomData, mem::MaybeUninit, ptr};
 
-use ndarray::{ArrayViewD, ArrayViewMut, IxDyn, ShapeBuilder};
+use ndarray::{ArrayView, ArrayViewD, ArrayViewMut, Axis, Dimension, IxDyn, ShapeBuilder};
 
 use crate::{c_link, et_c, et_rs_c, Error, Result};
 
@@ -127,8 +127,8 @@ impl Scalar for bool {
 
 pub struct Tensor<'a>(pub(crate) et_c::Tensor, PhantomData<&'a ()>);
 impl<'a> Tensor<'a> {
- pub fn new(tensor_impl: &'a mut TensorImpl<'a>) -> Self {
- let impl_ = &mut tensor_impl.0;
+ pub fn new(tensor_impl: &'a TensorImpl<'a>) -> Self {
+ let impl_ = &tensor_impl.0 as *const _ as *mut _;
  Self(et_c::Tensor { impl_ }, PhantomData)
  }
 
@@ -252,16 +252,20 @@ pub struct TensorImpl<'a>(et_c::TensorImpl, PhantomData<&'a ()>);
 impl<'a> TensorImpl<'a> {
  pub fn new<S: Scalar>(
  sizes: &'a [SizesType],
- data: &mut [S],
- data_order: &'a [DimOrderType],
+ data: &'a [S],
+ data_order: Option<&'a [DimOrderType]>,
  strides: &'a [StridesType],
  ) -> Self {
  let dim = sizes.len();
- assert_eq!(dim, data_order.len());
+ if let Some(data_order) = &data_order {
+ assert_eq!(dim, data_order.len());
+ }
  assert_eq!(dim, strides.len());
  let sizes = sizes as *const _ as *mut SizesType;
- let data = data.as_mut_ptr() as *mut _;
- let dim_order = data_order as *const _ as *mut DimOrderType;
+ let data = data as *const _ as *mut _;
+ let dim_order = data_order
+ .map(|p| p as *const _ as *mut DimOrderType)
+ .unwrap_or(ptr::null_mut());
  let strides = strides as *const _ as *mut StridesType;
  let impl_ = unsafe {
  et_c::TensorImpl::new(
@@ -276,4 +280,44 @@ impl<'a> TensorImpl<'a> {
  };
  Self(impl_, PhantomData)
  }
+
+ pub fn from_array<S: Scalar, D: Dimension>(array: ArrayView<'a, S, D>) -> impl AsRef<Self> {
+ struct Wrapper<'a, S: Scalar> {
+ sizes: Vec<SizesType>,
+ strides: Vec<StridesType>,
+ tensor_impl: MaybeUninit<TensorImpl<'a>>,
+ _phantom: PhantomData<S>,
+ }
+ impl<'a, S: Scalar> AsRef<TensorImpl<'a>> for Wrapper<'a, S> {
+ fn as_ref(&self) -> &TensorImpl<'a> {
+ unsafe { self.tensor_impl.assume_init_ref() }
+ }
+ }
+
+ let mut wrapper = Wrapper::<'a, S> {
+ sizes: array
+ .shape()
+ .iter()
+ .map(|&size| size as SizesType)
+ .collect(),
+ strides: (0..array.ndim())
+ .map(|d| array.stride_of(Axis(d)) as StridesType)
+ .collect(),
+ tensor_impl: MaybeUninit::uninit(),
+ _phantom: PhantomData,
+ };
+ let impl_ = unsafe {
+ et_c::TensorImpl::new(
+ S::TYPE.into_c_scalar_type(),
+ array.ndim() as isize,
+ wrapper.sizes.as_slice() as *const _ as *mut SizesType,
+ array.as_ptr() as *mut _,
+ ptr::null_mut(),
+ wrapper.strides.as_slice() as *const _ as *mut StridesType,
+ et_c::TensorShapeDynamism::STATIC,
+ )
+ };
+ wrapper.tensor_impl.write(Self(impl_, PhantomData));
+ wrapper
+ }
 }