ruprim/reduce/components/instructions/
maxabs.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 absolute value.
11#[derive(Debug, RudaType, Clone)]
12pub struct MaxAbs;
13
14impl ReduceFamily for MaxAbs {
15    type Instruction<P: ReducePrecision> = Self;
16    type Config = ();
17}
18
19#[ruda]
20impl<P: ReducePrecision> ReduceInstruction<P> for MaxAbs {
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        MaxAbs {}
34    }
35
36    fn null_input(_this: &Self) -> Vector<P::EI, P::SI> {
37        Vector::empty().fill(P::EI::from_int(0))
38    }
39
40    fn null_accumulator(_this: &Self) -> Accumulator<P> {
41        Accumulator::<P> {
42            elements: Value::new_single(Vector::empty().fill(P::EA::from_int(0))),
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(Vector::abs(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_abs = Vector::cast_from(Vector::abs(item.elements));
65                select_many(
66                    accumulator_item.greater_than(item_abs),
67                    accumulator_item,
68                    item_abs,
69                )
70            }
71        };
72
73        accumulator.elements.assign(&Value::new_single(elements));
74    }
75
76    fn fuse_accumulators(_this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
77        let accumulator_item = accumulator.elements.item();
78        let other_item = other.elements.item();
79
80        accumulator.elements.assign(&Value::new_single(select_many(
81            accumulator_item.greater_than(other_item),
82            accumulator_item,
83            other_item,
84        )));
85    }
86
87    fn plane_reduce_inplace(_this: &Self, accumulator: &mut Accumulator<P>) {
88        let acc_item = accumulator.elements.item();
89        let candidate_item = Vector::cast_from(plane_max(Vector::abs(acc_item)));
90        let max = select_many(
91            acc_item.greater_than(candidate_item),
92            acc_item,
93            candidate_item,
94        );
95        accumulator.elements.assign(&Value::new_single(max));
96    }
97
98    fn to_output_parallel<Out: Numeric>(
99        _this: &Self,
100        accumulator: Accumulator<P>,
101        _shape_axis_reduce: usize,
102    ) -> Value<Out> {
103        let mut max = P::EA::from_int(0);
104        let accumulator = accumulator.elements.item();
105        #[unroll]
106        for k in 0..accumulator.size() {
107            let candidate = accumulator[k];
108            max = select(candidate > max, candidate, max);
109        }
110        Value::new_single(Out::cast_from(max))
111    }
112
113    fn to_output_perpendicular<Out: Numeric>(
114        _this: &Self,
115        accumulator: Accumulator<P>,
116        _shape_axis_reduce: usize,
117    ) -> Value<Vector<Out, P::SI>> {
118        Value::new_single(Vector::cast_from(accumulator.elements.item()))
119    }
120}
ruprim/reduce/components/instructions/maxabs.rs

ruprim/reduce/components/instructions/
maxabs.rs
