ruprim/reduce/components/instructions/
argmax.rs

1use ruda_kernel::dsl as kernel_dsl;
2use super::{ArgAccumulator, ReduceFamily, ReduceInstruction, lowest_coordinate_matching};
3use crate::reduce::components::{
4    instructions::{Accumulator, AccumulatorFormat, Item, ReduceRequirements, ReduceStep, Value},
5    precision::ReducePrecision,
6};
7use ruda_kernel::dsl::prelude::*;
8
9/// Compute the coordinate of the maximum item returning the smallest coordinate in case of equality.
10#[derive(Debug, RudaType, Clone)]
11pub struct ArgMax {}
12
13#[ruda]
14impl ArgMax {
15    /// Compare two pairs of items and coordinates and return a new pair
16    /// where each element in the vectors is the maximal item with its coordinate.
17    /// In case of equality, the lowest coordinate is selected.
18    pub fn choose_argmax<T: Numeric, N: Size>(
19        items0: Vector<T, N>,
20        coordinates0: Vector<u32, N>,
21        items1: Vector<T, N>,
22        coordinates1: Vector<u32, N>,
23    ) -> (Vector<T, N>, Vector<u32, N>) {
24        let to_keep = select_many(
25            items0.equal(items1),
26            coordinates0.less_than(coordinates1),
27            items0.greater_than(items1),
28        );
29        let items = select_many(to_keep, items0, items1);
30        let coordinates = select_many(to_keep, coordinates0, coordinates1);
31        (items, coordinates)
32    }
33}
34
35impl ReduceFamily for ArgMax {
36    type Instruction<P: ReducePrecision> = Self;
37    type Config = ();
38}
39
40#[ruda]
41impl<P: ReducePrecision> ReduceInstruction<P> for ArgMax {
42    type SharedAccumulator = ArgAccumulator<P>;
43    type Config = ();
44
45    fn requirements(_this: &Self) -> ReduceRequirements {
46        ReduceRequirements { coordinates: true }
47    }
48
49    fn accumulator_format(_this: &Self) -> comptime_type!(AccumulatorFormat) {
50        AccumulatorFormat::Single
51    }
52
53    fn from_config(_config: Self::Config) -> Self {
54        ArgMax {}
55    }
56
57    fn null_input(_this: &Self) -> Vector<P::EI, P::SI> {
58        Vector::new(P::EI::min_value())
59    }
60
61    fn null_accumulator(_this: &Self) -> Accumulator<P> {
62        Accumulator::<P> {
63            elements: Value::new_single(Vector::new(P::EA::min_value())),
64            args: Value::new_single(Vector::new(u32::MAX)),
65        }
66    }
67
68    fn reduce(
69        _this: &Self,
70        accumulator: &mut Accumulator<P>,
71        item: Item<P>,
72        #[comptime] reduce_step: ReduceStep,
73    ) {
74        let coordinate = item.args.item();
75        let item = item.elements;
76
77        let (candidate_item, candidate_coordinate) = match reduce_step {
78            ReduceStep::Plane => {
79                let candidate_item = plane_max(item);
80                let candidate_coordinate =
81                    lowest_coordinate_matching(candidate_item, item, coordinate);
82                (candidate_item, candidate_coordinate)
83            }
84            ReduceStep::Identity => (item, coordinate),
85        };
86
87        let (elements, args) = Self::choose_argmax(
88            Vector::cast_from(candidate_item),
89            candidate_coordinate,
90            accumulator.elements.item(),
91            accumulator.args.item(),
92        );
93
94        accumulator.elements.assign(&Value::new_single(elements));
95        accumulator.args.assign(&Value::new_single(args));
96    }
97
98    fn plane_reduce_inplace(_this: &Self, accumulator: &mut Accumulator<P>) {
99        let acc_item = accumulator.elements.item();
100        let coordinate = accumulator.args.item();
101
102        let candidate_item = plane_max(acc_item);
103        let candidate_coordinate = lowest_coordinate_matching(candidate_item, acc_item, coordinate);
104
105        let (elements, args) = Self::choose_argmax(
106            accumulator.elements.item(),
107            accumulator.args.item(),
108            Vector::cast_from(candidate_item),
109            candidate_coordinate,
110        );
111
112        accumulator.elements.assign(&Value::new_single(elements));
113        accumulator.args.assign(&Value::new_single(args));
114    }
115
116    fn fuse_accumulators(_this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
117        let (elements, args) = Self::choose_argmax(
118            accumulator.elements.item(),
119            accumulator.args.item(),
120            other.elements.item(),
121            other.args.item(),
122        );
123
124        accumulator.elements.assign(&Value::new_single(elements));
125        accumulator.args.assign(&Value::new_single(args));
126    }
127
128    fn to_output_parallel<Out: Numeric>(
129        _this: &Self,
130        accumulator: Accumulator<P>,
131        _shape_axis_reduce: usize,
132    ) -> Value<Out> {
133        let vector_size = accumulator.elements.item().size().comptime();
134        let value = if vector_size > 1 {
135            let mut max = P::EA::min_value();
136            let mut coordinate = u32::MAX.runtime();
137            #[unroll]
138            for k in 0..vector_size {
139                let acc_element = accumulator.elements.item()[k];
140                let acc_coordinate = accumulator.args.item()[k];
141                if acc_element == max && acc_coordinate < coordinate {
142                    coordinate = acc_coordinate;
143                } else if acc_element > max {
144                    max = acc_element;
145                    coordinate = acc_coordinate;
146                }
147            }
148            Out::cast_from(coordinate)
149        } else {
150            Out::cast_from(accumulator.args.item())
151        };
152        Value::new_single(value)
153    }
154
155    fn to_output_perpendicular<Out: Numeric>(
156        _this: &Self,
157        accumulator: Accumulator<P>,
158        _shape_axis_reduce: usize,
159    ) -> Value<Vector<Out, P::SI>> {
160        Value::new_single(Vector::cast_from(accumulator.args.item()))
161    }
162}
ruprim/reduce/components/instructions/argmax.rs

ruprim/reduce/components/instructions/
argmax.rs
