diff --git a/src-tauri/src/transcription/npu.rs b/src-tauri/src/transcription/npu.rs index 9fde9cb..b8de977 100644 --- a/src-tauri/src/transcription/npu.rs +++ b/src-tauri/src/transcription/npu.rs @@ -1,12 +1,15 @@ -//! NPU transcriber: Whisper ONNX on the Intel NPU via ONNX Runtime + OpenVINO -//! (T3.4, ADR-0004). Same `Transcriber` trait and `TranscriptSegment` output as -//! the whisper.cpp path, so callers never branch on engine. +//! ONNX transcriber: Whisper ONNX via ONNX Runtime, on either the Intel NPU +//! (OpenVINO EP, T3.4) or a DX12 GPU (DirectML EP, P2 — the non-Vulkan path for +//! AMD/Intel). One engine, EP chosen by the requested `BackendId`; same +//! `Transcriber` trait and `TranscriptSegment` output as the whisper.cpp path, +//! so callers never branch on engine. The ONNX model artifacts are identical +//! across EPs (an EP is a device backend, not a different graph). //! //! Split of work, validated by the T3.4 spike (encoder ~3.6× faster on NPU): -//! - **encoder** (fixed `[1,80,3000]` shape) runs on the **NPU** via OpenVINO — -//! the expensive graph, and the shape NPUs want; +//! - **encoder** (fixed `[1,80,3000]` shape) runs on the **accelerator** (NPU or +//! GPU) — the expensive graph, and the fixed shape accelerators want; //! - **decoder** (dynamic, autoregressive) runs greedy on the **CPU** EP — a -//! dynamic KV loop is a poor NPU fit and the cheap half anyway. +//! dynamic KV loop is a poor accelerator fit and the cheap half anyway. //! //! ponytail: greedy, no-KV-cache decode (re-feeds the full token prefix each //! step). Correct and simple; windows are short so the token count is small. @@ -14,7 +17,9 @@ use super::{AudioWindow, SegmentSink, Transcriber, TrxError}; use crate::models::{BackendId, TranscriptSegment}; -use ort::execution_providers::OpenVINOExecutionProvider; +use ort::execution_providers::{ + DirectMLExecutionProvider, ExecutionProviderDispatch, OpenVINOExecutionProvider, +}; use ort::session::Session; use ort::value::Tensor; use std::collections::HashMap; @@ -25,7 +30,7 @@ use std::sync::Mutex; /// Safety cap on generated tokens per 30 s window (Whisper's own max is 448). const MAX_NEW_TOKENS: usize = 224; -pub struct OnnxNpuTranscriber { +pub struct OnnxTranscriber { encoder: Mutex, decoder: Mutex, enc_input: String, @@ -44,15 +49,21 @@ pub struct OnnxNpuTranscriber { next_id: AtomicU64, } -impl OnnxNpuTranscriber { - /// Points `ort` (load-dynamic) at the on-demand-downloaded runtime and puts - /// its DLLs on the search path, unless the caller already set `ORT_DYLIB_PATH` - /// (the test harness does, to target a dev runtime). Idempotent. - fn ensure_runtime_env() { +impl OnnxTranscriber { + /// Points `ort` (load-dynamic) at the given on-demand-downloaded runtime DLL + /// and puts its sibling DLLs on the search path, unless the caller already + /// set `ORT_DYLIB_PATH` (the test harness does, to target a dev runtime). + /// Idempotent. + /// + /// `ORT_DYLIB_PATH` is process-global and `ort` dlopens it once, so a single + /// process run must not mix EPs from different runtime bundles (OpenVINO vs + /// DirectML). That's fine: `best()` picks one backend per session. + // ponytail: one runtime per process. If backend hot-swapping between NPU and + // DirectML in one run is ever needed, that's a re-init-ort problem to solve then. + fn ensure_runtime_env(dll: std::path::PathBuf) { if std::env::var_os("ORT_DYLIB_PATH").is_some() { return; } - let dll = crate::paths::npu_runtime_dll(); std::env::set_var("ORT_DYLIB_PATH", &dll); if let Some(dir) = dll.parent() { let path = std::env::var_os("PATH").unwrap_or_default(); @@ -185,32 +196,52 @@ impl OnnxNpuTranscriber { } } -impl Transcriber for OnnxNpuTranscriber { - /// `model` is the directory holding the ONNX artifacts (see `onnx_models`). - /// `backend` is expected to be `Npu`; a non-NPU value is rejected so the - /// dispatcher can fall back to whisper.cpp rather than us guessing. +impl Transcriber for OnnxTranscriber { + /// `model` is the directory holding the ONNX artifacts (see `onnx_models`) — + /// the same artifacts for every EP. `backend` selects the accelerator: + /// `Npu` → OpenVINO EP; `Amd`/`Intel` → DirectML EP (the non-Vulkan GPU + /// path). Any other value is rejected so the dispatcher can fall back to + /// whisper.cpp rather than us guessing. fn load(model: &Path, backend: BackendId) -> Result { - if backend != BackendId::Npu { - return Err(TrxError::Load(format!( - "OnnxNpuTranscriber only serves the NPU backend, got {backend:?}" - ))); - } - Self::ensure_runtime_env(); + // Pick the runtime bundle + encoder EP for the requested accelerator. + // error_on_failure makes a failed accelerator registration LOUD (Err) + // instead of a silent CPU fallback, so the dispatcher can cleanly drop + // to the whisper.cpp path (T3.4 spike lesson). + let (runtime_dll, encoder_ep): (std::path::PathBuf, ExecutionProviderDispatch) = + match backend { + BackendId::Npu => ( + crate::paths::npu_runtime_dll(), + OpenVINOExecutionProvider::default() + .with_device_type("NPU") + .build() + .error_on_failure(), + ), + // ponytail: DirectML device_id 0 = the default DX12 adapter. + // Map the DXGI adapter index here if a multi-GPU box ever needs + // to pin a specific AMD/Intel card. + BackendId::Amd | BackendId::Intel => ( + crate::paths::directml_runtime_dll(), + DirectMLExecutionProvider::default() + .with_device_id(0) + .build() + .error_on_failure(), + ), + other => { + return Err(TrxError::Load(format!( + "OnnxTranscriber serves NPU/AMD/Intel, got {other:?}" + ))) + } + }; + Self::ensure_runtime_env(runtime_dll); let encoder_path = model.join("encoder_model.onnx"); let decoder_path = model.join("decoder_model.onnx"); let tokenizer_path = model.join("tokenizer.json"); let gen_config_path = model.join("generation_config.json"); - // Encoder pinned to the NPU. error_on_failure makes a failed NPU - // registration LOUD (Err) instead of a silent CPU fallback, so the - // dispatcher can cleanly drop to the whisper.cpp path (T3.4 spike lesson). let encoder = Session::builder() .map_err(|e| TrxError::Load(e.to_string()))? - .with_execution_providers([OpenVINOExecutionProvider::default() - .with_device_type("NPU") - .build() - .error_on_failure()]) + .with_execution_providers([encoder_ep]) .map_err(|e| TrxError::Load(e.to_string()))? .commit_from_file(&encoder_path) .map_err(|e| TrxError::Load(e.to_string()))?; @@ -434,7 +465,7 @@ mod tests { }; let wav = std::env::var("WA_NPU_TEST_WAV").expect("set WA_NPU_TEST_WAV"); let t0 = std::time::Instant::now(); - let t = OnnxNpuTranscriber::load(Path::new(&model_dir), BackendId::Npu) + let t = OnnxTranscriber::load(Path::new(&model_dir), BackendId::Npu) .expect("load NPU transcriber"); let load_ms = t0.elapsed().as_millis(); let t1 = std::time::Instant::now(); @@ -455,4 +486,45 @@ mod tests { ); } } + + /// Real DirectML GPU inference (P2), opt-in. Point `ORT_DYLIB_PATH` at a + /// DirectML-EP `onnxruntime.dll` (or stage `runtime\directml\`), reuse the + /// same ONNX model dir as the NPU spike, and run on the Intel/AMD GPU: + /// ORT_DYLIB_PATH=…directml\onnxruntime.dll WA_DML_MODEL_DIR=… \ + /// WA_DML_TEST_WAV=… cargo test --release --features npu \ + /// directml_transcribes -- --ignored --nocapture + #[test] + #[ignore = "requires a DirectML runtime + model + wav; run manually on a GPU"] + fn directml_transcribes_speech() { + let Some(model_dir) = std::env::var_os("WA_DML_MODEL_DIR") else { + eprintln!("skip: set WA_DML_MODEL_DIR"); + return; + }; + let wav = std::env::var("WA_DML_TEST_WAV").expect("set WA_DML_TEST_WAV"); + // Intel and AMD both take the DirectML EP; either BackendId exercises it. + let backend = match std::env::var("WA_DML_BACKEND").as_deref() { + Ok("amd") => BackendId::Amd, + _ => BackendId::Intel, + }; + let t0 = std::time::Instant::now(); + let t = OnnxTranscriber::load(Path::new(&model_dir), backend).expect("load DirectML"); + let load_ms = t0.elapsed().as_millis(); + let t1 = std::time::Instant::now(); + let segs = t.transcribe_file(Path::new(&wav)).expect("transcribe"); + let infer_ms = t1.elapsed().as_millis(); + let text = segs + .iter() + .map(|s| s.text.as_str()) + .collect::>() + .join(" ") + .to_lowercase(); + eprintln!("[spike] backend={backend:?} load={load_ms}ms infer={infer_ms}ms text={text:?}"); + assert!(!text.trim().is_empty(), "transcript was empty"); + if let Ok(expect) = std::env::var("WA_DML_EXPECT") { + assert!( + text.contains(&expect.to_lowercase()), + "transcript {text:?} missing expected {expect:?}" + ); + } + } }