tract_linalg/frame/
element_wise.rs

1use std::fmt::Debug;
2use std::marker::PhantomData;
3
4use tract_data::TractResult;
5
6use crate::LADatum;
7
8use super::element_wise_helper::map_slice_with_alignment;
9
10macro_rules! ew_impl_wrap {
11    ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $run: item) => {
12        paste! {
13            #[derive(Copy, Clone, Debug)]
14            #[allow(non_camel_case_types)]
15            pub struct $func;
16
17            impl crate::frame::element_wise::ElementWiseKer<$ti, $params> for $func {
18                #[inline(always)]
19                fn name() -> &'static str {
20                    stringify!($func)
21                }
22                #[inline(always)]
23                fn nr() -> usize {
24                    $nr
25                }
26                #[inline(always)]
27                fn alignment_items() -> usize {
28                    $alignment_items
29                }
30                $run
31            }
32        }
33    };
34}
35
36/// Define an f16 element-wise kernel for cores without native f16 arithmetic by
37/// round-tripping through an existing f32 kernel: convert each `CHUNK`-sized f16
38/// slice into an aligned f32 scratch, run the f32 kernel in place, convert back.
39///
40/// Callers supply the `unsafe` f16<->f32 conversion fns (their target-feature
41/// gating, if any, lives on those fns — this macro is architecture-agnostic), the
42/// f32 kernel to reuse, the f32-scratch `CHUNK`, and the scratch alignment (must
43/// satisfy the f32 kernel's input-alignment contract, since `run` is called
44/// directly, bypassing `map_slice_with_alignment`). The remaining arguments match
45/// `ew_impl_wrap!`.
46///
47/// `CHUNK` must be a multiple of `nr`: the f32 kernel steps `nr` lanes with no
48/// tail, and each chunk length passed to it is a multiple of `nr` only because
49/// both `CHUNK` and every buffer length are.
50///
51/// The param arm converts the f16-side param into the f32 kernel's param via
52/// `$pname => $pconv` (e.g. `f16, alpha => alpha.to_f32()`), computed once per call.
53macro_rules! ew_impl_f16_via_f32 {
54    ($func:ident, $nr:expr, $alignment_items:expr, $chunk:expr, $scratch_align:literal,
55     $cvt_in:path, $cvt_out:path, $f32_kernel:ty) => {
56        ew_impl_f16_via_f32!(@build $func, $nr, $alignment_items, $chunk, $scratch_align,
57            $cvt_in, $cvt_out, $f32_kernel, (), _params, ());
58    };
59    ($func:ident, $nr:expr, $alignment_items:expr, $chunk:expr, $scratch_align:literal,
60     $cvt_in:path, $cvt_out:path, $f32_kernel:ty, $params:ty, $pname:ident => $pconv:expr) => {
61        ew_impl_f16_via_f32!(@build $func, $nr, $alignment_items, $chunk, $scratch_align,
62            $cvt_in, $cvt_out, $f32_kernel, $params, $pname, $pconv);
63    };
64    (@build $func:ident, $nr:expr, $alignment_items:expr, $chunk:expr, $scratch_align:literal,
65     $cvt_in:path, $cvt_out:path, $f32_kernel:ty, $params:ty, $pname:ident, $pconv:expr) => {
66        ew_impl_wrap!(
67            f16, $func, $nr, $alignment_items, $params,
68            #[inline(never)]
69            fn run(buf: &mut [f16], $pname: $params) {
70                const _: () = assert!(
71                    $chunk % $nr == 0,
72                    "CHUNK must be a multiple of nr; the f32 kernel steps nr lanes with no tail"
73                );
74                #[repr(C, align($scratch_align))]
75                struct AlignedScratch([f32; $chunk]);
76                debug_assert!(buf.len() % Self::nr() == 0);
77                debug_assert!(buf.as_ptr() as usize % Self::alignment_bytes() == 0);
78                if buf.is_empty() {
79                    return;
80                }
81                let f32_params = $pconv;
82                let mut scratch = std::mem::MaybeUninit::<AlignedScratch>::uninit();
83                // SAFETY: f32 has no invalid bit patterns, and every `s[..n]` element is
84                // written by `$cvt_in` before the f32 kernel or `$cvt_out` reads it, so the
85                // scratch never needs zero-initialising.
86                let s = unsafe { &mut (*scratch.as_mut_ptr()).0 };
87                let mut i = 0;
88                while i < buf.len() {
89                    let n = ($chunk).min(buf.len() - i);
90                    unsafe { $cvt_in(&buf[i..i + n], &mut s[..n]) };
91                    <$f32_kernel>::run(&mut s[..n], f32_params);
92                    unsafe { $cvt_out(&s[..n], &mut buf[i..i + n]) };
93                    i += n;
94                }
95            }
96        );
97    };
98}
99
100macro_rules! ew_impl {
101    ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr) => {
102        paste! {
103            mod [<sys_ $func>] {
104                #[allow(unused_imports)]
105                use tract_data::prelude::f16;
106                extern_kernel!(fn $func(ptr: *mut $ti, count: usize) -> ());
107            }
108            ew_impl_wrap!($ti, $func, $nr, $alignment_items, (),
109                #[inline(never)]
110                fn run(buf: &mut [$ti], _params: ()) {
111                    unsafe { [<sys_ $func>]::$func(buf.as_mut_ptr(), buf.len()) }
112                }
113            );
114        }
115    };
116    ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty) => {
117        paste! {
118            mod [<sys_ $func>] {
119                #[allow(unused_imports)]
120                use tract_data::prelude::f16;
121                extern_kernel!(fn $func(ptr: *mut $ti, count: usize, params: $params) -> ());
122            }
123            ew_impl_wrap!($ti, $func, $nr, $alignment_items, $params,
124                #[inline(never)]
125                fn run(buf: &mut [$ti], params: $params) {
126                    unsafe { [<sys_ $func>]::$func(buf.as_mut_ptr(), buf.len(), params) }
127                }
128            );
129        }
130    };
131}
132
133pub trait ElementWise<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
134where
135    Params: Copy + Send + Sync + Debug + 'static + Default,
136    T: Copy + Debug + PartialEq + Send + Sync,
137{
138    fn name(&self) -> &'static str;
139    fn run(&self, vec: &mut [T]) -> TractResult<()> {
140        self.run_with_params(vec, Params::default())
141    }
142    fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<()>;
143}
144
145dyn_clone::clone_trait_object!(<T, Params> ElementWise<T, Params> where T: Copy, Params: Copy);
146
147#[derive(Debug, Clone, new)]
148pub struct ElementWiseImpl<K, T, Params = ()>
149where
150    T: LADatum,
151    Params: Copy + Send + Sync + Debug + 'static + Default,
152    K: ElementWiseKer<T, Params> + Clone,
153{
154    phantom: PhantomData<(K, T, Params)>,
155}
156
157impl<K, T, Params> ElementWise<T, Params> for ElementWiseImpl<K, T, Params>
158where
159    T: LADatum,
160    Params: Copy + Send + Sync + Debug + 'static + Default,
161    K: ElementWiseKer<T, Params> + Clone,
162{
163    fn name(&self) -> &'static str {
164        K::name()
165    }
166    fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<()> {
167        map_slice_with_alignment(vec, |data| K::run(data, params), K::nr(), K::alignment_bytes())
168    }
169}
170
171pub trait ElementWiseKer<T, Params = ()>:
172    Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
173where
174    Params: Copy + Send + Sync + Debug + 'static + Default,
175    T: LADatum,
176{
177    fn name() -> &'static str;
178    fn alignment_bytes() -> usize {
179        Self::alignment_items() * T::datum_type().size_of()
180    }
181    fn alignment_items() -> usize;
182    fn nr() -> usize;
183    fn run(vec: &mut [T], params: Params);
184    fn ew() -> Box<dyn ElementWise<T, Params>> {
185        Box::new(ElementWiseImpl::<Self, T, Params>::new())
186    }
187}
188
189#[cfg(test)]
190pub mod test {
191    use crate::{LADatum, frame::element_wise::*};
192    use num_traits::AsPrimitive;
193    use proptest::test_runner::{TestCaseError, TestCaseResult};
194    use tract_data::internal::*;
195
196    /// Every finite `f16`, or a 1/4096 grid of `[-30, 30]` for wider types.
197    ///
198    /// The grid samples where the `f16` set is enumerated, but it reaches past every input
199    /// clamp the f32 kernels apply, so no input outside its bounds takes a path it has not
200    /// already exercised.
201    ///
202    /// The step has to stay fine because these invariants break on value-specific
203    /// rounding, not over a contiguous region: the Tanh kernels leave `[-1, 1]` only
204    /// inside a band about 0.3 wide, and a 1/256 grid steps clean over
205    /// `arm64simd_tanh_f32_4n`'s share of it.
206    fn invariant_sweep<T: LADatum>() -> Vec<T>
207    where
208        f32: AsPrimitive<T>,
209    {
210        if T::datum_type() == f16::datum_type() {
211            let all: Vec<f16> =
212                (0..=u16::MAX).map(f16::from_bits).filter(|x| x.is_finite()).collect();
213            let all = tensor1(&all).cast_to::<T>().unwrap().into_owned();
214            return all.try_as_plain().unwrap().as_slice::<T>().unwrap().to_vec();
215        }
216        (-30 * 4096..=30 * 4096).map(|i| (i as f32 / 4096.).as_()).collect()
217    }
218
219    /// Assert `invariant` holds of every `(input, output)` pair a kernel produces over
220    /// [`invariant_sweep`], reporting `expected` on the first pair that breaks it.
221    ///
222    /// The accuracy tests cannot stand in for this on the saturating tails: there the true
223    /// value is smaller than the rounding error of the kernels' own arithmetic, so an
224    /// output that violates the range or the sign still compares close to the reference.
225    pub fn test_element_wise_invariant<K: ElementWiseKer<T>, T: LADatum>(
226        expected: &str,
227        invariant: impl Fn(T, T) -> bool,
228    ) -> TestCaseResult
229    where
230        f32: AsPrimitive<T>,
231    {
232        crate::setup_test_logger();
233        let values = invariant_sweep::<T>();
234        let mut found = values.clone();
235        K::ew().run(&mut found).unwrap();
236        for (x, y) in values.iter().zip(found.iter()) {
237            proptest::prop_assert!(
238                invariant(*x, *y),
239                "{}({x:?}) returned {y:?}, expected {expected}",
240                K::name()
241            );
242        }
243        Ok(())
244    }
245
246    pub fn test_element_wise<K: ElementWiseKer<T, ()>, T: LADatum, F: Fn(T) -> T>(
247        values: &[T],
248        reference: F,
249    ) -> TestCaseResult {
250        test_element_wise_params::<K, T, F, ()>(values, reference, ())
251    }
252
253    pub fn test_element_wise_params<
254        K: ElementWiseKer<T, Params>,
255        T: LADatum,
256        F: Fn(T) -> T,
257        Params,
258    >(
259        values: &[T],
260        reference: F,
261        params: Params,
262    ) -> TestCaseResult
263    where
264        Params: Copy + Send + Sync + Debug + 'static + Default,
265    {
266        crate::setup_test_logger();
267        let op = ElementWiseImpl::<K, T, Params>::new();
268        let mut values = values.to_vec();
269        while values.len() < K::nr() {
270            values.push(T::zero());
271        }
272        let expected = values.iter().copied().map(reference).collect::<Vec<_>>();
273        let mut found = values;
274        op.run_with_params(&mut found, params).unwrap();
275        tensor1(&found)
276            .close_enough(&tensor1(&expected), true)
277            .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
278        Ok(())
279    }
280}
tract_linalg/frame/element_wise.rs

tract_linalg/frame/
element_wise.rs
