Add an intra_threads option to Ortex.load

lib/ortex.ex

@@ -18,14 +18,23 @@ defmodule Ortex do
   `config.exs` where `EXECUTION_PROVIDERS` is a list of strings of which execution providers
   to enable.
 
+  ## Options
+
+    * `:intra_threads` - the number of threads ONNX Runtime spreads a single
+      inference across. By default ONNX Runtime uses one thread per physical
+      core, so one busy model can occupy most of the machine; a smaller pool
+      trades per-inference latency for cores left free for other work, such
+      as a second model. Must be a positive integer.
+
   ## 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, intra_threads: 2)
 
   """
-  defdelegate load(path, eps \\ [:cpu], opt \\ 3), to: Ortex.Model
+  defdelegate load(path, eps \\ [:cpu], opt \\ 3, options \\ []), to: Ortex.Model
 
   @doc """
   Run a forward pass through a model.

lib/ortex/model.ex

@@ -22,8 +22,11 @@ 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, options \\ []) do
+    options = Keyword.validate!(options, intra_threads: nil)
+    intra_threads = validate_intra_threads!(options[:intra_threads])
+
+    case Ortex.Native.init(path, eps, opt, intra_threads) do
       {:error, msg} ->
         raise msg
 
@@ -32,6 +35,16 @@ defmodule Ortex.Model do
     end
   end
 
+  defp validate_intra_threads!(nil), do: nil
+
+  defp validate_intra_threads!(threads) when is_integer(threads) and threads > 0,
+    do: threads
+
+  defp validate_intra_threads!(threads) do
+    raise ArgumentError,
+          "expected :intra_threads to be a positive integer, got: #{inspect(threads)}"
+  end
+
   @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, _intra_threads),
     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,
+    intra_threads: Option<usize>,
 ) -> 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, intra_threads)
         .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, None)?;
 //! let (inputs, outputs) = show(model)?;
 //! ```
 
@@ -31,16 +31,26 @@ 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.
+///
+/// `intra_threads` caps the thread pool one inference is spread across. `None`
+/// leaves the choice to ONNX Runtime, which defaults to one thread per physical
+/// core and honours `ORT_INTRA_OP_NUM_THREADS` when set in the OS environment.
 pub fn init(
     model_path: String,
     eps: Vec<ExecutionProviderDispatch>,
     opt: i32,
+    intra_threads: Option<usize>,
 ) -> Result<OrtexModel, Error> {
     // TODO: send tracing logs to erlang/elixir _somehow_
     // tracing_subscriber::fmt::init();
 
-    let session = Session::builder()?
-        .with_optimization_level(map_opt_level(opt))?
+    let mut builder = Session::builder()?.with_optimization_level(map_opt_level(opt))?;
+
+    if let Some(threads) = intra_threads {
+        builder = builder.with_intra_threads(threads)?;
+    }
+
+    let session = builder
         .with_execution_providers(eps)?
         .commit_from_file(model_path)?;
 

test/ortex_test.exs

@@ -23,6 +23,36 @@ defmodule OrtexTest do
     assert result |> Nx.backend_transfer() |> Nx.argmax(axis: 1) == Nx.tensor([499])
   end
 
+  describe "intra_threads" do
+    test "a session capped at one intra-op thread still runs the model" do
+      model = Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_threads: 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 "is not required" do
+      assert %Ortex.Model{} = Ortex.load("./models/tinymodel.onnx", [:cpu], 3, [])
+    end
+
+    test "must be a positive integer" do
+      for bad <- [0, -1, 2.0, "4"] do
+        assert_raise ArgumentError, ~r/intra_threads/, fn ->
+          Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_threads: bad)
+        end
+      end
+    end
+
+    test "an unknown option is refused rather than ignored" do
+      assert_raise ArgumentError, ~r/intra_thread/, fn ->
+        Ortex.load("./models/tinymodel.onnx", [:cpu], 3, intra_thread: 2)
+      end
+    end
+  end
+
   test "Nx.Serving with tinymodel" do
     model = Ortex.load("./models/tinymodel.onnx")