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
+    }
 }