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}