ruprim/reduce/components/instructions/
base.rs

1use ruda_kernel::dsl as kernel_dsl;
2use crate::reduce::components::{instructions::lowest_coordinate_matching, precision::ReducePrecision};
3use ruda_kernel::dsl::prelude::*;
4
5pub trait ReduceFamily: Send + Sync + 'static + std::fmt::Debug {
6    type Instruction<P: ReducePrecision>: ReduceInstruction<P, Config = Self::Config>;
7    type Config: RudaComptime + Send + Sync;
8}
9
10#[derive(RudaType, Clone, Copy)]
11/// Whether we keep track of coordinates of items
12pub struct ReduceRequirements {
13    #[ruda(comptime)]
14    pub coordinates: bool,
15}
16
17#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, RudaType)]
18pub enum AccumulatorFormat {
19    Multiple(usize),
20    Single,
21}
22
23impl AccumulatorFormat {
24    pub fn len(&self) -> usize {
25        match self {
26            AccumulatorFormat::Multiple(k) => *k,
27            AccumulatorFormat::Single => 1,
28        }
29    }
30
31    pub fn is_empty(&self) -> bool {
32        self.len() == 0
33    }
34}
35
36#[derive(RudaType)]
37/// Whether the accumulator has zero, one or more vectors
38pub enum Value<X: RudaPrimitive> {
39    Multiple(Array<X>),
40    /// Wrap the item to be able to modify it as a field
41    Single(ValueWrapper<X>),
42    None,
43}
44
45#[derive(RudaType)]
46/// Wrap the item to be able to modify it as a field
47pub struct ValueWrapper<X: RudaPrimitive> {
48    val: X,
49}
50
51#[ruda]
52impl<X: RudaPrimitive> ValueWrapper<X> {
53    pub fn unwrap(&self) -> X {
54        self.val
55    }
56}
57
58#[ruda]
59impl<X: RudaPrimitive> Value<X> {
60    pub fn new_single(val: X) -> Value<X> {
61        Value::new_Single(ValueWrapper::<X> { val })
62    }
63
64    pub fn item(&self) -> X {
65        match self {
66            Value::Multiple(_) => panic!("Tried item on Multiple"),
67            Value::Single(item) => item.val,
68            Value::None => panic!("Tried item on None"),
69        }
70    }
71
72    pub fn multiple(&self) -> &Array<X> {
73        match self {
74            Value::Multiple(array) => array,
75            Value::Single(_) => panic!("Tried multiple on Single"),
76            Value::None => panic!("Tried multiple on None"),
77        }
78    }
79
80    pub fn multiple_mut(&mut self) -> &mut Array<X> {
81        match self {
82            Value::Multiple(array) => array,
83            Value::Single(_) => panic!("Tried multiple on Single"),
84            Value::None => panic!("Tried multiple on None"),
85        }
86    }
87
88    pub fn assign(&mut self, other: &Value<X>) {
89        match (self, other) {
90            (Value::Multiple(this), Value::Multiple(other)) => {
91                for i in 0..this.len() {
92                    this[i] = other[i];
93                }
94            }
95            (Value::Single(this), Value::Single(other)) => {
96                this.val = other.val;
97            }
98            (Value::None, Value::None) => {}
99            _ => panic!("Tried assigning different accumulator kinds"),
100        }
101    }
102}
103
104#[ruda]
105pub fn plane_topk_insert<N: Numeric, S: Size>(
106    elements: &mut Array<Vector<N, S>>,
107    coordinates: &mut Value<Vector<u32, S>>,
108    item: Vector<N, S>,
109    coord: &Value<Vector<u32, S>>,
110    #[comptime] k: usize,
111    #[comptime] has_coords: bool,
112) {
113    let mut local_best_val = item;
114    let lane_id = Vector::new(UNIT_POS_X);
115
116    let mut local_best_coord = if has_coords {
117        coord.item()
118    } else {
119        Vector::new(u32::MAX)
120    };
121
122    #[unroll]
123    for _i in 0..k {
124        let winning_val = plane_max(local_best_val);
125
126        let winning_coord = if has_coords {
127            lowest_coordinate_matching(winning_val, local_best_val, local_best_coord)
128        } else {
129            let is_match = local_best_val.equal(winning_val);
130            let claim = select_many(is_match, lane_id, Vector::new(u32::MAX));
131            plane_min(claim)
132        };
133
134        let mut insert_val = winning_val;
135        let mut insert_coord = winning_coord;
136
137        if has_coords {
138            let coordinates = coordinates.multiple_mut();
139            #[unroll]
140            for j in 0..k {
141                let to_keep = select_many(
142                    elements[j].equal(insert_val),
143                    coordinates[j].less_than(insert_coord),
144                    elements[j].greater_than(insert_val),
145                );
146
147                let next_val = select_many(to_keep, insert_val, elements[j]);
148                elements[j] = select_many(to_keep, elements[j], insert_val);
149                insert_val = next_val;
150
151                let next_coord = select_many(to_keep, insert_coord, coordinates[j]);
152                coordinates[j] = select_many(to_keep, coordinates[j], insert_coord);
153                insert_coord = next_coord;
154            }
155        } else {
156            #[unroll]
157            for j in 0..k {
158                let to_keep = elements[j].greater_than(insert_val);
159                let next_val = select_many(to_keep, insert_val, elements[j]);
160                elements[j] = select_many(to_keep, elements[j], insert_val);
161                insert_val = next_val;
162            }
163        }
164
165        // Winner masking logic
166        let is_winner = if has_coords {
167            local_best_val
168                .equal(winning_val)
169                .and(local_best_coord.equal(winning_coord))
170        } else {
171            lane_id.equal(winning_coord)
172        };
173
174        local_best_val = select_many(is_winner, Vector::new(N::min_value()), local_best_val);
175        if has_coords {
176            local_best_coord = select_many(is_winner, Vector::new(u32::MAX), local_best_coord);
177        }
178    }
179}
180
181#[ruda]
182pub fn plane_topk_merge<N: Numeric, S: Size>(
183    elements: &mut Array<Vector<N, S>>,
184    coordinates: &mut Value<Vector<u32, S>>,
185    #[comptime] k: usize,
186    #[comptime] has_coords: bool,
187) {
188    let mut final_elements = Array::new(k);
189    let mut final_coords = Array::new(k);
190    let mut cursor = Vector::new(0u32);
191    let lane_id = Vector::new(UNIT_POS_X);
192
193    #[unroll]
194    for i in 0..k {
195        let mut local_val = Vector::new(N::min_value());
196        let mut local_coord = Vector::new(u32::MAX);
197
198        #[unroll]
199        for j in 0..k {
200            let is_pointed = cursor.equal(Vector::new(j as u32));
201            local_val = select_many(is_pointed, elements[j], local_val);
202            if has_coords {
203                let coords = coordinates.multiple_mut();
204                local_coord = select_many(is_pointed, coords[j], local_coord);
205            }
206        }
207
208        let winning_val = plane_max(local_val);
209        let winning_lane = if has_coords {
210            let best_c = lowest_coordinate_matching(winning_val, local_val, local_coord);
211            final_coords[i] = best_c;
212            let is_cand = local_val.equal(winning_val).and(local_coord.equal(best_c));
213            plane_min(select_many(is_cand, lane_id, Vector::new(u32::MAX)))
214        } else {
215            let is_cand = local_val.equal(winning_val);
216            plane_min(select_many(is_cand, lane_id, Vector::new(u32::MAX)))
217        };
218
219        final_elements[i] = winning_val;
220        let is_winner_thread = lane_id.equal(winning_lane);
221        cursor = select_many(is_winner_thread, cursor + Vector::new(1u32), cursor);
222    }
223
224    #[unroll]
225    for i in 0..k {
226        elements[i] = final_elements[i];
227        if has_coords {
228            let coords = coordinates.multiple_mut();
229            coords[i] = final_coords[i];
230        }
231    }
232}
233
234#[derive(RudaType)]
235/// Whether the accumulator has zero, one or more vectors
236/// This should be the same variant as AccumulatorKind for an instruction
237pub enum SharedAccumulatorKind<X: RudaPrimitive> {
238    Multiple(Sequence<SharedMemory<X>>),
239    Single(SharedMemory<X>),
240    None,
241}
242
243#[ruda]
244impl<X: RudaPrimitive> SharedAccumulatorKind<X> {
245    pub fn get(&self, i: usize) -> Value<X> {
246        match self {
247            SharedAccumulatorKind::Multiple(sequence) => {
248                let mut array = Array::new(sequence.len());
249                #[unroll]
250                for k_iter in 0..sequence.len() {
251                    array[k_iter] = sequence[k_iter][i];
252                }
253                Value::new_Multiple(array)
254            }
255            SharedAccumulatorKind::Single(shared_memory) => Value::new_single(shared_memory[i]),
256            SharedAccumulatorKind::None => Value::new_None(),
257        }
258    }
259
260    pub fn set(&mut self, i: usize, value: Value<X>) {
261        match self {
262            SharedAccumulatorKind::Multiple(sequence) =>
263            {
264                #[unroll]
265                for k_iter in 0..sequence.len() {
266                    let mut shared_acc = sequence[k_iter];
267                    shared_acc[i] = value.multiple()[k_iter];
268                }
269            }
270            SharedAccumulatorKind::Single(shared_memory) => shared_memory[i] = value.item(),
271            SharedAccumulatorKind::None => {}
272        }
273    }
274}
275
276/// An instruction for a reduce algorithm that works with [`Vector`].
277///
278/// See a provided implementation, such as [`Sum`](super::Sum) or [`ArgMax`](super::ArgMax) for an example how to implement
279/// this trait for a custom instruction.
280///
281/// A reduction works at three levels. First, it takes input data of type `In` and reduce them
282/// with their coordinate into an `AccumulatorItem`. Then, multiple `AccumulatorItem` are possibly fused
283/// together into a single accumulator that is converted to the expected output type.
284#[ruda]
285pub trait ReduceInstruction<P: ReducePrecision>:
286    Send + Sync + 'static + std::fmt::Debug + RudaType + Sized
287{
288    type Config: RudaComptime + Send + Sync;
289
290    /// When multiple agents are collaborating to reduce a single slice,
291    /// we need a share accumulator to store multiple `AccumulatorItem`.
292    /// This is most likely a `SharedMemory<Vector<T>>` or a struct or tuple of vectorized shared memories.
293    type SharedAccumulator: SharedAccumulator<P, Self>;
294
295    /// Requirements of the reduce.
296    fn requirements(this: &Self) -> ReduceRequirements;
297    fn accumulator_format(this: &Self) -> comptime_type!(AccumulatorFormat);
298
299    fn from_config(#[comptime] config: Self::Config) -> Self;
300    /// A input such that `Self::reduce(accumulator, Self::null_input(), coordinate, use_planes)`
301    /// is guaranteed to return `accumulator` unchanged for any choice of `coordinate`.
302    fn null_input(this: &Self) -> Vector<P::EI, P::SI>;
303
304    /// A accumulator such that `Self::fuse_accumulators(accumulator, Self::null_accumulator()` always returns
305    /// is guaranteed to return `accumulator` unchanged.
306    fn null_accumulator(this: &Self) -> Accumulator<P>;
307
308    /// If `ReduceStep` is `Plane`, reduce all the `item` and `coordinate` within the `accumulator`.
309    /// if `ReduceStep` is `Identity`, reduce the given `item` and `coordinate` into the accumulator.
310    fn reduce(
311        this: &Self,
312        accumulator: &mut Accumulator<P>,
313        item: Item<P>,
314        #[comptime] reduce_step: ReduceStep,
315    );
316
317    fn plane_reduce_inplace(this: &Self, accumulator: &mut Accumulator<P>);
318
319    /// Reduce a whole accumulator (other) in accumulator.
320    fn fuse_accumulators(this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>);
321
322    /// Reduce all elements of the accumulator into a single output element of type `Out`.
323    fn to_output_parallel<Out: Numeric>(
324        this: &Self,
325        accumulator: Accumulator<P>,
326        shape_axis_reduce: usize,
327    ) -> Value<Out>;
328
329    /// Convert each element of the accumulator into the expected output element of type `Out`.
330    fn to_output_perpendicular<Out: Numeric>(
331        this: &Self,
332        accumulator: Accumulator<P>,
333        shape_axis_reduce: usize,
334    ) -> Value<Vector<Out, P::SI>>;
335}
336
337#[derive(RudaType)]
338pub struct Item<P: ReducePrecision> {
339    pub elements: Vector<P::EI, P::SI>,
340    // Warning: should not be Multiple
341    pub args: Value<Vector<u32, P::SI>>,
342}
343
344#[derive(RudaType)]
345pub struct Accumulator<P: ReducePrecision> {
346    pub elements: Value<Vector<P::EA, P::SI>>,
347    pub args: Value<Vector<u32, P::SI>>,
348}
349
350/// A simple trait that abstract over a single or multiple shared memory.
351#[ruda]
352pub trait SharedAccumulator<P: ReducePrecision, I: ReduceInstruction<P>>:
353    RudaType + Send + Sync + 'static
354{
355    fn allocate(#[comptime] length: usize, #[comptime] _coordinate: bool, inst: &I) -> Self;
356
357    fn read(accumulator: &Self, index: usize) -> Accumulator<P>;
358
359    fn write(accumulator: &mut Self, index: usize, item: Accumulator<P>);
360}
361
362#[ruda]
363impl<P: ReducePrecision, I: ReduceInstruction<P>> SharedAccumulator<P, I>
364    for SharedMemory<Vector<P::EA, P::SI>>
365{
366    fn allocate(#[comptime] length: usize, #[comptime] _coordinate: bool, _inst: &I) -> Self {
367        SharedMemory::new(length)
368    }
369
370    fn read(accumulator: &Self, index: usize) -> Accumulator<P> {
371        Accumulator::<P> {
372            elements: Value::new_single(accumulator[index]),
373            args: Value::new_None(),
374        }
375    }
376
377    fn write(accumulator: &mut Self, index: usize, item: Accumulator<P>) {
378        accumulator[index] = item.elements.item();
379    }
380}
381
382/// A pair of shared memory used for [`ArgMax`](super::ArgMax) and [`ArgMin`](super::ArgMin).
383#[derive(RudaType)]
384pub struct ArgAccumulator<P: ReducePrecision> {
385    pub elements: SharedMemory<Vector<P::EA, P::SI>>,
386    pub args: SharedMemory<Vector<u32, P::SI>>,
387}
388
389/// For a single reduce step whether we need to do plane reduction
390#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
391pub enum ReduceStep {
392    /// Just keep the current value
393    Identity,
394    /// reduce across the plane
395    Plane,
396}
397
398#[ruda]
399impl<P: ReducePrecision, I: ReduceInstruction<P>> SharedAccumulator<P, I> for ArgAccumulator<P> {
400    fn allocate(#[comptime] length: usize, #[comptime] _coordinate: bool, _inst: &I) -> Self {
401        ArgAccumulator::<P> {
402            elements: SharedMemory::new(length),
403            args: SharedMemory::new(length),
404        }
405    }
406
407    fn read(accumulator: &Self, index: usize) -> Accumulator<P> {
408        Accumulator::<P> {
409            elements: Value::new_single(accumulator.elements[index]),
410            args: Value::new_single(accumulator.args[index]),
411        }
412    }
413
414    fn write(accumulator: &mut Self, index: usize, item: Accumulator<P>) {
415        accumulator.elements[index] = item.elements.item();
416        accumulator.args[index] = item.args.item();
417    }
418}
419
420#[ruda]
421pub fn reduce_inplace<P: ReducePrecision, R: ReduceInstruction<P>>(
422    inst: &R,
423    accumulator: &mut Accumulator<P>,
424    item: Item<P>,
425    #[comptime] reduce_step: ReduceStep,
426) {
427    R::reduce(inst, accumulator, item, reduce_step)
428}
429
430#[ruda]
431pub fn reduce_shared_inplace<P: ReducePrecision, R: ReduceInstruction<P>>(
432    inst: &R,
433    accumulator: &mut R::SharedAccumulator,
434    index: usize,
435    item: Item<P>,
436    #[comptime] reduce_step: ReduceStep,
437) {
438    let mut acc_item = R::SharedAccumulator::read(accumulator, index);
439    R::reduce(inst, &mut acc_item, item, reduce_step);
440    R::SharedAccumulator::write(accumulator, index, acc_item);
441}
442
443#[ruda]
444pub fn fuse_accumulator_inplace<P: ReducePrecision, R: ReduceInstruction<P>>(
445    inst: &R,
446    accumulator: &mut R::SharedAccumulator,
447    destination: usize,
448    origin: usize,
449) {
450    let mut acc = R::SharedAccumulator::read(accumulator, destination);
451    R::fuse_accumulators(
452        inst,
453        &mut acc,
454        &R::SharedAccumulator::read(accumulator, origin),
455    );
456    R::SharedAccumulator::write(accumulator, destination, acc);
457}
ruprim/reduce/components/instructions/base.rs

ruprim/reduce/components/instructions/
base.rs
