ruprim/reduce/components/instructions/
max.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// TODO Add to test framework.
10/// Return the item with the maximum value.
11#[derive(Debug, RudaType, Clone)]
12pub struct Max;
13
14impl ReduceFamily for Max {
15    type Instruction<P: ReducePrecision> = Self;
16    type Config = ();
17}
18
19#[ruda]
20impl<P: ReducePrecision> ReduceInstruction<P> for Max {
21    type SharedAccumulator = SharedMemory<Vector<P::EA, P::SI>>;
22    type Config = ();
23
24    fn requirements(_this: &Self) -> ReduceRequirements {
25        ReduceRequirements { coordinates: false }
26    }
27
28    fn accumulator_format(_this: &Self) -> comptime_type!(AccumulatorFormat) {
29        AccumulatorFormat::Single
30    }
31
32    fn from_config(_config: Self::Config) -> Self {
33        Max {}
34    }
35
36    fn null_input(_this: &Self) -> Vector<P::EI, P::SI> {
37        Vector::empty().fill(P::EI::min_value())
38    }
39
40    fn null_accumulator(_this: &Self) -> Accumulator<P> {
41        Accumulator::<P> {
42            elements: Value::new_single(Vector::empty().fill(P::EA::min_value())),
43            args: Value::new_None(),
44        }
45    }
46
47    fn reduce(
48        _this: &Self,
49        accumulator: &mut Accumulator<P>,
50        item: Item<P>,
51        #[comptime] reduce_step: ReduceStep,
52    ) {
53        let accumulator_item = accumulator.elements.item();
54        let elements = match reduce_step {
55            ReduceStep::Plane => {
56                let candidate_item = Vector::cast_from(plane_max(item.elements));
57                select_many(
58                    accumulator_item.greater_than(candidate_item),
59                    accumulator_item,
60                    candidate_item,
61                )
62            }
63            ReduceStep::Identity => {
64                let item = Vector::cast_from(item.elements);
65                select_many(accumulator_item.greater_than(item), accumulator_item, item)
66            }
67        };
68
69        accumulator.elements.assign(&Value::new_single(elements));
70    }
71
72    fn plane_reduce_inplace(_this: &Self, accumulator: &mut Accumulator<P>) {
73        let acc_item = accumulator.elements.item();
74        let candidate_item = Vector::cast_from(plane_max(acc_item));
75        let max = select_many(
76            acc_item.greater_than(candidate_item),
77            acc_item,
78            candidate_item,
79        );
80        accumulator.elements.assign(&Value::new_single(max));
81    }
82
83    fn fuse_accumulators(_this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
84        let accumulator_item = accumulator.elements.item();
85        let other_item = other.elements.item();
86
87        accumulator.elements.assign(&Value::new_single(select_many(
88            accumulator_item.greater_than(other_item),
89            accumulator_item,
90            other_item,
91        )));
92    }
93
94    fn to_output_parallel<Out: Numeric>(
95        _this: &Self,
96        accumulator: Accumulator<P>,
97        _shape_axis_reduce: usize,
98    ) -> Value<Out> {
99        let mut max = P::EA::min_value();
100        let accumulator = accumulator.elements.item();
101        #[unroll]
102        for k in 0..accumulator.size() {
103            let candidate = accumulator[k];
104            max = select(candidate > max, candidate, max);
105        }
106        Value::new_single(Out::cast_from(max))
107    }
108
109    fn to_output_perpendicular<Out: Numeric>(
110        _this: &Self,
111        accumulator: Accumulator<P>,
112        _shape_axis_reduce: usize,
113    ) -> Value<Vector<Out, P::SI>> {
114        Value::new_single(Vector::cast_from(accumulator.elements.item()))
115    }
116}
ruprim/reduce/components/instructions/max.rs

ruprim/reduce/components/instructions/
max.rs
