Enable trtx by default and load TensorRT from pip package if installed

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"]
@@ -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"] }
 

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",
@@ -46,4 +49,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"]

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::<String>()?,
+        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<Self> {
         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)?;