use super::BodyBuilder;
use crate::ir::{PtxType, Register};
impl BodyBuilder<'_> {
pub fn mov_imm_f64(&mut self, val: f64) -> Register {
let r = self.alloc_reg(PtxType::F64);
let bits = val.to_bits();
self.raw_ptx(&format!("mov.b64 {r}, 0D{bits:016X};"));
r
}
pub fn mov_imm_f32(&mut self, val: f32) -> Register {
let r = self.alloc_reg(PtxType::F32);
let bits = val.to_bits();
self.raw_ptx(&format!("mov.b32 {r}, 0F{bits:08X};"));
r
}
pub fn div_rn_f64(&mut self, a: &Register, b: &Register) -> Register {
let dst = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("div.rn.f64 {dst}, {a}, {b};"));
dst
}
pub fn div_rn_f32(&mut self, a: &Register, b: &Register) -> Register {
let dst = self.alloc_reg(PtxType::F32);
self.raw_ptx(&format!("div.rn.f32 {dst}, {a}, {b};"));
dst
}
fn mul_rn_f64(&mut self, a: &Register, b: &Register) -> Register {
let dst = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("mul.rn.f64 {dst}, {a}, {b};"));
dst
}
fn neg_f64(&mut self, x: &Register) -> Register {
let dst = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("neg.f64 {dst}, {x};"));
dst
}
pub fn exp_f64(&mut self, x: &Register) -> Register {
let log2e = self.mov_imm_f64(std::f64::consts::LOG2_E);
let kf = self.mul_rn_f64(x, &log2e);
let magic = self.mov_imm_f64(6_755_399_441_055_744.0);
let kf_plus = self.add_f64(kf, magic.clone());
let kf_rounded = self.sub_f64(kf_plus, magic);
let neg_ln2_hi = self.mov_imm_f64(-0.693_147_180_369_123_8);
let r = self.fma_f64(kf_rounded.clone(), neg_ln2_hi, x.clone());
let neg_ln2_lo = self.mov_imm_f64(-1.908_214_929_270_587_7e-10);
let r = self.fma_f64(kf_rounded.clone(), neg_ln2_lo, r);
let mut p = self.mov_imm_f64(1.0 / 5040.0);
let c6 = self.mov_imm_f64(1.0 / 720.0);
p = self.fma_f64(p, r.clone(), c6);
let c5 = self.mov_imm_f64(1.0 / 120.0);
p = self.fma_f64(p, r.clone(), c5);
let c4 = self.mov_imm_f64(1.0 / 24.0);
p = self.fma_f64(p, r.clone(), c4);
let c3 = self.mov_imm_f64(1.0 / 6.0);
p = self.fma_f64(p, r.clone(), c3);
let c2 = self.mov_imm_f64(0.5);
p = self.fma_f64(p, r.clone(), c2);
let c1 = self.mov_imm_f64(1.0);
p = self.fma_f64(p, r.clone(), c1);
let c0 = self.mov_imm_f64(1.0);
p = self.fma_f64(p, r, c0);
let ki = self.alloc_reg(PtxType::S32);
self.raw_ptx(&format!("cvt.rzi.s32.f64 {ki}, {kf_rounded};"));
self.raw_ptx(&format!("max.s32 {ki}, {ki}, -1075;"));
self.raw_ptx(&format!("min.s32 {ki}, {ki}, 1024;"));
let ki64 = self.alloc_reg(PtxType::S64);
self.raw_ptx(&format!("cvt.s64.s32 {ki64}, {ki};"));
let biased = self.alloc_reg(PtxType::S64);
self.raw_ptx(&format!("add.s64 {biased}, {ki64}, 1023;"));
let expbits = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!("shl.b64 {expbits}, {biased}, 52;"));
let two_k = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("mov.b64 {two_k}, {expbits};"));
let raw = self.mul_rn_f64(&p, &two_k);
let pos_inf = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("mov.b64 {pos_inf}, 0x7FF0000000000000;"));
let zero = self.mov_imm_f64(0.0);
let p_over = self.alloc_reg(PtxType::Pred);
self.raw_ptx(&format!("setp.gt.s32 {p_over}, {ki}, 1023;"));
let sat_over = self.selp(PtxType::F64, pos_inf, raw, p_over);
let p_under = self.alloc_reg(PtxType::Pred);
self.raw_ptx(&format!("setp.lt.s32 {p_under}, {ki}, -1022;"));
self.selp(PtxType::F64, zero, sat_over, p_under)
}
#[allow(clippy::many_single_char_names)]
pub fn log_f64(&mut self, x: &Register) -> Register {
let bits = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!("mov.b64 {bits}, {x};"));
let e_shifted = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!("shr.b64 {e_shifted}, {bits}, 52;"));
let e_raw = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!("and.b64 {e_raw}, {e_shifted}, 0x7FF;"));
let mant_frac = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!("and.b64 {mant_frac}, {bits}, 0x000FFFFFFFFFFFFF;"));
let mant_bits = self.alloc_reg(PtxType::B64);
self.raw_ptx(&format!(
"or.b64 {mant_bits}, {mant_frac}, 0x3FF0000000000000;"
));
let m = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("mov.b64 {m}, {mant_bits};"));
let e_i64 = self.alloc_reg(PtxType::S64);
self.raw_ptx(&format!("cvt.s64.u64 {e_i64}, {e_raw};"));
let e_unbiased = self.alloc_reg(PtxType::S64);
self.raw_ptx(&format!("sub.s64 {e_unbiased}, {e_i64}, 1023;"));
let ef = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("cvt.rn.f64.s64 {ef}, {e_unbiased};"));
let one = self.mov_imm_f64(1.0);
let f = self.sub_f64(m, one);
let two = self.mov_imm_f64(2.0);
let two_plus_f = self.add_f64(f.clone(), two);
let s = self.div_rn_f64(&f, &two_plus_f);
let z = self.mul_rn_f64(&s, &s);
let mut q = self.mov_imm_f64(1.0 / 9.0);
let c7 = self.mov_imm_f64(1.0 / 7.0);
q = self.fma_f64(q, z.clone(), c7);
let c5 = self.mov_imm_f64(1.0 / 5.0);
q = self.fma_f64(q, z.clone(), c5);
let c3 = self.mov_imm_f64(1.0 / 3.0);
q = self.fma_f64(q, z.clone(), c3);
let c1 = self.mov_imm_f64(1.0);
q = self.fma_f64(q, z, c1);
let two_s = self.add_f64(s.clone(), s);
let log_mant = self.mul_rn_f64(&two_s, &q);
let ln2 = self.mov_imm_f64(std::f64::consts::LN_2);
self.fma_f64(ef, ln2, log_mant)
}
pub fn erf_f64(&mut self, x: &Register) -> Register {
let ax = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("abs.f64 {ax}, {x};"));
let pconst = self.mov_imm_f64(0.327_591_1);
let one = self.mov_imm_f64(1.0);
let denom = self.fma_f64(pconst, ax.clone(), one.clone());
let t = self.div_rn_f64(&one, &denom);
let mut poly = self.mov_imm_f64(1.061_405_429);
let a4c = self.mov_imm_f64(-1.453_152_027);
poly = self.fma_f64(poly, t.clone(), a4c);
let a3c = self.mov_imm_f64(1.421_413_741);
poly = self.fma_f64(poly, t.clone(), a3c);
let a2c = self.mov_imm_f64(-0.284_496_736);
poly = self.fma_f64(poly, t.clone(), a2c);
let a1c = self.mov_imm_f64(0.254_829_592);
poly = self.fma_f64(poly, t.clone(), a1c);
poly = self.mul_rn_f64(&poly, &t);
let ax2 = self.mul_rn_f64(&ax, &ax);
let nax2 = self.neg_f64(&ax2);
let e = self.exp_f64(&nax2);
let pe = self.mul_rn_f64(&poly, &e);
let y = self.sub_f64(one, pe);
let neg_y = self.neg_f64(&y);
let zero = self.mov_imm_f64(0.0);
let p_neg = self.alloc_reg(PtxType::Pred);
self.raw_ptx(&format!("setp.lt.f64 {p_neg}, {x}, {zero};"));
self.selp(PtxType::F64, neg_y, y, p_neg)
}
pub fn tanh_f64(&mut self, x: &Register) -> Register {
let hi = self.mov_imm_f64(20.0);
let lo = self.mov_imm_f64(-20.0);
let xc = self.alloc_reg(PtxType::F64);
self.raw_ptx(&format!("min.f64 {xc}, {x}, {hi};"));
self.raw_ptx(&format!("max.f64 {xc}, {xc}, {lo};"));
let two_x = self.add_f64(xc.clone(), xc);
let e = self.exp_f64(&two_x);
let one = self.mov_imm_f64(1.0);
let num = self.sub_f64(e.clone(), one.clone());
let den = self.add_f64(e, one);
self.div_rn_f64(&num, &den)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::arch::SmVersion;
use crate::builder::KernelBuilder;
fn build_with_body<F>(f: F) -> String
where
F: FnOnce(&mut BodyBuilder<'_>) + 'static,
{
KernelBuilder::new("test")
.target(SmVersion::Sm80)
.param("a", PtxType::U64)
.param("n", PtxType::U32)
.body(f)
.build()
.expect("build should succeed")
}
#[test]
fn mov_imm_f64_emits_mov_b64() {
let ptx = build_with_body(|b| {
let _r = b.mov_imm_f64(1.5);
b.ret();
});
assert!(ptx.contains("mov.b64"), "expected mov.b64 in:\n{ptx}");
assert!(ptx.contains("0D"), "expected 0D literal in:\n{ptx}");
}
#[test]
fn mov_imm_f32_emits_mov_b32() {
let ptx = build_with_body(|b| {
let _r = b.mov_imm_f32(1.5_f32);
b.ret();
});
assert!(ptx.contains("mov.b32"), "expected mov.b32 in:\n{ptx}");
assert!(ptx.contains("0F"), "expected 0F literal in:\n{ptx}");
}
#[test]
fn div_rn_f64_emits_div() {
let ptx = build_with_body(|b| {
let a = b.mov_imm_f64(3.0);
let c = b.mov_imm_f64(2.0);
let _r = b.div_rn_f64(&a, &c);
b.ret();
});
assert!(ptx.contains("div.rn.f64"), "expected div.rn.f64 in:\n{ptx}");
}
#[test]
fn div_rn_f32_emits_div() {
let ptx = build_with_body(|b| {
let a = b.mov_imm_f32(3.0_f32);
let c = b.mov_imm_f32(2.0_f32);
let _r = b.div_rn_f32(&a, &c);
b.ret();
});
assert!(ptx.contains("div.rn.f32"), "expected div.rn.f32 in:\n{ptx}");
}
#[test]
fn exp_f64_emits_poly_and_2k() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.exp_f64(&x);
b.ret();
});
assert!(
ptx.contains("fma.rn.f64"),
"expected fma.rn.f64 (Taylor Horner) in:\n{ptx}"
);
assert!(ptx.contains("mov.b64"), "expected mov.b64 in:\n{ptx}");
assert!(
ptx.contains("shl.b64"),
"expected shl.b64 (2^k assembly) in:\n{ptx}"
);
assert!(
ptx.contains("cvt.rzi.s32.f64"),
"expected cvt.rzi.s32.f64 (k extraction) in:\n{ptx}"
);
}
#[test]
fn log_f64_emits_bit_split_and_poly() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(1.5);
let _r = b.log_f64(&x);
b.ret();
});
assert!(
ptx.contains("shr.b64"),
"expected shr.b64 (exponent extraction) in:\n{ptx}"
);
assert!(
ptx.contains("and.b64"),
"expected and.b64 (mantissa mask) in:\n{ptx}"
);
assert!(
ptx.contains("fma.rn.f64"),
"expected fma.rn.f64 (atanh Horner) in:\n{ptx}"
);
}
#[test]
fn erf_f64_emits_abs_and_exp() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.erf_f64(&x);
b.ret();
});
assert!(
ptx.contains("fma.rn.f64"),
"expected fma.rn.f64 (A&S Horner) in:\n{ptx}"
);
assert!(ptx.contains("abs.f64"), "expected abs.f64 (|x|) in:\n{ptx}");
assert!(
ptx.contains("shl.b64"),
"expected shl.b64 (exp 2^k assembly) in:\n{ptx}"
);
assert!(
ptx.contains("cvt.rzi.s32.f64"),
"expected cvt.rzi.s32.f64 (exp k extraction) in:\n{ptx}"
);
}
#[test]
fn tanh_f64_emits_div_and_exp() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.tanh_f64(&x);
b.ret();
});
assert!(
ptx.contains("div.rn.f64"),
"expected div.rn.f64 (final quotient) in:\n{ptx}"
);
assert!(
ptx.contains("shl.b64"),
"expected shl.b64 (exp 2^k assembly) in:\n{ptx}"
);
assert!(
ptx.contains("cvt.rzi.s32.f64"),
"expected cvt.rzi.s32.f64 (exp k extraction) in:\n{ptx}"
);
}
#[test]
fn no_approx_f64_emitted() {
let exp_ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.exp_f64(&x);
b.ret();
});
assert!(
!exp_ptx.contains(".approx.f64"),
"exp_f64 must not emit .approx.f64:\n{exp_ptx}"
);
let log_ptx = build_with_body(|b| {
let x = b.mov_imm_f64(1.5);
let _r = b.log_f64(&x);
b.ret();
});
assert!(
!log_ptx.contains(".approx.f64"),
"log_f64 must not emit .approx.f64:\n{log_ptx}"
);
let erf_ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.erf_f64(&x);
b.ret();
});
assert!(
!erf_ptx.contains(".approx.f64"),
"erf_f64 must not emit .approx.f64:\n{erf_ptx}"
);
let tanh_ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.tanh_f64(&x);
b.ret();
});
assert!(
!tanh_ptx.contains(".approx.f64"),
"tanh_f64 must not emit .approx.f64:\n{tanh_ptx}"
);
}
#[test]
fn exp_f64_emits_overflow_saturation() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.exp_f64(&x);
b.ret();
});
assert!(
ptx.contains("max.s32"),
"exp_f64 must clamp k with max.s32:\n{ptx}"
);
assert!(
ptx.contains("min.s32"),
"exp_f64 must clamp k with min.s32:\n{ptx}"
);
assert!(
ptx.contains("setp.gt.s32") && ptx.contains("setp.lt.s32"),
"exp_f64 must saturate via over/under predicates:\n{ptx}"
);
assert!(
ptx.contains("selp.f64"),
"exp_f64 must select the saturated value with selp.f64:\n{ptx}"
);
}
#[test]
fn tanh_f64_clamps_input() {
let ptx = build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.tanh_f64(&x);
b.ret();
});
assert!(
ptx.contains("min.f64") && ptx.contains("max.f64"),
"tanh_f64 must clamp its input to [-20, 20]:\n{ptx}"
);
}
fn find_ptxas() -> Option<std::path::PathBuf> {
if let Ok(path) = std::env::var("PATH") {
for dir in std::env::split_paths(&path) {
let candidate = dir.join("ptxas");
if candidate.is_file() {
return Some(candidate);
}
}
}
let fallback = std::path::PathBuf::from("/usr/local/cuda/bin/ptxas");
fallback.is_file().then_some(fallback)
}
#[test]
fn math_f64_kernels_assemble_for_sm86() {
let Some(ptxas) = find_ptxas() else {
println!("skipping: ptxas not found on PATH");
return;
};
let kernels = [
(
"exp",
build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.exp_f64(&x);
b.ret();
}),
),
(
"erf",
build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.erf_f64(&x);
b.ret();
}),
),
(
"tanh",
build_with_body(|b| {
let x = b.mov_imm_f64(0.5);
let _r = b.tanh_f64(&x);
b.ret();
}),
),
];
for (name, ptx) in kernels {
let mut ptx_path = std::env::temp_dir();
ptx_path.push(format!(
"oxicuda_math_f64_{name}_{}.ptx",
std::process::id()
));
std::fs::write(&ptx_path, &ptx).expect("write PTX to temp file");
let output = std::process::Command::new(&ptxas)
.arg("-arch=sm_86")
.arg(&ptx_path)
.arg("-o")
.arg("/dev/null")
.output()
.expect("invoke ptxas");
let _ = std::fs::remove_file(&ptx_path);
assert!(
output.status.success(),
"ptxas rejected {name}_f64 kernel:\n{}\n--- PTX ---\n{ptx}",
String::from_utf8_lossy(&output.stderr),
);
}
}
}