1use std::mem::ManuallyDrop;
2use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
3
4use crate::ll::*;
5
6use crate::{Buffer, Context, Program};
7
8pub struct Kernel {
9 program: ManuallyDrop<Program>,
10 inner: ManuallyDrop<Arc<RwLock<ClKernel>>>,
11 _unconstructable: (),
12}
13
14unsafe impl Send for Kernel {}
15unsafe impl Sync for Kernel {}
16
17impl Drop for Kernel {
18 fn drop(&mut self) {
19 unsafe {
20 ManuallyDrop::drop(&mut self.inner);
21 ManuallyDrop::drop(&mut self.program);
22 }
23 }
24}
25
26impl Clone for Kernel {
27 fn clone(&self) -> Kernel {
28 let inner_clone = self.read_lock().clone();
29 Kernel {
30 inner: ManuallyDrop::new(Arc::new(RwLock::new(inner_clone))),
31 program: ManuallyDrop::new(self.program().clone()),
32 _unconstructable: (),
33 }
34 }
35}
36
37impl Kernel {
38 pub unsafe fn new(kernel: ClKernel, program: Program) -> Kernel {
39 Kernel {
40 program: ManuallyDrop::new(program),
41 inner: ManuallyDrop::new(Arc::new(RwLock::new(kernel))),
42 _unconstructable: (),
43 }
44 }
45
46 pub fn create(program: &Program, name: &str) -> Output<Kernel> {
47 let ll_kernel = unsafe { ClKernel::create(program.low_level_program(), name) }?;
48 Ok(unsafe { Kernel::new(ll_kernel, program.clone()) })
49 }
50
51 pub unsafe fn set_arg<T>(&self, arg_index: usize, arg: &mut T) -> Output<()>
52 where
53 T: KernelArg,
54 {
55 self.write_lock().set_arg(arg_index, arg)
56 }
57
58 pub fn function_name(&self) -> Output<String> {
59 unsafe { self.read_lock().function_name() }
60 }
61
62 pub fn num_args(&self) -> Output<u32> {
63 unsafe { self.read_lock().num_args() }
64 }
65
66 pub fn reference_count(&self) -> Output<u32> {
67 unsafe { self.read_lock().reference_count() }
68 }
69
70 pub fn context(&self) -> &Context {
71 self.program().context()
72 }
73
74 pub fn program(&self) -> &Program {
75 &*self.program
76 }
77
78 pub fn attributes(&self) -> Output<String> {
79 unsafe { self.read_lock().attributes() }
80 }
81
82 pub fn read_lock(&self) -> RwLockReadGuard<ClKernel> {
83 self.inner.read().unwrap()
84 }
85
86 pub fn write_lock(&self) -> RwLockWriteGuard<ClKernel> {
87 self.inner.write().unwrap()
88 }
89}
90
91pub enum KernelOpArg<'a, T: ClNumber> {
92 Num(T),
93 Buffer(&'a Buffer),
94}
95
96pub trait ToKernelOpArg<'a, T: ClNumber> {
102 fn to_kernel_op_arg(&self) -> Output<KernelOpArg<'a, T>>;
103}
104
105impl<'a, T: ClNumber> ToKernelOpArg<'a, T> for T {
106 fn to_kernel_op_arg(&self) -> Output<KernelOpArg<'a, T>> {
107 Ok(KernelOpArg::Num(*self))
108 }
109}
110
111impl<'a, T: ClNumber> ToKernelOpArg<'a, T> for &'a Buffer {
112 fn to_kernel_op_arg(&self) -> Output<KernelOpArg<'a, T>> {
113 self.number_type().type_check(T::number_type())?;
114 Ok(KernelOpArg::Buffer(self))
115 }
116}
117
118impl<'a, T: ClNumber> KernelOpArg<'a, T> {
119 pub fn into_buffer(self) -> Output<&'a Buffer> {
120 if let KernelOpArg::Buffer(buffer) = self {
121 Ok(buffer)
122 } else {
123 Err(KernelError::KernelOpArgWasNotMem.into())
124 }
125 }
126
127 pub fn into_num(self) -> Output<T> {
128 if let KernelOpArg::Num(num) = self {
129 Ok(num)
130 } else {
131 Err(KernelError::KernelOpArgWasNotMem.into())
132 }
133 }
134}
135
136pub struct KernelOperation<'a, T: ClNumber + KernelArg> {
137 _name: String,
138 _args: Vec<KernelOpArg<'a, T>>,
139 _work: Option<Work>,
140 _returning: Option<usize>,
141 pub command_queue_opts: Option<CommandQueueOptions>,
142}
143
144impl<'a, T: ClNumber + KernelArg> KernelOperation<'a, T> {
145 pub fn new(name: &str) -> KernelOperation<T> {
146 KernelOperation {
147 _name: name.to_owned(),
148 _args: vec![],
149 _work: None,
150 _returning: None,
151 command_queue_opts: None,
152 }
153 }
154
155 pub fn name(&self) -> &str {
156 &self._name[..]
157 }
158
159 pub fn command_queue_opts(&self) -> Option<CommandQueueOptions> {
160 self.command_queue_opts.clone()
161 }
162
163 pub fn args(&self) -> &[KernelOpArg<T>] {
164 &self._args[..]
165 }
166
167 pub fn mut_args(&mut self) -> &mut [KernelOpArg<'a, T>] {
168 &mut self._args[..]
169 }
170
171 pub fn with_dims<D: Into<Dims>>(mut self, dims: D) -> KernelOperation<'a, T> {
172 self._work = Some(Work::new(dims.into()));
173 self
174 }
175
176 pub fn with_work<W: Into<Work>>(mut self, work: W) -> KernelOperation<'a, T> {
177 self._work = Some(work.into());
178 self
179 }
180
181 pub fn add_arg<A: Into<KernelOpArg<'a, T>>>(mut self, arg: A) -> KernelOperation<'a, T> {
182 self._args.push(arg.into());
183 self
184 }
185
186 pub fn with_command_queue_options(
187 mut self,
188 opts: CommandQueueOptions,
189 ) -> KernelOperation<'a, T> {
190 self.command_queue_opts = Some(opts);
191 self
192 }
193
194 pub fn with_returning_arg(mut self, arg_index: usize) -> KernelOperation<'a, T> {
195 self._returning = Some(arg_index);
196 self
197 }
198
199 pub fn argc(&self) -> usize {
200 self._args.len()
201 }
202
203 #[inline]
204 pub fn work(&self) -> Output<Work> {
205 self._work.clone().ok_or(KernelError::WorkIsRequired.into())
206 }
207}