speakrs 0.6.0

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};

use sha2::{Digest, Sha256};

use cudarc::driver::{CudaFunction, CudaModule};

use super::{ComputeCapability, CudaError, PtxTier};

/// Embeds a variant under exactly the GPU targets for which it is the best shipped PTX
///
/// The kernel check verifies these masks so tier-only builds include each area's
/// fallback without also including variants they do not need
macro_rules! tier_ptx {
    ([$($feature:literal),+], $stem:literal, [$($arch:literal),+]) => {{
        #[cfg(any($(feature = $feature),+))]
        const CUBINS: &[EmbeddedCubin] = &[$(EmbeddedCubin {
            arch: ComputeCapability::new($arch / 10, $arch % 10),
            bytes: shipped_cubin(concat!($stem, ".sm_", stringify!($arch), ".cubin")),
        }),+];
        #[cfg(any($(feature = $feature),+))]
        let ptx = Some(EmbeddedPtx {
            text: const { shipped_ptx(concat!($stem, ".ptx")) },
            cubins: CUBINS,
        });
        #[cfg(not(any($(feature = $feature),+)))]
        let ptx = None;
        ptx
    }};
}

/// A shipped PTX file, from the `speakrs-ptx-<tier>` crate that carries its tier
///
/// Each tier ships in its own crate so every published crate stays under the crates.io
/// upload limit. Call this only in const context outside tests: a run-time call would
/// link every file of every enabled tier crate into the binary
#[cfg(any(
    feature = "cuda-sm75",
    feature = "cuda-sm80",
    feature = "cuda-sm90",
    feature = "cuda-sm120"
))]
pub(crate) const fn shipped_ptx(path: &str) -> &'static str {
    if let Some(text) = speakrs_ptx_sm75::ptx(path) {
        return text;
    }

    #[cfg(any(feature = "cuda-sm80", feature = "cuda-sm90", feature = "cuda-sm120"))]
    if let Some(text) = speakrs_ptx_sm80::ptx(path) {
        return text;
    }

    #[cfg(feature = "cuda-sm120")]
    if let Some(text) = speakrs_ptx_sm120::ptx(path) {
        return text;
    }

    panic!("no enabled speakrs-ptx crate ships this PTX file")
}

/// A shipped cubin, from the crate that carries its PTX tier; const context only, like
/// [`shipped_ptx`]
#[cfg(any(
    feature = "cuda-sm75",
    feature = "cuda-sm80",
    feature = "cuda-sm90",
    feature = "cuda-sm120"
))]
const fn shipped_cubin(path: &str) -> &'static [u8] {
    if let Some(bytes) = speakrs_ptx_sm75::cubin(path) {
        return bytes;
    }

    #[cfg(any(feature = "cuda-sm80", feature = "cuda-sm90", feature = "cuda-sm120"))]
    if let Some(bytes) = speakrs_ptx_sm80::cubin(path) {
        return bytes;
    }

    #[cfg(feature = "cuda-sm120")]
    if let Some(bytes) = speakrs_ptx_sm120::cubin(path) {
        return bytes;
    }

    panic!("no enabled speakrs-ptx crate ships this cubin")
}

/// Set to `1` to disable cubins and use embedded PTX JIT for diagnosis
pub const FORCE_PTX_JIT_ENV: &str = "SPEAKRS_CUDA_FORCE_PTX_JIT";

/// Content identity of bytes embedded in the running binary
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ArtifactHash([u8; 32]);

impl ArtifactHash {
    /// Hash the actual artifact bytes, not a source-tree file or manifest claim
    pub fn of(bytes: &[u8]) -> Self {
        Self(Sha256::digest(bytes).into())
    }

    /// Read a canonical pinned hash; invalid pins fail at compile time in constants
    pub const fn from_hex(text: &str) -> Self {
        const fn digit(byte: u8) -> u8 {
            match byte {
                b'0'..=b'9' => byte - b'0',
                b'a'..=b'f' => byte - b'a' + 10,
                _ => panic!("hash must use lowercase hexadecimal"),
            }
        }
        assert!(
            text.len() == 64,
            "hash must contain 64 hexadecimal characters"
        );
        let mut bytes = [0; 32];
        let mut index = 0;
        while index < bytes.len() {
            bytes[index] =
                digit(text.as_bytes()[index * 2]) * 16 + digit(text.as_bytes()[index * 2 + 1]);
            index += 1;
        }
        Self(bytes)
    }
}

impl ArtifactHash {
    /// Byte equality usable in compile-time table validation
    pub(crate) const fn const_eq(self, other: Self) -> bool {
        let mut index = 0;
        while index < self.0.len() {
            if self.0[index] != other.0[index] {
                return false;
            }
            index += 1;
        }
        true
    }
}

impl std::fmt::Display for ArtifactHash {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        for byte in self.0 {
            write!(f, "{byte:02x}")?;
        }
        Ok(())
    }
}

/// Immutable bytes whose content identity can be reused across plans and runtimes
struct EmbeddedArtifact(&'static [u8]);

impl EmbeddedArtifact {
    fn sha256(self) -> ArtifactHash {
        static HASHES: OnceLock<Mutex<HashMap<(usize, usize), ArtifactHash>>> = OnceLock::new();
        let mut hashes = HASHES
            .get_or_init(Mutex::default)
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        // static slices cannot change or reuse their address; length distinguishes subviews
        *hashes
            .entry((self.0.as_ptr() as usize, self.0.len()))
            .or_insert_with(|| ArtifactHash::of(self.0))
    }
}

/// The exact artifact accepted by the driver for a kernel area
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum LoadedArtifact {
    /// Ready SASS compiled for this exact device capability
    Cubin {
        /// Exact cubin architecture, never an older compatible architecture
        arch: ComputeCapability,
        /// Hash of the binary bytes handed to the driver
        sha256: ArtifactHash,
    },
    /// Embedded PTX compiled by the device driver
    PtxJit {
        /// Hash of the PTX text handed to the driver
        sha256: ArtifactHash,
    },
}

/// One embedded binary, available only alongside its source PTX variant
#[derive(Debug, Clone, Copy)]
pub(super) struct EmbeddedCubin {
    pub arch: ComputeCapability,
    pub bytes: &'static [u8],
}

impl EmbeddedCubin {
    pub(super) fn sha256(self) -> ArtifactHash {
        EmbeddedArtifact(self.bytes).sha256()
    }
}

/// The PTX fallback and exact-architecture binaries for one tier
#[derive(Debug, Clone, Copy)]
pub(super) struct EmbeddedPtx {
    pub text: &'static str,
    pub cubins: &'static [EmbeddedCubin],
}

impl EmbeddedPtx {
    pub(super) fn sha256(self) -> ArtifactHash {
        EmbeddedArtifact(self.text.as_bytes()).sha256()
    }

    pub fn cubin(self, device: ComputeCapability) -> Option<EmbeddedCubin> {
        self.cubins
            .iter()
            .copied()
            .find(|cubin| cubin.arch == device)
    }
}

/// A complete module identity, resolved before the driver sees any bytes
///
/// Loading, the runtime's module cache and selection tokens compare whole
/// requests: the area, the PTX tier whose bytes are loaded, and the exact artifact
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub(crate) struct ModuleRequest {
    area: KernelModule,
    tier: PtxTier,
    artifact: LoadedArtifact,
}

impl ModuleRequest {
    /// A request whose cubin, if any, can run the tier; invalid constants fail to compile
    pub(crate) const fn new(area: KernelModule, tier: PtxTier, artifact: LoadedArtifact) -> Self {
        if let LoadedArtifact::Cubin { arch, .. } = artifact {
            let minimum = tier.min_capability();
            assert!(
                arch.major > minimum.major
                    || arch.major == minimum.major && arch.minor >= minimum.minor,
                "a cubin cannot be older than its PTX tier"
            );
        }
        Self {
            area,
            tier,
            artifact,
        }
    }

    pub(crate) const fn area(self) -> KernelModule {
        self.area
    }

    pub(crate) const fn tier(self) -> PtxTier {
        self.tier
    }

    pub(crate) const fn artifact(self) -> LoadedArtifact {
        self.artifact
    }

    /// Reject a different cached execution identity without replacing its module
    pub(crate) fn check_cached(self, cached: Self) -> Result<(), CudaError> {
        if self != cached {
            return Err(CudaError::ArtifactUnavailable {
                module: self.area.name(),
                artifact: self.artifact,
            });
        }
        Ok(())
    }

    /// Resolve a forced JIT diagnostic without changing the selected tier
    pub(crate) fn diagnostic_request(self, force_jit: bool) -> Result<Self, CudaError> {
        if !force_jit {
            return Ok(self);
        }
        let embedded =
            self.area
                .variants()
                .embedded(self.tier)
                .ok_or(CudaError::ArtifactUnavailable {
                    module: self.area.name(),
                    artifact: self.artifact,
                })?;
        Ok(self.ptx_jit(ArtifactHash::of(embedded.text.as_bytes())))
    }

    /// The same area and tier, loaded through driver JIT of the tier's embedded PTX
    pub(crate) const fn ptx_jit(self, sha256: ArtifactHash) -> Self {
        Self::new(self.area, self.tier, LoadedArtifact::PtxJit { sha256 })
    }
}

/// Failure to load the requested artifact, without substituting another artifact
#[derive(Debug, PartialEq, Eq)]
pub(super) enum ArtifactLoadError<E> {
    Unavailable,
    Driver(E),
}

/// Load exactly the requested bytes; driver rejection never tries the other format
pub(super) fn load_artifact<T, E>(
    requested: LoadedArtifact,
    cubin: Option<EmbeddedCubin>,
    ptx_sha256: ArtifactHash,
    load_cubin: impl FnOnce(&[u8]) -> Result<T, E>,
    load_ptx: impl FnOnce() -> Result<T, E>,
) -> Result<(T, LoadedArtifact), ArtifactLoadError<E>> {
    let loaded = match requested {
        LoadedArtifact::PtxJit { sha256 } if sha256 == ptx_sha256 => load_ptx(),
        LoadedArtifact::Cubin { arch, sha256 } => {
            let Some(cubin) = cubin.filter(|cubin| cubin.arch == arch && cubin.sha256() == sha256)
            else {
                return Err(ArtifactLoadError::Unavailable);
            };
            load_cubin(cubin.bytes)
        }
        _ => return Err(ArtifactLoadError::Unavailable),
    };
    loaded
        .map(|loaded| (loaded, requested))
        .map_err(ArtifactLoadError::Driver)
}

/// A cuda-oxide kernel area, embedded as committed PTX
///
/// `cargo xtask cuda-kernels build` regenerates `ptx/<area>.<tier>.ptx` in the
/// `speakrs-ptx-<tier>` crate on a GPU box,
/// and `cargo xtask cuda-kernels check` fails when that PTX is stale. Each GPU tier
/// feature embeds the best shipped variant of each area for that target. The `cuda`
/// feature enables all four targets and embeds all shipped variants
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum KernelModule {
    /// Toolchain probe kernels, which only the tier-dispatch tests load
    #[cfg(test)]
    Probe,
    /// Filterbank front-end kernels
    Fbank,
    /// ResNet34 multi-mask embedding kernels
    Embedding,
    /// Segmentation kernels
    Segmentation,
    /// Candidate kernels for the ResNet convolutions, kept apart from the Library-owned
    /// areas so module identities cannot overlap
    // the candidate areas are unused until a candidate plan loads one
    #[allow(dead_code)]
    Resnet,
    /// Candidate kernels for the LSTM stack
    #[allow(dead_code)]
    Lstm,
    /// Candidate kernels for the Sinc producer
    #[allow(dead_code)]
    Sincnet,
    /// The LSTM stack with its own input projections, which needs no library
    LstmProj,
    /// Record-owned filterbank DFT producer, separate from always-on Fbank
    FbankDft,
    /// Record-owned segmentation dense operators
    Segdense,
    /// Wide trunk convolutions: the stem, the 1x1 shortcuts, the strided layers from
    /// 64 channels and the 128- and 256-channel layers
    Wideconv,
}

impl KernelModule {
    /// The area name, which is also the PTX file stem
    pub const fn name(self) -> &'static str {
        match self {
            #[cfg(test)]
            Self::Probe => "probe",
            Self::Fbank => "fbank",
            Self::Embedding => "embedding",
            Self::Segmentation => "segmentation",
            Self::Resnet => "resnet",
            Self::Lstm => "lstm",
            Self::Sincnet => "sincnet",
            Self::LstmProj => "lstmproj",
            Self::FbankDft => "fbankdft",
            Self::Segdense => "segdense",
            Self::Wideconv => "wideconv",
        }
    }

    /// The build metadata embedded beside this area's artifact bytes
    pub(crate) const fn manifest(self) -> &'static str {
        match self {
            #[cfg(test)]
            Self::Probe => include_str!("ptx/probe.manifest"),
            Self::Fbank => include_str!("ptx/fbank.manifest"),
            Self::Embedding => include_str!("ptx/embedding.manifest"),
            Self::Segmentation => include_str!("ptx/segmentation.manifest"),
            Self::Resnet => include_str!("ptx/resnet.manifest"),
            Self::Lstm => include_str!("ptx/lstm.manifest"),
            Self::Sincnet => include_str!("ptx/sincnet.manifest"),
            Self::Segdense => include_str!("ptx/segdense.manifest"),
            Self::FbankDft => include_str!("ptx/fbankdft.manifest"),
            Self::Wideconv => include_str!("ptx/wideconv.manifest"),
            Self::LstmProj => include_str!("ptx/lstmproj.manifest"),
        }
    }

    /// The PTX variants embedded in this build
    pub const fn variants(self) -> AreaPtx {
        match self {
            #[cfg(test)]
            Self::Probe => AreaPtx {
                sm75: tier_ptx!(["cuda-sm75"], "ptx/probe.sm75", [75, 80, 86, 89, 90, 120]),
                sm80: tier_ptx!(
                    ["cuda-sm80", "cuda-sm90", "cuda-sm120"],
                    "ptx/probe.sm80",
                    [80, 86, 89, 90, 120]
                ),
                sm90: None,
                sm120: None,
            },
            Self::Fbank => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/fbank.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            Self::Embedding => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/embedding.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            Self::Segmentation => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/segmentation.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            // the capability 12.0 binding pins the sm75 bytes, so every target embeds
            // them; the sm80 tier adds the TF32 tensor-core entries
            Self::Resnet => AreaPtx {
                sm75: tier_ptx!(
                    ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                    "ptx/resnet.sm75",
                    [75, 80, 86, 89, 90, 120]
                ),
                sm80: tier_ptx!(
                    ["cuda-sm80", "cuda-sm90", "cuda-sm120"],
                    "ptx/resnet.sm80",
                    [80, 86, 89, 90, 120]
                ),
                sm90: None,
                sm120: None,
            },
            Self::Lstm => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/lstm.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            Self::Sincnet => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/sincnet.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            // the sm80 tier adds tensor-core kernels; on cc 8.0 and newer it replaces
            // the sm75 variant, whose tensor-core entries are trapping stubs
            Self::Segdense => AreaPtx {
                sm75: tier_ptx!(
                    ["cuda-sm75"],
                    "ptx/segdense.sm75",
                    [75, 80, 86, 89, 90, 120]
                ),
                sm80: tier_ptx!(
                    ["cuda-sm80", "cuda-sm90", "cuda-sm120"],
                    "ptx/segdense.sm80",
                    [80, 86, 89, 90, 120]
                ),
                sm90: None,
                sm120: None,
            },
            Self::FbankDft => AreaPtx::baseline(tier_ptx!(
                ["cuda-sm75", "cuda-sm80", "cuda-sm90", "cuda-sm120"],
                "ptx/fbankdft.sm75",
                [75, 80, 86, 89, 90, 120]
            )),
            Self::LstmProj => AreaPtx {
                sm75: tier_ptx!(
                    ["cuda-sm75"],
                    "ptx/lstmproj.sm75",
                    [75, 80, 86, 89, 90, 120]
                ),
                sm80: tier_ptx!(
                    ["cuda-sm80", "cuda-sm90"],
                    "ptx/lstmproj.sm80",
                    [80, 86, 89, 90, 120]
                ),
                sm90: None,
                sm120: tier_ptx!(["cuda-sm120"], "ptx/lstmproj.sm120", [120]),
            },
            Self::Wideconv => AreaPtx {
                sm75: tier_ptx!(
                    ["cuda-sm75"],
                    "ptx/wideconv.sm75",
                    [75, 80, 86, 89, 90, 120]
                ),
                sm80: tier_ptx!(
                    ["cuda-sm80", "cuda-sm90", "cuda-sm120"],
                    "ptx/wideconv.sm80",
                    [80, 86, 89, 90, 120]
                ),
                sm90: None,
                sm120: None,
            },
        }
    }
}

/// The PTX variants of one area embedded by this build
///
/// A tier feature may embed an older variant when that is the best the area ships
///
/// Every variant exports the same kernels with the same parameters, which
/// `cargo xtask cuda-kernels check` verifies, so the host can load any of them
#[derive(Debug, Clone, Copy)]
pub struct AreaPtx {
    sm75: Option<EmbeddedPtx>,
    sm80: Option<EmbeddedPtx>,
    sm90: Option<EmbeddedPtx>,
    sm120: Option<EmbeddedPtx>,
}

impl AreaPtx {
    /// Embedded variants without cubins, for host-only tests such as a fake higher tier
    #[cfg(all(test, feature = "_cuda-libraries"))]
    pub(crate) fn fixture(variants: &[(PtxTier, &'static str)]) -> Self {
        let mut area = Self::baseline(None);
        for (tier, text) in variants {
            let ptx = Some(EmbeddedPtx { text, cubins: &[] });
            match tier {
                PtxTier::Sm75 => area.sm75 = ptx,
                PtxTier::Sm80 => area.sm80 = ptx,
                PtxTier::Sm90 => area.sm90 = ptx,
                PtxTier::Sm120 => area.sm120 = ptx,
            }
        }
        area
    }

    const fn baseline(sm75: Option<EmbeddedPtx>) -> Self {
        Self {
            sm75,
            sm80: None,
            sm90: None,
            sm120: None,
        }
    }

    /// The embedded variants, lowest tier first
    pub fn iter(&self) -> impl Iterator<Item = (PtxTier, &'static str)> {
        [
            (PtxTier::Sm75, self.sm75),
            (PtxTier::Sm80, self.sm80),
            (PtxTier::Sm90, self.sm90),
            (PtxTier::Sm120, self.sm120),
        ]
        .into_iter()
        .filter_map(|(tier, ptx)| Some((tier, ptx?.text)))
    }

    /// The embedded bytes for one specific tier
    pub(super) fn embedded(&self, tier: PtxTier) -> Option<EmbeddedPtx> {
        match tier {
            PtxTier::Sm75 => self.sm75,
            PtxTier::Sm80 => self.sm80,
            PtxTier::Sm90 => self.sm90,
            PtxTier::Sm120 => self.sm120,
        }
    }

    /// The highest runnable embedded tier, without hashing its artifact
    pub(crate) fn runnable_tier(
        self,
        limit: PtxTier,
        device: ComputeCapability,
    ) -> Option<PtxTier> {
        [PtxTier::Sm120, PtxTier::Sm90, PtxTier::Sm80, PtxTier::Sm75]
            .into_iter()
            .find(|tier| {
                *tier <= limit && tier.min_capability() <= device && self.embedded(*tier).is_some()
            })
    }

    /// Resolve EmbeddedExact: the best runnable tier and this device's cubin or PTX
    pub(crate) fn driver_request(
        self,
        area: KernelModule,
        limit: PtxTier,
        device: ComputeCapability,
    ) -> Option<ModuleRequest> {
        let tier = self.runnable_tier(limit, device)?;
        let ptx = self.embedded(tier)?;
        let artifact = ptx.cubin(device).map_or_else(
            || LoadedArtifact::PtxJit {
                sha256: ptx.sha256(),
            },
            |cubin| LoadedArtifact::Cubin {
                arch: device,
                sha256: cubin.sha256(),
            },
        );
        Some(ModuleRequest::new(area, tier, artifact))
    }

    /// The highest embedded variant at or below `limit`
    ///
    /// Production never resolves a tier this way: a binding names its tier. Only tests
    /// and explicit development requests ask for the best embedded variant
    #[cfg(test)]
    pub fn select(&self, limit: PtxTier) -> Option<(PtxTier, &'static str)> {
        self.iter().filter(|(tier, _)| *tier <= limit).last()
    }

    /// Resolve a runnable variant or name the target feature that is missing
    #[cfg(test)]
    pub(super) fn resolve(
        &self,
        area: KernelModule,
        limit: PtxTier,
        device: ComputeCapability,
    ) -> Result<(PtxTier, &'static str), CudaError> {
        self.select(limit)
            .filter(|(tier, _)| tier.min_capability() <= device)
            .ok_or(CudaError::AreaTierNotCompiledIn {
                area: area.name(),
                tier: limit,
                device,
                feature: limit.feature(),
            })
    }
}

/// A kernel module loaded into a CUDA context
#[derive(Debug, Clone)]
pub struct LoadedKernels {
    request: ModuleRequest,
    inner: Arc<CudaModule>,
    ptx_sha256: ArtifactHash,
}

impl LoadedKernels {
    pub(super) fn new(
        request: ModuleRequest,
        inner: Arc<CudaModule>,
        ptx_sha256: ArtifactHash,
    ) -> Self {
        Self {
            request,
            inner,
            ptx_sha256,
        }
    }

    /// The complete identity the driver accepted
    pub(crate) fn request(&self) -> ModuleRequest {
        self.request
    }

    /// The PTX tier of the variant that was loaded
    pub fn tier(&self) -> PtxTier {
        self.request.tier
    }

    /// Identity of the artifact that the driver successfully loaded
    #[cfg(test)]
    pub fn artifact(&self) -> LoadedArtifact {
        self.request.artifact
    }

    /// Identity of the fallback PTX text embedded alongside the loaded artifact
    pub fn ptx_sha256(&self) -> ArtifactHash {
        self.ptx_sha256
    }

    /// Looks up a kernel by its PTX entry name, which is the Rust function name
    pub fn function(&self, kernel: &str) -> Result<CudaFunction, CudaError> {
        self.inner
            .load_function(kernel)
            .map_err(|source| CudaError::KernelMissing {
                module: self.request.area.name(),
                kernel: kernel.to_string(),
                source,
            })
    }
}

#[cfg(test)]
mod tests {
    use super::{AreaPtx, ComputeCapability, CudaError, KernelModule, PtxTier};

    #[test]
    fn embedded_hash_cache_preserves_content_and_subview_identity() {
        use super::{ArtifactHash, EmbeddedArtifact};

        static FIRST: &[u8] = b"first artifact";
        static SECOND: &[u8] = b"other artifact";
        for bytes in [FIRST, SECOND, &FIRST[..5], FIRST, SECOND] {
            assert_eq!(EmbeddedArtifact(bytes).sha256(), ArtifactHash::of(bytes));
        }
        assert_ne!(
            EmbeddedArtifact(FIRST).sha256(),
            EmbeddedArtifact(SECOND).sha256()
        );
        assert_ne!(
            EmbeddedArtifact(FIRST).sha256(),
            EmbeddedArtifact(&FIRST[..5]).sha256()
        );
    }

    #[test]
    #[ignore = "GPU artifact proof; run under the shared GPU flock"]
    fn driver_artifacts_resolve_every_ptx_entry() -> Result<(), CudaError> {
        use super::{ArtifactHash, LoadedArtifact, LoadedKernels, load_artifact};
        use crate::inference::cuda::CudaRuntime;
        use cudarc::nvrtc::Ptx;
        use std::fs::{OpenOptions, TryLockError};

        // the proof process must be serialized, including context creation and drops
        let lock = OpenOptions::new()
            .read(true)
            .write(true)
            .open("/workspace/gpu-bench.lock")
            .expect("proof runs under the shared GPU lock");
        assert!(matches!(lock.try_lock(), Err(TryLockError::WouldBlock)));
        let runtime = CudaRuntime::new(0)?;
        let device = runtime.compute_capability();
        let force_jit =
            std::env::var_os(super::FORCE_PTX_JIT_ENV).is_some_and(|value| value == "1");
        for area in [
            KernelModule::Probe,
            KernelModule::Fbank,
            KernelModule::Embedding,
            KernelModule::Segmentation,
            KernelModule::Resnet,
            KernelModule::Lstm,
            KernelModule::Sincnet,
            KernelModule::Segdense,
            KernelModule::FbankDft,
        ] {
            for (tier, ptx) in area.variants().iter() {
                if tier.min_capability() > device {
                    continue;
                }
                let embedded = area.variants().embedded(tier).expect("embedded variant");
                let cubin = embedded
                    .cubin(device)
                    .expect("proof device has exact cubins");
                let cubin_key = LoadedArtifact::Cubin {
                    arch: device,
                    sha256: ArtifactHash::of(cubin.bytes),
                };
                let ptx_hash = ArtifactHash::of(ptx.as_bytes());
                let (module, artifact) = load_artifact(
                    if force_jit {
                        LoadedArtifact::PtxJit { sha256: ptx_hash }
                    } else {
                        cubin_key
                    },
                    Some(cubin),
                    ptx_hash,
                    |bytes| {
                        runtime
                            .context()
                            .load_module(Ptx::from_binary(bytes.to_vec()))
                    },
                    || runtime.context().load_module(Ptx::from_src(ptx)),
                )
                .map_err(|error| match error {
                    super::ArtifactLoadError::Driver(source) => CudaError::Driver(source),
                    super::ArtifactLoadError::Unavailable => CudaError::ArtifactUnavailable {
                        module: area.name(),
                        artifact: cubin_key,
                    },
                })?;
                if force_jit {
                    assert_eq!(artifact, LoadedArtifact::PtxJit { sha256: ptx_hash });
                    assert_ne!(artifact, cubin_key);
                } else {
                    // this proof must fail, not silently accept a rejected cubin's JIT
                    assert_eq!(artifact, cubin_key);
                }
                let loaded = LoadedKernels::new(
                    super::ModuleRequest::new(area, tier, artifact),
                    module,
                    ptx_hash,
                );
                let mut entries = 0;
                for line in ptx.lines() {
                    let Some(entry) = line.trim().strip_prefix(".visible .entry ") else {
                        continue;
                    };
                    let name = entry.split(['(', ' ', '\t']).next().expect("entry name");
                    loaded.function(name)?;
                    entries += 1;
                }
                assert!(entries > 0, "proof requires declared entry points");
                println!(
                    "artifact_proof area={} tier={tier} device={device} artifact={artifact:?} embedded_ptx_sha256={ptx_hash} entries={entries}",
                    area.name()
                );
            }
            // a record-owned area without a binding on this device has no production load
            if runtime.production_module(area)?.is_some() {
                let loaded = runtime.load_kernels(area)?;
                assert_eq!(loaded.artifact(), runtime.load_kernels(area)?.artifact());
            }
        }
        runtime.synchronize()
    }

    #[test]
    fn pinned_loading_never_substitutes_artifacts() {
        use super::{
            ArtifactHash, ArtifactLoadError, EmbeddedCubin, LoadedArtifact, load_artifact,
        };
        let hash = ArtifactHash::of(b"ptx");
        let jit = LoadedArtifact::PtxJit { sha256: hash };
        for device in [ComputeCapability::new(8, 9), ComputeCapability::new(12, 0)] {
            let cubin = EmbeddedCubin {
                arch: device,
                bytes: b"cubin",
            };
            let binary = LoadedArtifact::Cubin {
                arch: device,
                sha256: ArtifactHash::of(cubin.bytes),
            };
            let result = load_artifact(
                binary,
                Some(cubin),
                hash,
                |bytes| Ok::<_, &str>(bytes.to_vec()),
                || panic!("cubin request cannot JIT"),
            )
            .unwrap();
            assert_eq!(result, (cubin.bytes.to_vec(), binary));
            assert_eq!(
                load_artifact::<(), _>(
                    binary,
                    Some(cubin),
                    hash,
                    |_| Err("driver refusal"),
                    || panic!("refused cubin must not JIT")
                ),
                Err(ArtifactLoadError::Driver("driver refusal"))
            );
            assert_eq!(
                load_artifact::<(), &str>(
                    binary,
                    None,
                    hash,
                    |_| panic!("missing cubin"),
                    || panic!("missing cubin must not JIT")
                ),
                Err(ArtifactLoadError::Unavailable)
            );
            assert_eq!(
                load_artifact(
                    jit,
                    Some(cubin),
                    hash,
                    |_| panic!("PTX request cannot load cubin"),
                    || Ok::<_, &str>("jit")
                ),
                Ok(("jit", jit))
            );
            assert_eq!(
                load_artifact::<(), _>(
                    jit,
                    Some(cubin),
                    hash,
                    |_| panic!("refused PTX cannot load cubin"),
                    || Err("PTX rejected")
                ),
                Err(ArtifactLoadError::Driver("PTX rejected"))
            );
        }
        let wrong = LoadedArtifact::PtxJit {
            sha256: ArtifactHash::of(b"stale ptx"),
        };
        assert_eq!(
            load_artifact::<(), &str>(
                wrong,
                None,
                hash,
                |_| panic!("wrong hash"),
                || panic!("wrong hash")
            ),
            Err(ArtifactLoadError::Unavailable)
        );
    }

    #[test]
    fn unknown_device_never_uses_an_older_cubin() {
        let Some(tier) = KernelModule::Fbank
            .variants()
            .iter()
            .next()
            .map(|(tier, _)| tier)
        else {
            return;
        };
        let embedded = KernelModule::Fbank.variants().embedded(tier).unwrap();
        for device in [ComputeCapability::new(8, 9), ComputeCapability::new(12, 0)] {
            assert_eq!(embedded.cubin(device).unwrap().arch, device);
        }
        assert!(embedded.cubin(ComputeCapability::new(12, 1)).is_none());
        assert!(embedded.cubin(ComputeCapability::new(10, 0)).is_none());
    }

    #[test]
    fn gpu_targets_embed_each_areas_best_shipped_variant() {
        let production = KernelModule::Fbank.variants();
        assert_eq!(
            production.iter().map(|(tier, _)| tier).collect::<Vec<_>>(),
            if PtxTier::compiled_in().next().is_some() {
                vec![PtxTier::Sm75]
            } else {
                vec![]
            }
        );
        let probe = KernelModule::Probe.variants();
        assert_eq!(probe.sm75.is_some(), cfg!(feature = "cuda-sm75"));
        assert_eq!(
            probe.sm80.is_some(),
            cfg!(any(
                feature = "cuda-sm80",
                feature = "cuda-sm90",
                feature = "cuda-sm120"
            ))
        );
        assert_eq!(PtxTier::Sm75.is_compiled_in(), cfg!(feature = "cuda-sm75"));
        let selected = probe.select(PtxTier::Sm120).map(|(tier, _)| tier);
        assert_eq!(
            selected,
            if probe.sm80.is_some() {
                Some(PtxTier::Sm80)
            } else if probe.sm75.is_some() {
                Some(PtxTier::Sm75)
            } else {
                None
            }
        );
    }

    #[test]
    fn missing_area_variant_names_area_device_and_target_feature() {
        let area = AreaPtx {
            sm75: None,
            sm80: Some(super::EmbeddedPtx {
                text: "test PTX",
                cubins: &[],
            }),
            sm90: None,
            sm120: None,
        };
        let device = ComputeCapability::new(7, 5);
        let error = area
            .resolve(KernelModule::Probe, PtxTier::Sm75, device)
            .unwrap_err();
        assert!(matches!(error, CudaError::AreaTierNotCompiledIn {
            area: "probe", tier: PtxTier::Sm75, device: actual, feature: "cuda-sm75"
        } if actual == device));
        assert!(
            error.to_string().contains("probe")
                && error.to_string().contains("7.5")
                && error.to_string().contains("cuda-sm75")
        );
    }
}