rave_tensorrt/

tensorrt.rs

//! TensorRT inference backend — ORT + TensorRTExecutionProvider + IO Binding.
//!
//! # Zero-copy contract
//!
//! Input and output tensors are bound to CUDA device pointers via ORT's
//! IO Binding API.  At no point does frame data touch host memory.
//!
//! # Execution provider policy
//!
//! **Only TensorRT EP is permitted.**  CPU EP is explicitly disabled.
//! After session creation, the provider list is validated.  If ORT falls
//! back to CPU for any graph node, `initialize()` returns an error.
//!
//! # CUDA stream ordering
//!
//! ORT creates its own internal CUDA stream for TensorRT EP execution.
//! We cannot inject `GpuContext::inference_stream` because `cudarc::CudaStream`
//! does not expose its raw `CUstream` handle (`pub(crate)` field).
//!
//! Correctness is maintained because `session.run_with_binding()` is
//! **synchronous** — ORT blocks the calling thread until all GPU kernels
//! on its internal stream complete.  Therefore:
//!
//! 1. Output buffer is fully written when `run_with_binding()` returns.
//! 2. CUDA global memory coherency guarantees visibility to any subsequent
//!    reader on any stream after this synchronization point.
//! 3. No additional inter-stream event is needed.
//!
//! # Output ring serialization
//!
//! `OutputRing` owns N pre-allocated device buffers.  `acquire()` checks
//! `Arc::strong_count == 1` before returning a slot, guaranteeing no
//! concurrent reader.  Ring size must be ≥ `downstream_channel_capacity + 2`.

#[cfg(target_os = "linux")]
use std::collections::HashSet;
use std::env;
#[cfg(target_os = "linux")]
use std::ffi::{CStr, CString, c_char, c_void};
#[cfg(target_os = "linux")]
use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, OnceLock};
#[cfg(target_os = "linux")]
use std::time::UNIX_EPOCH;

use async_trait::async_trait;
use cudarc::driver::{DevicePtr, DeviceSlice};
use tokio::sync::Mutex;
use tracing::{debug, info, warn};

use ort::session::Session;
use ort::sys as ort_sys;
use ort::value::Value as OrtValue;

use rave_core::backend::{ModelMetadata, UpscaleBackend};
use rave_core::context::GpuContext;
use rave_core::error::{EngineError, Result};
use rave_core::types::{GpuTexture, PixelFormat};

use ort::execution_providers::{CUDAExecutionProvider, TensorRTExecutionProvider};

// Helper to create tensor from device memory using C API (Zero-Copy)
unsafe fn create_tensor_from_device_memory(
    ptr: *mut std::ffi::c_void,
    bytes: usize,
    shape: &[i64],
    elem_type: ort::tensor::TensorElementType,
) -> Result<OrtValue> {
    let api = ort::api();

    // Create MemoryInfo for CUDA
    let mut mem_info_ptr: *mut ort_sys::OrtMemoryInfo = std::ptr::null_mut();
    let name = std::ffi::CString::new("Cuda").unwrap();
    let status = unsafe {
        (api.CreateMemoryInfo)(
            name.as_ptr(),
            ort_sys::OrtAllocatorType::OrtArenaAllocator,
            0, // device id
            ort_sys::OrtMemType::OrtMemTypeDefault,
            &mut mem_info_ptr,
        )
    };
    if !status.0.is_null() {
        unsafe { (api.ReleaseStatus)(status.0) };
        return Err(EngineError::Inference(ort::Error::new(
            "Failed to create MemoryInfo",
        )));
    }

    // Create Tensor
    let mut ort_value_ptr: *mut ort_sys::OrtValue = std::ptr::null_mut();
    // SAFETY: mem_info_ptr is a valid OrtMemoryInfo created above; ptr is a
    // caller-supplied device pointer of at least `bytes` bytes; shape is a
    // valid non-empty slice of i64 dimension values; ort_value_ptr is an
    // initialized null pointer used as the output parameter.
    let status = unsafe {
        (api.CreateTensorWithDataAsOrtValue)(
            mem_info_ptr,
            ptr,
            bytes as _,
            shape.as_ptr(),
            shape.len() as _,
            elem_type.into(),
            &mut ort_value_ptr,
        )
    };

    // ORT API contract: CreateTensorWithDataAsOrtValue does NOT take ownership
    // of mem_info — it copies what it needs internally. We release our handle
    // immediately after tensor creation, before any error check, as required.
    unsafe { (api.ReleaseMemoryInfo)(mem_info_ptr) };

    if !status.0.is_null() {
        unsafe { (api.ReleaseStatus)(status.0) };
        return Err(EngineError::Inference(ort::Error::new(
            "Failed to create Tensor",
        )));
    }

    // Wrap in OrtValue
    Ok(unsafe {
        ort::value::Value::<ort::value::DynValueTypeMarker>::from_ptr(
            std::ptr::NonNull::new(ort_value_ptr).unwrap(),
            None,
        )
    })
}

// ─── Precision policy ───────────────────────────────────────────────────────

/// TensorRT precision policy — controls EP optimization flags.
#[derive(Clone, Debug, Default)]
pub enum PrecisionPolicy {
    /// FP32 only — maximum accuracy, baseline performance.
    Fp32,
    /// FP16 mixed precision — 2× throughput on Tensor Cores.
    #[default]
    Fp16,
    /// INT8 quantized with calibration table — 4× throughput.
    /// Requires a pre-generated calibration table path.
    Int8 { calibration_table: PathBuf },
}

// ─── Batch config ──────────────────────────────────────────────────────────

/// Batch inference configuration.
#[derive(Clone, Debug)]
pub struct BatchConfig {
    /// Maximum batch size for pipelined inference.
    /// Must be ≤ model’s max dynamic batch axis.
    pub max_batch: usize,
    /// Collect at most this many frames before dispatching a batch,
    /// even if `max_batch` is not reached (latency bound).
    pub latency_deadline_us: u64,
}

impl Default for BatchConfig {
    fn default() -> Self {
        Self {
            max_batch: 1,
            latency_deadline_us: 8_000, // 8ms — half a 60fps frame
        }
    }
}

/// Validate a [`BatchConfig`], returning an error if `max_batch > 1`.
///
/// Micro-batching is not yet implemented; this enforces the `max_batch = 1` requirement.
pub fn validate_batch_config(cfg: &BatchConfig) -> Result<()> {
    if cfg.max_batch > 1 {
        return Err(EngineError::InvariantViolation(
            "micro-batching is not implemented; max_batch must be 1 (set max_batch=1)".into(),
        ));
    }
    Ok(())
}
// ─── Inference metrics ───────────────────────────────────────────────────────

/// Atomic counters for inference stage observability.
#[derive(Debug)]
pub struct InferenceMetrics {
    /// Total frames inferred.
    pub frames_inferred: AtomicU64,
    /// Cumulative inference time in microseconds (for avg latency).
    pub total_inference_us: AtomicU64,
    /// Peak single-frame inference time in microseconds.
    pub peak_inference_us: AtomicU64,
}

impl InferenceMetrics {
    /// Create a new zeroed [`InferenceMetrics`].
    pub const fn new() -> Self {
        Self {
            frames_inferred: AtomicU64::new(0),
            total_inference_us: AtomicU64::new(0),
            peak_inference_us: AtomicU64::new(0),
        }
    }

    /// Record a single frame's inference latency in microseconds.
    pub fn record(&self, elapsed_us: u64) {
        self.frames_inferred.fetch_add(1, Ordering::Relaxed);
        self.total_inference_us
            .fetch_add(elapsed_us, Ordering::Relaxed);
        self.peak_inference_us
            .fetch_max(elapsed_us, Ordering::Relaxed);
    }

    /// Return a point-in-time snapshot of all inference counters.
    pub fn snapshot(&self) -> InferenceMetricsSnapshot {
        let frames = self.frames_inferred.load(Ordering::Relaxed);
        let total = self.total_inference_us.load(Ordering::Relaxed);
        let peak = self.peak_inference_us.load(Ordering::Relaxed);
        InferenceMetricsSnapshot {
            frames_inferred: frames,
            avg_inference_us: if frames > 0 { total / frames } else { 0 },
            peak_inference_us: peak,
        }
    }
}

impl Default for InferenceMetrics {
    fn default() -> Self {
        Self::new()
    }
}

/// Snapshot of inference metrics for reporting.
#[derive(Clone, Debug)]
pub struct InferenceMetricsSnapshot {
    /// Total frames inferred.
    pub frames_inferred: u64,
    /// Average inference latency in microseconds.
    pub avg_inference_us: u64,
    /// Peak single-frame inference latency in microseconds.
    pub peak_inference_us: u64,
}

// ─── Ring metrics ────────────────────────────────────────────────────────────

/// A point-in-time snapshot of [`RingMetrics`] counters.
#[derive(Debug, Clone, Copy)]
pub struct RingMetricsSnapshot {
    /// Number of times a slot was acquired and was already free (strong_count == 1).
    pub reuse: u64,
    /// Number of times `acquire()` found a slot still held downstream (strong_count > 1).
    pub contention: u64,
    /// Number of times a slot was used for the very first time (not a reuse).
    pub first_use: u64,
}

/// Atomic counters for output ring buffer activity.
#[derive(Debug)]
pub struct RingMetrics {
    /// Successful slot reuses (slot was free, strong_count == 1).
    pub slot_reuse_count: AtomicU64,
    /// Times `acquire()` found a slot still held downstream (strong_count > 1).
    pub slot_contention_events: AtomicU64,
    /// Times a slot was acquired but it was the first use (not a reuse).
    pub slot_first_use_count: AtomicU64,
}

impl RingMetrics {
    /// Create a new zeroed [`RingMetrics`].
    pub const fn new() -> Self {
        Self {
            slot_reuse_count: AtomicU64::new(0),
            slot_contention_events: AtomicU64::new(0),
            slot_first_use_count: AtomicU64::new(0),
        }
    }

    /// Returns a point-in-time snapshot of all ring counters.
    pub fn snapshot(&self) -> RingMetricsSnapshot {
        RingMetricsSnapshot {
            reuse: self.slot_reuse_count.load(Ordering::Relaxed),
            contention: self.slot_contention_events.load(Ordering::Relaxed),
            first_use: self.slot_first_use_count.load(Ordering::Relaxed),
        }
    }
}

impl Default for RingMetrics {
    fn default() -> Self {
        Self::new()
    }
}

// ─── Output ring buffer ─────────────────────────────────────────────────────

/// Fixed-size ring of pre-allocated device buffers for inference output.
pub struct OutputRing {
    slots: Vec<Arc<cudarc::driver::CudaSlice<u8>>>,
    cursor: usize,
    /// Size of each slot in bytes (`3 × out_w × out_h × sizeof(f32)`).
    pub slot_bytes: usize,
    /// Input dimensions `(width, height)` used to allocate the current slots.
    pub alloc_dims: (u32, u32),
    /// Whether each slot has been used at least once (for first-use tracking).
    used: Vec<bool>,
    /// Atomic counters tracking reuse, contention, and first-use events.
    pub metrics: RingMetrics,
}

impl OutputRing {
    /// Allocate `count` output buffers.
    ///
    /// `min_slots` is the enforced minimum (`downstream_capacity + 2`).
    /// Returns error if `count < min_slots`.
    pub fn new(
        ctx: &GpuContext,
        in_w: u32,
        in_h: u32,
        scale: u32,
        count: usize,
        min_slots: usize,
    ) -> Result<Self> {
        if count < min_slots {
            return Err(EngineError::DimensionMismatch(format!(
                "OutputRing: ring_size ({count}) < required minimum ({min_slots}). \
                 Ring must be ≥ downstream_channel_capacity + 2."
            )));
        }
        if count < 2 {
            return Err(EngineError::DimensionMismatch(
                "OutputRing: ring_size must be ≥ 2 for double-buffering".into(),
            ));
        }

        let out_w = (in_w * scale) as usize;
        let out_h = (in_h * scale) as usize;
        let slot_bytes = 3 * out_w * out_h * std::mem::size_of::<f32>();

        let slots = (0..count)
            .map(|_| ctx.alloc(slot_bytes).map(Arc::new))
            .collect::<Result<Vec<_>>>()?;

        debug!(count, slot_bytes, out_w, out_h, "Output ring allocated");

        Ok(Self {
            slots,
            cursor: 0,
            slot_bytes,
            alloc_dims: (in_w, in_h),
            used: vec![false; count],
            metrics: RingMetrics::new(),
        })
    }

    /// Acquire the next ring slot for writing.
    ///
    /// # Serialization invariant
    ///
    /// Asserts `Arc::strong_count == 1` before returning.  If downstream
    /// still holds a reference, returns error and increments contention counter.
    pub fn acquire(&mut self) -> Result<Arc<cudarc::driver::CudaSlice<u8>>> {
        let slot = &self.slots[self.cursor];
        let sc = Arc::strong_count(slot);

        if sc != 1 {
            self.metrics
                .slot_contention_events
                .fetch_add(1, Ordering::Relaxed);
            return Err(EngineError::BufferTooSmall {
                need: self.slot_bytes,
                have: 0,
            });
        }

        // Debug assertion — belt-and-suspenders check.
        debug_assert_eq!(
            Arc::strong_count(slot),
            1,
            "OutputRing: slot {} strong_count must be 1 before reuse, got {}",
            self.cursor,
            sc
        );

        if self.used[self.cursor] {
            self.metrics
                .slot_reuse_count
                .fetch_add(1, Ordering::Relaxed);
        } else {
            self.used[self.cursor] = true;
            self.metrics
                .slot_first_use_count
                .fetch_add(1, Ordering::Relaxed);
        }

        let cloned = Arc::clone(slot);
        self.cursor = (self.cursor + 1) % self.slots.len();
        Ok(cloned)
    }

    /// Return `true` if the current slot dimensions differ from `(in_w, in_h)`.
    pub fn needs_realloc(&self, in_w: u32, in_h: u32) -> bool {
        self.alloc_dims != (in_w, in_h)
    }

    /// Reallocate all slots.  All must have `strong_count == 1`.
    pub fn reallocate(&mut self, ctx: &GpuContext, in_w: u32, in_h: u32, scale: u32) -> Result<()> {
        for (i, slot) in self.slots.iter().enumerate() {
            let sc = Arc::strong_count(slot);
            if sc != 1 {
                return Err(EngineError::DimensionMismatch(format!(
                    "Cannot reallocate ring: slot {} still in use (strong_count={})",
                    i, sc,
                )));
            }
        }

        // Free old slots — decrement VRAM accounting.
        for _slot in &self.slots {
            ctx.vram_dec(self.slot_bytes);
        }

        let count = self.slots.len();
        let out_w = (in_w * scale) as usize;
        let out_h = (in_h * scale) as usize;
        let slot_bytes = 3 * out_w * out_h * std::mem::size_of::<f32>();

        self.slots = (0..count)
            .map(|_| ctx.alloc(slot_bytes).map(Arc::new))
            .collect::<Result<Vec<_>>>()?;
        self.cursor = 0;
        self.slot_bytes = slot_bytes;
        self.alloc_dims = (in_w, in_h);
        self.used = vec![false; count];

        debug!(count, slot_bytes, out_w, out_h, "Output ring reallocated");
        Ok(())
    }

    /// Total number of slots.
    pub fn len(&self) -> usize {
        self.slots.len()
    }

    /// Return `true` if the ring has no slots (should not happen in normal operation).
    pub fn is_empty(&self) -> bool {
        self.slots.is_empty()
    }
}

// ─── Inference state ─────────────────────────────────────────────────────────

struct InferenceState {
    session: Session,
    ring: Option<OutputRing>,
}

/// Resolve ORT tensor element type from our PixelFormat.
fn ort_element_type(format: PixelFormat) -> ort::tensor::TensorElementType {
    match format {
        PixelFormat::RgbPlanarF16 => ort::tensor::TensorElementType::Float16,
        _ => ort::tensor::TensorElementType::Float32,
    }
}

// ─── Backend ─────────────────────────────────────────────────────────────────

/// TensorRT/CUDA ORT inference backend.
///
/// Implements [`UpscaleBackend`](rave_core::backend::UpscaleBackend) using
/// ONNX Runtime with a TensorRT or CUDA execution provider.  Output buffers
/// are managed via a fixed-size [`OutputRing`] to avoid per-frame allocation.
pub struct TensorRtBackend {
    model_path: PathBuf,
    ctx: Arc<GpuContext>,
    device_id: i32,
    ring_size: usize,
    min_ring_slots: usize,
    meta: OnceLock<ModelMetadata>,
    selected_provider: OnceLock<String>,
    state: Mutex<Option<InferenceState>>,
    /// Atomic inference latency and frame count metrics.
    pub inference_metrics: InferenceMetrics,
    /// Precision policy used when building the TensorRT EP session.
    pub precision_policy: PrecisionPolicy,
    /// Batch configuration (must have `max_batch = 1` — batching not yet implemented).
    pub batch_config: BatchConfig,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
enum OrtEpMode {
    Auto,
    TensorRtOnly,
    CudaOnly,
}

#[cfg(target_os = "linux")]
unsafe extern "C" {
    fn dlopen(filename: *const c_char, flags: i32) -> *mut c_void;
    fn dlerror() -> *const c_char;
}

#[cfg(target_os = "linux")]
const RTLD_NOW: i32 = 2;
#[cfg(all(target_os = "linux", test))]
const RTLD_LOCAL: i32 = 0;
#[cfg(target_os = "linux")]
const RTLD_GLOBAL: i32 = 0x100;

#[derive(Clone, Copy)]
enum OrtProviderKind {
    Cuda,
    TensorRt,
}

#[cfg(target_os = "linux")]
type ProviderCandidate = (PathBuf, &'static str);

impl OrtProviderKind {
    #[cfg(target_os = "linux")]
    fn soname(self) -> &'static str {
        match self {
            OrtProviderKind::Cuda => "libonnxruntime_providers_cuda.so",
            OrtProviderKind::TensorRt => "libonnxruntime_providers_tensorrt.so",
        }
    }

    #[cfg(target_os = "linux")]
    fn label(self) -> &'static str {
        match self {
            OrtProviderKind::Cuda => "providers_cuda",
            OrtProviderKind::TensorRt => "providers_tensorrt",
        }
    }
}

impl TensorRtBackend {
    /// Create a new backend instance.
    ///
    /// # Parameters
    ///
    /// - `ring_size`: number of output ring slots to pre-allocate.
    /// - `downstream_capacity`: the bounded channel capacity between inference
    ///   and the encoder.  Ring size is validated ≥ `downstream_capacity + 2`.
    pub fn new(
        model_path: PathBuf,
        ctx: Arc<GpuContext>,
        device_id: i32,
        ring_size: usize,
        downstream_capacity: usize,
    ) -> Self {
        Self::with_precision(
            model_path,
            ctx,
            device_id,
            ring_size,
            downstream_capacity,
            PrecisionPolicy::default(),
            BatchConfig::default(),
        )
    }

    /// Create with explicit precision policy and batch config.
    pub fn with_precision(
        model_path: PathBuf,
        ctx: Arc<GpuContext>,
        device_id: i32,
        ring_size: usize,
        downstream_capacity: usize,
        precision_policy: PrecisionPolicy,
        batch_config: BatchConfig,
    ) -> Self {
        let min_ring_slots = downstream_capacity + 2;
        assert!(
            ring_size >= min_ring_slots,
            "ring_size ({ring_size}) must be ≥ downstream_capacity + 2 ({min_ring_slots})"
        );
        Self {
            model_path,
            ctx,
            device_id,
            ring_size,
            min_ring_slots,
            meta: OnceLock::new(),
            selected_provider: OnceLock::new(),
            state: Mutex::new(None),
            inference_metrics: InferenceMetrics::new(),
            precision_policy,
            batch_config,
        }
    }

    /// Get or create cached ORT MemoryInfo (avoids per-frame allocation).
    /// Create a CUDA memory info structure.
    fn mem_info(&self) -> Result<ort::memory::MemoryInfo> {
        ort::memory::MemoryInfo::new(
            ort::memory::AllocationDevice::CUDA,
            0,
            ort::memory::AllocatorType::Device,
            ort::memory::MemoryType::Default,
        )
        .map_err(|e| EngineError::ModelMetadata(format!("MemoryInfo: {e}")))
    }

    /// Access ring metrics (if initialized).
    pub async fn ring_metrics(&self) -> Option<RingMetricsSnapshot> {
        let guard = self.state.lock().await;
        guard
            .as_ref()
            .and_then(|s| s.ring.as_ref())
            .map(|r| r.metrics.snapshot())
    }

    /// Active ORT execution provider selected during initialization.
    pub fn selected_provider(&self) -> Option<&str> {
        self.selected_provider.get().map(String::as_str)
    }

    fn infer_scale_from_name(name: &str) -> Option<u32> {
        let lower = name.to_ascii_lowercase();
        let bytes = lower.as_bytes();
        for i in 0..bytes.len() {
            if bytes[i] != b'x' {
                continue;
            }
            if i > 0 && bytes[i - 1].is_ascii_digit() {
                let d = (bytes[i - 1] - b'0') as u32;
                if d >= 2 {
                    return Some(d);
                }
            }
            if i + 1 < bytes.len() && bytes[i + 1].is_ascii_digit() {
                let d = (bytes[i + 1] - b'0') as u32;
                if d >= 2 {
                    return Some(d);
                }
            }
        }
        None
    }

    fn extract_metadata(session: &Session) -> Result<ModelMetadata> {
        let inputs = session.inputs();
        let outputs = session.outputs();

        if inputs.is_empty() || outputs.is_empty() {
            return Err(EngineError::ModelMetadata(
                "Model must have at least one input and one output tensor".into(),
            ));
        }

        // ... imports ...

        // ...

        // ... in extract_metadata ...
        // ...
        let input_info = &inputs[0];
        let output_info = &outputs[0];
        let input_name = input_info.name().to_string();
        let output_name = output_info.name().to_string();

        // ...

        let input_dims = match input_info.dtype() {
            ort::value::ValueType::Tensor { shape, .. } => shape.clone(),
            other => {
                return Err(EngineError::ModelMetadata(format!(
                    "Expected tensor input, got {:?}",
                    other
                )));
            }
        };

        let output_dims = match output_info.dtype() {
            ort::value::ValueType::Tensor { shape, .. } => shape.clone(),
            other => {
                return Err(EngineError::ModelMetadata(format!(
                    "Expected tensor output, got {:?}",
                    other
                )));
            }
        };

        if input_dims.len() != 4 || output_dims.len() != 4 {
            return Err(EngineError::ModelMetadata(format!(
                "Expected 4D tensors (NCHW), got input={}D output={}D",
                input_dims.len(),
                output_dims.len()
            )));
        }

        let input_channels = input_dims[1] as u32;

        let ih = input_dims[2];
        let iw = input_dims[3];
        let oh = output_dims[2];
        let ow = output_dims[3];

        let name = session
            .metadata()
            .map(|m| m.name().unwrap_or("unknown".to_string()))
            .unwrap_or_else(|_| "unknown".to_string());

        let scale = if ih > 0 && oh > 0 && iw > 0 && ow > 0 {
            (oh / ih) as u32
        } else if let Some(inferred) = Self::infer_scale_from_name(&name) {
            warn!(
                model_name = %name,
                inferred_scale = inferred,
                "Dynamic spatial axes — inferring upscale scale from model name"
            );
            inferred
        } else {
            warn!(
                model_name = %name,
                "Dynamic spatial axes — unable to infer scale from metadata; defaulting to scale=2"
            );
            2
        };

        let min_input_hw = (
            if ih > 0 { ih as u32 } else { 1 },
            if iw > 0 { iw as u32 } else { 1 },
        );

        let max_input_hw = (
            if ih > 0 { ih as u32 } else { u32::MAX },
            if iw > 0 { iw as u32 } else { u32::MAX },
        );

        Ok(ModelMetadata {
            name,
            scale,
            input_name,
            output_name,
            input_channels,
            min_input_hw,
            max_input_hw,
        })
    }

    fn ort_ep_mode() -> OrtEpMode {
        match env::var("RAVE_ORT_TENSORRT")
            .unwrap_or_else(|_| "auto".to_string())
            .to_lowercase()
            .as_str()
        {
            "0" | "off" | "false" | "cuda" | "cuda-only" => OrtEpMode::CudaOnly,
            "1" | "on" | "true" | "trt" | "trt-only" => OrtEpMode::TensorRtOnly,
            _ => OrtEpMode::Auto,
        }
    }

    #[cfg(target_os = "linux")]
    fn is_wsl2() -> bool {
        std::fs::read_to_string("/proc/sys/kernel/osrelease")
            .map(|s| s.to_ascii_lowercase().contains("microsoft"))
            .unwrap_or(false)
    }

    #[cfg(not(target_os = "linux"))]
    fn is_wsl2() -> bool {
        false
    }

    #[cfg(target_os = "linux")]
    fn ort_provider_cache_dirs_newest_first() -> Vec<PathBuf> {
        let mut dirs = Vec::<(u128, PathBuf)>::new();
        let Some(home) = env::var_os("HOME") else {
            return Vec::new();
        };
        let base = PathBuf::from(home).join(".cache/ort.pyke.io/dfbin");
        let Ok(triples) = std::fs::read_dir(base) else {
            return Vec::new();
        };

        for triple in triples.flatten() {
            let triple_path = triple.path();
            if !triple_path.is_dir() {
                continue;
            }
            let Ok(hashes) = std::fs::read_dir(triple_path) else {
                continue;
            };
            for hash in hashes.flatten() {
                let path = hash.path();
                if !path.is_dir() {
                    continue;
                }
                let modified = std::fs::metadata(&path)
                    .and_then(|m| m.modified())
                    .ok()
                    .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
                    .map(|d| d.as_nanos())
                    .unwrap_or(0);
                dirs.push((modified, path));
            }
        }

        dirs.sort_by(|a, b| b.cmp(a));
        dirs.into_iter().map(|(_, path)| path).collect()
    }

    #[cfg(target_os = "linux")]
    fn provider_dir_candidates(
        ort_dylib_path: Option<PathBuf>,
        ort_lib_location: Option<PathBuf>,
        exe_dir: Option<PathBuf>,
        cache_dirs: Vec<PathBuf>,
        on_wsl: bool,
    ) -> Vec<ProviderCandidate> {
        let mut dirs = Vec::<ProviderCandidate>::new();
        if let Some(dir) = ort_dylib_path {
            dirs.push((dir, "ORT_DYLIB_PATH"));
        }
        if let Some(dir) = ort_lib_location {
            dirs.push((dir, "ORT_LIB_LOCATION"));
        }
        if on_wsl {
            for dir in cache_dirs {
                dirs.push((dir, "ort_cache_newest"));
            }
            if let Some(dir) = exe_dir {
                dirs.push((dir.clone(), "exe_dir"));
                dirs.push((dir.join("deps"), "exe_dir/deps"));
            }
        } else {
            if let Some(dir) = exe_dir {
                dirs.push((dir.clone(), "exe_dir"));
                dirs.push((dir.join("deps"), "exe_dir/deps"));
            }
            for dir in cache_dirs {
                dirs.push((dir, "ort_cache_newest"));
            }
        }

        let mut uniq = HashSet::<PathBuf>::new();
        dirs.retain(|(p, _)| uniq.insert(p.clone()));
        dirs
    }

    #[cfg(target_os = "linux")]
    fn ort_provider_search_inputs() -> (
        Option<PathBuf>,
        Option<PathBuf>,
        Option<PathBuf>,
        Vec<PathBuf>,
        bool,
    ) {
        let ort_dylib_path = env::var_os("ORT_DYLIB_PATH").map(PathBuf::from);
        let ort_lib_location = env::var_os("ORT_LIB_LOCATION").map(PathBuf::from);
        let exe_dir = env::current_exe()
            .ok()
            .and_then(|exe| exe.parent().map(|dir| dir.to_path_buf()));
        let cache_dirs = Self::ort_provider_cache_dirs_newest_first();
        let on_wsl = Self::is_wsl2();
        (
            ort_dylib_path,
            ort_lib_location,
            exe_dir,
            cache_dirs,
            on_wsl,
        )
    }

    #[cfg(all(target_os = "linux", test))]
    fn ort_provider_search_dirs() -> Vec<ProviderCandidate> {
        let (ort_dylib_path, ort_lib_location, exe_dir, cache_dirs, on_wsl) =
            Self::ort_provider_search_inputs();
        Self::provider_dir_candidates(
            ort_dylib_path,
            ort_lib_location,
            exe_dir,
            cache_dirs,
            on_wsl,
        )
    }

    #[cfg(target_os = "linux")]
    fn render_provider_candidates(candidates: &[ProviderCandidate]) -> String {
        if candidates.is_empty() {
            return "none".into();
        }
        candidates
            .iter()
            .enumerate()
            .map(|(idx, (dir, source))| format!("{}:{source}:{}", idx + 1, dir.display()))
            .collect::<Vec<_>>()
            .join(", ")
    }

    #[cfg(all(target_os = "linux", test))]
    fn ort_provider_candidates(lib_name: &str) -> Vec<PathBuf> {
        Self::ort_provider_search_dirs()
            .into_iter()
            .map(|(dir, _)| dir.join(lib_name))
            .collect()
    }

    #[cfg(target_os = "linux")]
    fn configure_ort_loader_path(dir: &Path) {
        // SAFETY: Process env mutation is done at startup during backend init.
        unsafe { env::set_var("ORT_DYLIB_PATH", dir) };
        // SAFETY: Process env mutation is done at startup during backend init.
        unsafe { env::set_var("ORT_LIB_LOCATION", dir) };
    }

    #[cfg(target_os = "linux")]
    fn resolve_provider_dir(
        kind: OrtProviderKind,
        candidates: &[ProviderCandidate],
        ort_dylib_path: Option<&Path>,
        ort_lib_location: Option<&Path>,
        command: &str,
        mode: &str,
        ep_mode: &str,
    ) -> Result<(PathBuf, &'static str)> {
        let provider = kind.soname();
        for (dir, source) in candidates {
            let shared = dir.join("libonnxruntime_providers_shared.so");
            let provider_path = dir.join(provider);
            if shared.is_file() && provider_path.is_file() {
                info!(
                    path = %shared.display(),
                    source = *source,
                    "ORT providers_shared resolved"
                );
                info!(
                    provider = kind.label(),
                    path = %provider_path.display(),
                    source = *source,
                    "ORT provider resolved"
                );
                return Ok((dir.clone(), *source));
            }
        }

        let render_path = |v: Option<&Path>| -> String {
            v.map(|p| p.display().to_string())
                .unwrap_or_else(|| "<unset>".to_string())
        };
        let checked = Self::render_provider_candidates(candidates);
        Err(EngineError::ModelMetadata(format!(
            "ORT provider directory resolution failed: command={command} mode={mode} ep_mode={ep_mode} provider={} \
reason=no candidate directory contained required pair \
(libonnxruntime_providers_shared.so + {}). \
overrides={{ORT_DYLIB_PATH={}, ORT_LIB_LOCATION={}}}. \
candidates_checked=[{}]. \
Set ORT_DYLIB_PATH or ORT_LIB_LOCATION to a valid provider directory and re-run scripts/test_ort_provider_load.sh.",
            kind.label(),
            provider,
            render_path(ort_dylib_path),
            render_path(ort_lib_location),
            checked
        )))
    }

    #[cfg(target_os = "linux")]
    fn resolve_ort_provider_dir(kind: OrtProviderKind) -> Result<(PathBuf, &'static str)> {
        let (ort_dylib_path, ort_lib_location, exe_dir, cache_dirs, on_wsl) =
            Self::ort_provider_search_inputs();
        let candidates = Self::provider_dir_candidates(
            ort_dylib_path.clone(),
            ort_lib_location.clone(),
            exe_dir,
            cache_dirs,
            on_wsl,
        );
        let (dir, source) = Self::resolve_provider_dir(
            kind,
            &candidates,
            ort_dylib_path.as_deref(),
            ort_lib_location.as_deref(),
            "backend_initialize",
            "provider_preload",
            kind.label(),
        )?;
        Self::configure_ort_loader_path(&dir);
        Ok((dir, source))
    }

    #[cfg(target_os = "linux")]
    fn dlopen_path(path: &Path, flags: i32) -> Result<()> {
        let cpath = CString::new(path.to_string_lossy().as_bytes())
            .map_err(|_| EngineError::ModelMetadata("Invalid dlopen path".into()))?;
        // SAFETY: `dlopen` expects a valid, NUL-terminated C string and flags.
        let handle = unsafe { dlopen(cpath.as_ptr(), flags) };
        if handle.is_null() {
            // SAFETY: `dlerror` returns a thread-local pointer or null.
            let err = unsafe {
                let p = dlerror();
                if p.is_null() {
                    "unknown dlopen error".to_string()
                } else {
                    CStr::from_ptr(p).to_string_lossy().to_string()
                }
            };
            return Err(EngineError::ModelMetadata(format!(
                "dlopen failed for {}: {err}",
                path.display()
            )));
        }
        Ok(())
    }

    #[cfg(target_os = "linux")]
    fn preload_cuda_runtime_libs() {
        // Prefer explicit toolkit paths to avoid reliance on unrelated app installs.
        let mut roots = Vec::<PathBuf>::new();
        if let Some(dir) = env::var_os("RAVE_CUDA_RUNTIME_DIR") {
            roots.push(PathBuf::from(dir));
        }
        roots.push(PathBuf::from("/usr/local/cuda-12/targets/x86_64-linux/lib"));
        roots.push(PathBuf::from("/usr/local/cuda/lib64"));

        let libs = [
            "libcudart.so.12",
            "libcublasLt.so.12",
            "libcublas.so.12",
            "libnvrtc.so.12",
            "libcurand.so.10",
            "libcufft.so.11",
        ];

        for root in roots {
            if !root.is_dir() {
                continue;
            }
            let mut loaded = 0usize;
            for lib in libs {
                let p = root.join(lib);
                if !p.is_file() {
                    continue;
                }
                if Self::dlopen_path(&p, RTLD_NOW | RTLD_GLOBAL).is_ok() {
                    loaded += 1;
                }
            }
            if loaded > 0 {
                info!(root = %root.display(), loaded, "Preloaded CUDA runtime libraries for ORT");
                break;
            }
        }
    }

    /// When ORT is linked statically, TensorRT provider expects `Provider_GetHost`
    /// from `libonnxruntime_providers_shared.so` to already be present in the process.
    /// Preloading this bridge with `RTLD_GLOBAL` satisfies that symbol before TRT EP load.
    fn preload_ort_provider_pair(kind: OrtProviderKind) -> Result<()> {
        #[cfg(target_os = "linux")]
        {
            let (dir, source) = Self::resolve_ort_provider_dir(kind)?;
            let shared = dir.join("libonnxruntime_providers_shared.so");
            let provider = dir.join(kind.soname());

            Self::dlopen_path(&shared, RTLD_NOW | RTLD_GLOBAL).map_err(|e| {
                EngineError::ModelMetadata(format!(
                    "Failed loading providers_shared from {} ({e}). \
Ensure ORT_DYLIB_PATH/ORT_LIB_LOCATION points to a valid ORT cache dir.",
                    shared.display()
                ))
            })?;

            info!(
                path = %provider.display(),
                provider = kind.label(),
                "Skipping explicit provider dlopen; relying on startup LD_LIBRARY_PATH + ORT registration"
            );

            info!(
                source,
                dir = %dir.display(),
                provider = kind.label(),
                path = %provider.display(),
                "ORT provider pair prepared (providers_shared preloaded, provider path configured)"
            );
            Ok(())
        }

        #[cfg(not(target_os = "linux"))]
        {
            let _ = kind;
            Ok(())
        }
    }

    #[cfg(all(target_os = "linux", test))]
    fn preload_ort_provider_bridge() -> Result<()> {
        for path in Self::ort_provider_candidates("libonnxruntime_providers_shared.so") {
            if path.is_file() {
                return Self::dlopen_path(&path, RTLD_NOW | RTLD_GLOBAL);
            }
        }
        Err(EngineError::ModelMetadata(
            "Could not preload libonnxruntime_providers_shared.so from known search paths".into(),
        ))
    }

    #[cfg(all(not(target_os = "linux"), test))]
    #[allow(dead_code)]
    fn preload_ort_provider_bridge() -> Result<()> {
        Ok(())
    }

    #[cfg(not(target_os = "linux"))]
    fn preload_cuda_runtime_libs() {}

    fn build_trt_session(&self) -> Result<Session> {
        Self::preload_cuda_runtime_libs();
        Self::preload_ort_provider_pair(OrtProviderKind::TensorRt)?;

        let mut trt_ep = TensorRTExecutionProvider::default()
            .with_device_id(self.device_id)
            .with_engine_cache(true)
            .with_engine_cache_path(
                self.model_path
                    .parent()
                    .unwrap_or(&self.model_path)
                    .join("trt_cache")
                    .to_string_lossy()
                    .to_string(),
            );

        match &self.precision_policy {
            PrecisionPolicy::Fp32 => {
                info!("TRT precision: FP32 (no mixed precision)");
            }
            PrecisionPolicy::Fp16 => {
                trt_ep = trt_ep.with_fp16(true);
                info!("TRT precision: FP16 mixed precision");
            }
            PrecisionPolicy::Int8 { calibration_table } => {
                trt_ep = trt_ep.with_fp16(true).with_int8(true);
                info!(
                    table = %calibration_table.display(),
                    "TRT precision: INT8 with calibration table"
                );
            }
        }

        Session::builder()?
            .with_execution_providers([trt_ep.build().error_on_failure()])?
            .with_intra_threads(1)?
            .commit_from_file(&self.model_path)
            .map_err(Into::into)
    }

    fn build_cuda_session(&self) -> Result<Session> {
        Self::preload_cuda_runtime_libs();
        Self::preload_ort_provider_pair(OrtProviderKind::Cuda)?;
        let cuda_ep = CUDAExecutionProvider::default().with_device_id(self.device_id);
        Session::builder()?
            .with_execution_providers([cuda_ep.build().error_on_failure()])?
            .with_intra_threads(1)?
            .commit_from_file(&self.model_path)
            .map_err(Into::into)
    }

    fn pointer_identity_mismatch(
        input_ptr: u64,
        texture_ptr: u64,
        output_ptr: u64,
        ring_slot_ptr: u64,
    ) -> Option<String> {
        if input_ptr != texture_ptr {
            return Some(format!(
                "POINTER MISMATCH: IO-bound input (0x{input_ptr:016x}) != GpuTexture (0x{texture_ptr:016x})"
            ));
        }
        if output_ptr != ring_slot_ptr {
            return Some(format!(
                "POINTER MISMATCH: IO-bound output (0x{output_ptr:016x}) != ring slot (0x{ring_slot_ptr:016x})"
            ));
        }
        None
    }

    /// Verify that IO-bound device pointers match the source GpuTexture
    /// and OutputRing slot pointers exactly (pointer identity check).
    ///
    /// Called by `run_io_bound` to audit that ORT IO Binding uses our
    /// device pointers without any host staging or reallocation.
    fn verify_pointer_identity(
        input_ptr: u64,
        output_ptr: u64,
        input_texture: &GpuTexture,
        ring_slot_ptr: u64,
    ) -> Result<()> {
        let texture_ptr = input_texture.device_ptr();
        debug!(
            input_ptr = format!("0x{:016x}", input_ptr),
            texture_ptr = format!("0x{:016x}", texture_ptr),
            output_ptr = format!("0x{:016x}", output_ptr),
            ring_slot_ptr = format!("0x{:016x}", ring_slot_ptr),
            "IO-binding pointer identity audit"
        );

        if let Some(message) =
            Self::pointer_identity_mismatch(input_ptr, texture_ptr, output_ptr, ring_slot_ptr)
        {
            debug_assert!(false, "{message}");
            #[cfg(feature = "audit-no-host-copies")]
            if rave_core::host_copy_audit::is_strict_mode() {
                return Err(EngineError::InvariantViolation(message));
            }
            rave_core::host_copy_violation!("inference", "{message}");
        }
        Ok(())
    }

    fn run_io_bound(
        session: &mut Session,
        meta: &ModelMetadata,
        input: &GpuTexture,
        output_ptr: u64,
        output_bytes: usize,
        _ctx: &GpuContext,
        _mem_info: &ort::memory::MemoryInfo, // Unused argument if we recreate it, or use it if passed.
    ) -> Result<()> {
        let in_w = input.width as i64;
        let in_h = input.height as i64;
        let out_w = in_w * meta.scale as i64;
        let out_h = in_h * meta.scale as i64;

        let input_shape: Vec<i64> = vec![1, meta.input_channels as i64, in_h, in_w];
        let output_shape: Vec<i64> = vec![1, meta.input_channels as i64, out_h, out_w];

        // Resolve element type from pixel format (F16 or F32).
        let elem_type = ort_element_type(input.format);
        let elem_bytes = input.format.element_bytes();
        let input_bytes = input.data.len();

        let expected = (output_shape.iter().product::<i64>() as usize) * elem_bytes;
        if output_bytes < expected {
            return Err(EngineError::BufferTooSmall {
                need: expected,
                have: output_bytes,
            });
        }

        let input_ptr = input.device_ptr();

        // Phase 7: Pointer identity audit — verify device pointers match.
        Self::verify_pointer_identity(input_ptr, output_ptr, input, output_ptr)?;

        let mut binding = session.create_binding()?;

        // Create fresh MemoryInfo for this opt (cheap, avoids Sync issues with cache)
        // (Removed: using FFI helper instead)

        // SAFETY — ORT IO Binding with raw device pointers:
        unsafe {
            // Create Input Tensor from raw device memory (zero-copy)
            let input_tensor = create_tensor_from_device_memory(
                input_ptr as *mut _,
                input_bytes,
                &input_shape,
                elem_type,
            )?;

            binding.bind_input(&meta.input_name, &input_tensor)?;

            // Create Output Tensor from raw device memory (zero-copy)
            let output_tensor = create_tensor_from_device_memory(
                output_ptr as *mut _,
                output_bytes,
                &output_shape,
                elem_type,
            )?;

            binding.bind_output(&meta.output_name, output_tensor)?;
        }

        // Synchronous execution — ORT blocks until all TensorRT kernels complete.
        session.run_binding(&binding)?;

        Ok(())
    }
}

#[cfg(test)]
mod pointer_audit_tests {
    use super::TensorRtBackend;

    #[test]
    fn pointer_identity_audit_accepts_matching_addresses() {
        let mismatch = TensorRtBackend::pointer_identity_mismatch(0x1000, 0x1000, 0x2000, 0x2000);
        assert!(mismatch.is_none());
    }

    #[test]
    fn pointer_identity_audit_reports_mismatch() {
        let mismatch = TensorRtBackend::pointer_identity_mismatch(0x1000, 0x1111, 0x2000, 0x2000)
            .expect("mismatch should be reported");
        assert!(mismatch.contains("IO-bound input"));
    }
}

#[cfg(all(test, target_os = "linux"))]
mod provider_resolution_tests {
    use super::{OrtProviderKind, TensorRtBackend};
    use std::env;
    use std::fs::{self, File};
    use std::path::{Path, PathBuf};
    use std::process;
    use std::sync::atomic::{AtomicU64, Ordering};

    static NEXT_TMP_ID: AtomicU64 = AtomicU64::new(0);

    struct TempDir {
        path: PathBuf,
    }

    impl TempDir {
        fn new(prefix: &str) -> Self {
            let unique = format!(
                "rave_provider_resolution_{prefix}_{}_{}",
                process::id(),
                NEXT_TMP_ID.fetch_add(1, Ordering::Relaxed)
            );
            let path = env::temp_dir().join(unique);
            fs::create_dir_all(&path).expect("create temp dir");
            Self { path }
        }

        fn join(&self, suffix: &str) -> PathBuf {
            self.path.join(suffix)
        }
    }

    impl Drop for TempDir {
        fn drop(&mut self) {
            let _ = fs::remove_dir_all(&self.path);
        }
    }

    fn write_provider_pair(dir: &Path, kind: OrtProviderKind) {
        fs::create_dir_all(dir).expect("create provider dir");
        File::create(dir.join("libonnxruntime_providers_shared.so")).expect("create shared lib");
        File::create(dir.join(kind.soname())).expect("create provider lib");
    }

    #[test]
    fn provider_dir_env_override_wins() {
        let tmp = TempDir::new("env_override");
        let override_dir = tmp.join("override");
        let exe_dir = tmp.join("exe");
        let cache_dir = tmp.join("cache");

        write_provider_pair(&override_dir, OrtProviderKind::TensorRt);
        write_provider_pair(&exe_dir, OrtProviderKind::TensorRt);
        write_provider_pair(&cache_dir, OrtProviderKind::TensorRt);

        let candidates = TensorRtBackend::provider_dir_candidates(
            Some(override_dir.clone()),
            None,
            Some(exe_dir),
            vec![cache_dir],
            false,
        );
        let (resolved, source) = TensorRtBackend::resolve_provider_dir(
            OrtProviderKind::TensorRt,
            &candidates,
            Some(&override_dir),
            None,
            "upscale",
            "direct",
            "tensorrt",
        )
        .expect("provider directory should resolve");

        assert_eq!(resolved, override_dir);
        assert_eq!(source, "ORT_DYLIB_PATH");
    }

    #[test]
    fn provider_dir_fallback_order_uses_first_existing_candidate() {
        let tmp = TempDir::new("fallback_order");
        let exe_dir = tmp.join("exe");
        let exe_deps_dir = exe_dir.join("deps");
        let cache_dir = tmp.join("cache");

        write_provider_pair(&exe_deps_dir, OrtProviderKind::Cuda);
        write_provider_pair(&cache_dir, OrtProviderKind::Cuda);

        let candidates = TensorRtBackend::provider_dir_candidates(
            None,
            None,
            Some(exe_dir.clone()),
            vec![cache_dir],
            false,
        );
        let (resolved, source) = TensorRtBackend::resolve_provider_dir(
            OrtProviderKind::Cuda,
            &candidates,
            None,
            None,
            "benchmark",
            "provider_preload",
            "cuda",
        )
        .expect("provider directory should resolve");

        assert_eq!(resolved, exe_deps_dir);
        assert_eq!(source, "exe_dir/deps");
    }

    #[test]
    fn provider_dir_failure_reports_checked_candidates_in_order() {
        let tmp = TempDir::new("failure_diagnostics");
        let override_dir = tmp.join("missing_override");
        let lib_dir = tmp.join("missing_lib");

        let candidates = TensorRtBackend::provider_dir_candidates(
            Some(override_dir.clone()),
            Some(lib_dir.clone()),
            None,
            vec![],
            false,
        );
        let err = TensorRtBackend::resolve_provider_dir(
            OrtProviderKind::TensorRt,
            &candidates,
            Some(&override_dir),
            Some(&lib_dir),
            "validate",
            "provider_preload",
            "tensorrt",
        )
        .expect_err("resolution should fail");
        let msg = err.to_string();

        assert!(msg.contains("command=validate"));
        assert!(msg.contains("mode=provider_preload"));
        assert!(msg.contains("ep_mode=tensorrt"));
        assert!(msg.contains("candidates_checked=[1:ORT_DYLIB_PATH:"));
        let idx_override = msg
            .find("1:ORT_DYLIB_PATH:")
            .expect("missing ORT_DYLIB_PATH candidate");
        let idx_lib = msg
            .find("2:ORT_LIB_LOCATION:")
            .expect("missing ORT_LIB_LOCATION candidate");
        assert!(
            idx_override < idx_lib,
            "candidate order should be deterministic and preserve precedence"
        );
    }

    #[test]
    fn provider_candidate_rendering_is_deterministic() {
        let candidates = vec![
            (PathBuf::from("/tmp/a"), "ORT_DYLIB_PATH"),
            (PathBuf::from("/tmp/b"), "ORT_LIB_LOCATION"),
            (PathBuf::from("/tmp/c"), "exe_dir"),
        ];

        let rendered_a = TensorRtBackend::render_provider_candidates(&candidates);
        let rendered_b = TensorRtBackend::render_provider_candidates(&candidates);

        assert_eq!(rendered_a, rendered_b);
        assert_eq!(
            rendered_a,
            "1:ORT_DYLIB_PATH:/tmp/a, 2:ORT_LIB_LOCATION:/tmp/b, 3:exe_dir:/tmp/c"
        );
    }
}

#[cfg(all(test, target_os = "linux"))]
mod tests {
    use super::{RTLD_LOCAL, RTLD_NOW, TensorRtBackend};
    use rave_core::backend::UpscaleBackend;
    use std::env;
    use std::path::PathBuf;

    #[test]
    #[ignore = "requires ORT TensorRT provider libs on host"]
    fn providers_load_with_bridge_preloaded() {
        // Ensure bridge is globally visible before loading TensorRT provider.
        TensorRtBackend::preload_ort_provider_bridge().expect("failed to preload providers_shared");
        let trt_candidates =
            TensorRtBackend::ort_provider_candidates("libonnxruntime_providers_tensorrt.so");
        assert!(
            !trt_candidates.is_empty(),
            "no libonnxruntime_providers_tensorrt.so candidates found"
        );
        let path = trt_candidates[0].clone();
        TensorRtBackend::dlopen_path(&path, RTLD_NOW | RTLD_LOCAL)
            .expect("providers_tensorrt should dlopen after providers_shared preload");
    }

    #[test]
    #[ignore = "requires model + full ORT/TensorRT runtime"]
    fn ort_registers_tensorrt_ep_smoke() {
        let model = env::var("RAVE_TEST_ONNX_MODEL").expect("set RAVE_TEST_ONNX_MODEL");
        let backend = TensorRtBackend::new(
            PathBuf::from(model),
            rave_core::context::GpuContext::new(0).expect("cuda ctx"),
            0,
            6,
            4,
        );
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .expect("tokio runtime");
        rt.block_on(async move { backend.initialize().await })
            .expect("TensorRT EP registration should succeed");
    }
}

#[async_trait]
impl UpscaleBackend for TensorRtBackend {
    async fn initialize(&self) -> Result<()> {
        let mut guard = self.state.lock().await;
        if guard.is_some() {
            return Err(EngineError::ModelMetadata("Already initialized".into()));
        }

        validate_batch_config(&self.batch_config)?;

        let ep_mode = Self::ort_ep_mode();
        let on_wsl = Self::is_wsl2();
        info!(
            path = %self.model_path.display(),
            ?ep_mode,
            on_wsl,
            "Loading ONNX model with ORT execution provider policy"
        );

        let (session, active_provider) = match ep_mode {
            OrtEpMode::CudaOnly => {
                info!("RAVE_ORT_TENSORRT disables TensorRT EP; using CUDAExecutionProvider");
                (self.build_cuda_session()?, "CUDAExecutionProvider")
            }
            OrtEpMode::TensorRtOnly => (self.build_trt_session()?, "TensorrtExecutionProvider"),
            OrtEpMode::Auto => match self.build_trt_session() {
                Ok(session) => (session, "TensorrtExecutionProvider"),
                Err(e) => {
                    warn!(
                        error = %e,
                        on_wsl,
                        "TensorRT EP registration failed; falling back to CUDAExecutionProvider"
                    );
                    (self.build_cuda_session()?, "CUDAExecutionProvider")
                }
            },
        };
        info!(
            provider = active_provider,
            "ORT execution provider selected"
        );

        let metadata = Self::extract_metadata(&session)?;
        info!(
            name = %metadata.name,
            scale = metadata.scale,
            input = %metadata.input_name,
            output = %metadata.output_name,
            ring_size = self.ring_size,
            min_ring_slots = self.min_ring_slots,
            provider = active_provider,
            precision = ?self.precision_policy,
            max_batch = self.batch_config.max_batch,
            "Model loaded"
        );

        let _ = self.meta.set(metadata);
        let _ = self.selected_provider.set(active_provider.to_string());

        *guard = Some(InferenceState {
            session,
            ring: None,
        });

        Ok(())
    }

    async fn process(&self, input: GpuTexture) -> Result<GpuTexture> {
        // Accept both F32 and F16 planar RGB.
        match input.format {
            PixelFormat::RgbPlanarF32 | PixelFormat::RgbPlanarF16 => {}
            other => {
                return Err(EngineError::FormatMismatch {
                    expected: PixelFormat::RgbPlanarF32,
                    actual: other,
                });
            }
        }

        let meta = self.meta.get().ok_or(EngineError::NotInitialized)?;
        let mut guard = self.state.lock().await;
        let state = guard.as_mut().ok_or(EngineError::NotInitialized)?;

        // Lazy ring init / realloc.
        match &mut state.ring {
            Some(ring) if ring.needs_realloc(input.width, input.height) => {
                debug!(
                    old_w = ring.alloc_dims.0,
                    old_h = ring.alloc_dims.1,
                    new_w = input.width,
                    new_h = input.height,
                    "Reallocating output ring"
                );
                ring.reallocate(&self.ctx, input.width, input.height, meta.scale)?;
            }
            None => {
                debug!(
                    w = input.width,
                    h = input.height,
                    slots = self.ring_size,
                    "Lazily creating output ring"
                );
                state.ring = Some(OutputRing::new(
                    &self.ctx,
                    input.width,
                    input.height,
                    meta.scale,
                    self.ring_size,
                    self.min_ring_slots,
                )?);
            }
            Some(_) => {}
        }

        let ring = state.ring.as_mut().unwrap();

        // Debug-mode host allocation tracking.
        #[cfg(feature = "debug-alloc")]
        {
            rave_core::debug_alloc::reset();
            rave_core::debug_alloc::enable();
        }

        let output_arc = ring.acquire()?;
        let output_ptr = *(*output_arc).device_ptr();
        let output_bytes = ring.slot_bytes;

        // ── Inference with latency measurement ──
        let t_start = std::time::Instant::now();

        let mem_info = self.mem_info()?;
        Self::run_io_bound(
            &mut state.session,
            meta,
            &input,
            output_ptr,
            output_bytes,
            &self.ctx,
            &mem_info,
        )?;

        let elapsed_us = t_start.elapsed().as_micros() as u64;
        self.inference_metrics.record(elapsed_us);

        #[cfg(feature = "debug-alloc")]
        {
            rave_core::debug_alloc::disable();
            let host_allocs = rave_core::debug_alloc::count();
            debug_assert_eq!(
                host_allocs, 0,
                "VIOLATION: {host_allocs} host allocations during inference"
            );
        }

        let out_w = input.width * meta.scale;
        let out_h = input.height * meta.scale;
        let elem_bytes = input.format.element_bytes();

        Ok(GpuTexture {
            data: output_arc,
            width: out_w,
            height: out_h,
            pitch: (out_w as usize) * elem_bytes,
            format: input.format, // Preserve F32 or F16
        })
    }

    async fn shutdown(&self) -> Result<()> {
        let mut guard = self.state.lock().await;
        if let Some(state) = guard.take() {
            info!("Shutting down TensorRT backend");
            self.ctx.sync_all()?;

            // Report ring metrics.
            if let Some(ring) = &state.ring {
                let snap = ring.metrics.snapshot();
                info!(
                    reuse = snap.reuse,
                    contention = snap.contention,
                    first_use = snap.first_use,
                    "Final ring metrics"
                );
            }

            // Report inference metrics.
            let snap = self.inference_metrics.snapshot();
            info!(
                frames = snap.frames_inferred,
                avg_us = snap.avg_inference_us,
                peak_us = snap.peak_inference_us,
                precision = ?self.precision_policy,
                "Final inference metrics"
            );

            // Report VRAM.
            let (current, peak) = self.ctx.vram_usage();
            info!(
                current_mb = current / (1024 * 1024),
                peak_mb = peak / (1024 * 1024),
                "Final VRAM usage"
            );

            drop(state.ring);
            drop(state.session);
            debug!("TensorRT backend shutdown complete");
        }
        Ok(())
    }

    fn metadata(&self) -> Result<&ModelMetadata> {
        self.meta.get().ok_or(EngineError::NotInitialized)
    }
}

#[cfg(test)]
mod batch_config_tests {
    use super::{BatchConfig, validate_batch_config};

    #[test]
    fn batch_config_validator_accepts_single_frame() {
        let cfg = BatchConfig {
            max_batch: 1,
            latency_deadline_us: 8_000,
        };
        validate_batch_config(&cfg).expect("max_batch=1 should be accepted");
    }

    #[test]
    fn batch_config_validator_rejects_micro_batching() {
        let cfg = BatchConfig {
            max_batch: 2,
            latency_deadline_us: 8_000,
        };
        let err = validate_batch_config(&cfg).expect_err("max_batch>1 must be rejected");
        let msg = err.to_string();
        assert!(msg.contains("max_batch"));
        assert!(msg.contains("not implemented"));
    }
}

impl Drop for TensorRtBackend {
    fn drop(&mut self) {
        if let Ok(mut guard) = self.state.try_lock()
            && let Some(state) = guard.take()
        {
            let _ = self.ctx.sync_all();
            drop(state);
        }
    }
}