From 45c050b422d41d44b837eeec23dc04c62819fb2f Mon Sep 17 00:00:00 2001 From: Stephan Seitz Date: Wed, 7 Oct 2026 12:54:08 +0200 Subject: [PATCH 1/3] build: enable trtx-runtime feature by default --- Cargo.toml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 1a1a85b..57fbe2f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ crate-type = ["cdylib", "rlib"] name = "pywebnn" [features] -default = ["onnx-runtime", "dynamic-inputs"] +default = ["onnx-runtime", "trtx-runtime", "dynamic-inputs"] onnx-runtime = ["rustnn/onnx-runtime", "ort"] coreml-runtime = ["rustnn/coreml-runtime"] trtx-runtime = ["rustnn/trtx-runtime"] diff --git a/pyproject.toml b/pyproject.toml index ed7b38a..25ae986 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,4 +46,4 @@ Issues = "https://github.com/rustnn/pywebnn/issues" module-name = "webnn._rustnn" python-source = "python" # CoreML is opt-in because it requires Apple's Objective-C runtime. -features = ["pyo3/extension-module", "onnx-runtime", "dynamic-inputs"] +features = ["pyo3/extension-module", "onnx-runtime", "trtx-runtime", "dynamic-inputs"] From c9815b9d0468cb8d346efb5a32fea47c9f4e251a Mon Sep 17 00:00:00 2001 From: Stephan Seitz Date: Wed, 7 Oct 2026 12:59:57 +0200 Subject: [PATCH 2/3] deps: add pyproject extra `trtx` which installs `tensorrt_rtx_cu13_libs` --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 25ae986..89f6ec5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,9 @@ dependencies = [ ] [project.optional-dependencies] +trtx = [ + "tensorrt_rtx_cu13_libs==1.6.1.120", +] test = [ "pytest>=7.0", "pytest-asyncio", From 2d2ff8e3ec9b482b41bb8905ec16f2915188d3b8 Mon Sep 17 00:00:00 2001 From: Stephan Seitz Date: Wed, 7 Oct 2026 15:04:37 +0200 Subject: [PATCH 3/3] feat: load TRT libs from pip package if available --- Cargo.lock | 23 ++++------------------- Cargo.toml | 3 ++- src/python/context_state.rs | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 41 insertions(+), 20 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bcf4a1a..6d6566f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1571,21 +1571,6 @@ dependencies = [ "tempfile", ] -[[package]] -name = "ndarray" -version = "0.16.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "882ed72dce9365842bf196bdeedf5055305f11fc8c03dee7bb0194a6cad34841" -dependencies = [ - "matrixmultiply", - "num-complex", - "num-integer", - "num-traits", - "portable-atomic", - "portable-atomic-util", - "rawpointer", -] - [[package]] name = "ndarray" version = "0.17.2" @@ -1732,7 +1717,7 @@ checksum = "4336a1e2b38848325241c72889086886004e589b7c74f335e60a8e8db5138a0b" dependencies = [ "half", "libloading 0.9.0", - "ndarray 0.17.2", + "ndarray", "ort-sys", "smallvec", "tracing", @@ -2074,7 +2059,7 @@ version = "0.5.13" dependencies = [ "bytemuck", "half", - "ndarray 0.16.1", + "log", "ort", "pretty_env_logger", "pyo3", @@ -2301,7 +2286,7 @@ dependencies = [ [[package]] name = "rustnn" version = "0.5.11" -source = "git+https://github.com/rustnn/rustnn?branch=main#5f1a23b6591bf932383895aa656b08605835a0a2" +source = "git+https://github.com/theHamsta/rustnn?branch=bring-dynamically_load_tensorrt-back#818e79d225b15d72b3ff3a7c03ad2edfcf18b3c5" dependencies = [ "base64 0.23.1", "block", @@ -2317,7 +2302,7 @@ dependencies = [ "litert", "litert-sys", "log", - "ndarray 0.17.2", + "ndarray", "objc", "ort", "prost", diff --git a/Cargo.toml b/Cargo.toml index 57fbe2f..166d7ea 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,12 +26,13 @@ dynamic-inputs = ["rustnn/dynamic-inputs"] [dependencies] pyo3 = { version = "0.28", features = ["extension-module"] } bytemuck = "1.14" -rustnn = { git = "https://github.com/rustnn/rustnn", branch = "main" } +rustnn = { git = "https://github.com/theHamsta/rustnn", branch = "bring-dynamically_load_tensorrt-back" } serde_json = "1.0" webnn-graph = { git = "https://github.com/rustnn/webnn-graph", branch = "main" } half = "2.4" regex = "1.11" pretty_env_logger = "0.5" +log = "0.4" # Optional runtime dependencies ort = { version = "=2.0.0-rc.13", optional = true, features = ["half", "load-dynamic"] } diff --git a/src/python/context_state.rs b/src/python/context_state.rs index 66e40e5..ee52086 100644 --- a/src/python/context_state.rs +++ b/src/python/context_state.rs @@ -1,5 +1,7 @@ //! Internal rustnn `MLContext` state shared by pywebnn bindings. +#[cfg(feature = "trtx-runtime")] +use log::warn; use pyo3::exceptions::{PyRuntimeError, PyValueError}; use pyo3::prelude::*; use pyo3::types::{PyBytes, PyDict, PyTuple}; @@ -39,6 +41,35 @@ fn init_rust_logging_if_requested() { } } +/// Prefer the compatible optional Python runtime package over library discovery. +#[cfg(feature = "trtx-runtime")] +fn try_setting_tensorrt_libs(py: Python<'_>) -> PyResult<()> { + let metadata = py.import("importlib.metadata")?; + let version = match metadata.call_method1("version", ("tensorrt_rtx_cu13_libs",)) { + Ok(version) => version.extract::()?, + Err(err) if err.matches(py, &metadata.getattr("PackageNotFoundError")?)? => return Ok(()), + Err(err) => return Err(err), + }; + if version != "1.6.1.120" { + return Ok(()); + } + + let libs = py.import("tensorrt_rtx_libs")?; + let mut library_path: std::path::PathBuf = libs.getattr("__file__")?.extract()?; + let filename = if cfg!(target_os = "windows") { + "tensorrt_rtx_1_6.dll" + } else { + "libtensorrt_rtx.so.1" + }; + library_path.set_file_name(filename); + rustnn::backends::trtx::dynamically_load_tensorrt(Some(&library_path)).map_err(|err| { + PyRuntimeError::new_err(format!( + "Failed to load the packaged TensorRT RTX runtime at {}: {err}", + library_path.display() + )) + }) +} + pub(crate) fn map_rustnn_error(err: rustnn::error::Error) -> PyErr { PyRuntimeError::new_err(err.to_string()) } @@ -211,6 +242,10 @@ pub(crate) fn ml_tensor_descriptor( impl ContextState { pub fn new(options: MLContextOptions) -> PyResult { init_rust_logging_if_requested(); + #[cfg(feature = "trtx-runtime")] + let _ = Python::attach(try_setting_tensorrt_libs) + .map_err(|err| warn!("Failed to load TensorRT library from pip package: {err:?}")); + // SAFETY: `MLContext` and backend builders do not expose references tied to caller // stack frames; storing in this struct is the intended ownership model for pywebnn. let ml_context = MLContext::create(&options).map_err(map_rustnn_error)?;