cubecl_core/frontend/container/sequence/
launch.rs

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

cubecl_core/frontend/container/sequence/
launch.rs
