diff --git a/Cargo.toml b/Cargo.toml index 2e7488e..905a713 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,7 @@ tinympc-rs = { git = "https://github.com/3SE-Competitive-Robotics-Team/tinympc-r # for infer ndarray = "0.17.2" -ort = { version = "2.0.0-rc.12", features = ["tensorrt", "openvino", "coreml", "half"] } +ort = { version = "2.0.0-rc.12", features = ["tensorrt", "cuda", "openvino", "coreml", "half"] } half = "2.7.1" # error handle diff --git a/README.md b/README.md index 0d397d5..22b22dd 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ ## 🧠 特性 Highlights - 🏆 **面向比赛优化**：针对 `RoboMaster` 赛场需求深度定制，兼顾实时性与可靠性 -- 🦀 **全 Rust 实现**：零成本抽象，安全而强大。基于 `ort-rs` 实现跨平台端侧推理部署（支持几乎全部onnxruntime-EP，包括`TensorRT`, `OpenVino`, `RKNN`, `CANN` ... ） +- 🦀 **全 Rust 实现**：零成本抽象，安全而强大。基于 `ort-rs` 实现跨平台端侧推理部署（支持几乎全部onnxruntime-EP，包括`TensorRT`, `CUDA`, `OpenVino`, `RKNN`, `CANN` ... ） - 🚦 **多线程任务调度**：基于 `tokio` 打造全异步推理流，实时并发 - 🎯 **手搓 PnP 求解器（No OpenCV）**：高度特化 IPPE-PnP 模块，硬编码装甲板参数，省去平面化和各向同性归一化操作，解算速度超快 - ⚙️ **状态估计算法**：基于 `ESKF` 的状态估计器，内置装甲板选择与跟踪模块(Todo)，，使用双装甲板反向投影，机器人解算更加精确 @@ -62,6 +62,8 @@ 3. **设置环境变量** (每次打开终端都需要执行，建议自行持久化): ```bash # This sets the library path to find ONNX Runtime and OpenVINO native libraries. + # On linux aarch64, ort_ep = "auto" uses TensorRT when CUDA/TensorRT libraries are visible, + # CUDA when CUDA/cuDNN libraries are visible, otherwise CPU. Add your Jetson library paths here if needed. # Note: If your Python version is not 3.11, adjust the path accordingly. export LD_LIBRARY_PATH="$PWD/.venv/lib64/python3.11/site-packages/onnxruntime/capi:$PWD/.venv/lib64/python3.11/site-packages/openvino/libs:$LD_LIBRARY_PATH" ``` diff --git a/cfg/rbt_cfg.toml b/cfg/rbt_cfg.toml index 5192386..b1d8b2c 100644 --- a/cfg/rbt_cfg.toml +++ b/cfg/rbt_cfg.toml @@ -25,7 +25,9 @@ camera_img_width = 1280 camera_img_height = 720 infer_img_width = 640 infer_img_height = 640 -# auto: macOS -> CoreML, x86_64 Linux/Windows -> OpenVINO, other targets -> CPU. +# auto: macOS -> CoreML, linux aarch64 with CUDA + TensorRT -> TensorRT, +# linux aarch64 with CUDA + cuDNN -> CUDA, x86_64 Linux/Windows -> OpenVINO, +# other targets -> CPU. ort_ep = "auto" [detector_cfg.armor] diff --git a/docs/vivsionn-gap-checklist.md b/docs/vivsionn-gap-checklist.md index ffda047..e721783 100644 --- a/docs/vivsionn-gap-checklist.md +++ b/docs/vivsionn-gap-checklist.md @@ -18,7 +18,7 @@ This checklist tracks the functional gaps found while comparing this Rust worksp | [ ] P1 离线录制/回放 | 可录 `.avi + .csv`，强制 task mode 回放 | 缺主链路复盘工具 | 调现场问题很关键 | | [ ] P1 通信/MPC smoke 工具 | `testSerial`、`can_mpc_yaw_test` 工具链完整 | `comm_test` 基本空 | 接 CAN 后应尽快补 | | [ ] P2 显示/HUD/录制旁路 | MJPEG/Rerun/HUD/CSV/plot 脚本多 | 有 Rerun 和日志，但观测面较薄 | 影响调试效率 | -| [ ] P2 TRT/ROI/CUDA 性能路径 | TensorRT/CUDA 预处理更贴 Jetson | ORT + ONNX EP | 不是功能缺口，先看实测延迟 | +| [ ] P2 TRT/ROI/CUDA 性能路径 | TensorRT/CUDA 预处理更贴 Jetson | `auto` 在 linux aarch64 检测 CUDA + TensorRT runtime 后优先 TensorRT，只有 CUDA + cuDNN 时走 CUDA，否则 CPU fallback；其他平台继续走既有 CoreML/OpenVINO/CPU | TensorRT/CUDA EP 选择已接，ROI/CUDA 预处理仍未接 | | [x] 已基本对齐：模式路由 | `AutoShot`/`Outpost`/`Buff` 路由切换 | `RuntimeRouter`/`ModeContext` 已有 | 这块方向对 | | [x] 已基本对齐：yaw 发控主链 | yaw planner、二阶 MPC、shot phase | Rust 已迁主干 | 主要剩验证/调参 | | [x] 已基本对齐：CAN 协议格式 | `0x100`/`0x203`/`0x204` 协议 | `rbt_comm_frame.rs` 有实现和单测 | 协议定义不是短板 | diff --git a/lib/src/rbt_infra/rbt_ort_ep.rs b/lib/src/rbt_infra/rbt_ort_ep.rs index 3be93cd..6b11235 100644 --- a/lib/src/rbt_infra/rbt_ort_ep.rs +++ b/lib/src/rbt_infra/rbt_ort_ep.rs @@ -1,11 +1,12 @@ use crate::rbt_infra::rbt_err::{RbtError, RbtResult}; -use ort::{ep, session::builder::SessionBuilder}; +use ort::{ep, ep::ExecutionProviderDispatch, session::builder::SessionBuilder}; #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum OrtExecutionProvider { CoreML, OpenVINO, TensorRT, + CUDA, CPU, } @@ -15,6 +16,7 @@ impl OrtExecutionProvider { Self::CoreML => "CoreML", Self::OpenVINO => "OpenVINO", Self::TensorRT => "TensorRT", + Self::CUDA => "CUDA", Self::CPU => "CPU", } } @@ -29,6 +31,17 @@ fn default_execution_provider_for_target() -> OrtExecutionProvider { OrtExecutionProvider::CoreML } +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn default_execution_provider_for_target() -> OrtExecutionProvider { + if linux_aarch64_tensorrt_stack_available() { + OrtExecutionProvider::TensorRT + } else if linux_aarch64_cuda_stack_available() { + OrtExecutionProvider::CUDA + } else { + OrtExecutionProvider::CPU + } +} + #[cfg(all( not(target_os = "macos"), target_arch = "x86_64", @@ -40,6 +53,7 @@ fn default_execution_provider_for_target() -> OrtExecutionProvider { #[cfg(not(any( target_os = "macos", + all(target_os = "linux", target_arch = "aarch64"), all( target_arch = "x86_64", any(target_os = "linux", target_os = "windows") @@ -55,6 +69,7 @@ pub fn resolve_execution_provider(configured: &str) -> Option Some(OrtExecutionProvider::CoreML), "openvino" => Some(OrtExecutionProvider::OpenVINO), "tensorrt" => Some(OrtExecutionProvider::TensorRT), + "cuda" => Some(OrtExecutionProvider::CUDA), "cpu" => Some(OrtExecutionProvider::CPU), _ => None, } @@ -65,6 +80,11 @@ pub fn configure_session_builder( configured_ep: &str, engine_cache_path: &str, ) -> RbtResult<(SessionBuilder, OrtExecutionProvider)> { + let configured_ep = configured_ep.trim(); + if configured_ep.is_empty() || configured_ep.eq_ignore_ascii_case("auto") { + return configure_auto_session_builder(session_builder, engine_cache_path); + } + let provider = resolve_execution_provider(configured_ep) .ok_or_else(|| RbtError::UnsupportedExecutionProvider(configured_ep.to_string()))?; @@ -72,18 +92,10 @@ pub fn configure_session_builder( OrtExecutionProvider::CoreML => configure_coreml(session_builder, engine_cache_path)?, OrtExecutionProvider::OpenVINO => configure_openvino(session_builder)?, OrtExecutionProvider::TensorRT => { - session_builder.with_execution_providers([ep::TensorRTExecutionProvider::default() - .with_engine_cache(true) - .with_engine_cache_path(engine_cache_path) - .with_fp16(true) - .build() - .error_on_failure()])? - } - OrtExecutionProvider::CPU => { - session_builder.with_execution_providers([ep::CPUExecutionProvider::default() - .with_arena_allocator(true) - .build()])? + configure_tensorrt_strict(session_builder, engine_cache_path)? } + OrtExecutionProvider::CUDA => configure_cuda_strict(session_builder)?, + OrtExecutionProvider::CPU => configure_cpu(session_builder)?, }; Ok((session_builder, provider)) @@ -139,6 +151,205 @@ fn configure_openvino(_session_builder: SessionBuilder) -> RbtResult RbtResult { + Ok( + session_builder.with_execution_providers([ep::CPUExecutionProvider::default() + .with_arena_allocator(true) + .build()])?, + ) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn try_configure_execution_provider( + session_builder: SessionBuilder, + execution_provider: ExecutionProviderDispatch, +) -> Result { + match session_builder.with_execution_providers([execution_provider]) { + Ok(session_builder) => Ok(session_builder), + Err(err) => Err(err.recover()), + } +} + +fn tensorrt_execution_provider(engine_cache_path: &str) -> ExecutionProviderDispatch { + ep::TensorRTExecutionProvider::default() + .with_engine_cache(true) + .with_engine_cache_path(engine_cache_path) + .with_fp16(true) + .build() + .error_on_failure() +} + +#[cfg(any( + all( + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86_64") + ), + all(target_os = "windows", target_arch = "x86_64") +))] +fn cuda_execution_provider() -> ExecutionProviderDispatch { + ep::CUDAExecutionProvider::default() + .build() + .error_on_failure() +} + +fn configure_tensorrt_strict( + session_builder: SessionBuilder, + engine_cache_path: &str, +) -> RbtResult { + Ok(session_builder + .with_execution_providers([tensorrt_execution_provider(engine_cache_path)])?) +} + +#[cfg(any( + all( + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86_64") + ), + all(target_os = "windows", target_arch = "x86_64") +))] +fn configure_cuda_strict(session_builder: SessionBuilder) -> RbtResult { + Ok(session_builder.with_execution_providers([cuda_execution_provider()])?) +} + +#[cfg(not(any( + all( + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86_64") + ), + all(target_os = "windows", target_arch = "x86_64") +)))] +fn configure_cuda_strict(_session_builder: SessionBuilder) -> RbtResult { + Err(RbtError::UnsupportedExecutionProvider( + "CUDA is only enabled for aarch64/x86_64 Linux and x86_64 Windows targets".to_string(), + )) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn configure_auto_session_builder( + session_builder: SessionBuilder, + engine_cache_path: &str, +) -> RbtResult<(SessionBuilder, OrtExecutionProvider)> { + let mut session_builder = session_builder; + if linux_aarch64_tensorrt_stack_available() { + match try_configure_execution_provider( + session_builder, + tensorrt_execution_provider(engine_cache_path), + ) { + Ok(session_builder) => return Ok((session_builder, OrtExecutionProvider::TensorRT)), + Err(recovered) => { + session_builder = recovered; + } + } + } + + if linux_aarch64_cuda_stack_available() { + match try_configure_execution_provider(session_builder, cuda_execution_provider()) { + Ok(session_builder) => return Ok((session_builder, OrtExecutionProvider::CUDA)), + Err(recovered) => { + session_builder = recovered; + } + } + } + + Ok((configure_cpu(session_builder)?, OrtExecutionProvider::CPU)) +} + +#[cfg(not(all(target_os = "linux", target_arch = "aarch64")))] +fn configure_auto_session_builder( + session_builder: SessionBuilder, + engine_cache_path: &str, +) -> RbtResult<(SessionBuilder, OrtExecutionProvider)> { + let provider = default_execution_provider(); + let session_builder = match provider { + OrtExecutionProvider::CoreML => configure_coreml(session_builder, engine_cache_path)?, + OrtExecutionProvider::OpenVINO => configure_openvino(session_builder)?, + OrtExecutionProvider::TensorRT => { + configure_tensorrt_strict(session_builder, engine_cache_path)? + } + OrtExecutionProvider::CUDA => configure_cuda_strict(session_builder)?, + OrtExecutionProvider::CPU => configure_cpu(session_builder)?, + }; + Ok((session_builder, provider)) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn linux_aarch64_cuda_available() -> bool { + const CUDA_DIRS: &[&str] = &[ + "/usr/local/cuda", + "/usr/local/cuda/targets/aarch64-linux", + "/usr/lib/aarch64-linux-gnu/tegra", + ]; + const CUDA_LIB_DIRS: &[&str] = &[ + "/usr/lib/aarch64-linux-gnu", + "/usr/lib/aarch64-linux-gnu/tegra", + "/usr/local/cuda/lib64", + "/usr/local/cuda/targets/aarch64-linux/lib", + ]; + + CUDA_DIRS + .iter() + .any(|path| std::path::Path::new(path).exists()) + || CUDA_LIB_DIRS + .iter() + .any(|dir| directory_has_any_library(dir, &["libcuda.so", "libcudart.so"])) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn linux_aarch64_cudnn_available() -> bool { + const CUDNN_LIB_DIRS: &[&str] = &[ + "/usr/lib/aarch64-linux-gnu", + "/usr/lib/aarch64-linux-gnu/tegra", + "/usr/local/cuda/lib64", + "/usr/local/cuda/targets/aarch64-linux/lib", + ]; + + CUDNN_LIB_DIRS + .iter() + .any(|dir| directory_has_any_library(dir, &["libcudnn.so", "libcudnn.so.9"])) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn linux_aarch64_tensorrt_available() -> bool { + const TENSORRT_LIB_DIRS: &[&str] = &[ + "/usr/lib/aarch64-linux-gnu", + "/usr/lib/aarch64-linux-gnu/tegra", + "/usr/local/tensorrt/lib", + "/usr/local/TensorRT/lib", + "/opt/tensorrt/lib", + ]; + + TENSORRT_LIB_DIRS.iter().any(|dir| { + directory_has_any_library(dir, &["libnvinfer.so", "libnvinfer.so.10"]) + && directory_has_any_library(dir, &["libnvinfer_plugin.so", "libnvinfer_plugin.so.10"]) + }) +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn linux_aarch64_tensorrt_stack_available() -> bool { + linux_aarch64_cuda_stack_available() && linux_aarch64_tensorrt_available() +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn linux_aarch64_cuda_stack_available() -> bool { + linux_aarch64_cuda_available() && linux_aarch64_cudnn_available() +} + +#[cfg(all(target_os = "linux", target_arch = "aarch64"))] +fn directory_has_any_library(dir: &str, library_prefixes: &[&str]) -> bool { + let Ok(entries) = std::fs::read_dir(dir) else { + return false; + }; + entries.filter_map(Result::ok).any(|entry| { + let name = entry.file_name(); + let Some(name) = name.to_str() else { + return false; + }; + library_prefixes + .iter() + .any(|prefix| name == *prefix || name.starts_with(&format!("{prefix}."))) + }) +} + #[cfg(test)] mod tests { use super::*; @@ -169,12 +380,26 @@ mod tests { resolve_execution_provider("TensorRT"), Some(OrtExecutionProvider::TensorRT) ); + assert_eq!( + resolve_execution_provider("cuda"), + Some(OrtExecutionProvider::CUDA) + ); assert_eq!( resolve_execution_provider("cpu"), Some(OrtExecutionProvider::CPU) ); } + #[cfg(all(target_os = "linux", target_arch = "aarch64"))] + #[test] + fn linux_aarch64_auto_uses_accelerator_or_cpu() { + let provider = default_execution_provider(); + assert!(matches!( + provider, + OrtExecutionProvider::TensorRT | OrtExecutionProvider::CUDA | OrtExecutionProvider::CPU + )); + } + #[cfg(target_os = "macos")] #[test] fn macos_default_is_coreml() {