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:
@@ -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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user