oxiproj-core 0.1.2

//! PTX assembly toolkit for OxiProj GPU kernels.
//!
//! This module provides the building blocks shared by every OxiProj CUDA kernel:
//!
//! * [`PtxTarget`] / [`ptx_target_for_cc`] — map a device compute capability to
//!   the `.version` / `.target` header pair accepted by the JIT.
//! * [`ptx_header`] — the standard module preamble.
//! * [`f64_imm`] — format an `f64` as a PTX hexadecimal immediate.
//! * [`horner_f64`] — emit a numerically-stable Horner polynomial evaluation as
//!   PTX `fma.rn.f64` instructions (used by the transcendental device functions).
//!
//! The companion [`ptx_transc`](super::ptx_transc) module builds on this toolkit
//! to emit reusable `.func` device functions for the f64 transcendentals that PTX
//! lacks in hardware (only f32 `sin/cos/lg2/ex2.approx` exist). Those use
//! exact-rational (reciprocal-factorial / Cody–Waite) reductions over the
//! **bounded domains** that cartographic projections actually exercise, which
//! keeps the polynomials short while still reaching ~1e-12 relative accuracy.

/// PTX target triple derived from a device's compute capability.
#[derive(Debug, Clone)]
pub struct PtxTarget {
    /// The PTX ISA `.version` (`major`, `minor`).
    pub version: (u32, u32),
    /// The `.target` string, e.g. `"sm_86"`.
    pub sm: String,
}

/// Map a `(major, minor)` compute capability to a conservative PTX
/// `.version` / `.target` pair the in-driver JIT accepts.
///
/// Ported from the `oxicuda` driver's reference launch example; the mapping is
/// intentionally conservative (a lower `.version` than the maximum supported)
/// because OxiProj kernels use only classic instructions that have been stable
/// since PTX ISA 7.4.
#[must_use]
pub fn ptx_target_for_cc(major: i32, minor: i32) -> PtxTarget {
    let sm = format!("sm_{major}{minor}");
    let version = match (major, minor) {
        (7, 5) => (7, 4),
        (8, 0) | (8, 6) => (7, 5),
        (8, 9) => (8, 0),
        (9, 0 | 1) => (8, 0),
        (10, _) => (8, 5),
        (12, _) => (8, 7),
        _ => (7, 4),
    };
    PtxTarget { version, sm }
}

/// Build the standard PTX module header (`.version` / `.target` / `.address_size`).
#[must_use]
pub fn ptx_header(target: &PtxTarget) -> String {
    format!(
        ".version {}.{}\n.target {}\n.address_size 64\n",
        target.version.0, target.version.1, target.sm
    )
}

/// Format an `f64` as a PTX double-precision immediate (`0d` + 16 hex digits of
/// the IEEE-754 bit pattern), which is the only exact way to embed a constant in
/// PTX text.
#[must_use]
pub fn f64_imm(v: f64) -> String {
    format!("0d{:016X}", v.to_bits())
}

/// Emit PTX evaluating the polynomial `coeffs[0] + coeffs[1]·x + … ` at the
/// value held in register `%fd{x_reg}`, leaving the result in `%fd{out_reg}`.
///
/// Evaluation uses Horner's method with fused multiply-adds (`fma.rn.f64`) for
/// accuracy. `coeffs` is given in **ascending** power order. Scratch register
/// indices `tmp_base` and `tmp_base+? ` are not used — only `out_reg` and the
/// caller-supplied `x_reg` are touched, so the caller must ensure `out_reg` is
/// free. Returns the PTX snippet.
///
/// Requires `coeffs` to be non-empty.
#[must_use]
pub fn horner_f64(coeffs: &[f64], x_reg: u32, out_reg: u32) -> String {
    debug_assert!(
        !coeffs.is_empty(),
        "horner_f64 needs at least one coefficient"
    );
    let mut s = String::new();
    // Start from the highest-order coefficient.
    let last = coeffs.len() - 1;
    s.push_str(&format!(
        "\tmov.f64 \t%fd{out_reg}, {};\n",
        f64_imm(coeffs[last])
    ));
    for k in (0..last).rev() {
        // acc = acc * x + coeffs[k]
        s.push_str(&format!(
            "\tfma.rn.f64 \t%fd{out_reg}, %fd{out_reg}, %fd{x_reg}, {};\n",
            f64_imm(coeffs[k])
        ));
    }
    s
}

// Reusable f64 transcendental `.func` device functions (`OXP_EXP_F64`,
// `OXP_LN_F64`, `OXP_TAN_F64`, `OXP_ATAN_F64`, `OXP_ASINH_F64`, `OXP_SINH_F64`)
// are defined in `ptx_transc.rs` and used by the projection kernels.

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

    #[test]
    fn f64_imm_known_patterns() {
        assert_eq!(f64_imm(1.0), "0d3FF0000000000000");
        assert_eq!(f64_imm(2.0), "0d4000000000000000");
        assert_eq!(f64_imm(0.0), "0d0000000000000000");
    }

    #[test]
    fn ptx_target_ampere() {
        let t = ptx_target_for_cc(8, 6);
        assert_eq!(t.sm, "sm_86");
        assert_eq!(t.version, (7, 5));
    }

    #[test]
    fn horner_emits_fma_chain() {
        // p(x) = 1 + 2x + 3x²  → one mov + two fma
        let ptx = horner_f64(&[1.0, 2.0, 3.0], 0, 1);
        assert_eq!(ptx.matches("fma.rn.f64").count(), 2);
        assert_eq!(ptx.matches("mov.f64").count(), 1);
    }
}