Removed device from CopyToDevice for metal

jafioti · Apr 26, 2024 · c49507e · c49507e
1 parent 194f221
commit c49507e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 11 deletions.
diff --git a/crates/luminal_metal/src/lib.rs b/crates/luminal_metal/src/lib.rs
@@ -74,7 +74,7 @@ impl Data for MetalBuffer {
     }
 }
 
-pub trait MetalFloat: Copy + Debug + PartialEq + 'static {
+pub trait MetalFloat: Copy + Debug + PartialEq + 'static + Default {
     fn to_f32(self) -> f32;
     fn from_f32(a: f32) -> Self;
     fn is_f32() -> bool;

diff --git a/crates/luminal_metal/src/prim.rs b/crates/luminal_metal/src/prim.rs
@@ -52,16 +52,10 @@ impl<T: MetalFloat> Operator for MetalCopyToDevice<T> {
 }
 
 /// Copy a tensor from the GPU
-#[derive(Clone)]
-pub struct MetalCopyFromDevice<T>(Device, PhantomData<T>);
+#[derive(Clone, Default)]
+pub struct MetalCopyFromDevice<T>(PhantomData<T>);
 crate::debug_type!(MetalCopyFromDevice<T>);
 
-impl<T> MetalCopyFromDevice<T> {
-    pub fn new(dev: Device) -> Self {
-        Self(dev, Default::default())
-    }
-}
-
 impl<T: MetalFloat> Operator for MetalCopyFromDevice<T> {
     fn process(&mut self, mut inp: Vec<(InputTensor, ShapeTracker)>) -> Vec<Tensor> {
         if inp[0].0.borrowed().is::<Vec<f32>>() {
@@ -1620,7 +1614,7 @@ impl<T: MetalFloat + 'static> Compiler for PrimitiveCompiler<T> {
             {
                 let (input_order, output_order, shape) = edge_weight.as_data().unwrap();
                 let copy_from_node = graph
-                    .add_op(MetalCopyFromDevice::<T>::new(dev.clone()))
+                    .add_op(MetalCopyFromDevice::<T>::default())
                     .input(source, output_order, shape)
                     .finish();
                 graph.add_edge(
@@ -1663,7 +1657,7 @@ impl<T: MetalFloat + 'static> Compiler for PrimitiveCompiler<T> {
             } else {
                 // Create copy node
                 let copy_node = graph
-                    .add_op(MetalCopyFromDevice::<T>::new(dev.clone()))
+                    .add_op(MetalCopyFromDevice::<T>::default())
                     .input(output_node, 0, output_shape)
                     .finish();