cubecl_core/frontend/container/sequence/
launch.rs

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

cubecl_core/frontend/container/sequence/
launch.rs
