Impl ShallowCopy for LazyWrapper

elftausend · Dec 9, 2023 · 83680a0 · 83680a0
1 parent 9a81330
commit 83680a0
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 11 deletions.
diff --git a/src/modules/autograd.rs b/src/modules/autograd.rs
@@ -9,7 +9,7 @@ use core::cell::UnsafeCell;
 use crate::{
     pass_down_add_operation, pass_down_exec_now_module, register_buf, unregister_buf, AddGradFn,
     Alloc, Buffer, Device, HasId, Module, OnDropBuffer, OnNewBuffer, Parents, PtrConv, PtrType,
-    Retrieve, RunModule, Setup, Shape, TapeActions, WrappedData,
+    Retrieve, RunModule, Setup, Shape, TapeActions, WrappedData, ShallowCopy,
 };
 
 use super::{Cached, CachedModule};
@@ -50,6 +50,7 @@ impl<Mods> Autograd<Mods> {
     where
         T: 'static,
         D: Device + PtrConv + 'static,
+        // D::Data<T, S>: ShallowCopy,
         S: Shape,
     {
         let no_grads_pool = unsafe { &mut (*(self.grads.get())).no_grads_pool.cache };
@@ -67,6 +68,7 @@ impl<T, D, Mods> OnNewBuffer<T, D> for Autograd<Mods>
 where
     T: 'static,
     D: Alloc<T> + PtrConv + 'static,
+    // D::Data<T, S>: ShallowCopy,
     Mods: OnNewBuffer<T, D>,
 {
     #[inline]

diff --git a/src/modules/lazy.rs b/src/modules/lazy.rs
@@ -324,7 +324,7 @@ mod tests {
     #[cfg(feature = "cpu")]
     #[test]
     fn test_lazy_exec_with_range() {
-        use crate::{ExecNow, Run};
+        use crate::{ExecNow, Run, HostPtr};
 
         let device = CPU::<Lazy<Base>>::new();
         let mut out: Buffer<i32, _, ()> = device.retrieve(4, ());
@@ -357,7 +357,7 @@ mod tests {
     #[cfg(feature = "cpu")]
     #[test]
     fn test_lazy_exec_last_n() {
-        use crate::{ExecNow, Run};
+        use crate::{ExecNow, Run, HostPtr};
 
         let device = CPU::<Lazy<Base>>::new();
         let mut out: Buffer<i32, _, ()> = device.retrieve(4, ());

diff --git a/src/modules/lazy/wrapper.rs b/src/modules/lazy/wrapper.rs
@@ -1,16 +1,30 @@
 use core::{ops::{Deref, DerefMut}, marker::PhantomData};
 
-use crate::{HasId, PtrType, WrappedData, Lazy};
+use crate::{HasId, PtrType, WrappedData, Lazy, HostPtr, Id, ShallowCopy};
 
+#[derive(Debug, Default)]
 pub struct LazyWrapper<Data, T> {
     data: Option<Data>,
+    id: Option<Id>,
     _pd: PhantomData<T>
 }
 
+
+impl<Mods: WrappedData> WrappedData for Lazy<Mods> {
+    type Wrap<T, Base: HasId + PtrType> = LazyWrapper<Mods::Wrap<T, Base>, T>;
+    // type Wrap<T, Base: HasId + PtrType> = Mods::Wrap<T, Base>;
+
+    #[inline]
+    fn wrap_in_base<T, Base: HasId + PtrType>(&self, base: Base) -> Self::Wrap<T, Base> {
+        todo!()
+        // self.modules.wrap_in_base(base)
+    }
+}
+
 impl<Data: HasId, T> HasId for LazyWrapper<Data, T> {
     #[inline]
     fn id(&self) -> crate::Id {
-        self.data.as_ref().unwrap().id()
+        self.id.unwrap()
     }
 }
 
@@ -47,12 +61,25 @@ impl<Data: DerefMut<Target = [T]>, T> DerefMut for LazyWrapper<Data, T> {
     }
 }
 
-impl<Mods: WrappedData> WrappedData for Lazy<Mods> {
-    // type Wrap<T, Base: HasId + PtrType> = LazyWrapper<Mods::Wrap<T, Base>, T>;
-    type Wrap<T, Base: HasId + PtrType> = Mods::Wrap<T, Base>;
+impl<T, Data: HostPtr<T>> HostPtr<T> for LazyWrapper<Data, T> {
+    #[inline]
+    fn ptr(&self) -> *const T {
+        self.data.as_ref().unwrap().ptr()
+    }
 
     #[inline]
-    fn wrap_in_base<T, Base: HasId + PtrType>(&self, base: Base) -> Self::Wrap<T, Base> {
-        self.modules.wrap_in_base(base)
+    fn ptr_mut(&mut self) -> *mut T {
+        self.data.as_mut().unwrap().ptr_mut()
+    }
+}
+
+impl<Data: ShallowCopy, T> ShallowCopy for LazyWrapper<Data, T> {
+    #[inline]
+    unsafe fn shallow(&self) -> Self {
+        LazyWrapper {
+            id: self.id,
+            data: self.data.as_ref().map(|data| data.shallow()),
+            _pd: PhantomData
+        }
     }
 }
diff --git a/src/modules/mod.rs b/src/modules/mod.rs
@@ -28,7 +28,7 @@ mod fork;
 pub use fork::*;
 
 #[cfg(not(feature = "no-std"))]
-use crate::{flag::AllocFlag, Buffer, Device, HasId, HashLocation, Id, PtrConv, Shape, UniqueId};
+use crate::{flag::AllocFlag, Buffer, Device, HasId, HashLocation, Id, PtrConv, Shape, UniqueId, ShallowCopy};
 #[cfg(not(feature = "no-std"))]
 use core::{any::Any, hash::BuildHasher};
 
@@ -49,8 +49,11 @@ pub(crate) unsafe fn register_buf<T, D, S>(
 ) where
     T: 'static,
     D: Device + PtrConv + 'static,
+    // D::Data<T, S>: ShallowCopy,
     S: Shape,
 {
+
+    // buf.data
     let wrapped_data = D::convert::<T, S, T, S>(&buf.data, AllocFlag::Wrapper);
     let buf = Buffer {
         data: wrapped_data,