feat(transcription): generalize OnnxNpuTranscriber -> OnnxTranscriber with DirectML EP (P2)

load() now selects OpenVINO (NPU) or DirectML (AMD/Intel GPU) by BackendId over
the same ONNX artifacts; ensure_runtime_env takes the target runtime DLL. Adds
an opt-in directml_transcribes_speech spike for on-GPU validation.
This commit is contained in:
iamdoubz
2026-07-05 19:40:16 -05:00
parent 2e8bb730a3
commit 0f93f5947e
+104 -32
View File
@@ -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<Session>,
decoder: Mutex<Session>,
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<Self, TrxError> {
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::<Vec<_>>()
.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:?}"
);
}
}
}