ruda_kernel/dsl/frontend/
trigonometry.rs

1use ruda_core::ir::{ManagedVariable, Variable};
2
3use crate::dsl::prelude::*;
4
5define_scalar!(ElemA);
6define_size!(SizeA);
7
8/// Computes the hypotenuse of a right triangle given the lengths of the other two sides.
9///
10/// This function computes `sqrt(x² + y²)` in a numerically stable way that avoids
11/// overflow and underflow issues.
12#[ruda]
13pub fn hypot<F: Float, N: Size>(lhs: Vector<F, N>, rhs: Vector<F, N>) -> Vector<F, N> {
14    let one = Vector::new(F::from_int(1));
15    let a = lhs.abs();
16    let b = rhs.abs();
17    let max_val = max(a, b);
18    let max_val_is_zero = max_val.equal(Vector::new(F::from_int(0)));
19    let max_val_safe = select_many(max_val_is_zero, one, max_val);
20    let min_val = min(a, b);
21    let t = min_val / max_val_safe;
22
23    max_val * fma(t, t, one).sqrt()
24}
25
26#[allow(missing_docs)]
27pub fn expand_hypot(scope: &mut Scope, lhs: Variable, rhs: Variable, out: Variable) {
28    scope.register_type::<ElemA>(lhs.ty.storage_type());
29    scope.register_size::<SizeA>(lhs.vector_size());
30    let res = hypot::expand::<ElemA, SizeA>(
31        scope,
32        ManagedVariable::Plain(lhs).into(),
33        ManagedVariable::Plain(rhs).into(),
34    );
35    assign::expand_no_check(scope, res, ManagedVariable::Plain(out).into());
36}
37
38/// Computes the reciprocal of the hypotenuse of a right triangle given the lengths of the other two sides.
39///
40/// This function computes `1 / sqrt(x² + y²)` in a numerically stable way that avoids
41/// overflow and underflow issues.
42#[ruda]
43pub fn rhypot<F: Float, N: Size>(lhs: Vector<F, N>, rhs: Vector<F, N>) -> Vector<F, N> {
44    let one = Vector::new(F::from_int(1));
45    let a = lhs.abs();
46    let b = rhs.abs();
47    let max_val = max(a, b);
48    let max_val_is_zero = max_val.equal(Vector::new(F::from_int(0)));
49    let max_val_safe = select_many(max_val_is_zero, one, max_val);
50    let min_val = min(a, b);
51    let t = min_val / max_val_safe;
52
53    fma(t, t, one).inverse_sqrt() / max_val
54}
55
56#[allow(missing_docs)]
57pub fn expand_rhypot(scope: &mut Scope, lhs: Variable, rhs: Variable, out: Variable) {
58    scope.register_type::<ElemA>(lhs.ty.storage_type());
59    scope.register_size::<SizeA>(lhs.vector_size());
60    let res = rhypot::expand::<ElemA, SizeA>(
61        scope,
62        ManagedVariable::Plain(lhs).into(),
63        ManagedVariable::Plain(rhs).into(),
64    );
65    assign::expand_no_check(scope, res, ManagedVariable::Plain(out).into());
66}
ruda_kernel/dsl/frontend/trigonometry.rs

ruda_kernel/dsl/frontend/
trigonometry.rs
