Merge branch 'main' into webgl

Author: elftausend

src/cache/borrow_cache.rs

@@ -65,7 +65,7 @@ impl BorrowCache {
 
     pub fn add_or_get_mut<'a, T, D, S>(
         &mut self,
-        device: &D,
+        device: &'a D,
         id: Id,
         new_buf: &mut bool,
     ) -> &mut Buffer<'a, T, D, S>
@@ -75,7 +75,7 @@ impl BorrowCache {
         S: Shape,
     {
         self.add_buf_once::<T, D, S>(device, id, new_buf);
-        self.get_buf_mut(id).unwrap()
+        unsafe { self.get_buf_mut(id).unwrap() }
     }
 
     pub fn add_buf_once<T, D, S>(&mut self, device: &D, id: Id, new_buf: &mut bool)
@@ -91,7 +91,7 @@ impl BorrowCache {
         self.add_buf::<T, D, S>(device, id)
     }
 
-    pub fn add_buf<T, D, S>(&mut self, device: &D, id: Id)
+    pub fn add_buf<'a, T, D, S>(&'a mut self, device: &'a D, id: Id)
     where
         T: Unit + 'static,
         D: Alloc<T> + 'static,
@@ -123,7 +123,7 @@ impl BorrowCache {
     }
 
     #[inline]
-    pub fn get_buf<'a, T, D, S>(&self, id: Id) -> Result<&Buffer<'a, T, D, S>, CachingError>
+    pub unsafe fn get_buf<'a, T, D, S>(&self, id: Id) -> Result<&Buffer<'a, T, D, S>, CachingError>
     where
         T: Unit + 'static,
         D: Device + 'static,
@@ -137,7 +137,7 @@ impl BorrowCache {
     }
 
     #[inline]
-    pub fn get_buf_mut<'a, T, D, S>(
+    pub unsafe fn get_buf_mut<'a, T, D, S>(
         &mut self,
         id: Id,
     ) -> Result<&mut Buffer<'a, T, D, S>, CachingError>
@@ -161,16 +161,21 @@ impl BorrowCache {
 #[cfg(test)]
 mod tests {
 
-    /*#[test]
-    fn test_comp_error() {
-        let device = CPU::<Base>::new();
+    // #[test]
+    // #[cfg(feature = "cpu")]
+    // fn test_comp_error() {
+    //     use crate::{Base, BorrowCache, Id, CPU};
 
+    //     let mut cache = BorrowCache::default();
 
-        let a = {
-            let mut cache = BorrowingCache::default();
-            cache.add_or_get::<f32, CPU, ()>(&device, Id::new(10))
-        };
-    }*/
+
+    //     let a = {
+    //         let device = CPU::<Base>::new();
+    //         // drop(device);
+    //         let mut new_buf = false;
+    //         // cache.add_or_get::<f32, CPU, ()>(&device, Id { id: 0, len: 10}, &mut new_buf)
+    //     };
+    // }
 
     #[cfg(feature = "cpu")]
     #[test]
@@ -190,8 +195,8 @@ mod tests {
         cache.add_buf_once::<f32, _, ()>(&device, sid, &mut false);
         cache.add_buf_once::<f32, _, ()>(&device, tid, &mut false);
 
-        let a: &Buffer = cache.get_buf::<f32, _, ()>(fid).unwrap();
-        let b: &Buffer = cache.get_buf::<f32, _, ()>(fid).unwrap();
+        let a: &Buffer = unsafe { cache.get_buf::<f32, _, ()>(fid).unwrap() };
+        let b: &Buffer = unsafe { cache.get_buf::<f32, _, ()>(fid).unwrap() };
 
         assert_eq!(a.ptr, b.ptr);
     }

src/devices/opencl/cl_device.rs

@@ -3,7 +3,7 @@ use min_cl::CLDevice;
 use min_cl::api::{create_buffer, enqueue_full_copy_buffer, MemFlags};
 
 use super::{enqueue_kernel, AsClCvoidPtr, CLPtr};
-use crate::flag::AllocFlag;
+use crate::{flag::AllocFlag, opencl::KernelLaunch};
 use crate::{impl_device_traits, Shape, Unit};
 use crate::{
     pass_down_use_gpu_or_cpu, Alloc, Base, Buffer, Cached, CachedCPU, CloneBuf, Device,
@@ -160,7 +160,7 @@ impl<Mods> OpenCL<Mods> {
         lws: Option<[usize; 3]>,
         args: &[&dyn AsClCvoidPtr],
     ) -> crate::Result<()> {
-        enqueue_kernel(self, src, gws, lws, args)
+        self.device.launch_kernel(src, gws, lws, args)
     }
 }
 

src/devices/opencl/kernel_enqueue.rs

@@ -115,6 +115,29 @@ impl<T> AsClCvoidPtr for CLPtr<T> {
     }
 }
 
+pub trait KernelLaunch {
+    fn launch_kernel(
+        &self,
+        src: &str,
+        gws: [usize; 3],
+        lws: Option<[usize; 3]>,
+        args: &[&dyn AsClCvoidPtr],
+    ) -> crate::Result<()>;
+}
+
+impl KernelLaunch for CLDevice {
+    #[inline]
+    fn launch_kernel(
+        &self,
+        src: &str,
+        gws: [usize; 3],
+        lws: Option<[usize; 3]>,
+        args: &[&dyn AsClCvoidPtr],
+    ) -> crate::Result<()> {
+        enqueue_kernel(self, src, gws, lws, args)
+    }
+}
+
 /// Executes a cached OpenCL kernel.
 /// # Example
 ///

src/devices/untyped/dummy_cuda.rs

@@ -16,7 +16,7 @@ impl<Mods: OnDropBuffer> Device for CUDA<Mods> {
         }
 
         #[cfg(feature = "std")]
-        Err(crate::DeviceError::CPUDeviceNotAvailable.into())
+        Err(crate::DeviceError::CPUDeviceNotAvailable)
     }
 
     fn base_to_data<T: Unit, S: Shape>(&self, base: Self::Base<T, S>) -> Self::Data<T, S> {
@@ -30,15 +30,15 @@ impl<Mods: OnDropBuffer> Device for CUDA<Mods> {
         wrap
     }
 
-    fn data_as_wrap<'a, T: Unit, S: Shape>(
-        data: &'a Self::Data<T, S>,
-    ) -> &'a Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap<T: Unit, S: Shape>(
+        data: &Self::Data<T, S>,
+    ) -> &Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 
-    fn data_as_wrap_mut<'a, T: Unit, S: Shape>(
-        data: &'a mut Self::Data<T, S>,
-    ) -> &'a mut Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap_mut<T: Unit, S: Shape>(
+        data: &mut Self::Data<T, S>,
+    ) -> &mut Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 }

src/devices/untyped/untyped_device.rs

@@ -42,16 +42,16 @@ impl Device for Untyped {
     }
 
     #[inline]
-    fn data_as_wrap<'a, T: Unit, S: crate::Shape>(
-        data: &'a Self::Data<T, S>,
-    ) -> &'a Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap<T: Unit, S: crate::Shape>(
+        data: &Self::Data<T, S>,
+    ) -> &Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 
     #[inline]
-    fn data_as_wrap_mut<'a, T: Unit, S: crate::Shape>(
-        data: &'a mut Self::Data<T, S>,
-    ) -> &'a mut Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap_mut<T: Unit, S: crate::Shape>(
+        data: &mut Self::Data<T, S>,
+    ) -> &mut Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 

src/exec_on_cpu.rs

@@ -322,7 +322,19 @@ mod tests {
                 out
             }
         );
-        assert_eq!(a.read(), vec![2, 4, 6, 8])
+
+        assert_eq!(a.read(), vec![2, 4, 6, 8]);
+        let cpu = crate::CPU::<Base>::new();
+        let other_a: crate::Buffer<i32, crate::OpenCL> = cpu_exec!(
+            &device, &cpu, lhs, rhs; {
+                let mut out = cpu.retrieve(lhs.len(), (&lhs, &rhs)).unwrap();
+                for ((lhs, rhs), out) in lhs.iter().zip(rhs.iter()).zip(out.iter_mut()) {
+                    *out = lhs + rhs;
+                }
+                out
+            }
+        );
+        assert_eq!(other_a.read(), vec![2, 4, 6, 8]);
     }
 
     #[cfg(feature = "opencl")]
@@ -421,4 +433,51 @@ mod tests {
 
         Ok(())
     }
+
+    pub trait AddEw<T, D: crate::Device = Self>: crate::Device {
+        fn add(&self, lhs: &crate::Buffer<T, D>, rhs: &crate::Buffer<T, D>) -> crate::Buffer<T, D>;
+    }
+
+    impl<Mods, T> AddEw<T> for crate::CPU<Mods>
+    where
+        Mods: crate::hooks::OnDropBuffer + crate::Retrieve<Self, T> + 'static,
+        Self::Base<T, ()>: core::ops::Deref<Target = [T]>,
+        T: core::ops::Add<Output = T> + Copy,
+    {
+        fn add(
+            &self,
+            lhs: &crate::Buffer<T, Self>,
+            rhs: &crate::Buffer<T, Self>,
+        ) -> crate::Buffer<T, Self> {
+            use crate::Retriever;
+            let mut out = self.retrieve(lhs.len(), (lhs, rhs)).unwrap();
+            for idx in 0..lhs.len() {
+                out[idx] = lhs[idx] + rhs[idx]
+            }
+            out
+        }
+    }
+
+    #[cfg(feature = "opencl")]
+    #[test]
+    fn test_cpu_exec_macro() -> crate::Result<()> {
+        use crate::{prelude::chosen_cl_idx, Base, Cached, Device, OpenCL, CPU};
+
+        let device = OpenCL::<Cached<Base>>::new(chosen_cl_idx())?;
+        let cpu = CPU::<Cached<Base>>::new();
+
+        let lhs = device.buffer([1, 2, 3, 4, 5]);
+        let rhs = device.buffer([-1, -4, -1, -8, -1]);
+        
+        let out1 = crate::cpu_exec!(&device, &cpu, lhs, rhs; cpu.add(&lhs, &rhs));
+
+        let out = {             
+            let lhs = crate::Buffer::<_, _>::from(((&cpu), lhs. read_to_vec()));             
+            let rhs = crate::Buffer::<_, _>::from(((&cpu), rhs. read_to_vec()));
+            let cpu_out = cpu.add(&lhs, &rhs);
+            crate::Buffer::from((&device, cpu_out))         
+        };
+        assert_eq!(out1.read(), out.read());
+        Ok(())
+    }
 }

src/exec_on_cpu/cl_may_unified.rs

@@ -103,16 +103,21 @@ where
 /// This is way faster than [cpu_exec_reduce], as new memory is not allocated.
 ///
 /// `cpu_exec_binary_may_unified` can be used interchangeably with [cpu_exec_reduce].
-pub fn cpu_exec_reduce_may_unified<T, F>(device: &OpenCL, x: &Buffer<T, OpenCL>, f: F) -> T
+pub fn cpu_exec_reduce_may_unified<T, F, Mods>(
+    device: &min_cl::CLDevice,
+    x: &Buffer<T, OpenCL<Mods>>,
+    f: F,
+) -> T
 where
     T: Unit + Default + Clone,
     F: Fn(&CPU, &Buffer<T, CPU>) -> T,
+    Mods: OnDropBuffer + 'static,
 {
     let cpu = CPU::<crate::Base>::new();
 
     if device.unified_mem() {
         return f(&cpu, &unsafe {
-            Buffer::from_raw_host(x.data.host_ptr, x.len())
+            Buffer::from_raw_host(x.base().host_ptr, x.len())
         });
     }
     cpu_exec_reduce(x, f)

src/modules/autograd/gradients.rs

@@ -12,7 +12,7 @@ const INVALID_ID: &str = "A matching Buffer does not exist.";
 /// The cache is populated by `get_ref`, `get_like` or `get_mut_ref` calls.
 #[derive(Default)]
 pub struct Gradients {
-    pub grads_pool: BorrowCache,
+    pub(crate) grads_pool: BorrowCache,
     pub no_grads_pool: Buffers<Box<dyn BoxedShallowCopy>>,
     pub zero_grad_cbs: Vec<(Id, fn(&mut dyn Any))>,
     pub buf_requires_grad: HashMap<UniqueId, bool, BuildHasherDefault<NoHasher>>,
@@ -27,8 +27,6 @@ impl core::fmt::Debug for Gradients {
 }
 
 impl Gradients {
-    /// Clears the cache.
-    #[inline]
     pub fn zero_grad(&mut self) {
         for (id, cb) in &self.zero_grad_cbs {
             let grad_buf = self.grads_pool.cache.get_mut(id).unwrap();
@@ -42,6 +40,7 @@ impl Gradients {
         // self.grads_pool.cache.clear();
     }
 
+    #[inline]
     pub fn add_zero_grad_cb<T, D, S>(&mut self, id: &Id)
     where
         T: Unit + 'static,
@@ -57,18 +56,18 @@ impl Gradients {
 
     /// May get a reference to a gradient [`Buffer`].
     #[inline]
-    pub fn may_get_ref<'a, T, S, D>(&self, ident: Id) -> Result<&Buffer<'a, T, D, S>, CachingError>
+    pub(crate) unsafe fn may_get_ref<'a, T, S, D>(&self, ident: Id) -> Result<&Buffer<'a, T, D, S>, CachingError>
     where
         T: Unit + 'static,
         S: Shape,
         D: Alloc<T> + 'static,
     {
-        self.grads_pool.get_buf(ident)
+        unsafe { self.grads_pool.get_buf(ident) }
     }
 
     /// May get a mutable reference to a gradient [`Buffer`].
     #[inline]
-    pub fn may_get_mut<'a, T, S, D>(
+    pub(crate) unsafe fn may_get_mut<'a, T, S, D>(
         &mut self,
         id: Id,
     ) -> Result<&mut Buffer<'a, T, D, S>, CachingError>
@@ -77,7 +76,7 @@ impl Gradients {
         S: Shape,
         D: Alloc<T> + 'static,
     {
-        self.grads_pool.get_buf_mut(id)
+        unsafe { self.grads_pool.get_buf_mut(id) }
     }
 
     /// Returns a reference to a gradient [`Buffer`].
@@ -97,7 +96,7 @@ impl Gradients {
         if new_buf {
             self.add_zero_grad_cb::<T, D, S>(&id);
         }
-        self.grads_pool.get_buf(id).unwrap()
+        unsafe { self.grads_pool.get_buf(id).unwrap() }
     }
 
     /// Returns a mutable reference to a gradient [`Buffer`].
@@ -116,7 +115,7 @@ impl Gradients {
         if new_buf {
             self.add_zero_grad_cb::<T, D, S>(&id);
         }
-        self.grads_pool.get_buf_mut(id).unwrap()
+        unsafe { self.grads_pool.get_buf_mut(id).unwrap() }
     }
 
     /// Returns a reference to a gradient [`Buffer`] using information from `buf`.

src/modules/lazy.rs

@@ -752,7 +752,7 @@ mod tests {
                 .unwrap();
         }
 
-        if let Ok(_) = unsafe { device.run() } {
+        if unsafe { device.run() }.is_ok() {
             panic!()
         }
     }

src/modules/lazy/lazy_graph.rs

@@ -55,7 +55,7 @@ impl<B: AsAny, T> LazyGraph<B, T> {
         args: Args,
         op: fn(&mut Args) -> crate::Result<()>,
     ) -> Operation<B, T> {
-        // store ids and test if buffers are still in cache
+        // store ids and test if buffers are still in cache afterwards
         let arg_ids = args
             .maybe_ids()
             .into_iter()