Finished up with template tensor_ops

dfdx-core/src/tensor/webgpu/device.rs

@@ -37,8 +37,8 @@ impl core::ops::Deref for Buffer {
 }
 
 impl Buffer {
-    pub(crate) fn size(&self) -> u64 {
-        self.size as u64
+    pub(crate) fn size(&self) -> usize {
+        self.size
     }
 
     #[allow(unused)]
@@ -57,16 +57,16 @@ impl Buffer {
         let (sender, receiver) = std::sync::mpsc::channel();
         let buffer = dev.create_buffer(&BufferDescriptor {
             label: None,
-            size: self.size(),
+            size: self.size() as u64,
             usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
             mapped_at_creation: false,
         });
         {
             let mut encoder = dev.create_command_encoder(&Default::default());
-            encoder.copy_buffer_to_buffer(&self.data, 0, &buffer, 0, self.size());
+            encoder.copy_buffer_to_buffer(&self.data, 0, &buffer, 0, self.size() as u64);
             queue.submit(Some(encoder.finish()));
         }
-        let slice = buffer.slice(..self.size());
+        let slice = buffer.slice(..self.size() as u64);
         slice.map_async(wgpu::MapMode::Read, move |_| {
             sender.send(()).unwrap();
         });
@@ -190,11 +190,11 @@ impl<E> Clone for CachableBuffer<E> {
                 let mut encoder = self.dev.create_command_encoder(&Default::default());
                 let bfr = self.dev.create_buffer(&BufferDescriptor {
                     label: None,
-                    size: round_to_buffer_alignment(self.data.size()),
+                    size: round_to_buffer_alignment(self.data.size() as u64),
                     usage: BufferUsages::COPY_SRC | BufferUsages::COPY_DST,
                     mapped_at_creation: false,
                 });
-                encoder.copy_buffer_to_buffer(&self.data, 0, &bfr, 0, self.data.size());
+                encoder.copy_buffer_to_buffer(&self.data, 0, &bfr, 0, self.data.size() as u64);
                 (
                     encoder,
                     Buffer {
@@ -205,7 +205,7 @@ impl<E> Clone for CachableBuffer<E> {
             },
             |bfr| {
                 let mut encoder = self.dev.create_command_encoder(&Default::default());
-                encoder.copy_buffer_to_buffer(&self.data, 0, &bfr, 0, self.data.size());
+                encoder.copy_buffer_to_buffer(&self.data, 0, &bfr, 0, self.data.size() as u64);
                 (encoder, bfr)
             },
         );

dfdx-core/src/tensor_ops/abs/mod.rs

@@ -3,6 +3,9 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernel;
+
 use super::ops::{try_unary_op, UnaryKernel};
 use crate::{shapes::*, tensor::*};
 

dfdx-core/src/tensor_ops/abs/webgpu_kernel.rs

@@ -0,0 +1,28 @@
+use std::borrow::Cow;
+
+use crate::prelude::{ops::UnaryKernel, Dtype, Webgpu};
+
+impl<E: Dtype> UnaryKernel<super::AbsKernelOp, E> for Webgpu {
+    const BACKWARD_WITHOUT_INP: bool = false;
+
+    const BACKWARD_WITHOUT_DATA: bool = false;
+
+    fn forward<S: crate::prelude::Shape>(
+        &self,
+        op: super::AbsKernelOp,
+        inp: Cow<crate::prelude::Tensor<S, E, Self>>,
+    ) -> Result<crate::prelude::Tensor<S, E, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn backward<S: crate::prelude::Shape>(
+        &self,
+        op: super::AbsKernelOp,
+        inp: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_inp: &mut Self::Vec,
+        out: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_out: &Self::Vec,
+    ) -> Result<(), crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/accurate_gelu/mod.rs

@@ -3,6 +3,9 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernel;
+
 use super::ops::{try_unary_op, UnaryKernel};
 use crate::{shapes::*, tensor::*};
 

dfdx-core/src/tensor_ops/accurate_gelu/webgpu_kernel.rs

@@ -0,0 +1,28 @@
+use std::borrow::Cow;
+
+use crate::prelude::{ops::UnaryKernel, Dtype, Webgpu};
+
+impl<E: Dtype> UnaryKernel<super::AccurateGeLUKernelOp, E> for Webgpu {
+    const BACKWARD_WITHOUT_INP: bool = false;
+
+    const BACKWARD_WITHOUT_DATA: bool = false;
+
+    fn forward<S: crate::prelude::Shape>(
+        &self,
+        op: super::AccurateGeLUKernelOp,
+        inp: Cow<crate::prelude::Tensor<S, E, Self>>,
+    ) -> Result<crate::prelude::Tensor<S, E, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn backward<S: crate::prelude::Shape>(
+        &self,
+        op: super::AccurateGeLUKernelOp,
+        inp: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_inp: &mut Self::Vec,
+        out: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_out: &Self::Vec,
+    ) -> Result<(), crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/add/webgpu_kernel.rs

@@ -1,5 +1,4 @@
-extern crate alloc;
-use alloc::borrow::Cow;
+use std::borrow::Cow;
 
 use crate::prelude::{
     ops::{BinaryKernel, UnaryKernel},

dfdx-core/src/tensor_ops/axpy/mod.rs

@@ -7,6 +7,9 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernel;
+
 /// Elementwise `a * alpha + b * beta`.
 ///
 /// See [Tensor::axpy] for in place version.

dfdx-core/src/tensor_ops/axpy/webgpu_kernel.rs

@@ -0,0 +1,13 @@
+use crate::prelude::{Dtype, Webgpu};
+
+impl<E: Dtype> super::AxpyKernel<E> for Webgpu {
+    fn forward(
+        &self,
+        a: &mut Self::Vec,
+        alpha: E,
+        b: &Self::Vec,
+        beta: E,
+    ) -> Result<(), crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/bce/mod.rs

@@ -3,6 +3,9 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernel;
+
 use super::ops::{try_binary_op, BinaryKernel};
 use crate::{shapes::*, tensor::*};
 

dfdx-core/src/tensor_ops/bce/webgpu_kernel.rs

@@ -0,0 +1,27 @@
+use crate::prelude::{ops::BinaryKernel, Dtype, Webgpu};
+use std::borrow::Cow;
+
+impl<E: Dtype> BinaryKernel<super::BCEKernelOp, E> for Webgpu {
+    const BACKWARD_WITHOUT_DATA: bool = false;
+
+    fn forward<S: crate::prelude::Shape>(
+        &self,
+        op: super::BCEKernelOp,
+        lhs: Cow<crate::prelude::Tensor<S, E, Self>>,
+        rhs: Cow<crate::prelude::Tensor<S, E, Self>>,
+    ) -> Result<crate::prelude::Tensor<S, E, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn backward<S: crate::prelude::Shape>(
+        &self,
+        op: super::BCEKernelOp,
+        lhs: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_lhs: &mut Self::Vec,
+        rhs: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_rhs: &mut Self::Vec,
+        grad_out: &Self::Vec,
+    ) -> Result<(), crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/boolean/mod.rs

@@ -3,6 +3,9 @@ mod cpu_kernels;
 #[cfg(feature = "cuda")]
 mod cuda_kernels;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernels;
+
 use crate::{
     prelude::{OnesTensor, Tensor, ZerosTensor},
     shapes::*,

dfdx-core/src/tensor_ops/boolean/webgpu_kernels.rs

@@ -0,0 +1,34 @@
+use crate::prelude::Webgpu;
+
+impl super::BooleanKernel for Webgpu {
+    fn not<S: crate::prelude::Shape>(
+        &self,
+        inp: &crate::prelude::Tensor<S, bool, Self>,
+    ) -> Result<crate::prelude::Tensor<S, bool, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn and<S: crate::prelude::Shape>(
+        &self,
+        lhs: &crate::prelude::Tensor<S, bool, Self>,
+        rhs: &crate::prelude::Tensor<S, bool, Self>,
+    ) -> Result<crate::prelude::Tensor<S, bool, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn or<S: crate::prelude::Shape>(
+        &self,
+        lhs: &crate::prelude::Tensor<S, bool, Self>,
+        rhs: &crate::prelude::Tensor<S, bool, Self>,
+    ) -> Result<crate::prelude::Tensor<S, bool, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn xor<S: crate::prelude::Shape>(
+        &self,
+        lhs: &crate::prelude::Tensor<S, bool, Self>,
+        rhs: &crate::prelude::Tensor<S, bool, Self>,
+    ) -> Result<crate::prelude::Tensor<S, bool, Self>, crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/clamp/mod.rs

@@ -3,6 +3,9 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernel;
+
 use super::ops::{try_unary_op, UnaryKernel};
 use crate::{shapes::*, tensor::*};
 

dfdx-core/src/tensor_ops/clamp/webgpu_kernel.rs

@@ -0,0 +1,28 @@
+use std::borrow::Cow;
+
+use crate::prelude::{ops::UnaryKernel, Dtype, Webgpu};
+
+impl<E: Dtype> UnaryKernel<super::ClampKernelOp<E>, E> for Webgpu {
+    const BACKWARD_WITHOUT_INP: bool = false;
+
+    const BACKWARD_WITHOUT_DATA: bool = false;
+
+    fn forward<S: crate::prelude::Shape>(
+        &self,
+        op: super::ClampKernelOp<E>,
+        inp: Cow<crate::prelude::Tensor<S, E, Self>>,
+    ) -> Result<crate::prelude::Tensor<S, E, Self>, crate::prelude::Error> {
+        todo!()
+    }
+
+    fn backward<S: crate::prelude::Shape>(
+        &self,
+        op: super::ClampKernelOp<E>,
+        inp: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_inp: &mut Self::Vec,
+        out: &impl crate::prelude::Tensorlike<S, E, Self>,
+        grad_out: &Self::Vec,
+    ) -> Result<(), crate::prelude::Error> {
+        todo!()
+    }
+}

dfdx-core/src/tensor_ops/cmp/mod.rs

@@ -7,6 +7,9 @@ mod cpu_kernels;
 #[cfg(feature = "cuda")]
 mod cuda_kernels;
 
+#[cfg(feature = "webgpu")]
+mod webgpu_kernels;
+
 pub trait CmpKernel<Op, E>: Storage<E> + Storage<bool> {
     fn forward<S: Shape, T>(
         &self,