diff --git a/Cargo.toml b/Cargo.toml index df832b9..741fe1c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,10 +9,6 @@ license = "Apache-2.0" authors = ["Tarek Ziade", "Markus Tavenrath", "Stephan Seitz"] repository = "https://github.com/rustnn/trtx-rs" -[workspace.dependencies] -trtx-sys = { path = "trtx-sys" } -trtexec-rs = { path = "trtx" } - [profile.release] opt-level = 3 lto = true diff --git a/trtx/Cargo.toml b/trtx/Cargo.toml index d775264..83c93ad 100644 --- a/trtx/Cargo.toml +++ b/trtx/Cargo.toml @@ -14,13 +14,9 @@ categories = ["api-bindings", "science"] trtx-sys = { version = "0.8.1", path = "../trtx-sys", default-features = false } thiserror = "2.0" cxx = "1.0" -libc = "0.2" libloading = { version = "0.9", optional = true } autocxx = { version = "0.30" } -# cudarc for safe CUDA operations (required when real mode is enabled) -# Using cuda-12050 as fallback; CUDA 13.x should be compatible -cudarc = { version = "0.11", features = ["driver", "cuda-12050"] } log = "0.4.29" [features] @@ -48,6 +44,7 @@ v_1_6 = ["trtx-sys/v_1_6", "v_1_5"] enterprise = ["trtx-sys/enterprise"] [dev-dependencies] +cudarc = { version = "0.19.9", features = ["cuda-13020"] } # For examples and tests pretty_env_logger = "0.5" @@ -58,7 +55,3 @@ required-features = [] [[example]] name = "basic_workflow" required-features = ["onnxparser"] - -[[example]] -name = "rustnn_executor" -required-features = ["onnxparser"] diff --git a/trtx/examples/rustnn_executor.rs b/trtx/examples/rustnn_executor.rs deleted file mode 100644 index 97bc586..0000000 --- a/trtx/examples/rustnn_executor.rs +++ /dev/null @@ -1,100 +0,0 @@ -//! Example demonstrating rustnn-compatible executor API -//! -//! This example shows how to use trtx as an executor in the rustnn pattern: -//! - Load ONNX model bytes -//! - Execute with TensorRT -//! - Get output tensors -//! -//! Run with: cargo run --features mock --example rustnn_executor - -use std::error::Error; -use trtx::executor::{run_onnx_with_tensorrt, run_onnx_zeroed, TensorInput}; - -fn main() -> Result<(), Box> { - #[cfg(feature = "dlopen_tensorrt_rtx")] - trtx::dynamically_load_tensorrt(None::).unwrap(); - - println!("TensorRT-RTX Executor for rustnn"); - println!("==================================\n"); - - // Example 1: Execute with zero-filled inputs (for testing) - println!("1. Testing with zero-filled inputs..."); - - let dummy_onnx = create_dummy_onnx_model(); - - let input_descriptors = vec![("input".to_string(), vec![1, 3, 224, 224])]; - - match run_onnx_zeroed(&dummy_onnx, &input_descriptors) { - Ok(outputs) => { - println!(" ✓ Execution succeeded"); - println!(" Outputs:"); - for output in outputs { - println!( - " - {}: shape {:?}, {} values", - output.name, - output.shape, - output.data.len() - ); - } - } - Err(e) => { - println!(" ✗ Execution failed: {}", e); - println!(" (This is expected with a dummy model in mock mode)"); - } - } - - // Example 2: Execute with actual input data - println!("\n2. Testing with actual input data..."); - - let inputs = vec![TensorInput { - name: "input".to_string(), - shape: vec![1, 3, 224, 224], - data: create_sample_input(3 * 224 * 224), - }]; - - match run_onnx_with_tensorrt(&dummy_onnx, &inputs) { - Ok(outputs) => { - println!(" ✓ Execution succeeded"); - for output in outputs { - println!(" - {}: shape {:?}", output.name, output.shape); - println!( - " First 5 values: {:?}", - &output.data[..output.data.len().min(5)] - ); - } - } - Err(e) => { - println!(" ✗ Execution failed: {}", e); - println!(" (Expected with dummy model - use real ONNX for actual inference)"); - } - } - - println!("\n3. rustnn Integration Pattern"); - println!(" To use in rustnn, implement:"); - println!(" ```rust"); - println!(" #[cfg(feature = \"trtx-runtime\")]"); - println!(" pub fn run_trtx_with_inputs("); - println!(" model_bytes: &[u8],"); - println!(" inputs: &[TrtxInput],"); - println!(" ) -> Result> {{"); - println!(" trtx::run_onnx_with_tensorrt(model_bytes, inputs)"); - println!(" }}"); - println!(" ```"); - - println!("\n✓ Example completed"); - - Ok(()) -} - -/// Create a dummy ONNX model for demonstration -/// In real usage, you would load an actual ONNX file -fn create_dummy_onnx_model() -> Vec { - // This is just a placeholder - in real usage, load from file: - // std::fs::read("model.onnx")? - vec![0u8; 100] -} - -/// Create sample input data (random-ish values) -fn create_sample_input(size: usize) -> Vec { - (0..size).map(|i| (i as f32 * 0.001).sin()).collect() -} diff --git a/trtx/examples/tiny_network.rs b/trtx/examples/tiny_network.rs index 5eb4a42..bf7245f 100644 --- a/trtx/examples/tiny_network.rs +++ b/trtx/examples/tiny_network.rs @@ -9,12 +9,12 @@ //! Network architecture: //! Input [1, 3, 4, 4] -> ReLU -> Output [1, 3, 4, 4] +use cudarc::driver::{CudaContext, CudaSlice, DevicePtr, DevicePtrMut}; use trtx::builder::MemoryPoolType; -use trtx::cuda::{synchronize, DeviceBuffer}; use trtx::error::Result; use trtx::{ActivationType, Builder, DataType, Logger, Runtime}; -fn main() -> Result<()> { +fn main() -> std::result::Result<(), Box> { pretty_env_logger::init(); println!("=== Tiny Network Example ===\n"); @@ -67,45 +67,33 @@ fn main() -> Result<()> { println!(" Input shape: [1, 3, 4, 4] ({} elements)", input_size); println!(" First 8 input values: {:?}", &input_data[..8]); + let ctx = CudaContext::new(0)?; + let stream = ctx.new_stream()?; + // Allocate device memory - let mut input_device = DeviceBuffer::new(input_size * std::mem::size_of::())?; - let output_device = DeviceBuffer::new(output_size * std::mem::size_of::())?; - - // Copy input to device (convert f32 slice to bytes) - let input_bytes = unsafe { - std::slice::from_raw_parts( - input_data.as_ptr() as *const u8, - input_data.len() * std::mem::size_of::(), - ) - }; - input_device.copy_from_host(input_bytes)?; + let input_device = stream.clone_htod(&input_data)?; + let mut output_device: CudaSlice = stream.alloc_zeros(output_size)?; // 7. Set tensor addresses println!("\n7. Binding tensors..."); unsafe { - context.set_tensor_address("input", input_device.as_ptr())?; - context.set_tensor_address("output", output_device.as_ptr())?; + context + .set_input_tensor_address("input", input_device.device_ptr(&stream).0 as *const _)?; + context.set_tensor_address("output", output_device.device_ptr_mut(&stream).0 as *mut _)?; } // 8. Execute inference println!("8. Running inference..."); - let stream = trtx::cuda::default_stream(); unsafe { - context.enqueue_v3(stream)?; + context.enqueue_v3(stream.cu_stream() as *mut _)?; } - synchronize()?; + stream.synchronize()?; println!(" ✓ Inference completed"); // 9. Copy output back to host println!("\n9. Reading results..."); let mut output_data: Vec = vec![0.0; output_size]; - let output_bytes = unsafe { - std::slice::from_raw_parts_mut( - output_data.as_mut_ptr() as *mut u8, - output_data.len() * std::mem::size_of::(), - ) - }; - output_device.copy_to_host(output_bytes)?; + stream.memcpy_dtoh(&output_device, &mut output_data)?; println!(" Output shape: [1, 3, 4, 4] ({} elements)", output_size); println!(" First 8 output values: {:?}", &output_data[..8]); diff --git a/trtx/src/cuda.rs b/trtx/src/cuda.rs deleted file mode 100644 index effc1a9..0000000 --- a/trtx/src/cuda.rs +++ /dev/null @@ -1,111 +0,0 @@ -//! CUDA memory management utilities - -use crate::error::{Error, Result}; - -use cudarc::driver::{CudaDevice, CudaSlice, DevicePtr}; - -/// RAII wrapper for CUDA device memory -pub struct DeviceBuffer { - ptr: CudaSlice, - device: std::sync::Arc, - size: usize, -} - -impl DeviceBuffer { - pub fn new(size: usize) -> Result { - let device = CudaDevice::new(0) - .map_err(|e| Error::Cuda(format!("Failed to initialize CUDA device: {:?}", e)))?; - let ptr = device - .alloc_zeros::(size) - .map_err(|e| Error::Cuda(format!("Failed to allocate CUDA memory: {:?}", e)))?; - Ok(DeviceBuffer { ptr, device, size }) - } - - pub fn as_ptr(&self) -> *mut std::ffi::c_void { - *self.ptr.device_ptr() as *mut std::ffi::c_void - } - - pub fn size(&self) -> usize { - self.size - } - - pub fn copy_from_host(&mut self, data: &[u8]) -> Result<()> { - if data.len() > self.size { - return Err(Error::InvalidArgument( - "Data size exceeds buffer size".to_string(), - )); - } - self.device - .htod_copy_into(data.to_vec(), &mut self.ptr) - .map_err(|e| Error::Cuda(format!("Failed to copy to device: {:?}", e))) - } - - pub fn copy_to_host(&self, data: &mut [u8]) -> Result<()> { - if data.len() > self.size { - return Err(Error::InvalidArgument( - "Data size exceeds buffer size".to_string(), - )); - } - self.device - .dtoh_sync_copy_into(&self.ptr, data) - .map_err(|e| Error::Cuda(format!("Failed to copy from device: {:?}", e))) - } -} - -impl std::fmt::Debug for DeviceBuffer { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("DeviceBuffer") - .field("size", &self.size) - .finish_non_exhaustive() - } -} - -unsafe impl Send for DeviceBuffer {} - -/// Synchronize CUDA device -pub fn synchronize() -> Result<()> { - let device = CudaDevice::new(0) - .map_err(|e| Error::Cuda(format!("Failed to get CUDA device: {:?}", e)))?; - device - .synchronize() - .map_err(|e| Error::Cuda(format!("Failed to synchronize device: {:?}", e))) -} - -/// Get the default CUDA stream -pub fn default_stream() -> *mut std::ffi::c_void { - std::ptr::null_mut() -} - -#[deprecated = "use default_stream instead"] -pub fn get_default_stream() -> *mut std::ffi::c_void { - default_stream() -} - -#[cfg(test)] -#[cfg(not(feature = "mock_runtime"))] -mod tests { - use super::*; - - #[test] - fn test_device_buffer_allocation() { - let buffer = DeviceBuffer::new(1024); - assert!(buffer.is_ok()); - let buffer = buffer.unwrap(); - assert_eq!(buffer.size(), 1024); - } - - #[test] - fn test_device_buffer_copy() { - let mut buffer = DeviceBuffer::new(256).unwrap(); - let host_data = vec![42u8; 256]; - assert!(buffer.copy_from_host(&host_data).is_ok()); - let mut output = vec![0u8; 256]; - assert!(buffer.copy_to_host(&mut output).is_ok()); - assert_eq!(host_data, output); - } - - #[test] - fn test_synchronize() { - assert!(synchronize().is_ok()); - } -} diff --git a/trtx/src/execution_context.rs b/trtx/src/execution_context.rs index b1a7620..d73e5a0 100644 --- a/trtx/src/execution_context.rs +++ b/trtx/src/execution_context.rs @@ -563,9 +563,9 @@ impl<'engine> ExecutionContext<'engine> { #[cfg(test)] #[cfg(not(feature = "mock_runtime"))] mod tests { + use cudarc::driver::{CudaContext, CudaSlice, DevicePtr, DevicePtrMut}; + use crate::builder::{Builder, MemoryPoolType}; - use crate::cuda::{default_stream, synchronize, DeviceBuffer}; - use crate::error::Error; use crate::logger::Logger; use crate::{DataType, ElementWiseOperation, Runtime}; @@ -599,68 +599,36 @@ mod tests { .expect("execution context"); let elem_size = std::mem::size_of::(); - let mut a_buf = DeviceBuffer::new(elem_size).expect("buffer a"); - let mut b_buf = DeviceBuffer::new(elem_size).expect("buffer b"); - let c_buf = DeviceBuffer::new(elem_size).expect("buffer c"); - a_buf.copy_from_host(&2.0f32.to_le_bytes()).expect("copy a"); - b_buf.copy_from_host(&3.0f32.to_le_bytes()).expect("copy b"); + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.new_stream().unwrap(); + + let a_buf = stream.clone_htod(&2.0f32.to_le_bytes()).expect("copy a"); + let b_buf = stream.clone_htod(&3.0f32.to_le_bytes()).expect("copy b"); + let mut c_buf: CudaSlice = stream.alloc_zeros(elem_size).unwrap(); unsafe { context - .set_input_tensor_address("a", a_buf.as_ptr() as *const _) + .set_input_tensor_address("a", a_buf.device_ptr(&stream).0 as *const _) .expect("bind input a"); context - .set_input_tensor_address("b", b_buf.as_ptr() as *const _) + .set_input_tensor_address("b", b_buf.device_ptr(&stream).0 as *const _) .expect("bind input b"); context - .set_output_tensor_address("c", c_buf.as_ptr()) + .set_output_tensor_address("c", c_buf.device_ptr_mut(&stream).0 as *mut _) .expect("bind output c"); - context.enqueue_v3(default_stream()).expect("enqueue"); + context + .enqueue_v3(stream.cu_stream() as *mut _) + .expect("enqueue"); } - synchronize().expect("sync"); let mut out_bytes = [0u8; 4]; - c_buf.copy_to_host(&mut out_bytes).expect("copy c"); + stream.memcpy_dtoh(&c_buf, &mut out_bytes).unwrap(); + stream.synchronize().expect("sync"); + let c_val = f32::from_le_bytes(out_bytes); assert!( (c_val - 5.0f32).abs() < 1e-5, "expected a + b = 5.0, got {c_val}" ); } - - #[test] - fn set_input_output_tensor_address_invalid_name_fails() { - let logger = Logger::stderr().expect("logger"); - let engine_data = build_add_network(&logger).expect("build network"); - - let mut runtime = Runtime::new(&logger).expect("runtime"); - let mut engine = runtime - .deserialize_cuda_engine(&engine_data) - .expect("deserialize"); - let mut context = engine - .create_execution_context() - .expect("execution context"); - - let buf = DeviceBuffer::new(std::mem::size_of::()).expect("buffer"); - let ptr = buf.as_ptr(); - let const_ptr = ptr as *const std::ffi::c_void; - - unsafe { - let input_err = context - .set_input_tensor_address("not_a_tensor", const_ptr) - .unwrap_err(); - assert!( - matches!(input_err, Error::FailedToSetInputTensorAddress { .. }), - "unexpected error for invalid input name: {input_err:?}" - ); - - let output_err = context - .set_output_tensor_address("also_not_a_tensor", ptr) - .unwrap_err(); - assert!( - matches!(output_err, Error::FailedToSetOutputTensorAddress { .. }), - "unexpected error for invalid output name: {output_err:?}" - ); - } - } } diff --git a/trtx/src/executor.rs b/trtx/src/executor.rs deleted file mode 100644 index 1eb36d2..0000000 --- a/trtx/src/executor.rs +++ /dev/null @@ -1,253 +0,0 @@ -//! Executor module providing a rustnn-compatible interface. -//! -//! This module orchestrates [`crate::Builder`], [`crate::OnnxParser`], [`crate::Runtime`], and -//! [`crate::CudaEngine`], which wrap [`trtx_sys::nvinfer1::IBuilder`], [`trtx_sys::nvonnxparser::IParser`], -//! [`trtx_sys::nvinfer1::IRuntime`], and [`trtx_sys::nvinfer1::ICudaEngine`]. See the -//! [TensorRT for RTX C++ API](https://docs.nvidia.com/deeplearning/tensorrt-rtx/latest/_static/c-api/annotated.html). - -use crate::builder::network_flags; -use crate::error::Result; -use crate::{Builder, DeviceBuffer, OnnxParser}; -use crate::{Logger, Runtime}; - -/// Input descriptor for TensorRT execution -#[derive(Debug, Clone)] -pub struct TensorInput { - pub name: String, - pub shape: Vec, - pub data: Vec, -} - -/// Output descriptor from TensorRT execution -#[derive(Debug, Clone)] -pub struct TensorOutput { - pub name: String, - pub shape: Vec, - pub data: Vec, -} - -/// Execute an ONNX model with TensorRT using provided inputs -/// -/// This function follows the rustnn executor pattern: -/// 1. Parse ONNX model -/// 2. Build TensorRT engine -/// 3. Execute inference -/// 4. Return results -/// -/// # Arguments -/// -/// * `onnx_model_bytes` - ONNX model as byte slice -/// * `inputs` - Input tensors with names, shapes, and data -/// -/// # Returns -/// -/// Vector of output tensors with names, shapes, and computed data -pub fn run_onnx_with_tensorrt( - onnx_model_bytes: &[u8], - inputs: &[TensorInput], -) -> Result> { - // Create logger - let logger = Logger::stderr()?; - - // Build engine from ONNX - let engine_data = build_engine_from_onnx(&logger, onnx_model_bytes)?; - - // Execute inference - execute_engine(&logger, &engine_data, inputs) -} - -/// Build TensorRT engine from ONNX model -fn build_engine_from_onnx(logger: &Logger, onnx_bytes: &[u8]) -> Result> { - // Create builder - - use crate::builder::MemoryPoolType; - let mut builder = Builder::new(logger)?; - - // Create network with explicit batch - let network = builder.create_network(network_flags::EXPLICIT_BATCH)?; - - // Parse ONNX model - let mut parser = OnnxParser::new(network, logger)?; - parser.parse(onnx_bytes)?; - let network = parser.network_mut(); - - // Configure builder - let mut config = builder.create_config()?; - - // Set workspace memory (1GB) - config.set_memory_pool_limit(MemoryPoolType::kWORKSPACE, 1 << 30); - - // Build serialized engine - let memory = builder.build_serialized_network(network, &mut config)?; - - // This makes an extra copy since the `memory` depends on the lifetime of builder - Ok(memory.to_vec()) -} - -/// Execute TensorRT engine with inputs -fn execute_engine( - logger: &Logger, - engine_data: &[u8], - inputs: &[TensorInput], -) -> Result> { - // Create runtime and deserialize engine - let mut runtime = Runtime::new(logger)?; - let mut engine = runtime.deserialize_cuda_engine(engine_data)?; - let mut context = engine.create_execution_context()?; - - // Get tensor information - let num_tensors = engine.nb_io_tensors()?; - - // Prepare CUDA buffers for inputs and outputs - let mut device_buffers: Vec<(String, DeviceBuffer)> = Vec::new(); - let mut output_info: Vec<(String, Vec)> = Vec::new(); - - // Process each tensor - for i in 0..num_tensors { - let name = engine.io_tensor_name(i)?; - - // Check if this is an input or output - if let Some(input) = inputs.iter().find(|inp| inp.name == name) { - // Input tensor - validate shape matches engine expectations - let expected_shape_i64 = engine.tensor_shape(&name)?; - let expected_shape: Vec = - expected_shape_i64.iter().map(|&d| d as usize).collect(); - let expected_elements: usize = expected_shape.iter().product(); - let provided_elements: usize = input.shape.iter().product(); - - if provided_elements != expected_elements { - return Err(crate::Error::InvalidArgument(format!( - "Input tensor '{}' shape mismatch: expected {:?} ({} elements), got {:?} ({} elements)", - name, expected_shape, expected_elements, input.shape, provided_elements - ))); - } - - // Validate data length matches shape - if input.data.len() != provided_elements { - return Err(crate::Error::InvalidArgument(format!( - "Input tensor '{}' data length ({}) doesn't match shape {:?} ({} elements)", - name, - input.data.len(), - input.shape, - provided_elements - ))); - } - - // Allocate and copy data - let size_bytes = input.data.len() * std::mem::size_of::(); - let mut buffer = DeviceBuffer::new(size_bytes)?; - - // Copy input data to device - let input_bytes = - unsafe { std::slice::from_raw_parts(input.data.as_ptr() as *const u8, size_bytes) }; - buffer.copy_from_host(input_bytes)?; - - // Bind tensor address - unsafe { - context.set_tensor_address(&name, buffer.as_ptr())?; - } - - device_buffers.push((name.clone(), buffer)); - } else { - // Output tensor - query actual shape from engine - let shape_i64 = engine.tensor_shape(&name)?; - let shape: Vec = shape_i64.iter().map(|&d| d as usize).collect(); - - // Calculate actual buffer size needed - let num_elements: usize = shape.iter().product(); - let size_bytes = num_elements * std::mem::size_of::(); - let buffer = DeviceBuffer::new(size_bytes)?; - - unsafe { - context.set_tensor_address(&name, buffer.as_ptr())?; - } - - output_info.push((name.clone(), shape)); - device_buffers.push((name.clone(), buffer)); - } - } - - // Execute inference - unsafe { - context.enqueue_v3(crate::cuda::default_stream())?; - } - - // Synchronize to ensure completion - crate::cuda::synchronize()?; - - // Copy outputs back to host - let mut outputs = Vec::new(); - - for (name, shape) in output_info { - if let Some((_, buffer)) = device_buffers.iter().find(|(n, _)| n == &name) { - let size_bytes = shape.iter().product::() * std::mem::size_of::(); - let mut host_data: Vec = vec![0u8; size_bytes]; - - buffer.copy_to_host(host_data.as_mut_slice())?; - - // Convert bytes to f32 - let data: Vec = unsafe { - std::slice::from_raw_parts( - host_data.as_ptr() as *const f32, - size_bytes / std::mem::size_of::(), - ) - } - .to_vec(); - - outputs.push(TensorOutput { name, shape, data }); - } - } - - Ok(outputs) -} - -/// Simpler version: Execute with zero-filled inputs (useful for testing/validation) -pub fn run_onnx_zeroed( - onnx_model_bytes: &[u8], - input_descriptors: &[(String, Vec)], -) -> Result> { - // Create zero-filled inputs - let inputs: Vec = input_descriptors - .iter() - .map(|(name, shape)| { - let size: usize = shape.iter().product(); - TensorInput { - name: name.clone(), - shape: shape.clone(), - data: vec![0.0; size], - } - }) - .collect(); - - run_onnx_with_tensorrt(onnx_model_bytes, &inputs) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_tensor_input_creation() { - let input = TensorInput { - name: "input".to_string(), - shape: vec![1, 3, 224, 224], - data: vec![0.0; 3 * 224 * 224], - }; - - assert_eq!(input.name, "input"); - assert_eq!(input.shape, vec![1, 3, 224, 224]); - assert_eq!(input.data.len(), 3 * 224 * 224); - } - - #[test] - #[ignore] // Requires valid ONNX model - fn test_executor_basic() { - let dummy_onnx = vec![0u8; 100]; - let inputs = vec![("input".to_string(), vec![1, 3, 224, 224])]; - - let _result = run_onnx_zeroed(&dummy_onnx, &inputs); - // In mock mode, this should succeed - #[cfg(feature = "mock_runtime")] - assert!(_result.is_ok()); - } -} diff --git a/trtx/src/lib.rs b/trtx/src/lib.rs index 56b4c0b..1bb949a 100644 --- a/trtx/src/lib.rs +++ b/trtx/src/lib.rs @@ -135,13 +135,10 @@ pub mod axes; pub mod builder; pub mod builder_config; -pub mod cuda; pub mod cuda_engine; pub mod engine_inspector; pub mod error; pub mod execution_context; -#[cfg(feature = "onnxparser")] -pub mod executor; pub mod host_memory; pub mod interfaces; pub mod logger; @@ -159,12 +156,7 @@ pub mod tensor; // Re-export commonly used types pub use axes::Axes; pub use builder::{Builder, BuilderConfig}; -pub use cuda::{default_stream, synchronize, DeviceBuffer}; pub use error::{Error, Result}; -#[cfg(feature = "onnxparser")] -pub use executor::{run_onnx_with_tensorrt, run_onnx_zeroed}; -#[cfg(feature = "onnxparser")] -pub use executor::{TensorInput, TensorOutput}; #[cfg(feature = "dlopen_tensorrt_rtx")] use libloading::AsFilename; pub use logger::{LogHandler, Logger, Severity, StderrLogger}; diff --git a/trtx/src/runtime.rs b/trtx/src/runtime.rs index 7a71343..f6b6d04 100644 --- a/trtx/src/runtime.rs +++ b/trtx/src/runtime.rs @@ -243,10 +243,10 @@ impl<'runtime> Runtime<'runtime> { #[cfg(test)] #[cfg(not(feature = "mock_runtime"))] mod tests { + use std::ffi::c_void; use std::sync::{Arc, Mutex}; use crate::builder::{Builder, MemoryPoolType}; - use crate::cuda::{synchronize, DeviceBuffer}; use crate::interfaces::{ProcessDebugTensor, ProcessDebugTensorResult}; #[cfg(all(feature = "v_1_6", not(feature = "enterprise")))] use crate::interfaces::{ReadStreamV2, StreamReaderV2}; @@ -254,6 +254,7 @@ mod tests { #[cfg(all(feature = "v_1_6", not(feature = "enterprise")))] use crate::SeekPosition; use crate::{DataType, ElementWiseOperation, Runtime}; + use cudarc::driver::{CudaContext, CudaSlice, DevicePtrMut}; use trtx_sys::{Dims64, TensorLocation}; #[cfg(all(feature = "v_1_6", not(feature = "enterprise")))] @@ -288,7 +289,7 @@ mod tests { // recognized by cuPointerGetAttribute, so an attribute-query failure means host. let mut memory_type = 0_u32; let attribute_result = unsafe { - cudarc::driver::sys::lib().cuPointerGetAttribute( + cudarc::driver::sys::cuPointerGetAttribute( (&mut memory_type as *mut u32).cast(), cudarc::driver::sys::CUpointer_attribute::CU_POINTER_ATTRIBUTE_MEMORY_TYPE, destination as usize as cudarc::driver::sys::CUdeviceptr, @@ -489,6 +490,9 @@ mod tests { context.set_all_tensors_debug_state(true).unwrap(); context.set_unfused_tensors_debug_state(true).unwrap(); + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.new_stream().unwrap(); + // input: 1 channel 4x4, output: 4 channels 4x4 let input_elems = 4 * 4; let output_elems = 4 * 4 * 4; @@ -496,24 +500,31 @@ mod tests { let input_bytes: Vec = std::iter::repeat_n(1.0f32, input_elems) .flat_map(|v| v.to_le_bytes()) .collect(); - let mut input_device = DeviceBuffer::new(input_elems * elem_size).expect("input buffer"); - let output_device = DeviceBuffer::new(output_elems * elem_size).expect("output buffer"); - input_device - .copy_from_host(&input_bytes) - .expect("copy input"); + + let mut output_device: CudaSlice = stream + .alloc_zeros(output_elems * elem_size) + .expect("output buffer"); + + let mut input_device = stream.clone_htod(&input_bytes).expect("copy input"); unsafe { context - .set_tensor_address("input", input_device.as_ptr()) + .set_tensor_address( + "input", + input_device.device_ptr_mut(&stream).0 as *mut c_void, + ) .expect("set input"); context - .set_tensor_address("conv_out_2", output_device.as_ptr()) + .set_tensor_address( + "conv_out_2", + output_device.device_ptr_mut(&stream).0 as *mut c_void, + ) .expect("set output"); context - .enqueue_v3(crate::cuda::default_stream()) + .enqueue_v3(stream.cu_stream() as *mut c_void) .expect("enqueue"); } - synchronize().expect("sync"); + stream.synchronize().expect("sync"); let seen = seen.lock().unwrap(); assert!( @@ -551,29 +562,37 @@ mod tests { context.set_all_tensors_debug_state(true).unwrap(); context.set_unfused_tensors_debug_state(true).unwrap(); - let elem_size = std::mem::size_of::(); - let mut input_device = DeviceBuffer::new(elem_size).expect("input buffer"); - let output_device = DeviceBuffer::new(elem_size).expect("output buffer"); - input_device - .copy_from_host(&0.0f32.to_le_bytes()) + let ctx = CudaContext::new(0).expect("CUDA context"); + let stream = ctx.new_stream().expect("CUDA stream"); + let mut input_device = stream + .clone_htod(&0.0f32.to_le_bytes()) .expect("copy input"); + let mut output_device: CudaSlice = stream + .alloc_zeros(std::mem::size_of::()) + .expect("output buffer"); unsafe { context - .set_tensor_address("tensor_0", input_device.as_ptr()) + .set_tensor_address( + "tensor_0", + input_device.device_ptr_mut(&stream).0 as *mut c_void, + ) .expect("set input"); context - .set_tensor_address("tensor_4", output_device.as_ptr()) + .set_tensor_address( + "tensor_4", + output_device.device_ptr_mut(&stream).0 as *mut c_void, + ) .expect("set output"); context - .enqueue_v3(crate::cuda::default_stream()) + .enqueue_v3(stream.cu_stream() as *mut c_void) .expect("enqueue"); } - synchronize().expect("sync"); + stream.synchronize().expect("sync"); let mut output_bytes = [0u8; 4]; - output_device - .copy_to_host(&mut output_bytes) + stream + .memcpy_dtoh(&output_device, &mut output_bytes) .expect("copy output"); let output_val = f32::from_le_bytes(output_bytes); assert!( @@ -653,14 +672,12 @@ mod tests { })) .expect("stream reader"); - let cuda_device = cudarc::driver::CudaDevice::new(0).expect("CUDA device"); - let cuda_stream = cuda_device - .fork_default_stream() - .expect("non-blocking CUDA stream"); - let stream = cuda_stream.stream.cast(); + let ctx = CudaContext::new(0).unwrap(); + let stream = ctx.new_stream().unwrap(); + unsafe { engine - .load_weights_async(stream_reader.as_mut(), stream) + .load_weights_async(stream_reader.as_mut(), stream.cu_stream() as *mut _) .expect("load weights asynchronously"); } assert!(engine.weights_loaded()); @@ -673,10 +690,7 @@ mod tests { "TensorRT did not request any engine weights in device memory" ); - unsafe { - cudarc::driver::result::stream::synchronize(cuda_stream.stream) - .expect("synchronize asynchronous weight load"); - } + stream.synchronize().unwrap(); } #[test]