CPU implementation for upsample-nearest2d. (huggingface#339)

philparzer · Aug 7, 2023 · d0d7010 · d0d7010
1 parent fc265d9
commit d0d7010
Showing 5 changed files with 61 additions and 7 deletions.
diff --git a/candle-core/src/backend.rs b/candle-core/src/backend.rs
@@ -38,6 +38,7 @@ pub trait BackendStorage: Sized {
     ) -> Result<Self>;
 
     fn avg_pool2d(&self, _: &Layout, _: (usize, usize), _: (usize, usize)) -> Result<Self>;
+    fn upsample_nearest2d(&self, _: &Layout, _: usize, _: usize) -> Result<Self>;
 
     fn gather(&self, _: &Layout, _: &Self, _: &Layout, _: usize) -> Result<Self>;
     fn scatter_add(

diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs
@@ -672,6 +672,43 @@ impl Map1 for AvgPool2D {
     }
 }
 
+struct UpsampleNearest2D(usize, usize);
+
+impl Map1 for UpsampleNearest2D {
+    fn f<T: WithDType>(&self, src: &[T], layout: &Layout) -> Result<Vec<T>> {
+        // TODO: Specialized implementation for the case 2*h, 2*w?
+        let (dst_h, dst_w) = (self.0, self.1);
+        let (b_sz, c, src_h, src_w) = layout.shape().dims4()?;
+        let stride = layout.stride();
+        let (stride_h, stride_w) = (stride[2], stride[3]);
+        let src_index = layout.start_offset();
+        let scale_h = src_h as f64 / dst_h as f64;
+        let scale_w = src_w as f64 / dst_w as f64;
+        let mut dst = vec![T::zero(); b_sz * c * dst_h * dst_w];
+        let src_h_idxs = (0..src_h)
+            .map(|h_idx| usize::min(src_h - 1, (h_idx as f64 * scale_h) as usize))
+            .collect::<Vec<_>>();
+        let src_w_idxs = (0..src_w)
+            .map(|w_idx| usize::min(src_w - 1, (w_idx as f64 * scale_w) as usize))
+            .collect::<Vec<_>>();
+        for b_idx in 0..b_sz {
+            let dst = &mut dst[b_idx * c * dst_h * dst_w..];
+            let src_index = src_index + b_idx * stride[0];
+            for c_idx in 0..c {
+                let dst = &mut dst[c_idx * dst_h * dst_w..];
+                let src_index = src_index + c_idx * stride[1];
+                for (h_idx, src_h_idx) in src_h_idxs.iter().enumerate() {
+                    for (w_idx, src_w_idx) in src_w_idxs.iter().enumerate() {
+                        let src_index = src_index + src_h_idx * stride_h + src_w_idx * stride_w;
+                        dst[h_idx * dst_w + w_idx] = src[src_index]
+                    }
+                }
+            }
+        }
+        Ok(dst)
+    }
+}
+
 struct Gather<'a, I: IntDType> {
     ids: &'a [I],
     ids_l: &'a Layout,
@@ -1577,6 +1614,10 @@ impl BackendStorage for CpuStorage {
         AvgPool2D(kernel_size, stride).map(self, layout)
     }
 
+    fn upsample_nearest2d(&self, layout: &Layout, h: usize, w: usize) -> Result<Self> {
+        UpsampleNearest2D(h, w).map(self, layout)
+    }
+
     fn elu(&self, layout: &Layout, alpha: f64) -> Result<Self> {
         // TODO: Have some generic map for functions that apply on num_traits::Float elements.
         match self {

diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs
@@ -1385,6 +1385,10 @@ impl BackendStorage for CudaStorage {
         todo!()
     }
 
+    fn upsample_nearest2d(&self, _: &Layout, _: usize, _: usize) -> Result<Self> {
+        todo!()
+    }
+
     fn index_select(&self, ids: &Self, l: &Layout, ids_l: &Layout, dim: usize) -> Result<Self> {
         let device = self.device().clone();
         let slice = IndexSelect(ids, ids_l, dim).map(&self.slice, &device, l)?;

diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs
@@ -123,6 +123,10 @@ impl crate::backend::BackendStorage for CudaStorage {
     fn avg_pool2d(&self, _: &Layout, _: (usize, usize), _: (usize, usize)) -> Result<Self> {
         Err(Error::NotCompiledWithCudaSupport)
     }
+
+    fn upsample_nearest2d(&self, _: &Layout, _: usize, _: usize) -> Result<Self> {
+        Err(Error::NotCompiledWithCudaSupport)
+    }
 }
 
 impl crate::backend::BackendDevice for CudaDevice {

diff --git a/candle-core/src/storage.rs b/candle-core/src/storage.rs
@@ -284,13 +284,17 @@ impl Storage {
         }
     }
 
-    pub(crate) fn upsample_nearest2d(
-        &self,
-        _layout: &Layout,
-        _h: usize,
-        _w: usize,
-    ) -> Result<Self> {
-        todo!()
+    pub(crate) fn upsample_nearest2d(&self, layout: &Layout, h: usize, w: usize) -> Result<Self> {
+        match self {
+            Storage::Cpu(storage) => {
+                let storage = storage.upsample_nearest2d(layout, h, w)?;
+                Ok(Self::Cpu(storage))
+            }
+            Self::Cuda(storage) => {
+                let storage = storage.upsample_nearest2d(layout, h, w)?;
+                Ok(Self::Cuda(storage))
+            }
+        }
     }
 
     pub(crate) fn where_cond(