keyhog_scanner/

gpu.rs

//! GPU-accelerated batch inference for the MoE classifier via wgpu compute shaders.
//!
//! Processes N feature vectors in a single GPU dispatch, achieving ~10-100x
//! throughput over CPU for large batches. Falls back to CPU when no GPU is
//! available or for batches smaller than the crossover threshold.
//!
//! Architecture mirrors ml_scorer.rs exactly:
//! - Gate: Linear(55→6) + softmax
//! - 6 experts: Linear(55→32)+ReLU → Linear(32→16)+ReLU → Linear(16→1)
//! - Output: sigmoid(weighted sum of expert logits)
//!
//! ## Feature-gating in the lean build
//!
//! Every entry point that would touch wgpu / vyre-driver-wgpu directly is
//! wrapped in `#[cfg(feature = "gpu")]`. With the `gpu` feature off (the
//! `cargo install keyhog --no-default-features --features ci` path), the
//! GPU drivers aren't linked at all, the probe functions report "no GPU
//! available" without ever calling into wgpu, and the self-test functions
//! return a "not available in this build" `Err` instead of panicking.
//! The CPU MoE path in `ml_scorer.rs` is the entire scoring story under
//! that profile.

// Both submodules lean on the wgpu device/queue + bytemuck cast helpers.
// They only exist in `gpu`-on builds; the public API in this module
// short-circuits to "no GPU" via the `cfg` arms below when off.
// Submodules live in `gpu/` (native resolution), matching the `foo.rs` + `foo/`
// layout used across the workspace. Module names (gpu_shader/backend/policy) are
// unchanged; only the files moved (and gpu_moe_backend.rs/gpu_env.rs were
// renamed to match their module names).
#[cfg(feature = "gpu")]
mod adapter_probe;
#[cfg(feature = "gpu")]
mod backend;
#[cfg(feature = "gpu")]
pub(crate) mod gpu_shader;

mod policy;
pub use policy::*;
mod self_test;
pub use self_test::*;

#[cfg(feature = "gpu")]
pub(crate) use adapter_probe::{
    gpu_adapter_device_identity, gpu_adapter_probe, is_software_adapter,
};

/// Split timers: accumulated wall time in feature extraction vs MoE scoring
/// across all batch ML inference calls. Only the SCORING fraction is
/// GPU-offloadable; feature extraction is inherent per-candidate CPU work. This
/// is the data that decides whether moving the MoE to a unified GPU batch is
/// worth the recall cost of reordering finalization.
static MOE_FEATURE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
static MOE_SCORE_NS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);

/// Gated by the unified scanner profile switch and dumped as part of
/// [`crate::profile_dump`].
#[cfg(feature = "ml")]
fn ml_split_prof_enabled() -> bool {
    crate::scan_profile::enabled()
}

/// Print + reset the feature-vs-score split. Folded into the unified profiler:
/// called from [`crate::profile_dump`] (early-returns when no data).
pub(crate) fn ml_split_profile_dump() {
    use std::sync::atomic::Ordering::Relaxed;
    let f = MOE_FEATURE_NS.swap(0, Relaxed) as f64 / 1e6;
    let s = MOE_SCORE_NS.swap(0, Relaxed) as f64 / 1e6;
    if f == 0.0 && s == 0.0 {
        return;
    }
    eprintln!(
        "=== ML split: feature_extract={f:.1}ms moe_score={s:.1}ms (score = {:.1}% of ML compute; \
only this fraction is GPU-offloadable) ===",
        100.0 * s / (f + s).max(1e-9),
    );
}

pub(crate) fn ml_split_profile_reset() {
    use std::sync::atomic::Ordering::Relaxed;
    MOE_FEATURE_NS.store(0, Relaxed);
    MOE_SCORE_NS.store(0, Relaxed);
}

#[cfg(all(test, feature = "ml", feature = "multiline"))]
pub(crate) fn batch_ml_inference<T: crate::ml_scorer::MlScoreInput>(
    candidates: &[T],
    config: &crate::types::ScannerConfig,
) -> Vec<f64> {
    batch_ml_inference_with_timeout(
        candidates,
        config,
        std::time::Duration::from_millis(
            crate::scanner_config::ScannerTuningConfig::GPU_MOE_TIMEOUT_MS_DEFAULT,
        ),
    )
}

#[cfg(feature = "ml")]
pub(crate) fn batch_ml_inference_with_timeout<T: crate::ml_scorer::MlScoreInput>(
    candidates: &[T],
    config: &crate::types::ScannerConfig,
    gpu_moe_timeout: std::time::Duration,
) -> Vec<f64> {
    if candidates.is_empty() {
        return Vec::new();
    }

    #[cfg(feature = "ml")]
    {
        use rayon::prelude::*;
        #[cfg(not(feature = "gpu"))]
        let _ = gpu_moe_timeout; // LAW10: cfg-only GPU timeout marker; ML CPU scoring ignores GPU dispatch timeout by construction
        let prof = ml_split_prof_enabled();

        // Single-chunk and windowed scans commonly produce only a handful of
        // candidates. Coalesced scans aggregate pending rows across chunks
        // before entering here, but any batch below the measured GPU crossover
        // still avoids rayon split/join and GPU dispatch overhead through one
        // fused serial feature-and-score loop.
        if candidates.len() < crate::ml_scorer::GPU_BATCH_THRESHOLD {
            // Small-batch fused serial path (the ~99% case).
            let t = prof.then(std::time::Instant::now);
            let scores = crate::ml_scorer::score_input_batch_serial(candidates, config);
            if let Some(t) = t {
                // Fused loop: attribute the whole cost to feature+score combined
                // under MOE_SCORE_NS (kept separate from the large-batch split).
                MOE_SCORE_NS.fetch_add(
                    t.elapsed().as_nanos() as u64,
                    std::sync::atomic::Ordering::Relaxed,
                );
            }
            return scores;
        }

        // Large batch: parallel feature extraction, then GPU (or parallel CPU).
        let t_feat = prof.then(std::time::Instant::now);
        let features: Vec<[f32; crate::ml_scorer::NUM_FEATURES]> = candidates
            .par_iter()
            .map(|candidate| candidate.ml_features(config))
            .collect();
        if let Some(t) = t_feat {
            MOE_FEATURE_NS.fetch_add(
                t.elapsed().as_nanos() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
        }

        let t_score = prof.then(std::time::Instant::now);
        let score_features_on_cpu =
            || crate::ml_scorer::score_precomputed_batch_on_cpu(candidates, &features);
        let scores = {
            #[cfg(feature = "gpu")]
            {
                match backend::batch_score_features(&features, gpu_moe_timeout) {
                    Some(mut scores) if scores.len() == candidates.len() => {
                        crate::confidence::policy::apply_empty_candidate_score_policy(
                            candidates.iter().map(|candidate| candidate.ml_text()),
                            &mut scores,
                        );
                        scores
                    }
                    Some(scores) => {
                        // Defense in depth. `batch_score_features` OWNS the length
                        // invariant (backend.rs degrades + returns `None` when the
                        // GPU readback count != batch_size == features.len()), and
                        // this caller builds `features` one-per-candidate, so a
                        // `Some` whose length differs from `candidates` cannot occur
                        // via the real backend, this arm is unreachable today. Keep
                        // it fail-LOUD instead of a silent CPU fallback (Law 10): if
                        // a future backend change ever breaks that contract, route
                        // the degrade through the SAME owner as every other MoE
                        // dispatch failure (hard-fail under --require-gpu, one-shot
                        // eprintln otherwise) rather than a second hand-rolled warn.
                        debug_assert_eq!(
                            scores.len(),
                            candidates.len(),
                            "backend::batch_score_features must return one score per input"
                        );
                        backend::moe_runtime_degrade(&format!(
                            "caller-side score count mismatch: backend returned {} scores for {} candidates",
                            scores.len(),
                            candidates.len()
                        ));
                        score_features_on_cpu()
                    }
                    // `None` here is a genuine GPU dispatch failure that
                    // `batch_score_features` ALREADY degraded loudly (below-threshold
                    // `None` cannot occur: this branch only runs for large batches).
                    None => score_features_on_cpu(),
                }
            }
            #[cfg(not(feature = "gpu"))]
            {
                score_features_on_cpu()
            }
        };
        if let Some(t) = t_score {
            MOE_SCORE_NS.fetch_add(
                t.elapsed().as_nanos() as u64,
                std::sync::atomic::Ordering::Relaxed,
            );
        }
        scores
    }

    #[cfg(not(feature = "ml"))]
    {
        let _ = candidates; // LAW10: unused-binding marker (signature/borrowck/cfg/compile-time assert); no runtime effect, not a fallback
        let _ = config; // LAW10: unused-binding marker (signature/borrowck/cfg/compile-time assert); no runtime effect, not a fallback
        let _ = gpu_moe_timeout; // LAW10: unused-binding marker (signature/borrowck/cfg/compile-time assert); no runtime effect, not a fallback
        Vec::new()
    }
}

/// Return `true` when GPU scoring support is available in this build/runtime.
///
/// Honors the resolved runtime policy before touching the adapter path. A
/// caller asking after `--no-gpu` must get the same cheap "not available"
/// answer as `gpu_probe()` instead of triggering a wgpu adapter probe.
///
/// # Examples
///
/// ```rust
/// use keyhog_scanner::gpu::gpu_available;
/// let _ = gpu_available();
/// ```
pub fn gpu_available() -> bool {
    gpu_probe().available
}