Fix opencl compilation

Cargo.toml

@@ -49,7 +49,7 @@ min-cl = { path="../min-cl", optional=true }
 
 [features]
 # default = ["cpu", "autograd", "macro"]
-default = ["graph", "stack", "autograd", "cpu", "fork", "lazy", ]
+default = ["graph", "stack", "autograd", "cpu", "fork", "lazy", "opencl"]
 
 cpu = []
 opencl = ["dep:min-cl", "cpu", "cached"]

examples/modules.rs

@@ -101,7 +101,7 @@ impl<T, S, Mods> ElementWise<T, Self, S> for custos::OpenCL<Mods>
 where
     T: Add<Output = T> + Copy + CDatatype + Default,
     S: Shape,
-    Mods: Retrieve<Self, T> + AddOperation<T, Self> + UseGpuOrCpu,
+    Mods: Retrieve<Self, T> + AddOperation + UseGpuOrCpu + 'static,
 {
     fn add(
         &self,
@@ -110,12 +110,12 @@ where
     ) -> custos::Result<Buffer<T, Self, S>> {
         let mut out = self.retrieve(lhs.len(), (lhs, rhs));
 
-        self.add_op((lhs, rhs), Some(&mut out), |out, (lhs, rhs)| {
+        self.add_op((lhs, rhs, &mut out), |(lhs, rhs, out)| {
             let dev = lhs.device();
-            let out = out.as_mut().unwrap();
+            let out = &mut **out;
             #[cfg(unified_cl)]
             {
-                let cpu_out = unsafe { &mut *(*out as *mut Buffer<_, OpenCL<Mods>, _>) };
+                let cpu_out = unsafe { &mut *(out as *mut Buffer<_, OpenCL<Mods>, _>) };
                 dev.use_cpu_or_gpu(
                     (file!(), line!(), column!()).into(),
                     &[lhs.len()],

src/devices/opencl/cl_device.rs

@@ -7,7 +7,7 @@ use crate::flag::AllocFlag;
 use crate::{
     impl_buffer_hook_traits, impl_retriever, pass_down_optimize_mem_graph,
     pass_down_use_gpu_or_cpu, Alloc, Base, Buffer, Cached, CachedCPU, CloneBuf, Device, Module,
-    OnDropBuffer, Setup, CPU,
+    OnDropBuffer, Setup, CPU, pass_down_tape_actions,
 };
 use crate::{PtrConv, Shape};
 
@@ -281,19 +281,8 @@ impl<Mods> crate::LazySetup for OpenCL<Mods> {}
 
 #[cfg(feature = "lazy")]
 impl<Mods> crate::LazyRun for OpenCL<Mods> {}
-
-#[cfg(feature = "autograd")]
-impl<Mods: crate::TapeActions> crate::TapeActions for OpenCL<Mods> {
-    #[inline]
-    fn tape(&self) -> Option<core::cell::Ref<crate::Tape>> {
-        self.modules.tape()
-    }
-
-    #[inline]
-    fn tape_mut(&self) -> Option<core::cell::RefMut<crate::Tape>> {
-        self.modules.tape_mut()
-    }
-}
+
+pass_down_tape_actions!(OpenCL);
 
 #[cfg(test)]
 mod tests {

src/devices/opencl/ops.rs

@@ -182,7 +182,7 @@ impl<T, S, Mods> ApplyFunction<T, S> for OpenCL<Mods>
 where
     T: CDatatype + Number,
     S: Shape,
-    Mods: AddOperation<T, Self> + Retrieve<Self, T> + UseGpuOrCpu + 'static,
+    Mods: AddOperation + Retrieve<Self, T> + UseGpuOrCpu + 'static,
 {
     #[inline]
     fn apply_fn<F>(
@@ -195,9 +195,10 @@ where
     {
         let mut out = self.retrieve(buf.len(), buf);
 
-        self.add_op((buf, f.no_id()), Some(&mut out), |out, (buf, f)| {
+        self.add_op((buf, f.no_id(), &mut out), |(buf, f, out)| {
             let dev = buf.device();
-            let out: &mut Buffer<'_, T, OpenCL<Mods>, S> = out.as_mut().unwrap();
+            // let out: &mut Buffer<'_, T, OpenCL<Mods>, S> = out.as_mut().unwrap();
+            let out = &mut **out;
             #[cfg(unified_cl)]
             {
                 let cpu_out = unsafe { &mut *(out as *mut Buffer<_, OpenCL<Mods>, _>) };
@@ -254,7 +255,7 @@ where
     Ok(())
 }
 
-impl<T, S, Mods: OnDropBuffer + AddOperation<T, Self>> UnaryGrad<T, S> for OpenCL<Mods>
+impl<T, S, Mods: OnDropBuffer + AddOperation + 'static> UnaryGrad<T, S> for OpenCL<Mods>
 where
     T: CDatatype + Number,
     S: Shape,
@@ -269,10 +270,9 @@ where
     ) where
         F: ToCLSource,
     {
-        self.add_op::<S, _, 4>(
+        self.add_op::<_, 4>(
             (lhs, lhs_grad.buf_no_id(), out, lhs_grad_fn.no_id()),
-            None,
-            move |_, (lhs, lhs_grad, out, lhs_grad_fn)| {
+            move |(lhs, lhs_grad, out, lhs_grad_fn)| {
                 try_cl_add_unary_grad(lhs.device(), lhs, &mut **lhs_grad, out, **lhs_grad_fn)
             },
         )

src/modules/autograd.rs

@@ -401,7 +401,7 @@ mod tests {
 
     #[test]
     fn test_grad_new_api() {
-        use crate::{AddGradFn, AddOperation, Lazy};
+        use crate::AddGradFn;
 
         let device = CPU::<Autograd<Base>>::new();