Merge branch 'main' into cycle-detection

Author: elftausend

src/buffer.rs

@@ -12,7 +12,7 @@ use crate::CPU;
 use crate::{
     flag::AllocFlag, shape::Shape, Alloc, Base, ClearBuf, CloneBuf, Device, DevicelessAble, HasId,
     IsShapeIndep, OnDropBuffer, OnNewBuffer, PtrType, Read, ReplaceBuf, ShallowCopy, Unit,
-    WrappedData, WriteBuf, ZeroGrad,
+    WrappedCopy, WrappedData, WriteBuf, ZeroGrad,
 };
 
 pub use self::num::Num;
@@ -479,11 +479,14 @@ impl<'a, T: Unit, D: Device, S: Shape> Buffer<'a, T, D, S> {
     pub fn to_dims<O: Shape>(self) -> Buffer<'a, T, D, O>
     where
         D: crate::ToDim<T, S, O>,
-        D::Data<T, S>: ShallowCopy,
+        D::Data<T, S>: WrappedCopy<Base = D::Base<T, S>>,
+        D::Base<T, S>: ShallowCopy,
     {
+        let base = unsafe { (*self).shallow() };
+        let data = self.data.wrapped_copy(base);
         let buf = ManuallyDrop::new(self);
 
-        let mut data = buf.device().to_dim(unsafe { buf.data.shallow() });
+        let mut data = buf.device().to_dim(data);
         unsafe { data.set_flag(AllocFlag::None) };
 
         Buffer {

src/devices/cpu/cpu_ptr.rs

@@ -7,7 +7,7 @@ use core::{
 
 use std::alloc::handle_alloc_error;
 
-use crate::{flag::AllocFlag, HasId, HostPtr, Id, PtrType, ShallowCopy};
+use crate::{flag::AllocFlag, HasId, HostPtr, Id, PtrType, ShallowCopy, WrappedCopy};
 
 /// The pointer used for `CPU` [`Buffer`](crate::Buffer)s
 #[derive(Debug)]
@@ -229,6 +229,15 @@ impl<T> PtrType for CPUPtr<T> {
     }
 }
 
+impl<T> WrappedCopy for CPUPtr<T> {
+    type Base = Self;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        to_wrap
+    }
+}
+
 impl<T> ShallowCopy for CPUPtr<T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {

src/devices/cuda/cuda_ptr.rs

@@ -1,5 +1,5 @@
 use super::api::{cu_read, cufree, cumalloc, CudaResult};
-use crate::{flag::AllocFlag, HasId, Id, PtrType, ShallowCopy};
+use crate::{flag::AllocFlag, HasId, Id, PtrType, ShallowCopy, WrappedCopy};
 use core::marker::PhantomData;
 
 /// The pointer used for `CUDA` [`Buffer`](crate::Buffer)s
@@ -76,6 +76,15 @@ impl<T> Drop for CUDAPtr<T> {
     }
 }
 
+impl<T> WrappedCopy for CUDAPtr<T> {
+    type Base = Self;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        to_wrap
+    }
+}
+
 impl<T> ShallowCopy for CUDAPtr<T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {

src/devices/opencl/cl_ptr.rs

@@ -8,7 +8,7 @@ use crate::HostPtr;
 
 use min_cl::api::release_mem_object;
 
-use crate::{flag::AllocFlag, HasId, Id, PtrType, ShallowCopy};
+use crate::{flag::AllocFlag, HasId, Id, PtrType, ShallowCopy, WrappedCopy};
 
 /// The pointer used for `OpenCL` [`Buffer`](crate::Buffer)s
 #[derive(Debug, PartialEq, Eq)]
@@ -59,6 +59,15 @@ impl<T> CLPtr<T> {
     }
 }
 
+impl<T> WrappedCopy for CLPtr<T> {
+    type Base = Self;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        to_wrap
+    }
+}
+
 impl<T> ShallowCopy for CLPtr<T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {

src/devices/stack_array.rs

@@ -1,6 +1,6 @@
 use core::ops::{Deref, DerefMut};
 
-use crate::{shape::Shape, HasId, HostPtr, PtrType, ShallowCopy};
+use crate::{shape::Shape, HasId, HostPtr, PtrType, ShallowCopy, WrappedCopy};
 
 /// A possibly multi-dimensional array allocated on the stack.
 /// It uses `S:`[`Shape`] to get the type of the array.
@@ -137,6 +137,16 @@ impl<S: Shape, T> HostPtr<T> for StackArray<S, T> {
     }
 }
 
+
+impl<S: Shape, T> WrappedCopy for StackArray<S, T> {
+    type Base = Self;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        to_wrap
+    }
+}
+
 impl<S: Shape, T> ShallowCopy for StackArray<S, T>
 where
     S::ARR<T>: Copy,

src/devices/vulkan/vk_array.rs

@@ -9,7 +9,7 @@ use core::{
 };
 use std::rc::Rc;
 
-use crate::{flag::AllocFlag, HasId, HostPtr, PtrType, ShallowCopy};
+use crate::{flag::AllocFlag, HasId, HostPtr, PtrType, ShallowCopy, WrappedCopy};
 
 use super::{context::Context, submit_and_wait};
 
@@ -228,6 +228,15 @@ impl<T> VkArray<T> {
     }
 }
 
+impl<T> WrappedCopy for VkArray<T> {
+    type Base = Self;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        to_wrap
+    }
+}
+
 impl<T> ShallowCopy for VkArray<T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {

src/lib.rs

@@ -180,6 +180,11 @@ pub trait Unit {} // useful for Sync and Send or 'static
 
 impl<T> Unit for T {}
 
+pub trait WrappedCopy {
+    type Base;
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self;
+}
+
 /// Used to shallow-copy a pointer. Use is discouraged.
 pub trait ShallowCopy {
     /// # Safety

src/modules/autograd/wrapper.rs

@@ -1,6 +1,6 @@
 use core::marker::PhantomData;
 
-use crate::{flag::AllocFlag, Autograd, HasId, PtrType, ShallowCopy, WrappedData};
+use crate::{flag::AllocFlag, Autograd, HasId, PtrType, ShallowCopy, WrappedCopy, WrappedData};
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
 pub struct ReqGradWrapper<Data, T> {
@@ -74,6 +74,22 @@ impl<Data: PtrType, T> PtrType for ReqGradWrapper<Data, T> {
     }
 }
 
+impl<Data, T> WrappedCopy for ReqGradWrapper<Data, T>
+where
+    Data: WrappedCopy<Base = T>,
+{
+    type Base = T;
+
+    #[inline]
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        Self {
+            requires_grad: self.requires_grad,
+            data: self.data.wrapped_copy(to_wrap),
+            _pd: PhantomData,
+        }
+    }
+}
+
 impl<Data, T> ShallowCopy for ReqGradWrapper<Data, T>
 where
     Data: ShallowCopy,

src/modules/lazy/wrapper.rs

@@ -6,7 +6,9 @@ use core::{
     ops::{Deref, DerefMut},
 };
 
-use crate::{flag::AllocFlag, HasId, HostPtr, Lazy, PtrType, ShallowCopy, WrappedData};
+use crate::{
+    flag::AllocFlag, HasId, HostPtr, Lazy, PtrType, ShallowCopy, WrappedCopy, WrappedData,
+};
 
 #[derive(Debug, Default)]
 pub struct LazyWrapper<Data, T> {
@@ -102,6 +104,24 @@ impl<T, Data: HostPtr<T>> HostPtr<T> for LazyWrapper<Data, T> {
     }
 }
 
+impl<Data, T> WrappedCopy for LazyWrapper<Data, T>
+where
+    Data: WrappedCopy<Base = T>,
+{
+    type Base = T;
+
+    fn wrapped_copy(&self, to_wrap: Self::Base) -> Self {
+        LazyWrapper {
+            maybe_data: match &self.maybe_data {
+                MaybeData::Data(data) => MaybeData::Data(data.wrapped_copy(to_wrap)),
+                MaybeData::Id(id) => MaybeData::Id(*id),
+                MaybeData::None => unimplemented!(),
+            },
+            _pd: PhantomData,
+        }
+    }
+}
+
 impl<Data: ShallowCopy, T> ShallowCopy for LazyWrapper<Data, T> {
     #[inline]
     unsafe fn shallow(&self) -> Self {