ruprim/reduce/components/instructions/
prod.rs

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

ruprim/reduce/components/instructions/
prod.rs
