diff --git a/README.md b/README.md index 53c54fc..671fbd1 100644 --- a/README.md +++ b/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). diff --git a/lib/ortex.ex b/lib/ortex.ex index 51b0e45..6c3ce99 100644 --- a/lib/ortex.ex +++ b/lib/ortex.ex @@ -18,6 +18,19 @@ 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") @@ -25,7 +38,7 @@ defmodule Ortex do 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. diff --git a/lib/ortex/model.ex b/lib/ortex/model.ex index d8e550f..67ff760 100644 --- a/lib/ortex/model.ex +++ b/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 diff --git a/lib/ortex/native.ex b/lib/ortex/native.ex index ee16c07..2059ab6 100644 --- a/lib/ortex/native.ex +++ b/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) diff --git a/lib/ortex/session_options.ex b/lib/ortex/session_options.ex new file mode 100644 index 0000000..1dd54d1 --- /dev/null +++ b/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 diff --git a/native/ortex/src/lib.rs b/native/ortex/src/lib.rs index 0523105..ff13ea0 100644 --- a/native/ortex/src/lib.rs +++ b/native/ortex/src/lib.rs @@ -24,7 +24,21 @@ fn init( opt: i32, ) -> NifResult> { 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, + opt: i32, + options: model::SessionOptions, +) -> NifResult> { + 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, diff --git a/native/ortex/src/model.rs b/native/ortex/src/model.rs index f5e04b3..f80c2e4 100644 --- a/native/ortex/src/model.rs +++ b/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, + pub inter_op_num_threads: Option, + pub execution_mode: Option, + pub intra_op_allow_spinning: Option, + pub inter_op_allow_spinning: Option, +} + +/// 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, opt: i32, + options: SessionOptions, ) -> Result { // 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) diff --git a/test/session_options_test.exs b/test/session_options_test.exs new file mode 100644 index 0000000..3040383 --- /dev/null +++ b/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