opencl_core/
kernel.rs

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
96// pub enum ReturnArg<T: ClNumber> {
97//     Num(T),
98//     Buffer(Buffer<T>),
99// }
100
101pub 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}