tokitai-operator 0.1.0

//! Hardware abstraction: `ComputeHardware` trait and capability types.
//!
//! `ComputeHardware` is the abstract trait implemented by every
//! backend. It exposes `target()`, `device_capabilities()`, and the
//! fingerprint chain used by the math-bound audit trace. The
//! `HardwareTarget` enum (CpuScalar, GpuScaffold, RocmHip) is the
//! source-of-truth for the "what backends exist" question.
//!
//! `DeviceCapabilities` records the cap-fingerprint that the
//! audit-trace flow binds to. The cap-fingerprint is stable across
//! the same `(target, device_kind)` pair.
//!
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use crate::backend::BackendCapabilities;
use crate::object::Representation;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DeviceKind {
    Cpu,
    Gpu,
    Accelerator,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum MemorySpace {
    Host,
    Device,
    Unified,
    PinnedHost,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct HardwareTarget {
    pub id: String,
    pub kind: DeviceKind,
    pub memory_space: MemorySpace,
}

impl HardwareTarget {
    pub fn cpu_scalar() -> Self {
        Self {
            id: "cpu_scalar".to_string(),
            kind: DeviceKind::Cpu,
            memory_space: MemorySpace::Host,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DeviceCapabilities {
    pub target: HardwareTarget,
    pub backend: BackendCapabilities,
    pub exact: bool,
    pub deterministic: bool,
    pub supported_representations: Vec<String>,
    pub supported_domains: Vec<String>,
    pub semantic_degradations: Vec<String>,
}

impl DeviceCapabilities {
    pub fn from_backend(target: HardwareTarget, backend: BackendCapabilities) -> Self {
        Self {
            exact: backend.exact,
            deterministic: backend.deterministic,
            supported_representations: backend.supported_representations.clone(),
            supported_domains: backend.supported_domains.clone(),
            semantic_degradations: backend.semantic_degradations.clone(),
            target,
            backend,
        }
    }

    pub fn cpu_scalar() -> Self {
        let mut capabilities = Self::from_backend(
            HardwareTarget::cpu_scalar(),
            BackendCapabilities::cpu_scalar(),
        );
        capabilities.supported_domains.extend([
            "padic:fixed_precision".to_string(),
            "sheaf:finite_site".to_string(),
        ]);
        capabilities.supported_domains.sort();
        capabilities.supported_domains.dedup();
        capabilities.backend.supported_domains = capabilities.supported_domains.clone();
        capabilities
    }

    pub fn supports_representation(&self, representation: &Representation) -> bool {
        self.supported_representations
            .iter()
            .any(|item| item == &representation.id().0)
    }

    pub fn supports_domain(&self, domain: &str) -> bool {
        self.supported_domains.iter().any(|item| item == domain)
    }

    pub fn fingerprint(&self) -> String {
        let mut hasher = DefaultHasher::new();
        self.target.hash(&mut hasher);
        self.exact.hash(&mut hasher);
        self.deterministic.hash(&mut hasher);
        self.supported_representations.hash(&mut hasher);
        self.supported_domains.hash(&mut hasher);
        self.semantic_degradations.hash(&mut hasher);
        format!("hwcap-{:016x}", hasher.finish())
    }
}

pub trait ComputeHardware {
    fn target(&self) -> HardwareTarget;
    fn device_capabilities(&self) -> DeviceCapabilities;

    fn capability_fingerprint(&self) -> String {
        self.device_capabilities().fingerprint()
    }

    fn backend_capabilities(&self) -> BackendCapabilities {
        self.device_capabilities().backend
    }
}