Make sure EPs are supported

Author: [email protected]

mix.exs

@@ -31,7 +31,7 @@ defmodule Ortex.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:rustler, "~> 0.33"},
+      {:rustler, "~> 0.36.2"},
       {:nx, "~> 0.10"},
       {:tokenizers, "~> 0.5", only: :dev},
       {:ex_doc, "~> 0.38", only: :dev, runtime: false}

native/ortex/Cargo.lock

@@ -13,9 +13,9 @@ crate-type = ["cdylib"]
 resolver = "2"
 
 [dependencies]
-rustler = "0.33"
+rustler = "0.36.2"
 ort-sys = { version = "=2.0.0-rc.10", default-features = false }
-ort = { version = "2.0.0-rc.10", features = ["half"] }
+ort = { version = "2.0.0-rc.10", features = ["half", "cuda"] }
 ndarray = "0.16.1"
 half = "2.6.0"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }

native/ortex/Cargo.toml

@@ -104,21 +104,13 @@ pub fn concatenate<'a>(
     Ok(ResourceArc::new(concatted))
 }
 
+pub fn on_load(env: Env) -> bool {
+    env.register::<OrtexModel>().is_ok() && env.register::<OrtexTensor>().is_ok()
+}
+
 rustler::init!(
     "Elixir.Ortex.Native",
-    [
-        run,
-        init,
-        from_binary,
-        to_binary,
-        show_session,
-        slice,
-        reshape,
-        concatenate
-    ],
     load = |env: Env, _term: Term| -> bool {
-        rustler::resource!(OrtexModel, env);
-        rustler::resource!(OrtexTensor, env);
-        true
+        on_load(env)
     }
 );

native/ortex/src/lib.rs

@@ -17,6 +17,7 @@ use ort::execution_providers::ExecutionProviderDispatch;
 use ort::session::Session;
 use ort::Error;
 use rustler::Atom;
+use rustler::Resource;
 use rustler::ResourceArc;
 use std::error::Error as StdError;
 use std::sync::Mutex;
@@ -26,6 +27,7 @@ use std::sync::Mutex;
 pub struct OrtexModel {
     pub session: Mutex<ort::session::Session>,
 }
+impl Resource for OrtexModel {}
 
 // Since we're only using the session for inference and
 // inference is threadsafe, this Sync is safe. Additionally,

native/ortex/src/model.rs

@@ -5,6 +5,7 @@ use ndarray::{ArrayBase, ArrayView, Data, IxDyn, IxDynImpl, ViewRepr};
 use ort::value::Value;
 use ort::Error;
 use rustler::Atom;
+use rustler::Resource;
 use rustler::ResourceArc;
 use std::error::Error as StdError;
 
@@ -204,6 +205,8 @@ impl OrtexTensor {
     }
 }
 
+impl Resource for OrtexTensor {}
+
 fn slice_array<'a, T, D>(
     array: &'a Array<T, D>,
     slice_specs: &'a Vec<(isize, Option<isize>, isize)>,
@@ -441,7 +444,7 @@ macro_rules! concatenate {
     // `typ` is the actual datatype, `ort_tensor_kind` is the OrtexTensor variant
     ($tensors:expr, $axis:expr, $typ:ty, $ort_tensor_kind:ident) => {{
         type ArrayType<'a> = ArrayBase<ViewRepr<&'a $typ>, Dim<IxDynImpl>>;
-        fn filter(tensor: &OrtexTensor) -> Option<ArrayType> {
+        fn filter<'a>(tensor: &'a OrtexTensor) -> Option<ArrayType<'a>> {
             match tensor {
                 OrtexTensor::$ort_tensor_kind(x) => Some(x.view()),
                 _ => None,

native/ortex/src/tensor.rs

@@ -15,7 +15,7 @@ use ort::execution_providers::ExecutionProviderDispatch;
 use ort::session::builder::GraphOptimizationLevel;
 
 /// A faster (unsafe) way of creating an Array from an Erlang binary
-fn initialize_from_raw_ptr<T>(ptr: *const T, shape: &[Ix]) -> ArrayViewMut<T, IxDyn> {
+fn initialize_from_raw_ptr<T>(ptr: *const T, shape: &[Ix]) -> ArrayViewMut<'_, T, IxDyn> {
     let array = unsafe { ArrayViewMut::from_shape_ptr(shape, ptr as *mut T) };
     array
 }
@@ -96,7 +96,7 @@ pub fn map_eps(env: rustler::env::Env, eps: Vec<Atom>) -> Vec<ExecutionProviderD
     eps.iter()
         .map(|e| match &e.to_term(env).atom_to_string().unwrap()[..] {
             CPU => ort::execution_providers::cpu::CPUExecutionProvider::default().build(),
-            CUDA => ort::execution_providers::cuda::CUDAExecutionProvider::default().build(),
+            CUDA => ort::execution_providers::cuda::CUDAExecutionProvider::default().build().error_on_failure(),
             TENSORRT => {
                 ort::execution_providers::tensorrt::TensorRTExecutionProvider::default().build()
             }