1use ruda_kernel::dsl as kernel_dsl;
2use ruda_kernel::dsl::prelude::*;
3use ruda_kernel::library::tensor::r#virtual::VirtualTensor;
4use ruda_kernel::library::tensor::r#virtual::VirtualTensorOperations;
5use ruda_kernel::library::tensor::r#virtual::VirtualTensorOperationsExpand;
6use ruda_kernel::dsl::unexpanded;
7use std::marker::PhantomData;
8
9pub trait ReduceDType {
10 type In: Numeric;
11 type SizeIn: Size;
12 type Out: Numeric;
13 type SizeOut: Size;
14}
15
16impl<In: Numeric, SizeIn: Size, Out: Numeric, SizeOut: Size> ReduceDType
17 for ((In, SizeIn), (Out, SizeOut))
18{
19 type In = In;
20 type SizeIn = SizeIn;
21 type Out = Out;
22 type SizeOut = SizeOut;
23}
24
25pub trait NumericVector {
26 type T: Numeric;
27 type N: Size;
28}
29
30impl<T: Numeric, N: Size> NumericVector for (T, N) {
31 type T = T;
32 type N = N;
33}
34
35#[ruda]
36#[allow(dead_code)]
37pub trait ReduceArgs: Send + Sync + 'static + Clone {
38 type Input<E: Numeric, S: Size>: LaunchArg + RudaType;
39
40 type Output<E: Numeric, S: Size>: LaunchArg + RudaType;
41
42 type State<P: ReduceDType>: RudaType;
43
44 fn init_state<P: ReduceDType>(
45 input: &Self::Input<P::In, P::SizeIn>,
46 output: &mut Self::Output<P::Out, P::SizeOut>,
47 ) -> Self::State<P>;
48
49 fn read_input<P: ReduceDType>(state: &Self::State<P>, index: usize)
50 -> Vector<P::In, P::SizeIn>;
51 fn read_output<P: ReduceDType>(
52 state: &Self::State<P>,
53 index: usize,
54 ) -> Vector<P::Out, P::SizeOut>;
55
56 fn write_output<P: ReduceDType>(
57 state: &mut Self::State<P>,
58 index: usize,
59 value: Vector<P::Out, P::SizeOut>,
60 );
61
62 fn len_input<P: ReduceDType>(state: &Self::State<P>) -> usize;
63 fn len_output<P: ReduceDType>(state: &Self::State<P>) -> usize;
64
65 fn buffer_len_input<P: ReduceDType>(state: &Self::State<P>) -> usize;
66 fn buffer_len_output<P: ReduceDType>(state: &Self::State<P>) -> usize;
67
68 fn rank_input<P: ReduceDType>(state: &Self::State<P>) -> usize;
69 fn rank_output<P: ReduceDType>(state: &Self::State<P>) -> usize;
70
71 fn shape_input<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize;
72 fn shape_output<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize;
73
74 fn stride_input<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize;
75 fn stride_output<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize;
76
77 fn vector_size_input<P: ReduceDType>(state: &Self::State<P>) -> comptime_type!(VectorSize);
78 fn vector_size_output<P: ReduceDType>(state: &Self::State<P>) -> comptime_type!(VectorSize);
79}
80
81#[ruda]
82pub fn init_tensors<RA: ReduceArgs, In: Numeric, InSize: Size, Out: Numeric, OutSize: Size>(
83 input: &RA::Input<In, InSize>,
84 output: &mut RA::Output<Out, OutSize>,
85) -> (
86 VirtualTensor<In, InSize>,
87 VirtualTensor<Out, OutSize, ReadWrite>,
88) {
89 let mut state = RA::init_state::<((In, InSize), (Out, OutSize))>(input, output);
90
91 let input = TensorArg::new_input(&state);
92 let mut output = TensorArg::new_output(&mut state);
93
94 let input = VirtualTensor::<In, InSize>::new::<
95 TensorArg<((In, InSize), (Out, OutSize)), RA, Input>,
96 >(&input);
97 let output = VirtualTensor::<Out, OutSize, ReadWrite>::new::<
98 TensorArg<((In, InSize), (Out, OutSize)), RA, Output>,
99 >(&mut output);
100
101 (input, output)
102}
103
104#[derive(Clone)]
105pub struct TensorArgs;
106
107#[ruda]
108impl ReduceArgs for TensorArgs {
109 type Input<EG: Numeric, N: Size> = Tensor<Vector<EG, N>>;
110 type Output<EG: Numeric, N: Size> = Tensor<Vector<EG, N>>;
111 type State<P: ReduceDType> = (
112 *const Tensor<Vector<P::In, P::SizeIn>>,
113 *mut Tensor<Vector<P::Out, P::SizeOut>>,
114 );
115
116 fn init_state<P: ReduceDType>(
117 input: &Self::Input<P::In, P::SizeIn>,
118 output: &mut Self::Output<P::Out, P::SizeOut>,
119 ) -> Self::State<P> {
120 (input, output)
121 }
122
123 fn read_input<P: ReduceDType>(
124 state: &Self::State<P>,
125 index: usize,
126 ) -> Vector<P::In, P::SizeIn> {
127 unsafe { (*state.0)[index] }
128 }
129
130 fn read_output<P: ReduceDType>(
131 state: &Self::State<P>,
132 index: usize,
133 ) -> Vector<P::Out, P::SizeOut> {
134 unsafe { (*state.1)[index] }
135 }
136
137 fn write_output<P: ReduceDType>(
138 state: &mut Self::State<P>,
139 index: usize,
140 value: Vector<P::Out, P::SizeOut>,
141 ) {
142 unsafe { (*state.1)[index] = value }
143 }
144
145 fn buffer_len_input<P: ReduceDType>(state: &Self::State<P>) -> usize {
146 unsafe { (*state.0).buffer_len() }
147 }
148
149 fn buffer_len_output<P: ReduceDType>(state: &Self::State<P>) -> usize {
150 unsafe { (*state.1).buffer_len() }
151 }
152
153 fn len_input<P: ReduceDType>(state: &Self::State<P>) -> usize {
154 unsafe { (*state.0).len() }
155 }
156
157 fn len_output<P: ReduceDType>(state: &Self::State<P>) -> usize {
158 unsafe { (*state.1).len() }
159 }
160 fn rank_input<P: ReduceDType>(state: &Self::State<P>) -> usize {
161 unsafe { (*state.0).rank() }
162 }
163
164 fn rank_output<P: ReduceDType>(state: &Self::State<P>) -> usize {
165 unsafe { (*state.1).rank() }
166 }
167
168 fn shape_input<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize {
169 unsafe { (*state.0).shape(dim) }
170 }
171
172 fn shape_output<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize {
173 unsafe { (*state.1).shape(dim) }
174 }
175
176 fn stride_input<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize {
177 unsafe { (*state.0).stride(dim) }
178 }
179
180 fn stride_output<P: ReduceDType>(state: &Self::State<P>, dim: usize) -> usize {
181 unsafe { (*state.1).stride(dim) }
182 }
183
184 fn vector_size_input<P: ReduceDType>(state: &Self::State<P>) -> comptime_type!(VectorSize) {
185 unsafe { (*state.0).vector_size() }
186 }
187
188 fn vector_size_output<P: ReduceDType>(state: &Self::State<P>) -> comptime_type!(VectorSize) {
189 unsafe { (*state.1).vector_size() }
190 }
191}
192
193pub struct Input;
194pub struct Output;
195
196pub struct TensorArg<P: ReduceDType, RA: ReduceArgs, Tag> {
197 _state: *mut RA::State<P>,
198 tag: PhantomData<Tag>,
199}
200
201pub struct TensorArgExpand<P: ReduceDType, RA: ReduceArgs, Tag> {
202 state: <RA::State<P> as RudaType>::ExpandType,
203 tag: PhantomData<Tag>,
204}
205
206impl<P: ReduceDType, RA: ReduceArgs> TensorArg<P, RA, Input> {
207 pub fn new_input(_state: &RA::State<P>) -> Self {
208 unexpanded!()
209 }
210 pub fn __expand_new_input(
211 _scope: &mut Scope,
212 state: <RA::State<P> as RudaType>::ExpandType,
213 ) -> TensorArgExpand<P, RA, Input> {
214 TensorArgExpand {
215 state,
216 tag: PhantomData,
217 }
218 }
219}
220
221impl<P: ReduceDType, RA: ReduceArgs> TensorArg<P, RA, Output> {
222 pub fn new_output(_state: &mut RA::State<P>) -> Self {
223 unexpanded!()
224 }
225 pub fn __expand_new_output(
226 _scope: &mut Scope,
227 state: <RA::State<P> as RudaType>::ExpandType,
228 ) -> TensorArgExpand<P, RA, Output> {
229 TensorArgExpand {
230 state,
231 tag: PhantomData,
232 }
233 }
234}
235
236impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperations<P::Out, P::SizeOut>
237 for TensorArg<P, RA, Output>
238{
239}
240impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperations<P::In, P::SizeIn>
241 for TensorArg<P, RA, Input>
242{
243}
244
245impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperationsExpand<P::In, P::SizeIn>
246 for TensorArgExpand<P, RA, Input>
247{
248 fn __expand_read_method(
249 &self,
250 scope: &mut Scope,
251 index: NativeExpand<usize>,
252 ) -> NativeExpand<Vector<P::In, P::SizeIn>> {
253 RA::__expand_read_input(scope, self.state.clone(), index)
254 }
255
256 fn __expand_write_method(
257 &self,
258 _scope: &mut Scope,
259 _index: NativeExpand<usize>,
260 _value: NativeExpand<Vector<P::In, P::SizeIn>>,
261 ) {
262 unreachable!("Can't write to input")
263 }
264
265 fn __expand_shape_method(
266 &self,
267 scope: &mut Scope,
268 axis: NativeExpand<usize>,
269 ) -> NativeExpand<usize> {
270 RA::__expand_shape_input(scope, self.state.clone(), axis)
271 }
272
273 fn __expand_stride_method(
274 &self,
275 scope: &mut Scope,
276 axis: NativeExpand<usize>,
277 ) -> NativeExpand<usize> {
278 RA::__expand_stride_input(scope, self.state.clone(), axis)
279 }
280
281 fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
282 RA::__expand_rank_input(scope, self.state.clone())
283 }
284 fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
285 RA::__expand_len_input(scope, self.state.clone())
286 }
287 fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
288 RA::__expand_buffer_len_input(scope, self.state.clone())
289 }
290
291 fn __expand_read_window_method(
292 &self,
293 _context: &mut Scope,
294 _start: NativeExpand<usize>,
295 _end: NativeExpand<usize>,
296 ) -> SliceExpand<Vector<P::In, P::SizeIn>, ReadOnly> {
297 panic!("Unsupported")
298 }
299
300 fn __expand_as_tensor_map_method(
301 &self,
302 scope: &mut Scope,
303 ) -> ComptimeOptionExpand<TensorMap<P::In, Tiled>> {
304 ComptimeOption::__expand_new_None(scope)
305 }
306}
307
308impl<P: ReduceDType, RA: ReduceArgs> Vectorized for TensorArg<P, RA, Input> {}
309impl<P: ReduceDType, RA: ReduceArgs> VectorizedExpand for TensorArgExpand<P, RA, Input> {
310 fn vector_size(&self) -> usize {
311 let mut scope = Scope::root(false);
312 RA::__expand_vector_size_input(&mut scope, self.state.clone())
313 }
314}
315
316impl<P: ReduceDType, RA: ReduceArgs> VirtualTensorOperationsExpand<P::Out, P::SizeOut>
317 for TensorArgExpand<P, RA, Output>
318{
319 fn __expand_read_method(
320 &self,
321 scope: &mut Scope,
322 index: NativeExpand<usize>,
323 ) -> NativeExpand<Vector<P::Out, P::SizeOut>> {
324 RA::__expand_read_output(scope, self.state.clone(), index)
325 }
326
327 fn __expand_write_method(
328 &self,
329 scope: &mut Scope,
330 index: NativeExpand<usize>,
331 value: NativeExpand<Vector<P::Out, P::SizeOut>>,
332 ) {
333 RA::__expand_write_output(scope, self.state.clone(), index, value)
334 }
335
336 fn __expand_shape_method(
337 &self,
338 scope: &mut Scope,
339 axis: NativeExpand<usize>,
340 ) -> NativeExpand<usize> {
341 RA::__expand_shape_output(scope, self.state.clone(), axis)
342 }
343
344 fn __expand_stride_method(
345 &self,
346 scope: &mut Scope,
347 axis: NativeExpand<usize>,
348 ) -> NativeExpand<usize> {
349 RA::__expand_stride_output(scope, self.state.clone(), axis)
350 }
351
352 fn __expand_rank_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
353 RA::__expand_rank_output(scope, self.state.clone())
354 }
355
356 fn __expand_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
357 RA::__expand_len_output(scope, self.state.clone())
358 }
359 fn __expand_buffer_len_method(&self, scope: &mut Scope) -> NativeExpand<usize> {
360 RA::__expand_buffer_len_output(scope, self.state.clone())
361 }
362
363 fn __expand_read_window_method(
364 &self,
365 _context: &mut Scope,
366 _start: NativeExpand<usize>,
367 _end: NativeExpand<usize>,
368 ) -> SliceExpand<Vector<P::Out, P::SizeOut>, ReadOnly> {
369 panic!("Unsupported")
370 }
371
372 fn __expand_as_tensor_map_method(
373 &self,
374 scope: &mut Scope,
375 ) -> ComptimeOptionExpand<TensorMap<P::Out, Tiled>> {
376 ComptimeOption::__expand_new_None(scope)
377 }
378}
379
380impl<P: ReduceDType, RA: ReduceArgs> Vectorized for TensorArg<P, RA, Output> {}
381impl<P: ReduceDType, RA: ReduceArgs> VectorizedExpand for TensorArgExpand<P, RA, Output> {
382 fn vector_size(&self) -> usize {
383 let mut scope = Scope::root(false);
384 RA::__expand_vector_size_output(&mut scope, self.state.clone())
385 }
386}
387
388mod __tensor_arg {
389 use super::*;
390
391 impl<P: ReduceDType, RA: ReduceArgs, Tag> RudaType for TensorArg<P, RA, Tag> {
392 type ExpandType = TensorArgExpand<P, RA, Tag>;
393 }
394
395 impl<P: ReduceDType, RA: ReduceArgs, Tag> IntoMut for TensorArgExpand<P, RA, Tag> {
396 fn into_mut(self, _scope: &mut Scope) -> Self {
397 self
398 }
399 }
400
401 impl<P: ReduceDType, RA: ReduceArgs, Tag> RudaDebug for TensorArgExpand<P, RA, Tag> {}
402 impl<P: ReduceDType, RA: ReduceArgs, Tag> Clone for TensorArgExpand<P, RA, Tag> {
403 fn clone(&self) -> Self {
404 Self {
405 state: self.state.clone(),
406 tag: self.tag,
407 }
408 }
409 }
410}