Merge branch 'main' into lazy-graph-comb

Cargo.toml

@@ -50,7 +50,7 @@ min-cl = { git = "https://github.com/elftausend/min-cl", optional = true }
 
 [features]
 # default = ["cpu", "autograd", "macro"]
-default = ["cpu", "lazy", "static-api", "opencl", "graph", "autograd", "fork", "stack"]
+default = ["cpu", "lazy", "static-api", "graph", "autograd", "fork", "serde", "json"]
 
 cpu = []
 opencl = ["dep:min-cl", "cpu", "cached"]

build.rs

@@ -58,8 +58,6 @@ fn has_device_unified_mem() -> bool {
         .unified_mem
 }
 
-use std::path::{Path, PathBuf};
-
 // https://github.com/coreylowman/cudarc/blob/main/build.rs
 #[cfg(feature = "cuda")]
 fn link_cuda() {

src/buffer/num.rs

@@ -88,15 +88,15 @@ impl Device for () {
     }
 
     #[inline(always)]
-    fn data_as_wrap<'a, T, S: crate::Shape>(
-        data: &'a Self::Data<T, S>,
-    ) -> &'a Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap<T, S: crate::Shape>(
+        data: &Self::Data<T, S>,
+    ) -> &Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 
-    fn data_as_wrap_mut<'a, T, 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, S: crate::Shape>(
+        data: &mut Self::Data<T, S>,
+    ) -> &mut Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 }
@@ -125,14 +125,12 @@ impl WrappedData for () {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         wrap
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         wrap
     }
 }

src/devices.rs

@@ -55,12 +55,10 @@ pub trait Device: OnDropBuffer + Sized {
     // FIXME: probably a better way to realize these
     fn base_to_data<T, S: Shape>(&self, base: Self::Base<T, S>) -> Self::Data<T, S>;
     fn wrap_to_data<T, S: Shape>(&self, wrap: Self::Wrap<T, Self::Base<T, S>>) -> Self::Data<T, S>;
-    fn data_as_wrap<'a, T, S: Shape>(
-        data: &'a Self::Data<T, S>,
-    ) -> &'a Self::Wrap<T, Self::Base<T, S>>;
-    fn data_as_wrap_mut<'a, T, S: Shape>(
-        data: &'a mut Self::Data<T, S>,
-    ) -> &'a mut Self::Wrap<T, Self::Base<T, S>>;
+    fn data_as_wrap<T, S: Shape>(data: &Self::Data<T, S>) -> &Self::Wrap<T, Self::Base<T, S>>;
+    fn data_as_wrap_mut<T, S: Shape>(
+        data: &mut Self::Data<T, S>,
+    ) -> &mut Self::Wrap<T, Self::Base<T, S>>;
 
     /// Creates a new [`Buffer`] using `A`.
     ///
@@ -111,7 +109,7 @@ macro_rules! impl_device_traits {
         $crate::impl_wrapped_data!($device);
 
         #[cfg(feature = "graph")]
-        crate::pass_down_optimize_mem_graph!($device);
+        $crate::pass_down_optimize_mem_graph!($device);
 
         $crate::pass_down_grad_fn!($device);
         $crate::pass_down_tape_actions!($device);

src/devices/cpu/cpu_device.rs

@@ -54,16 +54,14 @@ impl<Mods: OnDropBuffer> Device for CPU<Mods> {
     }
 
     #[inline(always)]
-    fn data_as_wrap<'a, T, S: Shape>(
-        data: &'a Self::Data<T, S>,
-    ) -> &'a Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap<T, S: Shape>(data: &Self::Data<T, S>) -> &Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 
     #[inline(always)]
-    fn data_as_wrap_mut<'a, T, S: Shape>(
-        data: &'a mut Self::Data<T, S>,
-    ) -> &'a mut Self::Wrap<T, Self::Base<T, S>> {
+    fn data_as_wrap_mut<T, S: Shape>(
+        data: &mut Self::Data<T, S>,
+    ) -> &mut Self::Wrap<T, Self::Base<T, S>> {
         data
     }
 

src/devices/cpu/mod.rs

@@ -11,4 +11,3 @@ mod cpu_ptr;
 mod ops;
 
 pub use cpu_ptr::*;
-pub use ops::*;

src/devices/opencl/mod.rs

@@ -74,6 +74,12 @@ impl<T> HasId for CLPtr<T> {
     }
 }
 
+impl<T> CLPtr<T> {
+    pub fn len(&self) -> usize {
+        self.len
+    }
+}
+
 impl<T> ShallowCopy for CLPtr<T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {
@@ -116,6 +122,7 @@ impl<T> HostPtr<T> for CLPtr<T> {
     }
 }
 
+#[cfg(unified_cl)]
 impl<T> Deref for CLPtr<T> {
     type Target = [T];
 
@@ -125,6 +132,7 @@ impl<T> Deref for CLPtr<T> {
     }
 }
 
+#[cfg(unified_cl)]
 impl<T> DerefMut for CLPtr<T> {
     #[inline]
     fn deref_mut(&mut self) -> &mut Self::Target {

src/devices/opencl/ops.rs

@@ -211,7 +211,10 @@ where
                 Ok(())
             }
             #[cfg(not(unified_cl))]
-            try_cl_apply_fn_mut(dev, buf, out, f);
+            {
+                try_cl_apply_fn_mut(dev, buf, out, **f)?;
+                Ok(())
+            }
         })
         .unwrap();
 

src/exec_on_cpu.rs

@@ -221,7 +221,7 @@ macro_rules! to_cpu {
 macro_rules! to_raw_host {
     ($cpu:expr, $($t:ident),*) => {
         $(
-            let $t = &unsafe { $crate::Buffer::<_, _, ()>::from_raw_host_device($cpu, $t.data.host_ptr, $t.len()) };
+            let $t = &unsafe { $crate::Buffer::<_, _, ()>::from_raw_host_device($cpu, $t.base().host_ptr, $t.len()) };
         )*
     };
 }
@@ -233,7 +233,7 @@ macro_rules! to_raw_host_mut {
     ($cpu:expr, $($t:ident, $cpu_name:ident),*) => {
         $(
             #[allow(unused_mut)]
-            let mut $cpu_name = &mut unsafe { $crate::Buffer::<_, _, ()>::from_raw_host_device($cpu, $t.data.host_ptr, $t.len()) };
+            let mut $cpu_name = &mut unsafe { $crate::Buffer::<_, _, ()>::from_raw_host_device($cpu, $t.base().host_ptr, $t.len()) };
         )*
     };
 }

src/exec_on_cpu/cl_may_unified.rs

@@ -202,7 +202,7 @@ macro_rules! cl_cpu_exec_unified {
         let cpu = CPU::<Base>::new();
         if $device.unified_mem() {
 
-            $crate::to_raw_host!($crate::CPU::<$crate::CachedModule<$crate::Base, $crate::CPU>>, $($t),*);
+            $crate::to_raw_host!(&$device.cpu, $($t),*);
 
             #[cfg(not(feature = "realloc"))]
             {

src/features.rs

@@ -102,7 +102,7 @@ macro_rules! pass_down_grad_fn {
             fn add_grad_fn<Args: $crate::Parents<N> + $crate::UpdateArgs, const N: usize>(
                 &self,
                 args: Args,
-                op: fn(&mut Args) -> crate::Result<()>,
+                op: fn(&mut Args) -> $crate::Result<()>,
             ) {
                 self.modules.add_grad_fn(args, op)
             }
@@ -253,7 +253,7 @@ macro_rules! pass_down_add_operation {
             fn add_op<Args: $crate::Parents<N> + $crate::UpdateArgs, const N: usize>(
                 &self,
                 args: Args,
-                operation: fn(&mut Args) -> crate::Result<()>,
+                operation: fn(&mut Args) -> $crate::Result<()>,
             ) -> $crate::Result<()> {
                 self.modules.add_op(args, operation)
             }

src/hooks.rs

@@ -6,7 +6,7 @@ pub trait OnDropBuffer: WrappedData {
     fn on_drop_buffer<T, D: Device, S: Shape>(&self, _device: &D, _buf: &Buffer<T, D, S>) {}
 }
 
-pub trait OnNewBuffer<T, D: Device, S: Shape> {
+pub trait OnNewBuffer<T, D: Device, S: Shape = ()> {
     #[track_caller]
     fn on_new_buffer(&self, _device: &D, _new_buf: &Buffer<T, D, S>) {}
 }

src/id.rs

@@ -102,7 +102,7 @@ impl<T> UpdateArg for NoId<T> {
     #[inline]
     #[cfg(not(feature = "no-std"))]
     fn update_arg<B>(
-        &mut self,
+        _to_update: &mut Self,
         _id: Option<UniqueId>,
         _buffers: &mut crate::Buffers<B>,
     ) -> crate::Result<()> {
@@ -113,7 +113,7 @@ impl<T> UpdateArg for NoId<T> {
 impl<'a, T: 'static, D: Device + 'static, S: Shape + 'static> UpdateArg for &Buffer<'a, T, D, S> {
     #[cfg(not(feature = "no-std"))]
     fn update_arg<B: crate::AsAny>(
-        &mut self,
+        to_update: &mut Self,
         id: Option<UniqueId>,
         buffers: &mut crate::Buffers<B>,
     ) -> crate::Result<()> {
@@ -123,7 +123,7 @@ impl<'a, T: 'static, D: Device + 'static, S: Shape + 'static> UpdateArg for &Buf
             .get(&id.unwrap())
             .ok_or(DeviceError::InvalidLazyBuf)?;
         // let any = buf.as_any();
-        *self = unsafe { &*(buf.as_any() as *const dyn Any as *const Buffer<T, D, S>) };
+        *to_update = unsafe { &*(buf.as_any() as *const dyn Any as *const Buffer<T, D, S>) };
         //    *self = buffers.get(&self.id()).unwrap().downcast_ref().unwrap();
         Ok(())
     }
@@ -133,8 +133,8 @@ impl<'a, T: 'static, D: Device + 'static, S: Shape + 'static> UpdateArg
     for &mut Buffer<'a, T, D, S>
 {
     #[cfg(not(feature = "no-std"))]
-    fn update_arg<B>(
-        &mut self,
+    fn update_arg<B: crate::AsAny>(
+        to_update: &mut Self,
         id: Option<UniqueId>,
         buffers: &mut crate::Buffers<B>,
     ) -> crate::Result<()> {
@@ -143,7 +143,7 @@ impl<'a, T: 'static, D: Device + 'static, S: Shape + 'static> UpdateArg
         let buf = buffers
             .get_mut(&id.unwrap())
             .ok_or(DeviceError::InvalidLazyBuf)?;
-        // *self = unsafe { &mut *(&mut **buf as *mut dyn ShallowCopyable as *mut Buffer<T, D, S>) };
+        // *to_update = unsafe { &mut *(&mut **buf as *mut dyn Any as *mut Buffer<T, D, S>) };
         Ok(())
         //    *self = buffers.get(&self.id()).unwrap().downcast_ref().unwrap();
     }

src/modules/autograd.rs

@@ -29,14 +29,12 @@ impl<Mods: WrappedData> WrappedData for Autograd<Mods> {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         Mods::wrapped_as_base(wrap)
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         Mods::wrapped_as_base_mut(wrap)
     }
 }

src/modules/base.rs

@@ -15,14 +15,12 @@ impl WrappedData for Base {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         wrap
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         wrap
     }
 }

src/modules/cached.rs

@@ -34,14 +34,12 @@ impl<Mods: WrappedData, SD: Device> WrappedData for CachedModule<Mods, SD> {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         Mods::wrapped_as_base(wrap)
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         Mods::wrapped_as_base_mut(wrap)
     }
 }

src/modules/fork.rs

@@ -34,14 +34,12 @@ impl<Mods: WrappedData> WrappedData for Fork<Mods> {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         Mods::wrapped_as_base(wrap)
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         Mods::wrapped_as_base_mut(wrap)
     }
 }

src/modules/graph.rs

@@ -32,14 +32,12 @@ impl<Mods: WrappedData> WrappedData for Graph<Mods> {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         Mods::wrapped_as_base(wrap)
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         Mods::wrapped_as_base_mut(wrap)
     }
 }

src/modules/lazy.rs

@@ -82,9 +82,10 @@ impl<Mods: AddOperation> AddOperation for Lazy<Mods> {
         args: Args,
         operation: fn(&mut Args) -> crate::Result<()>,
     ) -> crate::Result<()> {
-        Ok(self.graph.try_borrow_mut()
-            .expect("already borrowed: BorrowMutError - is the inner operation trying to add an operation as well?")
-            .add_operation(args, operation))
+        self.graph.try_borrow_mut()
+        .expect("already borrowed: BorrowMutError - is the inner operation trying to add an operation as well?")
+        .add_operation(args, operation);
+        Ok(())
     }
 }
 

src/modules/lazy/wrapper.rs

@@ -26,14 +26,12 @@ impl<Mods: WrappedData> WrappedData for Lazy<Mods> {
     }
 
     #[inline]
-    fn wrapped_as_base<'a, T, Base: HasId + PtrType>(wrap: &'a Self::Wrap<T, Base>) -> &'a Base {
+    fn wrapped_as_base<T, Base: HasId + PtrType>(wrap: &Self::Wrap<T, Base>) -> &Base {
         Mods::wrapped_as_base(wrap.data.as_ref().expect(MISSING_DATA))
     }
 
     #[inline]
-    fn wrapped_as_base_mut<'a, T, Base: HasId + PtrType>(
-        wrap: &'a mut Self::Wrap<T, Base>,
-    ) -> &'a mut Base {
+    fn wrapped_as_base_mut<T, Base: HasId + PtrType>(wrap: &mut Self::Wrap<T, Base>) -> &mut Base {
         Mods::wrapped_as_base_mut(wrap.data.as_mut().expect(MISSING_DATA))
     }
 }
@@ -71,7 +69,7 @@ impl<Data: PtrType, T> PtrType for LazyWrapper<Data, T> {
     }
 }
 
-const MISSING_DATA: &'static str =
+const MISSING_DATA: &str =
     "This lazy buffer does not contain any data. Try with a buffer.replace() call.";
 
 impl<Data: Deref<Target = [T]>, T> Deref for LazyWrapper<Data, T> {

src/parents.rs

@@ -22,7 +22,7 @@ impl AllParents for () {}
 impl UpdateArg for () {
     #[cfg(not(feature = "no-std"))]
     fn update_arg<B>(
-        &mut self,
+        _to_update: &mut Self,
         _id: Option<crate::UniqueId>,
         _buffers: &mut crate::Buffers<B>,
     ) -> crate::Result<()> {
@@ -73,7 +73,7 @@ macro_rules! impl_parents {
                 let mut ids = ids.iter();
                 #[allow(non_snake_case)]
                 let ($($to_impl,)+) = self;
-                $($to_impl.update_arg(*ids.next().unwrap(), buffers)?;)*
+                $($to_impl::update_arg($to_impl, *ids.next().unwrap(), buffers)?;)*
                 Ok(())
             }
         }