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,
9    prelude::{ArgSettings, CompilationArg, LaunchArg},
10};
11
12use super::{Sequence, SequenceExpand};
13
14pub struct SequenceArg<'a, R: Runtime, T: LaunchArg> {
15    pub values: Vec<T::RuntimeArg<'a, R>>,
16}
17
18impl<R: Runtime, T: LaunchArg> Default for SequenceArg<'_, R, T> {
19    fn default() -> Self {
20        Self::new()
21    }
22}
23
24impl<'a, R: Runtime, T: LaunchArg> SequenceArg<'a, R, T> {
25    pub fn new() -> Self {
26        Self { values: Vec::new() }
27    }
28    pub fn push(&mut self, arg: T::RuntimeArg<'a, R>) {
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> CompilationArg for SequenceCompilationArg<C> {}
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<'a, R: Runtime> = SequenceArg<'a, R, C>;
69    type CompilationArg = SequenceCompilationArg<C>;
70
71    fn compilation_arg<R: Runtime>(runtime_arg: &Self::RuntimeArg<'_, R>) -> Self::CompilationArg {
72        SequenceCompilationArg {
73            values: runtime_arg
74                .values
75                .iter()
76                .map(|value| C::compilation_arg(value))
77                .collect(),
78        }
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, T: LaunchArg> ArgSettings<R> for SequenceArg<'_, R, T> {
107    fn register(&self, launcher: &mut crate::prelude::KernelLauncher<R>) {
108        self.values.iter().for_each(|arg| arg.register(launcher));
109    }
110}
111
112impl<'a, R: Runtime, E: LaunchArg> FromIterator<E::RuntimeArg<'a, R>> for SequenceArg<'a, R, E> {
113    fn from_iter<T: IntoIterator<Item = E::RuntimeArg<'a, R>>>(iter: T) -> Self {
114        SequenceArg {
115            values: iter.into_iter().collect(),
116        }
117    }
118}
cubecl_core/frontend/container/sequence/launch.rs

cubecl_core/frontend/container/sequence/
launch.rs
