Skip to main content

hermes_simd_core/ops/
unary.rs

1//! Single-operand SIMD elementwise operation strategies.
2//!
3//! `UnaryOp<T>` is a sealed ZST trait for single-operand transforms. Implementors
4//! define how a single vector is transformed (`apply`) and how a single scalar element
5//! is transformed (`apply_scalar`). Both paths are `#[inline(always)]`.
6
7use crate::kernel::SimdKernel;
8use crate::ops::elementwise::Clamp;
9use crate::scalar::{NumericElement, Scalar};
10
11// ---------------------------------------------------------------------------
12// UnaryOp — single-operand elementwise strategy
13// ---------------------------------------------------------------------------
14
15/// Sealed ZST trait for single-operand SIMD elementwise operations.
16///
17/// Implementors define how a single vector is transformed (`apply`) and how
18/// a single scalar element is transformed (`apply_scalar`). Both paths are
19/// `#[inline(always)]` — DCE eliminates unused strategies entirely.
20///
21/// # Zero-Cost Guarantee
22///
23/// Every `impl UnaryOp<T>` passes through to an `#[inline(always)]
24/// SimdKernel<T>` method. The ZST strategy parameter is erased at every
25/// monomorphization site: `size_of::<Abs>() == 0`.
26pub trait UnaryOp<T: Scalar>: crate::private::Sealed + Copy + 'static {
27    /// Apply the operation to a vector: `self.apply::<Arch>(v) -> result`.
28    ///
29    /// Takes `self` by value so `Clamp<T>` can access its bounds; for true ZST
30    /// strategies (`Abs`, `Neg`, `Sqrt`), `self` has size zero and the compiler
31    /// removes it entirely from the generated code.
32    ///
33    /// # Safety
34    /// Processor must support the target feature of `Arch`.
35    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector;
36
37    /// Apply the operation to a single scalar element.
38    ///
39    /// Used for the SIMD tail (elements that do not fill a complete vector).
40    /// Requires only `T: Scalar` — no unsafe, no vector loads or stores.
41    fn apply_scalar(self, a: T) -> T;
42}
43
44// ---------------------------------------------------------------------------
45// Concrete unary ZSTs
46// ---------------------------------------------------------------------------
47
48/// Elementwise absolute value: `|a[i]|`.
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct Abs;
51
52/// Elementwise negation: `-a[i]`.
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct Neg;
55
56/// Elementwise square root: `sqrt(a[i])`.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub struct Sqrt;
59
60/// Elementwise reciprocal square root: `1.0 / sqrt(a[i])`.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub struct RecipSqrt;
63
64/// Elementwise population count: count of set bits in each lane.
65#[derive(Debug, Clone, Copy, PartialEq, Eq)]
66pub struct Popcount;
67
68// ---------------------------------------------------------------------------
69// Sealing impls
70// ---------------------------------------------------------------------------
71
72impl crate::private::Sealed for Abs {}
73impl crate::private::Sealed for Neg {}
74impl crate::private::Sealed for Sqrt {}
75impl crate::private::Sealed for RecipSqrt {}
76impl crate::private::Sealed for Popcount {}
77
78// ---------------------------------------------------------------------------
79// UnaryOp impls
80// ---------------------------------------------------------------------------
81
82impl<T: Scalar> UnaryOp<T> for Abs {
83    #[inline(always)]
84    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
85        Arch::abs(v)
86    }
87    #[inline(always)]
88    fn apply_scalar(self, a: T) -> T {
89        a.abs()
90    }
91}
92
93impl<T: Scalar> UnaryOp<T> for Neg {
94    #[inline(always)]
95    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
96        Arch::neg(v)
97    }
98    #[inline(always)]
99    fn apply_scalar(self, a: T) -> T {
100        T::ZERO - a
101    }
102}
103
104impl<T: Scalar> UnaryOp<T> for Sqrt {
105    #[inline(always)]
106    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
107        Arch::sqrt(v)
108    }
109    #[inline(always)]
110    fn apply_scalar(self, a: T) -> T {
111        a.sqrt()
112    }
113}
114
115impl<T: Scalar> UnaryOp<T> for RecipSqrt {
116    #[inline(always)]
117    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
118        Arch::recip_sqrt(v)
119    }
120    #[inline(always)]
121    fn apply_scalar(self, a: T) -> T {
122        T::ONE / a.sqrt()
123    }
124}
125
126impl<T: Scalar + PartialOrd + NumericElement> UnaryOp<T> for Clamp<T> {
127    #[inline(always)]
128    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
129        // clamp(v, lo, hi) = max(lo, min(v, hi))
130        let lo_vec = Arch::splat(self.lo);
131        let hi_vec = Arch::splat(self.hi);
132        let clamped_hi = Arch::min(v, hi_vec);
133        Arch::max(clamped_hi, lo_vec)
134    }
135    #[inline(always)]
136    fn apply_scalar(self, a: T) -> T {
137        a.min_scalar(self.hi).max_scalar(self.lo)
138    }
139}
140
141impl<T: Scalar> UnaryOp<T> for Popcount {
142    #[inline(always)]
143    unsafe fn apply<Arch: SimdKernel<T>>(self, v: Arch::Vector) -> Arch::Vector {
144        Arch::popcount(v)
145    }
146    #[inline(always)]
147    fn apply_scalar(self, a: T) -> T {
148        T::cast_from(a.count_ones() as i32)
149    }
150}