gigastt-core 2.18.0

//! ANE encoder session: runs the GigaAM v3 Conformer encoder on the Apple
//! Neural Engine via a per-bucket fixed-shape `.mlpackage`, with a pad-up +
//! fill-floor selection policy and an ort encoder fallback.
//!
//! ISOLATION: all `objc2_core_ml` usage stays in `runtime/coreml/` (the model
//! handle here is produced by [`super::bridge`]; this file only holds it and
//! calls `bridge::predict_f32`). Gated `#[cfg(all(feature = "ane",
//! target_os = "macos"))]`.

use std::sync::Arc;

use objc2::rc::Retained;
use objc2_core_ml::MLModel;

use crate::runtime::{
    error::RuntimeError,
    session::RuntimeSession,
    tensor::{Shape, Tensor, TensorData, TensorView},
};

use super::bridge;

/// Encoder output channel dim (`[1, ENC_DIM, T']`, channels-first).
const ENC_DIM: usize = 768;
/// Mel feature bins (`[1, N_MELS, T]`, channels-first).
const N_MELS: usize = 64;

/// Minimum fraction of a bucket the real mel must fill for the **file-mode**
/// pad-up path to be trusted. Below this, the mask-free zero-padded encoder
/// output diverges enough from the unpadded baseline that a borderline token can
/// flip; clips that don't reach the floor fall back to the variable-length ort
/// encoder.
///
/// Calibrated by the pad-parity experiment (mask-free pad-up): at >= 50% fill the
/// padded ANE output stays at cosine >= 0.94 vs the ort baseline and WER tracks
/// the baseline; below 50% the raw output diverges and the transcript is no
/// longer trustworthy.
///
/// The cos >= 0.94 @ >= 50%-fill calibration was measured against the REAL FP16
/// `.mlpackage` (Float16 mel input, see [`super::bridge::predict_f32`]), so the
/// floor already absorbs the Float16 quantization of the mel input — not just the
/// zero-pad effect. A future reader must NOT "fix" this floor assuming the
/// encoder sees f32 input; the f16 round-trip is baked into the threshold.
const FILL_FLOOR: f64 = 0.5;

/// Fill floor for **streaming / low-latency** encoder calls.
///
/// Streaming windows are capped at ~2.5 s (≈249 mel frames), which underfills
/// the smallest shipped bucket (512 → file-mode floor 256) at ~48.6%. The file
/// floor would force every streaming window onto ort. A zero floor lets any
/// window that fits under the max bucket pad into the smallest eligible bucket
/// (typically 512) and run on ANE.
///
/// Streaming pad may underfill the bucket and is preferred for latency on Mac
/// ANE builds over the correct-but-slower ort fallback. File-mode keeps
/// [`FILL_FLOOR`] = 0.5 so the quality/cost tradeoff on long clips is unchanged.
const STREAMING_FILL_FLOOR: f64 = 0.0;

/// `Retained<MLModel>` is not auto-`Send`/`Sync` (it wraps an Objective-C
/// pointer), but Apple documents `-[MLModel predictionFromFeatures:error:]` as
/// thread-safe: a single loaded model may be used for prediction concurrently
/// from multiple threads. We only ever call prediction (never mutate the model)
/// once it is loaded, so sharing an `Arc<SharedModel>` across pooled encoder
/// sessions is sound.
///
/// SAFETY: the wrapped `MLModel` is immutable after load and Apple guarantees
/// `predictionFromFeatures:error:` is safe to call concurrently; no other
/// Objective-C method is invoked on it from this crate.
pub struct SharedModel(pub Retained<MLModel>);

// SAFETY: see the type-level doc above — prediction on a loaded MLModel is
// documented thread-safe and we never mutate it after load.
unsafe impl Send for SharedModel {}
// SAFETY: same justification — concurrent `&MLModel` prediction is supported.
unsafe impl Sync for SharedModel {}

/// One compiled bucket model paired with its mel-frame window size.
pub struct BucketModel {
    /// Padded mel window length `N` this package expects (`[1, 64, N]`).
    pub size: usize,
    /// Shared compiled model (loaded once, shared across all pool slots).
    pub model: Arc<SharedModel>,
}

/// Encoder time subsampling: two stride-2 conv layers (×4 total), each
/// `out = (in - 1) / 2 + 1` (integer division). Equivalent to `in.div_ceil(2)`
/// applied twice, which is exactly what the Candle `StridingSubsampling` and the
/// ONNX-derived bucket package compute, so the trim below stays aligned with the
/// frames the model actually emits.
pub fn calc_output_length(t: usize) -> usize {
    fn stride2(x: usize) -> usize {
        (x - 1) / 2 + 1
    }
    stride2(stride2(t))
}

/// Pick the smallest available bucket `N >= t` whose fill ratio `t / N` meets
/// `floor`. Returns `None` when no bucket satisfies both (caller falls back to
/// the variable-length ort encoder). `buckets` need not be sorted.
pub fn select_bucket(t: usize, buckets: &[usize], floor: f64) -> Option<usize> {
    buckets
        .iter()
        .copied()
        .filter(|&n| n >= t && t as f64 / n as f64 >= floor)
        .min()
}

/// Zero-pad a channels-first `[1, channels, t]` row-major mel buffer up to
/// `[1, channels, n]` (append `n - t` zero frames to each channel row). `t <= n`
/// is required; the layout is channel-major with the time axis contiguous-last,
/// so each of the `channels` rows is copied then zero-filled to length `n`.
pub fn pad_time(mel: &[f32], channels: usize, t: usize, n: usize) -> Vec<f32> {
    debug_assert!(t <= n, "pad target {n} must be >= source {t}");
    debug_assert_eq!(mel.len(), channels * t, "mel len mismatch in pad_time");
    let mut out = vec![0.0f32; channels * n];
    for c in 0..channels {
        let src = &mel[c * t..c * t + t];
        out[c * n..c * n + t].copy_from_slice(src);
    }
    out
}

/// Trim a channels-first `[1, channels, t_padded]` row-major buffer down to
/// `[1, channels, t_keep]` (keep the first `t_keep` frames of each channel row).
/// `t_keep <= t_padded` is required.
pub fn trim_time(out: &[f32], channels: usize, t_padded: usize, t_keep: usize) -> Vec<f32> {
    debug_assert!(t_keep <= t_padded, "trim {t_keep} must be <= {t_padded}");
    debug_assert_eq!(
        out.len(),
        channels * t_padded,
        "out len mismatch in trim_time"
    );
    let mut trimmed = Vec::with_capacity(channels * t_keep);
    for c in 0..channels {
        let row = &out[c * t_padded..c * t_padded + t_keep];
        trimmed.extend_from_slice(row);
    }
    trimmed
}

/// ANE-backed encoder session. Runs the encoder on the Neural Engine via a
/// per-bucket fixed-shape package when a clip pads-up into a bucket at the
/// active fill floor ([`FILL_FLOOR`] for file mode, [`STREAMING_FILL_FLOOR`] for
/// streaming); otherwise delegates to the ort encoder fallback.
pub struct AneEncoderSession {
    /// Available compiled bucket models, sorted ascending by `size`.
    buckets: Vec<BucketModel>,
    /// Variable-length ort encoder for clips outside the fill-floor / bucket range.
    ort_fallback: Box<dyn RuntimeSession>,
}

impl AneEncoderSession {
    pub fn new(mut buckets: Vec<BucketModel>, ort_fallback: Box<dyn RuntimeSession>) -> Self {
        buckets.sort_by_key(|b| b.size);
        Self {
            buckets,
            ort_fallback,
        }
    }

    fn bucket_sizes(&self) -> Vec<usize> {
        self.buckets.iter().map(|b| b.size).collect()
    }

    fn model_for(&self, size: usize) -> Option<&Arc<SharedModel>> {
        self.buckets
            .iter()
            .find(|b| b.size == size)
            .map(|b| &b.model)
    }

    /// Encoder contract (mirrors the ort ONNX encoder, which emits two outputs):
    /// `inputs[0] = mel [1, 64, T] F32`, `inputs[1] = length [1]` (ignored; T is
    /// read from the mel shape, which is authoritative on this path). Returns
    /// `[encoded [1, 768, T'] F32, encoded_len [1] I64]` so the engine's decode
    /// loop reads `encoder_outputs[1]` for the frame count exactly as for ort.
    ///
    /// `floor` is the pad fill floor: file mode uses [`FILL_FLOOR`], streaming
    /// uses [`STREAMING_FILL_FLOOR`].
    fn run_with_floor(&self, inputs: &[Tensor], floor: f64) -> Result<Vec<Tensor>, RuntimeError> {
        if inputs.is_empty() {
            return Err(RuntimeError::InvalidInputCount {
                expected: 2,
                got: inputs.len(),
            });
        }

        let mel_view: TensorView<'_> = inputs[0].view();
        let mel = mel_view.data().as_f32().ok_or_else(|| {
            RuntimeError::InferenceFailed("ANE encoder mel input is not f32".to_string())
        })?;
        let dims = mel_view.shape().dims();
        if dims.len() != 3 || dims[0] != 1 || dims[1] != N_MELS {
            return Err(RuntimeError::InferenceFailed(format!(
                "ANE encoder expects mel shape [1, {N_MELS}, T], got {dims:?}"
            )));
        }
        let t = dims[2];

        let sizes = self.bucket_sizes();
        match select_bucket(t, &sizes, floor) {
            Some(n) => {
                let model = self.model_for(n).ok_or_else(|| {
                    RuntimeError::InferenceFailed(format!("no compiled model for bucket {n}"))
                })?;
                let padded = pad_time(mel, N_MELS, t, n);
                let (out, out_shape) =
                    bridge::predict_f32(&model.0, "mel", &padded, &[1, N_MELS, n], "encoded")?;
                // The package emits [1, 768, T'_N] for the padded window N.
                if out_shape.len() != 3 || out_shape[0] != 1 || out_shape[1] != ENC_DIM {
                    return Err(RuntimeError::InferenceFailed(format!(
                        "ANE encoder output shape {out_shape:?} != [1, {ENC_DIM}, T']"
                    )));
                }
                let t_padded_prime = out_shape[2];
                let t_real_prime = calc_output_length(t);
                if t_real_prime > t_padded_prime {
                    return Err(RuntimeError::InferenceFailed(format!(
                        "computed encoder output length {t_real_prime} exceeds bucket output {t_padded_prime}"
                    )));
                }
                tracing::debug!(
                    t,
                    bucket = n,
                    floor,
                    t_real_prime,
                    t_padded_prime,
                    "ANE encoder path (bucketed pad-up)"
                );
                let trimmed = trim_time(&out, ENC_DIM, t_padded_prime, t_real_prime);
                Ok(vec![
                    Tensor::new(
                        Shape::new(vec![1, ENC_DIM, t_real_prime]),
                        TensorData::F32(trimmed),
                    )?,
                    Tensor::new(
                        Shape::new(vec![1]),
                        TensorData::I64(vec![t_real_prime as i64]),
                    )?,
                ])
            }
            None => {
                tracing::debug!(
                    t,
                    floor,
                    "ANE encoder path (ort fallback: no bucket within fill-floor)"
                );
                self.ort_fallback.run(inputs)
            }
        }
    }
}

impl RuntimeSession for AneEncoderSession {
    fn run(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, RuntimeError> {
        self.run_with_floor(inputs, FILL_FLOOR)
    }

    fn run_low_latency(&self, inputs: &[Tensor]) -> Result<Vec<Tensor>, RuntimeError> {
        self.run_with_floor(inputs, STREAMING_FILL_FLOOR)
    }

    fn is_ane_encoder(&self) -> bool {
        true
    }
}

#[cfg(test)]
mod tests;