tenferro_cpu/
analytic.rs

1use num_complex::{Complex32, Complex64};
2use num_traits::{One, Zero};
3use strided_kernel::{map_into, reduce, zip_map2_into, StridedView};
4use tenferro_core_ops::PrimitiveOpKind;
5
6use super::{tensor_from_array, typed_array_uninit_from_pool, typed_view, typed_view_from_view};
7use crate::buffer_pool::{BufferPool, PoolScalar};
8use tenferro_tensor::{
9    BackendId, CapabilityAxis, DType, Tensor, TensorRank, TensorRead, TensorScalar, TensorView,
10    TypedTensor, TypedTensorView,
11};
12
13trait UnaryAnalyticElem: Copy + Clone + One + Zero {
14    fn exp_elem(self) -> Self;
15    fn log_elem(self) -> Self;
16    fn sin_elem(self) -> Self;
17    fn cos_elem(self) -> Self;
18    fn tanh_elem(self) -> Self;
19    fn sqrt_elem(self) -> Self;
20    fn rsqrt_elem(self) -> Self;
21    fn expm1_elem(self) -> Self;
22    fn log1p_elem(self) -> Self;
23}
24
25trait PowElem: Copy + Clone + Zero {
26    fn pow_elem(self, exponent: Self) -> Self;
27}
28
29trait IntegerPowElem: PowElem + TensorScalar + PoolScalar + Send + Sync + 'static {
30    fn is_negative_exponent(self) -> bool;
31    fn wrapping_pow_nonnegative(self, exponent: Self) -> Self;
32}
33
34macro_rules! impl_real_analytic_elem {
35    ($ty:ty) => {
36        impl UnaryAnalyticElem for $ty {
37            fn exp_elem(self) -> Self {
38                self.exp()
39            }
40
41            fn log_elem(self) -> Self {
42                self.ln()
43            }
44
45            fn sin_elem(self) -> Self {
46                self.sin()
47            }
48
49            fn cos_elem(self) -> Self {
50                self.cos()
51            }
52
53            fn tanh_elem(self) -> Self {
54                self.tanh()
55            }
56
57            fn sqrt_elem(self) -> Self {
58                self.sqrt()
59            }
60
61            fn rsqrt_elem(self) -> Self {
62                Self::one() / self.sqrt()
63            }
64
65            fn expm1_elem(self) -> Self {
66                self.exp_m1()
67            }
68
69            fn log1p_elem(self) -> Self {
70                self.ln_1p()
71            }
72        }
73
74        impl PowElem for $ty {
75            fn pow_elem(self, exponent: Self) -> Self {
76                self.powf(exponent)
77            }
78        }
79    };
80}
81
82macro_rules! impl_complex_analytic_elem {
83    ($ty:ty) => {
84        impl UnaryAnalyticElem for $ty {
85            fn exp_elem(self) -> Self {
86                self.exp()
87            }
88
89            fn log_elem(self) -> Self {
90                self.ln()
91            }
92
93            fn sin_elem(self) -> Self {
94                self.sin()
95            }
96
97            fn cos_elem(self) -> Self {
98                self.cos()
99            }
100
101            fn tanh_elem(self) -> Self {
102                self.tanh()
103            }
104
105            fn sqrt_elem(self) -> Self {
106                self.sqrt()
107            }
108
109            fn rsqrt_elem(self) -> Self {
110                Self::one() / self.sqrt()
111            }
112
113            fn expm1_elem(self) -> Self {
114                self.exp() - Self::one()
115            }
116
117            fn log1p_elem(self) -> Self {
118                (self + Self::one()).ln()
119            }
120        }
121
122        impl PowElem for $ty {
123            fn pow_elem(self, exponent: Self) -> Self {
124                self.powc(exponent)
125            }
126        }
127    };
128}
129
130impl_real_analytic_elem!(f32);
131impl_real_analytic_elem!(f64);
132impl_complex_analytic_elem!(Complex32);
133impl_complex_analytic_elem!(Complex64);
134
135macro_rules! impl_integer_pow_elem {
136    ($ty:ty) => {
137        impl PowElem for $ty {
138            fn pow_elem(self, exponent: Self) -> Self {
139                self.wrapping_pow_nonnegative(exponent)
140            }
141        }
142
143        impl IntegerPowElem for $ty {
144            fn is_negative_exponent(self) -> bool {
145                self < 0
146            }
147
148            fn wrapping_pow_nonnegative(self, exponent: Self) -> Self {
149                let mut base = self;
150                let mut exp = exponent as u64;
151                let mut acc: Self = 1;
152                while exp != 0 {
153                    if exp & 1 == 1 {
154                        acc = acc.wrapping_mul(base);
155                    }
156                    exp >>= 1;
157                    if exp != 0 {
158                        base = base.wrapping_mul(base);
159                    }
160                }
161                acc
162            }
163        }
164    };
165}
166
167impl_integer_pow_elem!(i32);
168impl_integer_pow_elem!(i64);
169
170fn with_local_pool<T>(f: impl FnOnce(&mut BufferPool) -> T) -> T {
171    let mut buffers = BufferPool::new();
172    f(&mut buffers)
173}
174
175enum AnalyticReadView<'a> {
176    F32(TypedTensorView<'a, f32>),
177    F64(TypedTensorView<'a, f64>),
178    I32(TypedTensorView<'a, i32>),
179    I64(TypedTensorView<'a, i64>),
180    Bool,
181    C32(TypedTensorView<'a, Complex32>),
182    C64(TypedTensorView<'a, Complex64>),
183}
184
185fn read_as_analytic_view(input: TensorRead<'_>) -> AnalyticReadView<'_> {
186    match input {
187        TensorRead::Tensor(Tensor::F32(tensor)) => AnalyticReadView::F32(tensor.as_view()),
188        TensorRead::Tensor(Tensor::F64(tensor)) => AnalyticReadView::F64(tensor.as_view()),
189        TensorRead::Tensor(Tensor::I32(tensor)) => AnalyticReadView::I32(tensor.as_view()),
190        TensorRead::Tensor(Tensor::I64(tensor)) => AnalyticReadView::I64(tensor.as_view()),
191        TensorRead::Tensor(Tensor::Bool(_)) => AnalyticReadView::Bool,
192        TensorRead::Tensor(Tensor::C32(tensor)) => AnalyticReadView::C32(tensor.as_view()),
193        TensorRead::Tensor(Tensor::C64(tensor)) => AnalyticReadView::C64(tensor.as_view()),
194        TensorRead::View(TensorView::F32(view)) => AnalyticReadView::F32(view),
195        TensorRead::View(TensorView::F64(view)) => AnalyticReadView::F64(view),
196        TensorRead::View(TensorView::I32(view)) => AnalyticReadView::I32(view),
197        TensorRead::View(TensorView::I64(view)) => AnalyticReadView::I64(view),
198        TensorRead::View(TensorView::Bool(_)) => AnalyticReadView::Bool,
199        TensorRead::View(TensorView::C32(view)) => AnalyticReadView::C32(view),
200        TensorRead::View(TensorView::C64(view)) => AnalyticReadView::C64(view),
201    }
202}
203
204fn typed_unary_with_pool<T>(
205    op: &'static str,
206    buffers: &mut BufferPool,
207    input: &TypedTensor<T>,
208    f: impl Fn(T) -> T + Copy + Sync,
209) -> crate::Result<TypedTensor<T>>
210where
211    T: Copy + PoolScalar + 'static,
212{
213    // SAFETY: the following kernel overwrites every output element before any read.
214    let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
215    map_into(&mut out.view_mut(), &typed_view(op, input)?, f)
216        .map_err(|err| crate::Error::backend_failure(op, err))?;
217    Ok(tensor_from_array(out))
218}
219
220fn typed_unary_view_with_pool<T, R>(
221    op: &'static str,
222    buffers: &mut BufferPool,
223    input: &TypedTensorView<'_, T, R>,
224    f: impl Fn(T) -> T + Copy + Sync,
225) -> crate::Result<TypedTensor<T>>
226where
227    T: Copy + PoolScalar + 'static,
228    R: TensorRank,
229{
230    // SAFETY: the following kernel overwrites every output element before any read.
231    let mut out = unsafe { typed_array_uninit_from_pool(buffers, input.shape()) }?;
232    map_into(&mut out.view_mut(), &typed_view_from_view(op, input)?, f)
233        .map_err(|err| crate::Error::backend_failure(op, err))?;
234    Ok(tensor_from_array(out))
235}
236
237fn typed_unary_tensor_with_pool<T>(
238    op: &'static str,
239    buffers: &mut BufferPool,
240    input: &TypedTensor<T>,
241    f: impl Fn(T) -> T + Copy + Sync,
242) -> crate::Result<Tensor>
243where
244    T: Copy + PoolScalar + TensorScalar + 'static,
245{
246    let out = typed_unary_with_pool(op, buffers, input, f)?;
247    Ok(T::typed_tensor_into_tensor(out))
248}
249
250fn typed_unary_view_tensor_with_pool<T, R>(
251    op: &'static str,
252    buffers: &mut BufferPool,
253    input: &TypedTensorView<'_, T, R>,
254    f: impl Fn(T) -> T + Copy + Sync,
255) -> crate::Result<Tensor>
256where
257    T: Copy + PoolScalar + TensorScalar + 'static,
258    R: TensorRank,
259{
260    let out = typed_unary_view_with_pool(op, buffers, input, f)?;
261    Ok(T::typed_tensor_into_tensor(out))
262}
263
264fn require_cpu_capability(
265    op_kind: PrimitiveOpKind,
266    op: &'static str,
267    dtype: DType,
268    axis: CapabilityAxis,
269) -> crate::Result<()> {
270    let supported = crate::cpu_capabilities()
271        .iter()
272        .copied()
273        .find(|entry| entry.op == op_kind && entry.dtype == dtype)
274        .is_some_and(|entry| entry.axis(axis).is_supported());
275    if supported {
276        Ok(())
277    } else {
278        Err(crate::Error::unsupported_op_dtype(
279            op,
280            dtype,
281            BackendId::Cpu,
282        ))
283    }
284}
285
286fn strided_view_contains<T>(
287    op: &'static str,
288    view: &StridedView<'_, T>,
289    pred: impl Fn(T) -> bool + Copy + Sync,
290) -> crate::Result<bool>
291where
292    T: Copy + Send + Sync,
293{
294    reduce(view, pred, |lhs, rhs| lhs || rhs, false)
295        .map_err(|err| crate::Error::backend_failure(op, err))
296}
297
298fn ensure_nonnegative_integer_exponents<T>(
299    op: &'static str,
300    rhs: &StridedView<'_, T>,
301) -> crate::Result<()>
302where
303    T: IntegerPowElem,
304{
305    if strided_view_contains(op, rhs, |value| value.is_negative_exponent())? {
306        return Err(crate::Error::negative_integer_exponent(op, T::dtype()));
307    }
308    Ok(())
309}
310
311fn typed_pow_view_with_pool<T, L, R>(
312    op: &'static str,
313    buffers: &mut BufferPool,
314    lhs: &TypedTensorView<'_, T, L>,
315    rhs: &TypedTensorView<'_, T, R>,
316) -> crate::Result<TypedTensor<T>>
317where
318    T: PowElem + PoolScalar + 'static,
319    L: TensorRank,
320    R: TensorRank,
321{
322    let output_shape = if lhs.shape() == rhs.shape() {
323        lhs.shape()
324    } else if lhs.shape().is_empty() {
325        rhs.shape()
326    } else if rhs.shape().is_empty() {
327        lhs.shape()
328    } else {
329        return Err(crate::Error::ShapeMismatch {
330            op,
331            lhs: lhs.shape().to_vec(),
332            rhs: rhs.shape().to_vec(),
333        });
334    };
335    // SAFETY: the selected map kernel overwrites every output element.
336    let mut out = unsafe { typed_array_uninit_from_pool(buffers, output_shape) }?;
337    if lhs.shape() == rhs.shape() {
338        zip_map2_into(
339            &mut out.view_mut(),
340            &typed_view_from_view(op, lhs)?,
341            &typed_view_from_view(op, rhs)?,
342            |x, y| x.pow_elem(y),
343        )
344        .map_err(|err| crate::Error::backend_failure(op, err))?;
345    } else if lhs.shape().is_empty() {
346        let scalar = typed_view_from_view(op, lhs)?.get(&[]);
347        map_into(&mut out.view_mut(), &typed_view_from_view(op, rhs)?, |x| {
348            scalar.pow_elem(x)
349        })
350        .map_err(|err| crate::Error::backend_failure(op, err))?;
351    } else {
352        let scalar = typed_view_from_view(op, rhs)?.get(&[]);
353        map_into(&mut out.view_mut(), &typed_view_from_view(op, lhs)?, |x| {
354            x.pow_elem(scalar)
355        })
356        .map_err(|err| crate::Error::backend_failure(op, err))?;
357    }
358    Ok(tensor_from_array(out))
359}
360
361macro_rules! define_unary_analytic_dispatch {
362    ($dispatch_fn:ident, $dispatch_with_pool_fn:ident, $dispatch_read_with_pool_fn:ident, $op_kind:ident, $elem_fn:ident) => {
363        #[cfg(test)]
364        pub(crate) fn $dispatch_fn(input: &Tensor) -> crate::Result<Tensor> {
365            with_local_pool(|buffers| $dispatch_with_pool_fn(buffers, input))
366        }
367
368        pub(crate) fn $dispatch_with_pool_fn(
369            buffers: &mut BufferPool,
370            input: &Tensor,
371        ) -> crate::Result<Tensor> {
372            require_cpu_capability(
373                PrimitiveOpKind::$op_kind,
374                stringify!($dispatch_fn),
375                input.dtype(),
376                CapabilityAxis::OwnedResult,
377            )?;
378            tenferro_tensor::with_scalar!(
379                input,
380                float_complex,
381                backend = BackendId::Cpu,
382                op = stringify!($dispatch_fn),
383                |tensor| -> crate::Result<Tensor> {
384                    typed_unary_tensor_with_pool(stringify!($dispatch_fn), buffers, tensor, |x| {
385                        x.$elem_fn()
386                    })
387                }
388            )
389        }
390
391        pub(crate) fn $dispatch_read_with_pool_fn(
392            buffers: &mut BufferPool,
393            input: TensorRead<'_>,
394        ) -> crate::Result<Tensor> {
395            let dtype = input.dtype();
396            require_cpu_capability(
397                PrimitiveOpKind::$op_kind,
398                stringify!($dispatch_fn),
399                dtype,
400                CapabilityAxis::ReadInputs,
401            )?;
402            tenferro_tensor::with_scalar_read!(
403                input,
404                float_complex,
405                backend = BackendId::Cpu,
406                op = stringify!($dispatch_fn),
407                |view| -> crate::Result<Tensor> {
408                    typed_unary_view_tensor_with_pool(
409                        stringify!($dispatch_fn),
410                        buffers,
411                        &view,
412                        |x| x.$elem_fn(),
413                    )
414                }
415            )
416        }
417    };
418}
419
420define_unary_analytic_dispatch!(exp, exp_with_pool, exp_read_with_pool, Exp, exp_elem);
421define_unary_analytic_dispatch!(log, log_with_pool, log_read_with_pool, Log, log_elem);
422define_unary_analytic_dispatch!(sin, sin_with_pool, sin_read_with_pool, Sin, sin_elem);
423define_unary_analytic_dispatch!(cos, cos_with_pool, cos_read_with_pool, Cos, cos_elem);
424define_unary_analytic_dispatch!(tanh, tanh_with_pool, tanh_read_with_pool, Tanh, tanh_elem);
425define_unary_analytic_dispatch!(sqrt, sqrt_with_pool, sqrt_read_with_pool, Sqrt, sqrt_elem);
426define_unary_analytic_dispatch!(
427    rsqrt,
428    rsqrt_with_pool,
429    rsqrt_read_with_pool,
430    Rsqrt,
431    rsqrt_elem
432);
433define_unary_analytic_dispatch!(
434    expm1,
435    expm1_with_pool,
436    expm1_read_with_pool,
437    Expm1,
438    expm1_elem
439);
440define_unary_analytic_dispatch!(
441    log1p,
442    log1p_with_pool,
443    log1p_read_with_pool,
444    Log1p,
445    log1p_elem
446);
447
448pub fn pow(lhs: &Tensor, rhs: &Tensor) -> crate::Result<Tensor> {
449    with_local_pool(|buffers| pow_with_pool(buffers, lhs, rhs))
450}
451
452pub(crate) fn pow_with_pool(
453    buffers: &mut BufferPool,
454    lhs: &Tensor,
455    rhs: &Tensor,
456) -> crate::Result<Tensor> {
457    match (lhs, rhs) {
458        (Tensor::F32(a), Tensor::F32(b)) => Ok(Tensor::F32(typed_pow_with_pool(buffers, a, b)?)),
459        (Tensor::F64(a), Tensor::F64(b)) => Ok(Tensor::F64(typed_pow_with_pool(buffers, a, b)?)),
460        (Tensor::I32(a), Tensor::I32(b)) => {
461            Ok(Tensor::I32(typed_integer_pow_with_pool(buffers, a, b)?))
462        }
463        (Tensor::I64(a), Tensor::I64(b)) => {
464            Ok(Tensor::I64(typed_integer_pow_with_pool(buffers, a, b)?))
465        }
466        (Tensor::C32(a), Tensor::C32(b)) => Ok(Tensor::C32(typed_pow_with_pool(buffers, a, b)?)),
467        (Tensor::C64(a), Tensor::C64(b)) => Ok(Tensor::C64(typed_pow_with_pool(buffers, a, b)?)),
468        _ => Err(crate::Error::DTypeMismatch {
469            op: "pow",
470            lhs: lhs.dtype(),
471            rhs: rhs.dtype(),
472        }),
473    }
474}
475
476pub(crate) fn pow_read_with_pool(
477    buffers: &mut BufferPool,
478    lhs: TensorRead<'_>,
479    rhs: TensorRead<'_>,
480) -> crate::Result<Tensor> {
481    let lhs_dtype = lhs.dtype();
482    let rhs_dtype = rhs.dtype();
483    match (read_as_analytic_view(lhs), read_as_analytic_view(rhs)) {
484        (AnalyticReadView::F32(a), AnalyticReadView::F32(b)) => Ok(Tensor::F32(
485            typed_pow_view_with_pool("pow", buffers, &a, &b)?,
486        )),
487        (AnalyticReadView::F64(a), AnalyticReadView::F64(b)) => Ok(Tensor::F64(
488            typed_pow_view_with_pool("pow", buffers, &a, &b)?,
489        )),
490        (AnalyticReadView::I32(a), AnalyticReadView::I32(b)) => Ok(Tensor::I32(
491            typed_integer_pow_view_with_pool(buffers, &a, &b)?,
492        )),
493        (AnalyticReadView::I64(a), AnalyticReadView::I64(b)) => Ok(Tensor::I64(
494            typed_integer_pow_view_with_pool(buffers, &a, &b)?,
495        )),
496        (AnalyticReadView::C32(a), AnalyticReadView::C32(b)) => Ok(Tensor::C32(
497            typed_pow_view_with_pool("pow", buffers, &a, &b)?,
498        )),
499        (AnalyticReadView::C64(a), AnalyticReadView::C64(b)) => Ok(Tensor::C64(
500            typed_pow_view_with_pool("pow", buffers, &a, &b)?,
501        )),
502        _ => Err(crate::Error::DTypeMismatch {
503            op: "pow",
504            lhs: lhs_dtype,
505            rhs: rhs_dtype,
506        }),
507    }
508}
509
510fn typed_pow_with_pool<T>(
511    buffers: &mut BufferPool,
512    lhs: &TypedTensor<T>,
513    rhs: &TypedTensor<T>,
514) -> crate::Result<TypedTensor<T>>
515where
516    T: PowElem + PoolScalar,
517{
518    let output_shape = if lhs.shape() == rhs.shape() {
519        lhs.shape()
520    } else if lhs.shape().is_empty() {
521        rhs.shape()
522    } else if rhs.shape().is_empty() {
523        lhs.shape()
524    } else {
525        return Err(crate::Error::ShapeMismatch {
526            op: "pow",
527            lhs: lhs.shape().to_vec(),
528            rhs: rhs.shape().to_vec(),
529        });
530    };
531    // SAFETY: the selected map kernel overwrites every output element.
532    let mut out = unsafe { typed_array_uninit_from_pool(buffers, output_shape) }?;
533    if lhs.shape() == rhs.shape() {
534        zip_map2_into(
535            &mut out.view_mut(),
536            &typed_view("pow", lhs)?,
537            &typed_view("pow", rhs)?,
538            |x, y| x.pow_elem(y),
539        )
540        .map_err(|err| crate::Error::backend_failure("pow", err))?;
541    } else if lhs.shape().is_empty() {
542        let scalar = typed_view("pow", lhs)?.get(&[]);
543        map_into(&mut out.view_mut(), &typed_view("pow", rhs)?, |x| {
544            scalar.pow_elem(x)
545        })
546        .map_err(|err| crate::Error::backend_failure("pow", err))?;
547    } else {
548        let scalar = typed_view("pow", rhs)?.get(&[]);
549        map_into(&mut out.view_mut(), &typed_view("pow", lhs)?, |x| {
550            x.pow_elem(scalar)
551        })
552        .map_err(|err| crate::Error::backend_failure("pow", err))?;
553    }
554    Ok(tensor_from_array(out))
555}
556
557fn typed_integer_pow_with_pool<T>(
558    buffers: &mut BufferPool,
559    lhs: &TypedTensor<T>,
560    rhs: &TypedTensor<T>,
561) -> crate::Result<TypedTensor<T>>
562where
563    T: IntegerPowElem,
564{
565    let rhs_view = typed_view("pow", rhs)?;
566    ensure_nonnegative_integer_exponents("pow", &rhs_view)?;
567    typed_pow_with_pool(buffers, lhs, rhs)
568}
569
570fn typed_integer_pow_view_with_pool<T, L, R>(
571    buffers: &mut BufferPool,
572    lhs: &TypedTensorView<'_, T, L>,
573    rhs: &TypedTensorView<'_, T, R>,
574) -> crate::Result<TypedTensor<T>>
575where
576    T: IntegerPowElem,
577    L: TensorRank,
578    R: TensorRank,
579{
580    let rhs_view = typed_view_from_view("pow", rhs)?;
581    ensure_nonnegative_integer_exponents("pow", &rhs_view)?;
582    typed_pow_view_with_pool("pow", buffers, lhs, rhs)
583}
tenferro_cpu/analytic.rs

tenferro_cpu/
analytic.rs
