feat: auto-select TensorRT on linux aarch64

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

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"
     ```

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]

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` 有实现和单测 | 协议定义不是短板 |

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<OrtExecutionProvid
         "coreml" => Some(OrtExecutionProvider::CoreML),
         "openvino" => Some(OrtExecutionProvider::OpenVINO),
         "tensorrt" => Some(OrtExecutionProvider::TensorRT),
+        "cuda" => Some(OrtExecutionProvider::CUDA),
         "cpu" => Some(OrtExecutionProvider::CPU),
         _ => None,
     }
@@ -65,25 +80,22 @@ 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()))?;
 
     let session_builder = match provider {
         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<SessionBuil
     ))
 }
 
+fn configure_cpu(session_builder: SessionBuilder) -> RbtResult<SessionBuilder> {
+    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<SessionBuilder, SessionBuilder> {
+    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<SessionBuilder> {
+    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<SessionBuilder> {
+    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<SessionBuilder> {
+    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() {