oxiproj-core 0.1.2

//! f64 transcendental PTX device-function generators.
//!
//! PTX exposes only **f32** hardware transcendentals (`sin/cos/lg2/ex2.approx.f32`,
//! ~1e-7), which is far too coarse for geodetic f64. This module emits hand-rolled
//! **f64** `.func` device functions — `exp`, `ln`, `tan`, `atan`, `sinh`, `asinh` —
//! built from native f64 arithmetic (`fma.rn.f64`, `sqrt.rn.f64`, `div.rn.f64`,
//! bit manipulation) plus short polynomials. They are generated in Rust so every
//! constant is an exact `f64` ([`f64_imm`]) and the Horner chains are emitted by
//! [`horner_f64`], eliminating transcription error.
//!
//! Each function uses a classic reduction tuned to the **bounded domains** that
//! cartographic projections exercise, which keeps the polynomials short while
//! reaching ~1e-12 relative accuracy (verified against `std` on real hardware in
//! this crate's GPU tests):
//!
//! * `exp` — Cody–Waite `x = k·ln2 + r`, `eˣ = 2ᵏ·P(r)`, `2ᵏ` assembled into the
//!   double exponent field. Valid for normalised `2ᵏ` (`|k| ≲ 1000`).
//! * `ln`  — `x = m·2ᵉ` (`m ∈ [√½, √2)`), `ln = e·ln2 + 2·atanh((m−1)/(m+1))`.
//! * `tan` — octant fold to `[0, π/4]`, `tan = sin/cos` via `selp` (branch-free).
//! * `atan`— `|x|>1 ⇒ π/2 − atan(1/x)`, then a `π/12` sub-reduction and odd poly.
//! * `sinh`— `(eˣ − e⁻ˣ)/2` (calls `exp`).
//! * `asinh`— `sign(x)·ln(|x| + √(x²+1))` (calls `ln`; the `|x|` form avoids
//!   cancellation).
//!
//! # ABI
//!
//! Every function uses the param-space ABI `(.param .f64 p_ret) f (.param .f64
//! p_arg)`. Call them with [`call_unary`], which wraps the standard
//! `st.param / call.uni / ld.param` sequence. Concatenate the needed definitions
//! (callees before callers) ahead of a kernel's `.entry`; the [`merc_forward_funcs`]
//! / [`merc_inverse_funcs`] helpers do this for the Web Mercator kernels.

use super::ptx_math::{f64_imm, horner_f64};

// ── Cody–Waite split of ln2 and related constants ────────────────────────────
const LN2_HI: f64 = f64::from_bits(0x3FE6_2E42_FEE0_0000);
const LN2_LO: f64 = f64::from_bits(0x3DEA_39EF_3579_3C76);
const INV_LN2: f64 = std::f64::consts::LOG2_E; // 1/ln2
const SQRT2: f64 = std::f64::consts::SQRT_2;
const PI_2: f64 = std::f64::consts::FRAC_PI_2;
const PI_4: f64 = std::f64::consts::FRAC_PI_4;
const PI_6: f64 = std::f64::consts::FRAC_PI_6;
const INV_SQRT3: f64 = 0.577_350_269_189_625_7; // 1/√3
const TAN_PI_12: f64 = 0.267_949_192_431_122_7; // tan(15°) = 2 − √3
const HALF: &str = "0d3FE0000000000000"; // 0.5

/// `n!` as `f64` (`0! = 1`). Exact for the small `n` used here.
fn factorial(n: u64) -> f64 {
    (1..=n).map(|i| i as f64).product()
}

/// Emit a unary-`.func` call (param-space ABI): leaves `func(%fd{arg})` in
/// `%fd{out}`. `id` disambiguates the scoped `.param` names when one function
/// body (or kernel) issues several calls.
#[must_use]
pub fn call_unary(func: &str, arg: u32, out: u32, id: u32) -> String {
    let mut s = String::from("\t{\n");
    s.push_str(&format!("\t.param .f64 ci{id};\n\t.param .f64 co{id};\n"));
    s.push_str(&format!("\tst.param.f64 [ci{id}], %fd{arg};\n"));
    s.push_str(&format!("\tcall.uni (co{id}), {func}, (ci{id});\n"));
    s.push_str(&format!("\tld.param.f64 %fd{out}, [co{id}];\n\t}}\n"));
    s
}

/// Open a unary `.func` definition with the standard `(p_ret)(p_arg)` ABI.
fn func_open(name: &str) -> String {
    format!(".func (.param .f64 p_ret) {name} (.param .f64 p_arg)\n{{\n")
}

/// `exp(x)` device function. See module docs for the algorithm.
#[must_use]
pub fn emit_exp_f64() -> String {
    let coeffs: Vec<f64> = (0..=13).map(|k| 1.0 / factorial(k)).collect();
    let mut s = func_open("oxp_exp_f64");
    s.push_str("\t.reg .b32 %r<2>;\n\t.reg .b64 %rd<2>;\n\t.reg .f64 %fd<8>;\n");
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n");
    s.push_str(&format!("\tmul.f64 %fd1, %fd0, {};\n", f64_imm(INV_LN2)));
    s.push_str("\tcvt.rni.s32.f64 %r0, %fd1;\n"); // k = round(x/ln2)
    s.push_str("\tcvt.rn.f64.s32 %fd1, %r0;\n"); // kf
    s.push_str(&format!(
        "\tfma.rn.f64 %fd2, %fd1, {}, %fd0;\n",
        f64_imm(-LN2_HI)
    ));
    s.push_str(&format!(
        "\tfma.rn.f64 %fd2, %fd1, {}, %fd2;\n",
        f64_imm(-LN2_LO)
    ));
    s.push_str(&horner_f64(&coeffs, 2, 3)); // P(r) -> %fd3
    s.push_str("\tadd.s32 %r0, %r0, 1023;\n");
    s.push_str("\tcvt.s64.s32 %rd0, %r0;\n");
    s.push_str("\tshl.b64 %rd0, %rd0, 52;\n");
    s.push_str("\tmov.b64 %fd4, %rd0;\n"); // 2^k
    s.push_str("\tmul.f64 %fd5, %fd3, %fd4;\n");
    s.push_str("\tst.param.f64 [p_ret], %fd5;\n\tret;\n}\n");
    s
}

/// `ln(x)` device function for `x > 0`. See module docs for the algorithm.
#[must_use]
pub fn emit_ln_f64() -> String {
    // P(s²) = Σ s²ᵏ/(2k+1), so ln(m) = 2·s·P(s²).
    let coeffs: Vec<f64> = (0..8).map(|k| 1.0 / (2 * k + 1) as f64).collect();
    let mut s = func_open("oxp_ln_f64");
    s.push_str(
        "\t.reg .pred %p<1>;\n\t.reg .b32 %r<2>;\n\t.reg .b64 %rd<2>;\n\t.reg .f64 %fd<8>;\n",
    );
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n");
    s.push_str("\tmov.b64 %rd0, %fd0;\n");
    s.push_str("\tshr.u64 %rd1, %rd0, 52;\n");
    s.push_str("\tand.b64 %rd1, %rd1, 0x7FF;\n");
    s.push_str("\tcvt.u32.u64 %r0, %rd1;\n");
    s.push_str("\tsub.s32 %r0, %r0, 1023;\n"); // e
    s.push_str("\tand.b64 %rd0, %rd0, 0x000FFFFFFFFFFFFF;\n");
    s.push_str("\tor.b64 %rd0, %rd0, 0x3FF0000000000000;\n");
    s.push_str("\tmov.b64 %fd1, %rd0;\n"); // m ∈ [1,2)
    s.push_str(&format!("\tsetp.gt.f64 %p0, %fd1, {};\n", f64_imm(SQRT2)));
    s.push_str(&format!("\t@%p0 mul.f64 %fd1, %fd1, {HALF};\n"));
    s.push_str("\t@%p0 add.s32 %r0, %r0, 1;\n");
    s.push_str(&format!("\tsub.f64 %fd2, %fd1, {};\n", f64_imm(1.0)));
    s.push_str(&format!("\tadd.f64 %fd3, %fd1, {};\n", f64_imm(1.0)));
    s.push_str("\tdiv.rn.f64 %fd2, %fd2, %fd3;\n"); // s
    s.push_str("\tmul.f64 %fd3, %fd2, %fd2;\n"); // s²
    s.push_str(&horner_f64(&coeffs, 3, 4)); // P(s²) -> %fd4
    s.push_str("\tmul.f64 %fd5, %fd2, %fd4;\n");
    s.push_str("\tadd.f64 %fd5, %fd5, %fd5;\n"); // ln(m) = 2·s·P
    s.push_str("\tcvt.rn.f64.s32 %fd6, %r0;\n"); // e as f64
    s.push_str(&format!(
        "\tfma.rn.f64 %fd7, %fd6, {}, %fd5;\n",
        f64_imm(LN2_LO)
    ));
    s.push_str(&format!(
        "\tfma.rn.f64 %fd7, %fd6, {}, %fd7;\n",
        f64_imm(LN2_HI)
    ));
    s.push_str("\tst.param.f64 [p_ret], %fd7;\n\tret;\n}\n");
    s
}

/// `tan(x)` device function for `x ∈ (−π/2, π/2)`. Branch-free octant fold.
#[must_use]
pub fn emit_tan_f64() -> String {
    let sin_c: Vec<f64> = (0..8).map(|k| sign(k) / factorial(2 * k + 1)).collect();
    let cos_c: Vec<f64> = (0..9).map(|k| sign(k) / factorial(2 * k)).collect();
    let mut s = func_open("oxp_tan_f64");
    s.push_str("\t.reg .pred %p<1>;\n\t.reg .f64 %fd<16>;\n");
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n"); // phi
    s.push_str("\tabs.f64 %fd1, %fd0;\n"); // a = |phi|
    s.push_str(&format!("\tsetp.gt.f64 %p0, %fd1, {};\n", f64_imm(PI_4)));
    s.push_str(&format!("\tmov.f64 %fd2, {};\n", f64_imm(PI_2)));
    s.push_str("\tsub.f64 %fd3, %fd2, %fd1;\n"); // π/2 − a
    s.push_str("\tselp.f64 %fd2, %fd3, %fd1, %p0;\n"); // arg
    s.push_str("\tmul.f64 %fd4, %fd2, %fd2;\n"); // t²
    s.push_str(&horner_f64(&sin_c, 4, 5)); // Psin -> %fd5
    s.push_str("\tmul.f64 %fd6, %fd2, %fd5;\n"); // sa = arg·Psin
    s.push_str(&horner_f64(&cos_c, 4, 7)); // ca -> %fd7
    s.push_str("\tselp.f64 %fd8, %fd7, %fd6, %p0;\n"); // num = high?ca:sa
    s.push_str("\tselp.f64 %fd9, %fd6, %fd7, %p0;\n"); // den = high?sa:ca
    s.push_str("\tdiv.rn.f64 %fd10, %fd8, %fd9;\n");
    // copysign.f64 d, a, b = |b| with sign of a → sign from phi (%fd0).
    s.push_str("\tcopysign.f64 %fd11, %fd0, %fd10;\n");
    s.push_str("\tst.param.f64 [p_ret], %fd11;\n\tret;\n}\n");
    s
}

/// `atan(x)` device function for any finite `x`. Two-stage reduction.
#[must_use]
pub fn emit_atan_f64() -> String {
    let coeffs: Vec<f64> = (0..11).map(|k| sign(k) / (2 * k + 1) as f64).collect();
    let mut s = func_open("oxp_atan_f64");
    s.push_str("\t.reg .pred %p<2>;\n\t.reg .f64 %fd<20>;\n");
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n"); // x
    s.push_str("\tabs.f64 %fd1, %fd0;\n"); // ax
    s.push_str(&format!("\tsetp.gt.f64 %p0, %fd1, {};\n", f64_imm(1.0)));
    s.push_str("\trcp.rn.f64 %fd2, %fd1;\n");
    s.push_str("\tselp.f64 %fd3, %fd2, %fd1, %p0;\n"); // z0 ∈ [0,1]
    s.push_str(&format!(
        "\tsetp.gt.f64 %p1, %fd3, {};\n",
        f64_imm(TAN_PI_12)
    ));
    s.push_str(&format!("\tmov.f64 %fd4, {};\n", f64_imm(INV_SQRT3)));
    s.push_str("\tsub.f64 %fd5, %fd3, %fd4;\n");
    s.push_str(&format!(
        "\tfma.rn.f64 %fd6, %fd3, %fd4, {};\n",
        f64_imm(1.0)
    ));
    s.push_str("\tdiv.rn.f64 %fd5, %fd5, %fd6;\n"); // zr
    s.push_str("\tselp.f64 %fd7, %fd5, %fd3, %p1;\n"); // z
    s.push_str("\tmul.f64 %fd8, %fd7, %fd7;\n"); // z²
    s.push_str(&horner_f64(&coeffs, 8, 9)); // P(z²) -> %fd9
    s.push_str("\tmul.f64 %fd10, %fd7, %fd9;\n"); // z·P
    s.push_str(&format!("\tadd.f64 %fd11, %fd10, {};\n", f64_imm(PI_6)));
    s.push_str("\tselp.f64 %fd12, %fd11, %fd10, %p1;\n"); // atan(z0)
    s.push_str(&format!("\tmov.f64 %fd13, {};\n", f64_imm(PI_2)));
    s.push_str("\tsub.f64 %fd14, %fd13, %fd12;\n");
    s.push_str("\tselp.f64 %fd15, %fd14, %fd12, %p0;\n"); // atan(ax)
                                                          // copysign.f64 d, a, b = |b| with sign of a → sign from x (%fd0).
    s.push_str("\tcopysign.f64 %fd16, %fd0, %fd15;\n");
    s.push_str("\tst.param.f64 [p_ret], %fd16;\n\tret;\n}\n");
    s
}

/// `sinh(x)` device function. Depends on [`emit_exp_f64`].
#[must_use]
pub fn emit_sinh_f64() -> String {
    let mut s = func_open("oxp_sinh_f64");
    s.push_str("\t.reg .f64 %fd<8>;\n");
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n");
    s.push_str(&call_unary("oxp_exp_f64", 0, 1, 0)); // eˣ
    s.push_str("\tneg.f64 %fd2, %fd0;\n");
    s.push_str(&call_unary("oxp_exp_f64", 2, 3, 1)); // e⁻ˣ
    s.push_str("\tsub.f64 %fd4, %fd1, %fd3;\n");
    s.push_str(&format!("\tmul.f64 %fd5, %fd4, {HALF};\n"));
    s.push_str("\tst.param.f64 [p_ret], %fd5;\n\tret;\n}\n");
    s
}

/// `asinh(x)` device function. Depends on [`emit_ln_f64`].
#[must_use]
pub fn emit_asinh_f64() -> String {
    let mut s = func_open("oxp_asinh_f64");
    s.push_str("\t.reg .f64 %fd<6>;\n");
    s.push_str("\tld.param.f64 %fd0, [p_arg];\n");
    s.push_str("\tabs.f64 %fd1, %fd0;\n");
    s.push_str(&format!(
        "\tfma.rn.f64 %fd2, %fd1, %fd1, {};\n",
        f64_imm(1.0)
    ));
    s.push_str("\tsqrt.rn.f64 %fd2, %fd2;\n");
    s.push_str("\tadd.f64 %fd2, %fd1, %fd2;\n"); // |x| + √(x²+1)
    s.push_str(&call_unary("oxp_ln_f64", 2, 3, 0));
    // copysign.f64 d, a, b = |b| with sign of a → sign from x (%fd0).
    s.push_str("\tcopysign.f64 %fd4, %fd0, %fd3;\n");
    s.push_str("\tst.param.f64 [p_ret], %fd4;\n\tret;\n}\n");
    s
}

/// `(-1)^k` as `f64`.
fn sign(k: u64) -> f64 {
    if k.is_multiple_of(2) {
        1.0
    } else {
        -1.0
    }
}

/// Device functions required by the Web Mercator **forward** kernel, emitted in
/// dependency order (`ln` before `asinh`; `tan`).
#[must_use]
pub fn merc_forward_funcs() -> String {
    let mut s = emit_ln_f64();
    s.push_str(&emit_asinh_f64());
    s.push_str(&emit_tan_f64());
    s
}

/// Device functions required by the Web Mercator **inverse** kernel, emitted in
/// dependency order (`exp` before `sinh`; `atan`).
#[must_use]
pub fn merc_inverse_funcs() -> String {
    let mut s = emit_exp_f64();
    s.push_str(&emit_sinh_f64());
    s.push_str(&emit_atan_f64());
    s
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::gpu::{self, ptx_math, DeviceBuffer};

    /// A test `.entry` that applies one unary device function element-wise:
    /// `out[i] = f(in[i])`. `call` is the [`call_unary`] snippet producing the
    /// result in `%fd1` from `%fd0`.
    fn unary_test_entry(call: &str) -> String {
        let mut s = String::from(
            "\n.visible .entry oxp_unary_test(\n\t.param .u64 p_in,\n\t.param .u64 p_out,\n\t.param .u32 p_n\n)\n{\n",
        );
        s.push_str(
            "\t.reg .pred %p<1>;\n\t.reg .b32 %r<6>;\n\t.reg .b64 %rd<6>;\n\t.reg .f64 %fd<4>;\n",
        );
        s.push_str("\tmov.u32 %r0, %ntid.x;\n\tmov.u32 %r1, %ctaid.x;\n\tmov.u32 %r2, %tid.x;\n");
        s.push_str("\tmad.lo.u32 %r3, %r1, %r0, %r2;\n\tld.param.u32 %r4, [p_n];\n");
        s.push_str("\tsetp.ge.u32 %p0, %r3, %r4;\n\t@%p0 bra DONE;\n");
        s.push_str("\tmul.wide.u32 %rd0, %r3, 8;\n\tld.param.u64 %rd1, [p_in];\n\tadd.u64 %rd2, %rd1, %rd0;\n");
        s.push_str("\tld.param.u64 %rd3, [p_out];\n\tadd.u64 %rd4, %rd3, %rd0;\n");
        s.push_str("\tld.global.f64 %fd0, [%rd2];\n");
        s.push_str(call);
        s.push_str("\tst.global.f64 [%rd4], %fd1;\nDONE:\n\tret;\n}\n");
        s
    }

    fn run_unary(key: &'static str, prelude: String, func: &str, xs: &[f64]) -> Vec<f64> {
        let n = xs.len();
        let mut out = vec![0.0f64; n];
        let call = call_unary(func, 0, 1, 0);
        gpu::with_kernel(
            key,
            "oxp_unary_test",
            |target| {
                let mut p = ptx_math::ptx_header(target);
                p.push_str(&prelude);
                p.push_str(&unary_test_entry(&call));
                p
            },
            |kernel, stream| {
                let d_in = gpu::upload(xs, stream)?;
                let d_out = DeviceBuffer::<f64>::alloc(n)?;
                let args = (d_in.as_device_ptr(), d_out.as_device_ptr(), n as u32);
                gpu::launch_1d(kernel, stream, n as u32, &args)?;
                d_out.copy_to_host(&mut out)?;
                Ok(())
            },
        )
        .expect("unary test kernel launch");
        out
    }

    fn check(
        key: &'static str,
        prelude: String,
        func: &str,
        xs: &[f64],
        want: impl Fn(f64) -> f64,
        tol: f64,
    ) {
        let got = run_unary(key, prelude, func, xs);
        let mut max_rel = 0.0f64;
        for (&x, &g) in xs.iter().zip(got.iter()) {
            let w = want(x);
            let rel = (g - w).abs() / w.abs().max(1.0);
            max_rel = max_rel.max(rel);
            assert!(
                rel < tol,
                "{func}({x}): got {g}, want {w}, rel {rel:e} > {tol:e}"
            );
        }
        eprintln!(
            "{func}: max rel error {max_rel:e} over {} samples",
            xs.len()
        );
    }

    fn grid(lo: f64, hi: f64, n: usize) -> Vec<f64> {
        (0..=n)
            .map(|i| lo + (hi - lo) * i as f64 / n as f64)
            .collect()
    }

    #[test]
    fn exp_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        check(
            "t_exp",
            emit_exp_f64(),
            "oxp_exp_f64",
            &grid(-30.0, 30.0, 4000),
            f64::exp,
            1e-12,
        );
    }

    #[test]
    fn ln_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        check(
            "t_ln",
            emit_ln_f64(),
            "oxp_ln_f64",
            &grid(1e-3, 1.0e6, 4000),
            f64::ln,
            1e-12,
        );
    }

    #[test]
    fn tan_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        // Web Mercator latitude domain stays within ±~85°; sample to ±88°.
        check(
            "t_tan",
            emit_tan_f64(),
            "oxp_tan_f64",
            &grid(-1.535, 1.535, 4000),
            f64::tan,
            1e-11,
        );
    }

    #[test]
    fn atan_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        check(
            "t_atan",
            emit_atan_f64(),
            "oxp_atan_f64",
            &grid(-60.0, 60.0, 4000),
            f64::atan,
            1e-12,
        );
    }

    #[test]
    fn sinh_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        check(
            "t_sinh",
            merc_inverse_funcs(),
            "oxp_sinh_f64",
            &grid(-5.0, 5.0, 4000),
            f64::sinh,
            1e-12,
        );
    }

    #[test]
    fn asinh_matches_std() {
        if !gpu::cuda_available() {
            return;
        }
        check(
            "t_asinh",
            merc_forward_funcs(),
            "oxp_asinh_f64",
            &grid(-60.0, 60.0, 4000),
            f64::asinh,
            1e-12,
        );
    }
}