Add ability to set ONNX session options

lib/ortex.ex

@@ -18,14 +18,24 @@ defmodule Ortex do
   `config.exs` where `EXECUTION_PROVIDERS` is a list of strings of which execution providers
   to enable.
 
+  Session config entries can be passed as a fourth argument. These are handed
+  straight to ONNX Runtime as
+  [session configuration keys](https://github.com/microsoft/onnxruntime/blob/main/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h),
+  so any key ONNX Runtime understands works here. Values may be given as strings,
+  booleans (`true`/`false` become `"1"`/`"0"`), or numbers.
+
   ## Examples
 
       iex> Ortex.load("./models/tinymodel.onnx")
       iex> Ortex.load("./models/tinymodel.onnx", [:cuda, :cpu])
       iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 0)
+      iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 3,
+      ...>   "session.intra_op.allow_spinning": false,
+      ...>   "session.inter_op.allow_spinning": false
+      ...> )
 
   """
-  defdelegate load(path, eps \\ [:cpu], opt \\ 3), to: Ortex.Model
+  defdelegate load(path, eps \\ [:cpu], opt \\ 3, session_options \\ []), to: Ortex.Model
 
   @doc """
   Run a forward pass through a model.

lib/ortex/model.ex

@@ -22,8 +22,8 @@ defmodule Ortex.Model do
   defstruct [:reference]
 
   @doc false
-  def load(path, eps \\ [:cpu], opt \\ 3) do
-    case Ortex.Native.init(path, eps, opt) do
+  def load(path, eps \\ [:cpu], opt \\ 3, session_options \\ []) do
+    case Ortex.Native.init(path, eps, opt, normalize_session_options(session_options)) do
       {:error, msg} ->
         raise msg
 
@@ -32,6 +32,24 @@ defmodule Ortex.Model do
     end
   end
 
+  # ONNX Runtime session config entries are string key/value pairs. Accept the
+  # natural Elixir spellings (atoms, booleans, numbers) and stringify them.
+  defp normalize_session_options(session_options) do
+    Enum.map(session_options, fn {key, value} ->
+      {config_key(key), config_value(value)}
+    end)
+  end
+
+  defp config_key(key) when is_atom(key), do: Atom.to_string(key)
+  defp config_key(key) when is_binary(key), do: key
+
+  defp config_value(true), do: "1"
+  defp config_value(false), do: "0"
+  defp config_value(value) when is_binary(value), do: value
+  defp config_value(value) when is_atom(value), do: Atom.to_string(value)
+  defp config_value(value) when is_integer(value), do: Integer.to_string(value)
+  defp config_value(value) when is_float(value), do: Float.to_string(value)
+
   @doc false
   def run(%Ortex.Model{} = model, tensor) when not is_tuple(tensor) do
     run(model, {tensor})

lib/ortex/native.ex

@@ -25,7 +25,7 @@ defmodule Ortex.Native do
 
   # When loading a NIF module, dummy clauses for all NIF function are required.
   # NIF dummies usually just error out when called when the NIF is not loaded, as that should never normally happen.
-  def init(_model_path, _execution_providers, _optimization_level),
+  def init(_model_path, _execution_providers, _optimization_level, _session_options),
     do: :erlang.nif_error(:nif_not_loaded)
 
   def run(_model, _inputs), do: :erlang.nif_error(:nif_not_loaded)

native/ortex/src/lib.rs

@@ -22,9 +22,10 @@ fn init(
     model_path: String,
     eps: Vec<Atom>,
     opt: i32,
+    session_options: Vec<(String, String)>,
 ) -> NifResult<ResourceArc<model::OrtexModel>> {
     let eps = utils::map_eps(env, eps);
-    let model = model::init(model_path, eps, opt)
+    let model = model::init(model_path, eps, opt, session_options)
         .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?;
     Ok(ResourceArc::new(model))
 }

native/ortex/src/model.rs

@@ -4,7 +4,7 @@
 //! # Examples
 //!
 //! ```
-//! let model = init("./models/resnet50.onnx", vec![])?;
+//! let model = init("./models/resnet50.onnx", vec![], 3, vec![])?;
 //! let (inputs, outputs) = show(model)?;
 //! ```
 
@@ -30,19 +30,27 @@ pub struct OrtexModel {
 unsafe impl Sync for OrtexModel {}
 
 /// Creates a model given the path to the model and vector of execution providers.
-/// The execution providers are Atoms from Erlang/Elixir.
+/// The execution providers are Atoms from Erlang/Elixir. `session_options` are
+/// key/value pairs passed straight through to ONNX Runtime as session config
+/// entries, e.g. `("session.intra_op.allow_spinning", "0")`.
 pub fn init(
     model_path: String,
     eps: Vec<ExecutionProviderDispatch>,
     opt: i32,
+    session_options: Vec<(String, String)>,
 ) -> Result<OrtexModel, Error> {
     // TODO: send tracing logs to erlang/elixir _somehow_
     // tracing_subscriber::fmt::init();
 
-    let session = Session::builder()?
+    let mut builder = Session::builder()?
         .with_optimization_level(map_opt_level(opt))?
-        .with_execution_providers(eps)?
-        .commit_from_file(model_path)?;
+        .with_execution_providers(eps)?;
+
+    for (key, value) in session_options.iter() {
+        builder = builder.with_config_entry(key, value)?;
+    }
+
+    let session = builder.commit_from_file(model_path)?;
 
     let state = OrtexModel { session };
     Ok(state)

test/ortex_test.exs

@@ -23,6 +23,21 @@ defmodule OrtexTest do
     assert result |> Nx.backend_transfer() |> Nx.argmax(axis: 1) == Nx.tensor([499])
   end
 
+  test "session config entries" do
+    model =
+      Ortex.load("./models/tinymodel.onnx", [:cpu], 3,
+        "session.intra_op.allow_spinning": false,
+        "session.inter_op.allow_spinning": 0,
+        "session.disable_prepacking": "1"
+      )
+
+    {%Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}, %Nx.Tensor{shape: {1, 10}}} =
+      Ortex.run(model, {
+        Nx.broadcast(0, {1, 100}) |> Nx.as_type(:s32),
+        Nx.broadcast(0.0, {1, 100}) |> Nx.as_type(:f32)
+      })
+  end
+
   test "Nx.Serving with tinymodel" do
     model = Ortex.load("./models/tinymodel.onnx")