NVExecutionProvider

Author: decahedron1

Cargo.toml

@@ -84,6 +84,7 @@ cann = [ "ort-sys/cann" ]
 qnn = [ "ort-sys/qnn" ]
 webgpu = [ "ort-sys/webgpu" ]
 azure = [ "ort-sys/azure" ]
+nv = [ "ort-sys/nv" ]
 
 [dependencies]
 ort-sys = { version = "=2.0.0-rc.9", path = "ort-sys", default-features = false }

ort-sys/Cargo.toml

@@ -44,6 +44,7 @@ cann = []
 qnn = []
 webgpu = [ "dep:glob" ]
 azure = []
+nv = []
 
 [build-dependencies]
 ureq = { version = "3", optional = true, default-features = false, features = [ "native-tls", "socks-proxy" ] }

src/execution_providers/mod.rs

@@ -63,6 +63,8 @@ pub mod webgpu;
 pub use self::webgpu::WebGPUExecutionProvider;
 pub mod azure;
 pub use self::azure::AzureExecutionProvider;
+pub mod nv;
+pub use self::nv::NVExecutionProvider;
 
 /// ONNX Runtime works with different hardware acceleration libraries through its extensible **Execution Providers**
 /// (EP) framework to optimally execute the ONNX models on the hardware platform. This interface enables flexibility for

src/execution_providers/nv.rs

@@ -0,0 +1,52 @@
+use super::{ExecutionProvider, ExecutionProviderOptions, RegisterError};
+use crate::{error::Result, session::builder::SessionBuilder};
+
+#[derive(Debug, Default, Clone)]
+pub struct NVExecutionProvider {
+	options: ExecutionProviderOptions
+}
+
+super::impl_ep!(arbitrary; NVExecutionProvider);
+
+impl NVExecutionProvider {
+	pub fn with_device_id(mut self, device_id: u32) -> Self {
+		self.options.set("ep.nvtensorrtrtxexecutionprovider.device_id", device_id.to_string());
+		self
+	}
+
+	pub fn with_cuda_graph(mut self, enable: bool) -> Self {
+		self.options
+			.set("ep.nvtensorrtrtxexecutionprovider.nv_cuda_graph_enable", if enable { "1" } else { "0" });
+		self
+	}
+}
+
+impl ExecutionProvider for NVExecutionProvider {
+	fn as_str(&self) -> &'static str {
+		"NvTensorRTRTXExecutionProvider"
+	}
+
+	fn supported_by_platform(&self) -> bool {
+		cfg!(any(all(target_os = "linux", any(target_arch = "aarch64", target_arch = "x86_64")), all(target_os = "windows", target_arch = "x86_64")))
+	}
+
+	#[allow(unused, unreachable_code)]
+	fn register(&self, session_builder: &mut SessionBuilder) -> Result<(), RegisterError> {
+		#[cfg(any(feature = "load-dynamic", feature = "nv"))]
+		{
+			use crate::{AsPointer, ortsys};
+
+			let ffi_options = self.options.to_ffi();
+			ortsys![unsafe SessionOptionsAppendExecutionProvider(
+				session_builder.ptr_mut(),
+				c"NvTensorRtRtx".as_ptr().cast::<core::ffi::c_char>(),
+				ffi_options.key_ptrs(),
+				ffi_options.value_ptrs(),
+				ffi_options.len(),
+			)?];
+			return Ok(());
+		}
+
+		Err(RegisterError::MissingFeature)
+	}
+}

src/memory.rs

@@ -255,13 +255,12 @@ impl AllocationDevice {
 	pub const CANN: AllocationDevice = AllocationDevice("Cann\0");
 	pub const CANN_PINNED: AllocationDevice = AllocationDevice("CannPinned\0");
 	pub const DIRECTML: AllocationDevice = AllocationDevice("DML\0");
-	pub const DIRECTML_CPU: AllocationDevice = AllocationDevice("DML CPU\0");
 	pub const HIP: AllocationDevice = AllocationDevice("Hip\0");
 	pub const HIP_PINNED: AllocationDevice = AllocationDevice("HipPinned\0");
 	pub const OPENVINO_CPU: AllocationDevice = AllocationDevice("OpenVINO_CPU\0");
 	pub const OPENVINO_GPU: AllocationDevice = AllocationDevice("OpenVINO_GPU\0");
-	pub const XNNPACK: AllocationDevice = AllocationDevice("XnnpackExecutionProvider\0");
-	pub const TVM: AllocationDevice = AllocationDevice("TVM\0");
+	pub const QNN_HTP_SHARED: AllocationDevice = AllocationDevice("QnnHtpShared\0");
+	pub const WEBGPU_BUFFER: AllocationDevice = AllocationDevice("WebGPU_Buffer\0");
 
 	pub fn as_str(&self) -> &'static str {
 		&self.0[..self.0.len() - 1]

src/session/builder/editable.rs

@@ -1,4 +1,4 @@
-use alloc::sync::Arc;
+use alloc::{boxed::Box, sync::Arc};
 use core::{
 	ops::Deref,
 	ptr::{self, NonNull}

src/value/impl_tensor/copy.rs

@@ -43,7 +43,7 @@ fn ep_for_device(device: AllocationDevice, device_id: i32) -> Result<ep::Executi
 			.with_conv_max_workspace(false)
 			.with_conv_algorithm_search(ep::cuda::CuDNNConvAlgorithmSearch::Default)
 			.build(),
-		AllocationDevice::DIRECTML | AllocationDevice::DIRECTML_CPU => ep::DirectMLExecutionProvider::default().with_device_id(device_id).build(),
+		AllocationDevice::DIRECTML => ep::DirectMLExecutionProvider::default().with_device_id(device_id).build(),
 		AllocationDevice::CANN | AllocationDevice::CANN_PINNED => ep::CANNExecutionProvider::default()
 			.with_arena_extend_strategy(ep::ArenaExtendStrategy::SameAsRequested)
 			.with_cann_graph(false)
@@ -63,8 +63,6 @@ fn ep_for_device(device: AllocationDevice, device_id: i32) -> Result<ep::Executi
 			.with_exhaustive_conv_search(false)
 			.with_device_id(device_id)
 			.build(),
-		AllocationDevice::TVM => ep::TVMExecutionProvider::default().build(),
-		AllocationDevice::XNNPACK => ep::XNNPACKExecutionProvider::default().build(),
 		_ => return Err(crate::Error::new("Unsupported allocation device {device} for tensor copy target"))
 	})
 }