conv backward contiguous, map_index for remap

NVIDIA · Dec 15, 2020 · bacf460 · bacf460
1 parent 2114cfa
commit bacf460
Show file tree

Hide file tree

Showing 5 changed files with 30 additions and 19 deletions.
diff --git a/pybind/extern.hpp b/pybind/extern.hpp
@@ -65,7 +65,7 @@ ConvolutionForwardCPU(at::Tensor const &in_feat,                         //
 template <typename coordinate_type>
 std::pair<at::Tensor, at::Tensor>
 ConvolutionBackwardCPU(at::Tensor const &in_feat,                         //
-                       at::Tensor const &grad_out_feat,                   //
+                       at::Tensor &grad_out_feat,                         //
                        at::Tensor const &kernel,                          //
                        default_types::stride_type const &kernel_size,     //
                        default_types::stride_type const &kernel_stride,   //
@@ -95,7 +95,7 @@ template <typename coordinate_type,
           template <typename C> class TemplatedAllocator>
 std::pair<at::Tensor, at::Tensor> ConvolutionBackwardGPU(
     at::Tensor const &in_feat,                         //
-    at::Tensor const &grad_out_feat,                   //
+    at::Tensor &grad_out_feat,                         //
     at::Tensor const &kernel,                          //
     default_types::stride_type const &kernel_size,     //
     default_types::stride_type const &kernel_stride,   //

diff --git a/src/convolution_cpu.cpp b/src/convolution_cpu.cpp
@@ -107,7 +107,7 @@ ConvolutionForwardCPU(at::Tensor const &in_feat,                         //
 template <typename coordinate_type>
 std::pair<at::Tensor, at::Tensor>
 ConvolutionBackwardCPU(at::Tensor const &in_feat,                         //
-                       at::Tensor const &grad_out_feat,                   //
+                       at::Tensor &grad_out_feat,                         //
                        at::Tensor const &kernel,                          //
                        default_types::stride_type const &kernel_size,     //
                        default_types::stride_type const &kernel_stride,   //
@@ -119,7 +119,8 @@ ConvolutionBackwardCPU(at::Tensor const &in_feat,                         //
                        cpu_manager_type<coordinate_type> *p_map_manager) {
 
   ASSERT(in_feat.is_contiguous(), "in_feat must be contiguous");
-  ASSERT(grad_out_feat.is_contiguous(), "grad_out_feata must be contiguous");
+  // ASSERT(grad_out_feat.is_contiguous(), "grad_out_feata must be contiguous");
+  grad_out_feat = grad_out_feat.contiguous();
   ASSERT(kernel.is_contiguous(), "kernel must be contiguous");
 
   ASSERT(!in_feat.is_cuda(), "in_feat must be CPU");
@@ -184,7 +185,7 @@ template at::Tensor ConvolutionForwardCPU<default_types::dcoordinate_type>(
 template std::pair<at::Tensor, at::Tensor>
 ConvolutionBackwardCPU<default_types::dcoordinate_type>(
     at::Tensor const &in_feat,                         //
-    at::Tensor const &grad_out_feat,                   //
+    at::Tensor &grad_out_feat,                         //
     at::Tensor const &kernel,                          //
     default_types::stride_type const &kernel_size,     //
     default_types::stride_type const &kernel_stride,   //

diff --git a/src/convolution_gpu.cu b/src/convolution_gpu.cu
@@ -130,7 +130,7 @@ template <typename coordinate_type,
           template <typename C> class TemplatedAllocator>
 std::pair<at::Tensor, at::Tensor> ConvolutionBackwardGPU(
     at::Tensor const &in_feat,                         //
-    at::Tensor const &grad_out_feat,                   //
+    at::Tensor &grad_out_feat,                         //
     at::Tensor const &kernel,                          //
     default_types::stride_type const &kernel_size,     //
     default_types::stride_type const &kernel_stride,   //
@@ -142,7 +142,8 @@ std::pair<at::Tensor, at::Tensor> ConvolutionBackwardGPU(
     gpu_manager_type<coordinate_type, TemplatedAllocator> *p_map_manager) {
 
   ASSERT(in_feat.is_contiguous(), "in_feat must be contiguous");
-  ASSERT(grad_out_feat.is_contiguous(), "grad_out_feata must be contiguous");
+  // ASSERT(grad_out_feat.is_contiguous(), "grad_out_feata must be contiguous");
+  grad_out_feat = grad_out_feat.contiguous();
   ASSERT(kernel.is_contiguous(), "kernel must be contiguous");
 
   ASSERT(in_feat.is_cuda(), "in_feat must be CUDA");
@@ -239,7 +240,7 @@ template std::pair<at::Tensor, at::Tensor>
 ConvolutionBackwardGPU<default_types::dcoordinate_type,
                        detail::default_allocator>(
     at::Tensor const &in_feat,                         //
-    at::Tensor const &grad_out_feat,                   //
+    at::Tensor &grad_out_feat,                         //
     at::Tensor const &kernel,                          //
     default_types::stride_type const &kernel_size,     //
     default_types::stride_type const &kernel_stride,   //
@@ -256,7 +257,7 @@ template std::pair<at::Tensor, at::Tensor>
 ConvolutionBackwardGPU<default_types::dcoordinate_type,
                        detail::c10_allocator>(
     at::Tensor const &in_feat,                         //
-    at::Tensor const &grad_out_feat,                   //
+    at::Tensor &grad_out_feat,                         //
     at::Tensor const &kernel,                          //
     default_types::stride_type const &kernel_size,     //
     default_types::stride_type const &kernel_stride,   //

diff --git a/src/coordinate_map_gpu.cu b/src/coordinate_map_gpu.cu
@@ -173,11 +173,10 @@ void CoordinateMapGPU<coordinate_type, TemplatedAllocator>::insert(
                               // of successful insertions
   ) {
     thrust::counting_iterator<uint32_t> count_begin{0};
-    thrust::for_each(count_begin, count_begin + number_of_valid,
-                     detail::update_value<coordinate_type, map_type>{
-                         *m_map, const_coordinate_data(),
-                         thrust::raw_pointer_cast(m_valid_row_index.data()),
-                         m_coordinate_size});
+    thrust::for_each(
+        count_begin, count_begin + number_of_valid,
+        detail::update_value_with_offset<index_type, map_type>{
+            *m_map, thrust::raw_pointer_cast(m_valid_map_index.data())});
 
     size_type const num_threads = N;
     auto const num_blocks = GET_BLOCKS(num_threads, CUDA_NUM_THREADS);
@@ -361,11 +360,12 @@ CoordinateMapGPU<coordinate_type, TemplatedAllocator>::stride(
   LOG_DEBUG("Reduced to", number_of_valid);
 
   // remap values
-  thrust::for_each(count_begin, count_begin + number_of_valid,
-                   detail::update_value<coordinate_type, map_type>{
-                       *stride_map.m_map, stride_map.const_coordinate_data(),
-                       thrust::raw_pointer_cast(stride_valid_row_index.data()),
-                       m_coordinate_size});
+  thrust::for_each(
+      count_begin, count_begin + number_of_valid,
+      detail::update_value_with_offset<index_type, map_type>{
+          *stride_map.m_map,
+          thrust::raw_pointer_cast(stride_map.m_valid_map_index.data())});
+
   LOG_DEBUG("Stride remap done");
 
   return stride_map;

diff --git a/src/kernel_map.cuh b/src/kernel_map.cuh
@@ -260,6 +260,15 @@ public:
     return nmap;
   }
 
+  size_type max_size() const {
+    size_type nmap = 0;
+    for (auto const &k : m_kernel_size_map) {
+      if (k.second > nmap)
+        nmap = k.second;
+    }
+    return nmap;
+  }
+
   size_type size(index_type const kernel_index) const {
     auto const iter = m_kernel_size_map.find(kernel_index);
     if (iter == m_kernel_size_map.end())