ruprim/reduce/components/instructions/
argtopk.rs1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::comptime;
3use ruda_kernel::dsl::ruda;
4use ruda_kernel::dsl::frontend::RudaIndexMutExpand;
5use ruda_kernel::dsl::prelude::*;
6
7use crate::reduce::components::instructions::AccumulatorFormat;
8use crate::reduce::components::instructions::plane_topk_insert;
9use crate::reduce::components::instructions::plane_topk_merge;
10use crate::reduce::components::instructions::{Accumulator, Item, Value};
11use crate::reduce::{
12 ReduceFamily, ReduceInstruction, ReducePrecision,
13 components::instructions::{ReduceRequirements, ReduceStep, SharedAccumulator},
14};
15use ruda_kernel::dsl::frontend::Numeric;
16
17#[derive(Debug, RudaType, Clone)]
18pub struct ArgTopK {
19 #[ruda(comptime)]
20 pub k: usize,
21}
22
23impl ReduceFamily for ArgTopK {
24 type Instruction<P: ReducePrecision> = Self;
25 type Config = usize;
26}
27
28#[derive(RudaType)]
29pub struct ArgTopkAccumulator<E: Scalar, S: Size> {
30 pub elements: Array<Vector<E, S>>,
31 pub coordinates: Array<Vector<u32, S>>,
32}
33
34#[derive(RudaType)]
35pub struct ArgTopKSharedAccumulator<P: ReducePrecision> {
36 elements: Sequence<SharedMemory<Vector<P::EA, P::SI>>>,
37 args: Sequence<SharedMemory<Vector<u32, P::SI>>>,
38 #[ruda(comptime)]
39 k: usize,
40}
41
42#[ruda]
43impl<P: ReducePrecision> SharedAccumulator<P, ArgTopK> for ArgTopKSharedAccumulator<P> {
44 fn allocate(#[comptime] length: usize, #[comptime] _coordinate: bool, inst: &ArgTopK) -> Self {
45 let mut elements = Sequence::new();
46 let mut args = Sequence::new();
47 for _ in 0..inst.k {
48 elements.push(SharedMemory::new(length));
49 args.push(SharedMemory::new(length));
50 }
51 ArgTopKSharedAccumulator::<P> {
52 elements,
53 args,
54 k: inst.k,
55 }
56 }
57
58 fn read(accumulator: &Self, index: usize) -> Accumulator<P> {
59 let mut values = Array::new(accumulator.k);
60 let mut args = Array::new(accumulator.k);
61 #[unroll]
62 for i in 0..accumulator.k {
63 values[i] = accumulator.elements[i][index];
64 args[i] = accumulator.args[i][index];
65 }
66 Accumulator::<P> {
67 elements: Value::new_Multiple(values),
68 args: Value::new_Multiple(args),
69 }
70 }
71
72 fn write(accumulator: &mut Self, index: usize, item: Accumulator<P>) {
73 let values = item.elements.multiple();
74 let args = item.args.multiple();
75 #[unroll]
76 for i in 0..accumulator.k {
77 let values_acc = values[i];
78 let args_acc = args[i];
79
80 let mut shared_acc = accumulator.elements[i];
81 shared_acc[index] = values_acc;
82
83 let mut shared_arg_acc = accumulator.args[i];
84 shared_arg_acc[index] = args_acc;
85 }
86 }
87}
88
89#[ruda]
90impl<P: ReducePrecision> ReduceInstruction<P> for ArgTopK {
91 type SharedAccumulator = ArgTopKSharedAccumulator<P>;
92 type Config = usize;
93
94 fn requirements(_this: &Self) -> super::ReduceRequirements {
95 ReduceRequirements { coordinates: true }
96 }
97
98 fn accumulator_format(this: &Self) -> comptime_type!(AccumulatorFormat) {
99 comptime!(AccumulatorFormat::Multiple(this.k))
100 }
101
102 fn from_config(#[comptime] config: Self::Config) -> Self {
103 ArgTopK { k: config }
104 }
105
106 fn null_input(_this: &Self) -> Vector<P::EI, P::SI> {
107 Vector::empty().fill(P::EI::min_value())
108 }
109
110 fn null_accumulator(this: &Self) -> Accumulator<P> {
111 let mut elements = Array::new(comptime!(this.k));
112 let mut args = Array::new(comptime!(this.k));
113 #[unroll]
114 for i in 0..this.k {
115 elements[i] = Vector::new(P::EA::min_value());
116 args[i] = Vector::new(u32::MAX);
117 }
118
119 Accumulator::<P> {
120 elements: Value::new_Multiple(elements),
121 args: Value::new_Multiple(args),
122 }
123 }
124
125 fn reduce(
126 this: &Self,
127 accumulator: &mut Accumulator<P>,
128 item: Item<P>,
129 #[comptime] reduce_step: ReduceStep,
130 ) {
131 let elements = accumulator.elements.multiple_mut();
132
133 match reduce_step {
134 ReduceStep::Plane => {
135 plane_topk_insert::<P::EA, P::SI>(
136 elements,
137 &mut accumulator.args,
138 Vector::cast_from(item.elements),
139 &item.args,
140 this.k,
141 true,
142 );
143 }
144 ReduceStep::Identity => {
145 let coordinates = accumulator.args.multiple_mut();
146 let mut insert_val = Vector::cast_from(item.elements);
147 let mut insert_coord = item.args.item();
148
149 for j in 0..this.k {
150 let to_keep = select_many(
151 elements[j].equal(insert_val),
152 coordinates[j].less_than(insert_coord),
153 elements[j].greater_than(insert_val),
154 );
155 let best_value = select_many(to_keep, elements[j], insert_val);
156 let loser_value = select_many(to_keep, insert_val, elements[j]);
157 let best_coordinate = select_many(to_keep, coordinates[j], insert_coord);
158 let loser_coordinate = select_many(to_keep, insert_coord, coordinates[j]);
159
160 elements[j] = best_value;
161 coordinates[j] = best_coordinate;
162 insert_val = loser_value;
163 insert_coord = loser_coordinate;
164 }
165 }
166 };
167 }
168
169 fn plane_reduce_inplace(this: &Self, accumulator: &mut Accumulator<P>) {
170 plane_topk_merge::<P::EA, P::SI>(
171 accumulator.elements.multiple_mut(),
172 &mut accumulator.args,
173 this.k,
174 true,
175 );
176 }
177
178 fn fuse_accumulators(this: &Self, accumulator: &mut Accumulator<P>, other: &Accumulator<P>) {
179 let elements = accumulator.elements.multiple_mut();
180 let coordinates = accumulator.args.multiple_mut();
181 let other_elements = other.elements.multiple();
182 let other_coords = other.args.multiple();
183
184 for i in 0..this.k {
185 let mut insert_val = other_elements[i];
186 let mut insert_coord = other_coords[i];
187 for j in 0..this.k {
188 let to_keep = select_many(
189 elements[j].equal(insert_val),
190 coordinates[j].less_than(insert_coord),
191 elements[j].greater_than(insert_val),
192 );
193 let best_value = select_many(to_keep, elements[j], insert_val);
194 let best_coordinate = select_many(to_keep, coordinates[j], insert_coord);
195 let loser_value = select_many(to_keep, insert_val, elements[j]);
196 let loser_coordinate = select_many(to_keep, insert_coord, coordinates[j]);
197
198 elements[j] = best_value;
199 coordinates[j] = best_coordinate;
200 insert_val = loser_value;
201 insert_coord = loser_coordinate;
202 }
203 }
204 }
205
206 fn to_output_parallel<Out: Numeric>(
207 this: &Self,
208 accumulator: Accumulator<P>,
209 _shape_axis_reduce: usize,
210 ) -> Value<Out> {
211 let coords = accumulator.args.multiple();
212 let vals = accumulator.elements.multiple();
213 let vector_size = coords[0].size().comptime();
214
215 let mut topk_vals = Array::new(this.k);
216 let mut topk_coords = Array::new(this.k);
217
218 #[unroll]
219 for slot in 0..this.k {
220 topk_vals[slot] = P::EA::min_value();
221 topk_coords[slot] = u32::MAX;
222 }
223
224 #[unroll]
225 for i in 0..this.k {
226 #[unroll]
227 for j in 0..vector_size {
228 let mut value = vals[i][j];
229 let mut coordinate = coords[i][j];
230
231 #[unroll]
232 for slot in 0..this.k {
233 let current_value = topk_vals[slot];
234 let current_coordinate = topk_coords[slot];
235
236 let to_keep = select(
237 current_value == value,
238 current_coordinate < coordinate,
239 current_value > value,
240 );
241
242 topk_vals[slot] = select(to_keep, current_value, value);
243 topk_coords[slot] = select(to_keep, current_coordinate, coordinate);
244
245 value = select(to_keep, value, current_value);
246 coordinate = select(to_keep, coordinate, current_coordinate);
247 }
248 }
249 }
250
251 let mut out = Array::new(this.k);
252 #[unroll]
253 for i in 0..this.k {
254 out[i] = Out::cast_from(topk_coords[i]);
255 }
256 Value::new_Multiple(out)
257 }
258
259 fn to_output_perpendicular<Out: Numeric>(
260 this: &Self,
261 accumulator: Accumulator<P>,
262 _shape_axis_reduce: usize,
263 ) -> Value<Vector<Out, P::SI>> {
264 let acc_args = accumulator.args.multiple();
265 let mut output = Array::new(this.k);
266
267 #[unroll]
268 for i in 0..this.k {
269 output[i] = Vector::cast_from(acc_args[i]);
270 }
271
272 Value::new_Multiple(output)
273 }
274}