86 lines
3.9 KiB
Rust
86 lines
3.9 KiB
Rust
//! Multiplies by a splatted f32 (`VSplat` of an F32 value into `VFma` / `VX(FMul)`): arm64 folds
|
|
//! them into the by-element forms (`fmla vd.4s, vn.4s, vm.s[0]`, arm64 fold_lane_splats); every
|
|
//! negation form and operand position against Rust's `mul_add` / `*`, run on the host.
|
|
#![cfg(any(all(target_arch = "x86_64", unix), all(target_arch = "aarch64", unix)))]
|
|
use relax_backend::{exec::Code, rir::*};
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
type Entry = extern "sysv64" fn(*const f32, *const f32, *const f32, *mut f32);
|
|
#[cfg(target_arch = "aarch64")]
|
|
type Entry = extern "C" fn(*const f32, *const f32, *const f32, *mut f32);
|
|
|
|
#[cfg(target_arch = "x86_64")]
|
|
fn compile(f: &Func) -> Code {
|
|
let mut options = relax_session::Options::default();
|
|
options.cpu = relax_backend::x64::detect_cpu(&options);
|
|
let env = relax_backend::x64::Env { slots: SlotsPtr::NONE, tls_key: 0, tls_slow: 0, resume: 0, opt_regs: true, syms: 0, sym_class: 0, near: false };
|
|
let c = relax_backend::x64::compile(&options, f, &env).unwrap();
|
|
Code::new(&c.code).unwrap()
|
|
}
|
|
|
|
#[cfg(target_arch = "aarch64")]
|
|
fn compile(f: &Func) -> Code {
|
|
let options = relax_session::Options::default();
|
|
let env = relax_backend::arm64::Env { slots: SlotsPtr::NONE, tls_key: 0, tls_slow: 0, resume: 0, opt_regs: true, syms: 0, sym_class: 0, near: false };
|
|
let c = relax_backend::arm64::compile(&options, f, &env).unwrap();
|
|
Code::new(&c.code).unwrap()
|
|
}
|
|
|
|
/// fn(a, s, c, d): x = splat(*s); d[0..4] = fma(a, x, c); d[4..8] = fma(x, a, c); d[8..12] = a * x;
|
|
/// d[12..16] = x * a (the fma negations np / na)
|
|
fn func(np: bool, na: bool, nb: NanBits) -> Func {
|
|
let mut f = Func::new("lane_splat".into(), 0);
|
|
let (pa, ps, pc, pd) = (f.vreg(Cls::I), f.vreg(Cls::I), f.vreg(Cls::I), f.vreg(Cls::I));
|
|
f.params = vec![pa, ps, pc, pd];
|
|
f.rets = vec![];
|
|
f.block();
|
|
let s = f.vreg(Cls::F32);
|
|
let (x, a, c) = (f.vreg(Cls::V128), f.vreg(Cls::V128), f.vreg(Cls::V128));
|
|
let (d0, d1, d2, d3) = (f.vreg(Cls::V128), f.vreg(Cls::V128), f.vreg(Cls::V128), f.vreg(Cls::V128));
|
|
let insts = vec![
|
|
Inst::Load(Mem::F32, s, ps, 0),
|
|
Inst::VSplat(x, s),
|
|
Inst::VLoad(a, pa, 0),
|
|
Inst::VLoad(c, pc, 0),
|
|
Inst::VFma(np, na, d0, a, x, c),
|
|
Inst::VFma(np, na, d1, x, a, c),
|
|
Inst::VX(VX::FMul(VW::S, nb), d2, a, x),
|
|
Inst::VX(VX::FMul(VW::S, nb), d3, x, a),
|
|
Inst::VStore(pd, 0, d0),
|
|
Inst::VStore(pd, 16, d1),
|
|
Inst::VStore(pd, 32, d2),
|
|
Inst::VStore(pd, 48, d3),
|
|
];
|
|
for i in insts {
|
|
f.blocks[0].insts.push(i);
|
|
f.blocks[0].pos.push(0);
|
|
}
|
|
f.blocks[0].term = Term::Ret(vec![]);
|
|
f.verify().unwrap();
|
|
f
|
|
}
|
|
|
|
#[test]
|
|
fn splatted_multiplicands_match_lanewise_mul_add_and_mul() {
|
|
let a = [1.5f32, -2.25, 3.0e-3, 7.0];
|
|
let c = [0.5f32, 100.0, -0.0, 1.0e30];
|
|
for s in [3.0f32, -0.75, 1.0e-20] {
|
|
for (np, na) in [(false, false), (true, false), (false, true), (true, true)] {
|
|
for nb in [NanBits::Host, NanBits::Neon] {
|
|
let code = compile(&func(np, na, nb));
|
|
let entry: Entry = unsafe { std::mem::transmute(code.ptr()) };
|
|
let mut d = [0f32; 16];
|
|
entry(a.as_ptr(), &s, c.as_ptr(), d.as_mut_ptr());
|
|
for l in 0..4 {
|
|
let pa = if np { -a[l] } else { a[l] };
|
|
let ac = if na { -c[l] } else { c[l] };
|
|
let want = pa.mul_add(s, ac);
|
|
assert_eq!(d[l].to_bits(), want.to_bits(), "fma(a, x) lane {l} np {np} na {na} s {s}");
|
|
assert_eq!(d[4 + l].to_bits(), want.to_bits(), "fma(x, a) lane {l} np {np} na {na} s {s}");
|
|
assert_eq!(d[8 + l].to_bits(), (a[l] * s).to_bits(), "a * x lane {l}");
|
|
assert_eq!(d[12 + l].to_bits(), (s * a[l]).to_bits(), "x * a lane {l}");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|