ruprim/reduce/components/instructions/
argmin.rs

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

ruprim/reduce/components/instructions/
argmin.rs
