Add per-session ONNX Runtime threading options

README.md

@@ -62,3 +62,42 @@ end
 ```
 
 You will need [Rust](https://www.rust-lang.org/tools/install) for compilation to succeed.
+
+### Per-session CPU threading
+
+`Ortex.load/4` accepts a keyword list of ONNX Runtime session options:
+
+```elixir
+model = Ortex.load("model.onnx", [:cpu], 3,
+  intra_op_num_threads: 1,
+  inter_op_num_threads: 1,
+  execution_mode: :sequential,
+  intra_op_allow_spinning: false,
+  inter_op_allow_spinning: false
+)
+```
+
+Omitting the fourth argument, or passing `[]`, retains the existing behavior.
+Omitted individual settings leave ONNX Runtime defaults intact; `0` for either
+thread count lets ONNX Runtime choose. Inter-op parallelism applies only to
+`:parallel` execution. Settings belong to the loaded session, not the process or
+node. Choose them using measurements for your model and hardware; a smaller
+thread pool does not guarantee faster inference.
+
+These settings use the existing `ort` 2.0.0-rc.8 APIs and session config entries;
+no ONNX Runtime or Rust dependency upgrade is required. Build the Elixir code
+**and** native library from the same revision (`mix deps.compile ortex --force`
+for a dependency). A copied/precompiled old NIF cannot implement the new
+`init_with_options/4` entry point. Ortex normally compiles Rust from source;
+custom release pipelines must invalidate native caches, retain `Cargo.lock`,
+build for each target OS/architecture, and package the matching ONNX Runtime
+libraries as before. The Hex package includes the new Elixir source through its
+existing `lib` entry.
+
+`Ortex.run/2` still uses a `DirtyIo` NIF. Tensor transfers use other NIFs, including
+`DirtyCpu` operations. Timing the Elixir `run/2` call includes transfer and
+scheduler wait, not just ONNX graph execution. Scheduler classification deserves
+a separate investigation and is deliberately unchanged here.
+
+References: [ONNX Runtime threading](https://onnxruntime.ai/docs/performance/tune-performance/threading.html),
+[ERTS dirty NIFs](https://www.erlang.org/doc/apps/erts/erl_nif.html#dirty-nifs).

lib/ortex.ex

@@ -18,14 +18,27 @@ defmodule Ortex do
   `config.exs` where `EXECUTION_PROVIDERS` is a list of strings of which execution providers
   to enable.
 
+  The optional fourth argument configures this session only:
+
+    * `:intra_op_num_threads` / `:inter_op_num_threads` — non-negative integers
+      (0 lets ONNX Runtime choose). Inter-op threads apply to parallel execution.
+    * `:execution_mode` — `:sequential` or `:parallel`.
+    * `:intra_op_allow_spinning` / `:inter_op_allow_spinning` — booleans.
+
+  Omitted options preserve ONNX Runtime defaults. For example,
+  `Ortex.load(path, [:cpu], 3, intra_op_num_threads: 1, inter_op_num_threads: 1,
+  execution_mode: :sequential, intra_op_allow_spinning: false,
+  inter_op_allow_spinning: false)`. Settings cannot be changed after loading.
+  This API requires rebuilding the native library along with the Elixir code.
+
   ## Examples
 
       iex> Ortex.load("./models/tinymodel.onnx")
       iex> Ortex.load("./models/tinymodel.onnx", [:cuda, :cpu])
       iex> Ortex.load("./models/tinymodel.onnx", [:cpu], 0)
 
   """
-  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,15 @@ 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
+    options = Ortex.SessionOptions.normalize(session_options)
+
+    result =
+      if session_options == [],
+        do: Ortex.Native.init(path, eps, opt),
+        else: Ortex.Native.init_with_options(path, eps, opt, options)
+
+    case result do
       {:error, msg} ->
         raise msg
 

lib/ortex/native.ex

@@ -28,6 +28,9 @@ defmodule Ortex.Native do
   def init(_model_path, _execution_providers, _optimization_level),
     do: :erlang.nif_error(:nif_not_loaded)
 
+  def init_with_options(_model_path, _execution_providers, _optimization_level, _options),
+    do: :erlang.nif_error(:nif_not_loaded)
+
   def run(_model, _inputs), do: :erlang.nif_error(:nif_not_loaded)
   def from_binary(_bin, _shape, _type), do: :erlang.nif_error(:nif_not_loaded)
   def to_binary(_reference, _bits, _limit), do: :erlang.nif_error(:nif_not_loaded)

lib/ortex/session_options.ex

@@ -0,0 +1,37 @@
+defmodule Ortex.SessionOptions do
+  @moduledoc false
+
+  @defaults %{
+    intra_op_num_threads: nil,
+    inter_op_num_threads: nil,
+    execution_mode: nil,
+    intra_op_allow_spinning: nil,
+    inter_op_allow_spinning: nil
+  }
+
+  def normalize(options) do
+    unless Keyword.keyword?(options),
+      do: raise(ArgumentError, "session options must be a keyword list")
+
+    if length(Keyword.keys(options)) != length(Enum.uniq(Keyword.keys(options))) do
+      raise ArgumentError, "duplicate session options"
+    end
+
+    Enum.reduce(options, @defaults, fn {key, value}, acc ->
+      unless valid?(key, value),
+        do: raise(ArgumentError, "invalid session option #{inspect(key)}: #{inspect(value)}")
+
+      Map.put(acc, key, value)
+    end)
+  end
+
+  defp valid?(key, value) when key in [:intra_op_num_threads, :inter_op_num_threads],
+    do: is_integer(value) and value >= 0 and value <= 2_147_483_647
+
+  defp valid?(:execution_mode, value), do: value in [:sequential, :parallel]
+
+  defp valid?(key, value) when key in [:intra_op_allow_spinning, :inter_op_allow_spinning],
+    do: is_boolean(value)
+
+  defp valid?(_, _), do: false
+end

native/ortex/src/lib.rs

@@ -24,7 +24,21 @@ fn init(
     opt: i32,
 ) -> 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, model::SessionOptions::default())
+        .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?;
+    Ok(ResourceArc::new(model))
+}
+
+#[rustler::nif(schedule = "DirtyIo")]
+fn init_with_options(
+    env: Env,
+    model_path: String,
+    eps: Vec<Atom>,
+    opt: i32,
+    options: model::SessionOptions,
+) -> NifResult<ResourceArc<model::OrtexModel>> {
+    let eps = utils::map_eps(env, eps);
+    let model = model::init(model_path, eps, opt, options)
         .map_err(|e| rustler::Error::Term(Box::new(e.to_string())))?;
     Ok(ResourceArc::new(model))
 }
@@ -109,6 +123,7 @@ rustler::init!(
     [
         run,
         init,
+        init_with_options,
         from_binary,
         to_binary,
         show_session,

native/ortex/src/model.rs

@@ -17,8 +17,22 @@ use ort::{Error, ExecutionProviderDispatch, Session};
 use rustler::resource::ResourceArc;
 use rustler::Atom;
 
-/// Holds the model state which include onnxruntime session and environment. All
-/// are threadsafe so this can be called concurrently from the beam.
+#[derive(rustler::NifUnitEnum, Clone, Copy)]
+pub enum ExecutionMode {
+    Sequential,
+    Parallel,
+}
+
+#[derive(rustler::NifMap, Default)]
+pub struct SessionOptions {
+    pub intra_op_num_threads: Option<u32>,
+    pub inter_op_num_threads: Option<u32>,
+    pub execution_mode: Option<ExecutionMode>,
+    pub intra_op_allow_spinning: Option<bool>,
+    pub inter_op_allow_spinning: Option<bool>,
+}
+
+/// Holds the thread-safe ONNX Runtime session shared by BEAM callers.
 pub struct OrtexModel {
     pub session: ort::Session,
 }
@@ -35,14 +49,31 @@ pub fn init(
     model_path: String,
     eps: Vec<ExecutionProviderDispatch>,
     opt: i32,
+    options: SessionOptions,
 ) -> 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)?;
+
+    if let Some(threads) = options.intra_op_num_threads {
+        builder = builder.with_intra_threads(threads as usize)?;
+    }
+    if let Some(threads) = options.inter_op_num_threads {
+        builder = builder.with_inter_threads(threads as usize)?;
+    }
+    if let Some(mode) = options.execution_mode {
+        builder = builder.with_parallel_execution(matches!(mode, ExecutionMode::Parallel))?;
+    }
+    if let Some(spinning) = options.intra_op_allow_spinning {
+        builder = builder.with_intra_op_spinning(spinning)?;
+    }
+    if let Some(spinning) = options.inter_op_allow_spinning {
+        builder = builder.with_inter_op_spinning(spinning)?;
+    }
+    let session = builder.commit_from_file(model_path)?;
 
     let state = OrtexModel { session };
     Ok(state)

test/session_options_test.exs

@@ -0,0 +1,62 @@
+defmodule Ortex.SessionOptionsTest do
+  use ExUnit.Case, async: true
+
+  @tuned [
+    intra_op_num_threads: 1,
+    inter_op_num_threads: 1,
+    execution_mode: :sequential,
+    intra_op_allow_spinning: false,
+    inter_op_allow_spinning: false
+  ]
+
+  test "legacy arities, empty options, and configured sessions produce equivalent results" do
+    reference = run(Ortex.load("models/tinymodel.onnx"))
+
+    models = [
+      Ortex.load("models/tinymodel.onnx", [:cpu]),
+      Ortex.load("models/tinymodel.onnx", [:cpu], 3),
+      Ortex.load("models/tinymodel.onnx", [:cpu], 3, []),
+      Ortex.load("models/tinymodel.onnx", [:cpu], 3, @tuned),
+      Ortex.load(
+        "models/tinymodel.onnx",
+        [:cpu],
+        3,
+        Keyword.merge(@tuned,
+          execution_mode: :parallel,
+          intra_op_allow_spinning: true,
+          inter_op_allow_spinning: true
+        )
+      )
+    ]
+
+    for model <- models, do: assert(run(model) == reference)
+  end
+
+  test "each option is independently optional and zero thread counts retain runtime choice" do
+    for options <- Enum.map(@tuned, &[&1]) ++ [[intra_op_num_threads: 0, inter_op_num_threads: 0]] do
+      assert %Ortex.Model{} = Ortex.load("models/tinymodel.onnx", [:cpu], 3, options)
+    end
+  end
+
+  test "invalid options fail before native loading rather than silently using defaults" do
+    for options <- [
+          [unknown: 1],
+          [intra_op_num_threads: -1],
+          [inter_op_num_threads: 1.5],
+          [intra_op_num_threads: 2_147_483_648],
+          [execution_mode: :typo],
+          [intra_op_allow_spinning: 0],
+          [inter_op_allow_spinning: nil],
+          [intra_op_num_threads: 1, intra_op_num_threads: 2],
+          %{}
+        ] do
+      assert_raise ArgumentError, fn -> Ortex.load("missing.onnx", [:cpu], 3, options) end
+    end
+  end
+
+  defp run(model) do
+    Ortex.run(model, {Nx.broadcast(1, {1, 100}) |> Nx.as_type(:s32), Nx.broadcast(1.0, {1, 100})})
+    |> Tuple.to_list()
+    |> Enum.map(&Nx.to_binary/1)
+  end
+end