cubecl_core/frontend/container/sequence/
launch.rs

1use alloc::{rc::Rc, vec::Vec};
2use core::cell::RefCell;
3
4use cubecl_runtime::runtime::Runtime;
5use cubecl_zspace::SmallVec;
6
7use crate::{
8    compute::{KernelBuilder, KernelLauncher},
9    prelude::LaunchArg,
10};
11
12use super::{Sequence, SequenceExpand};
13
14pub struct SequenceArg<R: Runtime, T: LaunchArg> {
15    pub values: SmallVec<[T::RuntimeArg<R>; 5]>,
16}
17
18impl<R: Runtime, T: LaunchArg> Default for SequenceArg<R, T> {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl<R: Runtime, T: LaunchArg> SequenceArg<R, T> {
25    pub fn new() -> Self {
26        Self {
27            values: SmallVec::new(),
28        }
29    }
30    pub fn push(&mut self, arg: T::RuntimeArg<R>) {
31        self.values.push(arg);
32    }
33}
34
35pub struct SequenceCompilationArg<C: LaunchArg> {
36    pub values: SmallVec<[C::CompilationArg; 5]>,
37}
38
39impl<C: LaunchArg> Clone for SequenceCompilationArg<C> {
40    fn clone(&self) -> Self {
41        Self {
42            values: self.values.clone(),
43        }
44    }
45}
46
47impl<C: LaunchArg> core::hash::Hash for SequenceCompilationArg<C> {
48    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
49        self.values.hash(state)
50    }
51}
52
53impl<C: LaunchArg> core::cmp::PartialEq for SequenceCompilationArg<C> {
54    fn eq(&self, other: &Self) -> bool {
55        self.values.eq(&other.values)
56    }
57}
58
59impl<C: LaunchArg> core::fmt::Debug for SequenceCompilationArg<C> {
60    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
61        f.write_str("Sequence ")?;
62        self.values.fmt(f)
63    }
64}
65impl<C: LaunchArg> core::cmp::Eq for SequenceCompilationArg<C> {}
66
67impl<C: LaunchArg> LaunchArg for Sequence<C> {
68    type RuntimeArg<R: Runtime> = SequenceArg<R, C>;
69    type CompilationArg = SequenceCompilationArg<C>;
70
71    fn register<R: Runtime>(
72        arg: Self::RuntimeArg<R>,
73        launcher: &mut KernelLauncher<R>,
74    ) -> Self::CompilationArg {
75        arg.values
76            .into_iter()
77            .map(|arg| C::register(arg, launcher))
78            .collect()
79    }
80
81    fn expand(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> SequenceExpand<C> {
82        let values = arg
83            .values
84            .iter()
85            .map(|value| C::expand(value, builder))
86            .collect::<Vec<_>>();
87
88        SequenceExpand {
89            values: Rc::new(RefCell::new(values)),
90        }
91    }
92
93    fn expand_output(arg: &Self::CompilationArg, builder: &mut KernelBuilder) -> SequenceExpand<C> {
94        let values = arg
95            .values
96            .iter()
97            .map(|value| C::expand_output(value, builder))
98            .collect::<Vec<_>>();
99
100        SequenceExpand {
101            values: Rc::new(RefCell::new(values)),
102        }
103    }
104}
105
106impl<R: Runtime, E: LaunchArg> FromIterator<E::RuntimeArg<R>> for SequenceArg<R, E> {
107    fn from_iter<T: IntoIterator<Item = E::RuntimeArg<R>>>(iter: T) -> Self {
108        SequenceArg {
109            values: iter.into_iter().collect(),
110        }
111    }
112}
113
114impl<E: LaunchArg> FromIterator<E::CompilationArg> for SequenceCompilationArg<E> {
115    fn from_iter<T: IntoIterator<Item = E::CompilationArg>>(iter: T) -> Self {
116        Self {
117            values: iter.into_iter().collect(),
118        }
119    }
120}
cubecl_core/frontend/container/sequence/launch.rs

cubecl_core/frontend/container/sequence/
launch.rs
