From dd64465899f4b58628642b406c465d35ddfe8f79 Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Thu, 24 Aug 2023 10:16:37 +0100
Subject: [PATCH] Add a test for conv2d with padding + bugfix the random number
 generation on cuda. (#578)

* Add a test for conv2d with padding.

* Cosmetic changes.

* Bugfix the rand function on the cuda backend.
---
 candle-core/examples/cuda_basics.rs |  3 +++
 candle-core/src/cuda_backend.rs     | 17 +++++++++--------
 candle-core/tests/conv_tests.rs     | 12 ++++++++++++
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/candle-core/examples/cuda_basics.rs b/candle-core/examples/cuda_basics.rs
index ac435488..6a3aaacc 100644
--- a/candle-core/examples/cuda_basics.rs
+++ b/candle-core/examples/cuda_basics.rs
@@ -9,6 +9,9 @@ use candle_core::{Device, Tensor};
 
 fn main() -> Result<()> {
     let device = Device::new_cuda(0)?;
+    let t = Tensor::rand(-1f32, 1f32, 96, &device)?;
+    println!("{t}");
+
     let t = Tensor::randn(0f32, 1f32, (2, 4, 96, 96), &device)?;
     let w = Tensor::randn(0f32, 1f32, (320, 4, 3, 3), &device)?;
     let res = t.conv2d(&w, 1, 1, 1)?;
diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs
index 9809bb4a..fbe55266 100644
--- a/candle-core/src/cuda_backend.rs
+++ b/candle-core/src/cuda_backend.rs
@@ -295,10 +295,12 @@ impl BackendDevice for CudaDevice {
                 CudaStorageSlice::F64(data)
             }
         };
-        if lo != 0.0 || up != 1.0 {
+        let slice = if lo == 0. && up == 1.0 {
+            slice
+        } else {
             let layout = Layout::contiguous(shape);
-            Affine(up - lo, lo).map(&slice, self, &layout)?;
-        }
+            Affine(up - lo, lo).map(&slice, self, &layout)?
+        };
         Ok(CudaStorage {
             slice,
             device: self.clone(),
@@ -955,7 +957,7 @@ impl<'a> Map2 for Conv1D<'a> {
         } else if dims.len() == 2 {
             [&[1], dims, &[1], inp_l.stride(), k_l.dims(), k_l.stride()].concat()
         } else {
-            panic!("unexpected input shape for conv1d {dims:?}")
+            crate::bail!("unexpected input shape for conv1d {dims:?}")
         };
         let ds = dev.htod_copy(ds).w()?;
         let params = (el, l_out, p.stride, p.padding, &ds, inp, k, &out);
@@ -993,7 +995,7 @@ impl<'a> Map2 for Conv2D<'a> {
         let ds = if dims.len() == 4 {
             [dims, inp_l.stride(), k_l.dims(), k_l.stride()].concat()
         } else {
-            panic!("unexpected input shape for conv1d {dims:?}")
+            crate::bail!("unexpected input shape for conv2d {dims:?}")
         };
         let ds = dev.htod_copy(ds).w()?;
         let params = (el, out_w, out_h, p.stride, p.padding, &ds, inp, k, &out);
@@ -1030,7 +1032,7 @@ impl Map1 for Pool2D {
         let ds = if dims.len() == 4 {
             [dims, inp_l.stride()].concat()
         } else {
-            panic!("unexpected input shape for conv1d {dims:?}")
+            crate::bail!("unexpected input shape for pool {dims:?}")
         };
         let el = shape.elem_count();
         let out_w = (dims[2] - self.w_k) / self.w_stride + 1;
@@ -1076,7 +1078,7 @@ impl Map1 for UpsampleNearest2D {
         let ds = if dims.len() == 4 {
             [dims, inp_l.stride()].concat()
         } else {
-            panic!("unexpected input shape for conv1d {dims:?}")
+            crate::bail!("unexpected input shape for upsample {dims:?}")
         };
         let (out_w, out_h) = (self.0, self.1);
         let dst_el = out_w * out_h * dims[0] * dims[1];
@@ -1622,7 +1624,6 @@ impl BackendStorage for CudaStorage {
                     .map_err(crate::Error::wrap)?;
                 S::F16(out)
             }
-
             (S::F32(inp), S::F32(k)) => {
                 let inp = &inp.slice(inp_l.start_offset()..);
                 let k = &k.slice(kernel_l.start_offset()..);
diff --git a/candle-core/tests/conv_tests.rs b/candle-core/tests/conv_tests.rs
index d09fa344..1177c1cf 100644
--- a/candle-core/tests/conv_tests.rs
+++ b/candle-core/tests/conv_tests.rs
@@ -149,6 +149,18 @@ fn conv2d_small(dev: &Device) -> Result<()> {
         test_utils::to_vec1_round(&res.flatten_all()?, 4)?,
         [0.164, -0.0111, -0.1742, 2.6437, -2.0268, 1.1823, 3.2855, -1.0324, 0.2539]
     );
+    let res = t.conv2d(&w, 2, 1, 1)?;
+    assert_eq!(res.dims(), [1, 1, 7, 7]);
+    assert_eq!(
+        test_utils::to_vec1_round(&res.flatten_all()?, 4)?,
+        [
+            0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
+            0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1640, -0.0111, -0.1742, 0.0000, 0.0000,
+            0.0000, 0.0000, 2.6437, -2.0268, 1.1823, 0.0000, 0.0000, 0.0000, 0.0000, 3.2855,
+            -1.0324, 0.2539, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
+            0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000
+        ]
+    );
     Ok(())
 }