wai-quantum 0.4.4

//! The gate kernel for CUDA devices, and what names a CUDA member: the
//! kernel's text (`lanes64.cu`), its implementation id, a device's identity
//! and the mutants a conforming device must be caught running. No device is
//! touched here, so all of it is testable on any host; the CUDA member itself
//! is `quantum_fabric_cuda`.
//!
//! **Why the kernel is bit-exact.** Every CUDA device has 64-bit integer
//! arithmetic, so the kernel is a transcription of [`lanes::exact_row`] and
//! [`lanes::apply_words`] rather than of the 32-bit [`lanes`] steps a WGSL
//! kernel needs:
//!
//! - a gate entry and an amplitude are 32-bit values, so their product is at
//!   most `2^62` in size, and `fx(a, b) = (a·b + 2^29) >> 30` is exact in 64
//!   bits, with `>>` arithmetic as Rust's is on `i64`;
//! - a row's sum of four such terms is at most `2^34` in size, exact too;
//! - the lane keeps the sum's low 32 bits, as [`lanes::apply_words`] does,
//!   and a sum outside 32 bits records the gate's op-log index, the lowest
//!   winning ([`lanes::run_words`]).
//!
//! Each thread owns one amplitude pair and a launch is one gate, so no thread
//! reads what another writes within a launch, and the order threads run in
//! cannot change a bit. The member's probe checks the device's products and
//! shifts against [`lanes::mul_wide`] and [`lanes::fxmul_wide`] before the
//! device is a member, and its certificates check whole runs against the CPU's
//! pinned corpus digests.
//!
//! [`lanes::exact_row`]: super::lanes
//! [`lanes::apply_words`]: super::lanes::apply_words
//! [`lanes::run_words`]: super::lanes::run_words
//! [`lanes::mul_wide`]: super::lanes::mul_wide
//! [`lanes::fxmul_wide`]: super::lanes::fxmul_wide
//! [`lanes`]: super::lanes

use super::FabricClass;
use super::wgsl::{Detector, Mutant};

/// The kernel, `lanes64.cu`, exactly as it is compiled.
pub const SOURCE: &str = include_str!("lanes64.cu");

/// The implementation name the id starts with.
pub const IMPLEMENTATION: &str = "cuda-lanes64";

/// Version of the host-side layout: interleaved `int2` lanes, one launch per
/// gate with the gate's words as kernel parameters, the report word (the
/// lowest out-of-range op index, `u32::MAX` for none), the probe's `int2`
/// cases and `uint4` outputs, and the launch shape ([`BLOCK`],
/// [`GRID_MAX`]). Bumped on any change to them; it is part of the
/// implementation id.
pub const HOST_LAYOUT: u32 = 1;

/// Threads per block.
pub const BLOCK: u32 = 256;

/// Blocks per launch, at most; the kernel strides over the rest.
pub const GRID_MAX: u32 = 4096;

/// The options the runtime compiler is given, besides the target
/// architecture, which is the device's own: no fused multiply-add, though the
/// kernel has no floating point for it to touch. Fast math is off unless
/// asked for, and is not.
pub const COMPILE_OPTIONS: [&str; 2] = ["--std=c++17", "--fmad=false"];

const DOMAIN_IMPLEMENTATION: &[u8] = b"wai:quantum-fabric-cuda\x01";
const DOMAIN_MEMBER: &[u8] = b"wai:quantum-fabric-cuda-member\x01";

fn hex(bytes: &[u8]) -> String {
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

/// A length-prefixed byte string, so that no field reads as part of another.
fn put(o: &mut Vec<u8>, bytes: &[u8]) {
    o.extend_from_slice(&(bytes.len() as u64).to_be_bytes());
    o.extend_from_slice(bytes);
}

/// `"cuda-lanes64/<16 hex>"`: the first 8 bytes of
/// BLAKE3(`"wai:quantum-fabric-cuda\x01"` ‖ put([`SOURCE`]) ‖ put of each of
/// [`COMPILE_OPTIONS`] ‖ [`HOST_LAYOUT`] as a big-endian `u32`), in hex. Any
/// edit to the kernel's text, its compile options or the host layout changes
/// it. The target architecture is the device's and is part of the member id
/// instead.
pub fn implementation_id() -> String {
    implementation_id_of(SOURCE)
}

/// [`implementation_id`] for another kernel text, such as a mutant's.
pub fn implementation_id_of(source: &str) -> String {
    let mut o = DOMAIN_IMPLEMENTATION.to_vec();
    put(&mut o, source.as_bytes());
    for option in COMPILE_OPTIONS {
        put(&mut o, option.as_bytes());
    }
    o.extend_from_slice(&HOST_LAYOUT.to_be_bytes());
    format!("{IMPLEMENTATION}/{}", hex(&blake3::hash(&o).as_bytes()[..8]))
}

/// What a CUDA member is, as its id hashes it. The member fills it from the
/// driver; the fields are plain values here so that the id is testable
/// without a device.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DeviceIdentity {
    pub name: String,
    /// The device's UUID, in hex.
    pub uuid: String,
    /// The compute capability, `(major, minor)`: the architecture the kernel
    /// is compiled for.
    pub compute: (u32, u32),
    /// `"<domain>:<bus>:<device>"`, in hex.
    pub pci: String,
    /// The driver's CUDA version, as `cuDriverGetVersion` reports it
    /// (`major·1000 + minor·10`).
    pub driver: u32,
    /// The runtime compiler's version, `(major, minor)`.
    pub compiler: (u32, u32),
    /// The device's memory, bytes.
    pub memory: u64,
}

impl DeviceIdentity {
    /// Always [`FabricClass::Gpu`].
    pub fn class(&self) -> FabricClass {
        FabricClass::Gpu
    }

    /// The target architecture the kernel is compiled for:
    /// `"compute_<major><minor>"`.
    pub fn arch(&self) -> String {
        format!("compute_{}{}", self.compute.0, self.compute.1)
    }

    /// `"local/gpu-cuda-<d12>"`. `d12` is the first 12 hex digits of
    /// BLAKE3(`"wai:quantum-fabric-cuda-member\x01"` ‖ put(name) ‖ put(uuid) ‖
    /// compute major ‖ compute minor ‖ put(pci) ‖ driver ‖ compiler major ‖
    /// compiler minor ‖ memory), the numbers big-endian. It sorts after every
    /// `local/cpu-<arch>`, so with no energy figures
    /// [`place`](super::place) keeps the CPU reference.
    pub fn member_id(&self) -> String {
        let mut o = DOMAIN_MEMBER.to_vec();
        put(&mut o, self.name.as_bytes());
        put(&mut o, self.uuid.as_bytes());
        o.extend_from_slice(&self.compute.0.to_be_bytes());
        o.extend_from_slice(&self.compute.1.to_be_bytes());
        put(&mut o, self.pci.as_bytes());
        o.extend_from_slice(&self.driver.to_be_bytes());
        o.extend_from_slice(&self.compiler.0.to_be_bytes());
        o.extend_from_slice(&self.compiler.1.to_be_bytes());
        o.extend_from_slice(&self.memory.to_be_bytes());
        format!("local/gpu-cuda-{}", hex(&blake3::hash(&o).as_bytes()[..6]))
    }

    /// The key a runner declares in CI: `"gpu-cuda-sm<major><minor>"`, for
    /// instance `"gpu-cuda-sm89"`.
    pub fn key(&self) -> String {
        format!("gpu-cuda-sm{}{}", self.compute.0, self.compute.1)
    }
}

/// The mutants: textual edits of [`SOURCE`] that a conforming device must be
/// caught running, each by the detector named. Each `find` occurs exactly
/// once in [`SOURCE`].
pub const MUTANTS: [Mutant; 6] = [
    // The rounding constant is half what it should be.
    Mutant { name: "C1", find: "((i64)a * (i64)b + ((i64)1 << 29)) >> 30;", replace: "((i64)a * (i64)b + ((i64)1 << 28)) >> 30;", detector: Detector::Probe },
    // The product is shifted by one bit too few.
    Mutant { name: "C2", find: "((i64)a * (i64)b + ((i64)1 << 29)) >> 30;", replace: "((i64)a * (i64)b + ((i64)1 << 29)) >> 29;", detector: Detector::Probe },
    // The probe's product is taken in 32 bits.
    Mutant { name: "C3", find: "u64 prod = (u64)((i64)c.x * (i64)c.y);", replace: "u64 prod = (u64)(unsigned int)(c.x * c.y);", detector: Detector::Probe },
    // The control test is inverted.
    Mutant { name: "C4", find: "if ((i & (u64)ctrl) != (u64)ctrl) {", replace: "if ((i & (u64)ctrl) == (u64)ctrl) {", detector: Detector::Certificate },
    // A sign of row 0's real part is wrong, on the imaginary part of the
    // top-right entry, which `Y` has. (The top-left entry is real for every
    // gate of the set, so the same edit on its imaginary part changes no run.)
    Mutant { name: "C5", find: "fx(m2, a1.x) - fx(m3, a1.y);", replace: "fx(m2, a1.x) + fx(m3, a1.y);", detector: Detector::Certificate },
    // The range guard never fires.
    Mutant { name: "C6", find: "if (!(fits32(o0r) && fits32(o0i) && fits32(o1r) && fits32(o1i))) {", replace: "if (false) {", detector: Detector::Range },
];

/// [`SOURCE`] with `m` applied. Panics unless `m.find` occurs exactly once.
pub fn mutate(m: &Mutant) -> String {
    let n = SOURCE.matches(m.find).count();
    assert_eq!(n, 1, "mutant {}: `{}` occurs {n} times in the kernel", m.name, m.find);
    SOURCE.replacen(m.find, m.replace, 1)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn identity() -> DeviceIdentity {
        DeviceIdentity {
            name: "a device".into(),
            uuid: "00112233445566778899aabbccddeeff".into(),
            compute: (8, 9),
            pci: "0000:01:00".into(),
            driver: 13020,
            compiler: (12, 8),
            memory: 6 << 30,
        }
    }

    #[test]
    fn every_field_of_the_identity_is_in_the_member_id() {
        let base = identity();
        let id = base.member_id();
        assert!(id.starts_with("local/gpu-cuda-") && id.len() == "local/gpu-cuda-".len() + 12, "{id}");
        assert!(id.as_str() > "local/cpu-x86_64" && id.as_str() > "local/cpu-aarch64", "sorts after every CPU reference");
        let variants = [
            DeviceIdentity { name: "another".into(), ..identity() },
            DeviceIdentity { uuid: "ff".into(), ..identity() },
            DeviceIdentity { compute: (9, 0), ..identity() },
            DeviceIdentity { compute: (8, 6), ..identity() },
            DeviceIdentity { pci: "0000:02:00".into(), ..identity() },
            DeviceIdentity { driver: 12080, ..identity() },
            DeviceIdentity { compiler: (13, 0), ..identity() },
            DeviceIdentity { compiler: (12, 9), ..identity() },
            DeviceIdentity { memory: 8 << 30, ..identity() },
        ];
        for v in variants {
            assert_ne!(v.member_id(), id, "{v:?}");
        }
        assert_eq!(base.key(), "gpu-cuda-sm89");
        assert_eq!(base.arch(), "compute_89");
        assert_eq!(base.class(), FabricClass::Gpu);
    }

    #[test]
    fn the_implementation_id_follows_the_kernel_text() {
        let id = implementation_id();
        assert!(id.starts_with("cuda-lanes64/") && id.len() == "cuda-lanes64/".len() + 16, "{id}");
        for m in &MUTANTS {
            assert_ne!(implementation_id_of(&mutate(m)), id, "mutant {}", m.name);
        }
    }

    #[test]
    fn every_mutant_edits_the_kernel_once_and_names_distinct_edits() {
        let mut seen = std::collections::BTreeSet::new();
        for m in &MUTANTS {
            let edited = mutate(m);
            assert_ne!(edited, SOURCE, "mutant {} changes nothing", m.name);
            assert!(seen.insert(edited), "mutant {} duplicates another", m.name);
        }
        for d in [Detector::Probe, Detector::Certificate, Detector::Range] {
            assert!(MUTANTS.iter().any(|m| m.detector == d), "no mutant for {d:?}");
        }
    }
}