Remove custom wgpu unfold kernel since doesn't handle complex cases w…

…ithout performance benefits
tracel-ai · nathanielsimard · Oct 22, 2023 · Oct 15, 2023 · Oct 15, 2023 · Oct 15, 2023
commit 7d49d2cd3e96a313d435fd8485ede5a827f9cb34
diff --git a/burn-tensor/src/tensor/ops/modules/unfold.rs b/burn-tensor/src/tensor/ops/modules/unfold.rs
@@ -16,7 +16,7 @@ use super::{ConvOptions, UnfoldOptions};
 pub(crate) fn create_unfolding_weight<B: Backend>(
     in_channels: usize,
     kernel_size: [usize; 2],
-    devide: &B::Device,
+    device: &B::Device,
 ) -> B::TensorPrimitive<4> {
     let shape = Shape::new([
         in_channels * kernel_size[0] * kernel_size[1],
@@ -53,7 +53,7 @@ pub(crate) fn create_unfolding_weight<B: Backend>(
         }
     }
 
-    B::from_data(Data::new(weight, shape), devide)
+    B::from_data(Data::new(weight, shape), device)
 }
 
 /// Compute the unfold4d operation using the conv2d operations.

diff --git a/burn-wgpu/src/kernel/mod.rs b/burn-wgpu/src/kernel/mod.rs
@@ -25,8 +25,6 @@ pub mod matmul;
 pub mod pool;
 /// Pseudo-random number generator kernels
 pub mod prng;
-/// Unfold kernel
-pub mod unfold;
 
 pub(crate) use cat::*;
 pub(crate) use comparison::*;

diff --git a/burn-wgpu/src/kernel/unfold/mod.rs b/burn-wgpu/src/kernel/unfold/mod.rs
diff --git a/burn-wgpu/src/kernel/unfold/unfold4d.rs b/burn-wgpu/src/kernel/unfold/unfold4d.rs
diff --git a/burn-wgpu/src/ops/module_ops.rs b/burn-wgpu/src/ops/module_ops.rs
@@ -1,9 +1,5 @@
-use burn_tensor::{
-    ops::{
-        ConvOptions, ConvTransposeOptions, MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps,
-        TensorOps, UnfoldOptions,
-    },
-    Shape,
+use burn_tensor::ops::{
+    ConvOptions, ConvTransposeOptions, MaxPool2dBackward, MaxPool2dWithIndices, ModuleOps,
 };
 
 use crate::{
@@ -37,19 +33,6 @@ where
         kernel::conv::conv_transpose2d(x, weight, bias, options)
     }
 
-    fn unfold4d(
-        x: FloatTensor<Self, 4>,
-        kernel_size: [usize; 2],
-        options: UnfoldOptions,
-    ) -> FloatTensor<Self, 3> {
-        let unfold_intermediate = kernel::unfold::unfold4d(x, kernel_size, options);
-        let [batch_size, channels_out, out_height, out_width] = unfold_intermediate.shape.dims;
-        <WgpuBackend<G, F, I> as TensorOps<WgpuBackend<G, F, I>>>::reshape(
-            unfold_intermediate,
-            Shape::new([batch_size, channels_out, out_height * out_width]),
-        )
-    }
-
     fn avg_pool2d(
         x: FloatTensor<Self, 4>,
         kernel_size: [usize; 2],

diff --git a/burn-wgpu/src/template/unfold/fill_weights.wgsl b/burn-wgpu/src/template/unfold/fill_weights.wgsl