ruprim/reduce/components/instructions/
mean.rs

1use ruda_kernel::dsl as kernel_dsl;
2use super::{ReduceFamily, ReduceInstruction, ReduceRequirements, Sum};
3use crate::reduce::components::{
4    instructions::{Accumulator, AccumulatorFormat, Item, ReduceStep, Value},
5    precision::ReducePrecision,
6};
7use ruda_kernel::dsl::prelude::*;
8
9#[derive(Debug, RudaType, Clone)]
10pub struct Mean {
11    pub(crate) sum: Sum,
12}
13
14impl ReduceFamily for Mean {
15    type Instruction<P: ReducePrecision> = Self;
16    type Config = ();
17}
18
19#[ruda]
20fn null_input<P: ReducePrecision, SI: ReduceInstruction<P>>(sum: &SI) -> Vector<P::EI, P::SI> {
21    SI::null_input(sum)
22}
23
24#[ruda]
25impl<P: ReducePrecision> ReduceInstruction<P> for Mean {
26    type SharedAccumulator = SharedMemory<Vector<P::EA, P::SI>>;
27    type Config = ();
28
29    fn requirements(_this: &Self) -> ReduceRequirements {
30        ReduceRequirements { coordinates: false }
31    }
32
33    fn accumulator_format(_this: &Self) -> comptime_type!(AccumulatorFormat) {
34        AccumulatorFormat::Single
35    }
36
37    fn from_config(_config: Self::Config) -> Self {
38        Mean { sum: Sum {} }
39    }
40
41    fn null_input(this: &Self) -> Vector<P::EI, P::SI> {
42        <Sum as ReduceInstruction<P>>::null_input(&this.sum)
43    }
44
45    fn null_accumulator(this: &Self) -> Accumulator<P> {
46        <Sum as ReduceInstruction<P>>::null_accumulator(&this.sum)
47    }
48
49    fn reduce(
50        this: &Self,
51        accumulator: &mut Accumulator<P>,
52        item: Item<P>,
53        #[comptime] reduce_step: ReduceStep,
54    ) {
55        <Sum as ReduceInstruction<P>>::reduce(&this.sum, accumulator, item, reduce_step)
56    }
57
58    fn plane_reduce_inplace(this: &Self, accumulator: &mut Accumulator<P>) {
59        <Sum as ReduceInstruction<P>>::plane_reduce_inplace(&this.sum, accumulator)
60    }
61
62    fn fuse_accumulators(this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
63        <Sum as ReduceInstruction<P>>::fuse_accumulators(&this.sum, accumulator, other)
64    }
65
66    fn to_output_parallel<Out: Numeric>(
67        this: &Self,
68        accumulator: Accumulator<P>,
69        shape_axis_reduce: VectorSize,
70    ) -> Value<Out> {
71        let sum = <Sum as ReduceInstruction<P>>::to_output_parallel::<P::EA>(
72            &this.sum,
73            accumulator,
74            shape_axis_reduce,
75        )
76        .item();
77
78        let value = Out::cast_from(sum / P::EA::cast_from(shape_axis_reduce));
79        Value::new_single(value)
80    }
81
82    fn to_output_perpendicular<Out: Numeric>(
83        this: &Self,
84        accumulator: Accumulator<P>,
85        shape_axis_reduce: VectorSize,
86    ) -> Value<Vector<Out, P::SI>> {
87        let sum = <Sum as ReduceInstruction<P>>::to_output_perpendicular::<P::EA>(
88            &this.sum,
89            accumulator,
90            shape_axis_reduce,
91        )
92        .item();
93
94        let vector = Vector::cast_from(sum / Vector::cast_from(shape_axis_reduce));
95        Value::new_single(vector)
96    }
97}
ruprim/reduce/components/instructions/mean.rs

ruprim/reduce/components/instructions/
mean.rs
