[Major] Update convolution.

Author: kentang-mit

README.md

@@ -9,6 +9,10 @@
  year = {2020}
 }
 ```
+
+**[NEW!!]** We are releasing `torchsparse` v1.1, which has a significant speedup over the previous v1.0, please have a look!
+
+
 ## Overview
 
 We release `torchsparse`, a high-performance computing library for efficient 3D sparse convolution. This library aims at accelerating sparse computation in 3D, in particular the Sparse Convolution operation. 

torchsparse/src/convolution/convolution.cpp

@@ -24,7 +24,60 @@ void ConvolutionForwardGPU(at::Tensor in_feat, at::Tensor out_feat,
 
 
  int kernel_volume = kernel.size(0);
+ int in_buffer_size = 1;
+ bool flag = false;
+ // memory optimization
+ if(kernel_volume % 2 && out_nrows == in_feat.size(0)){
+ flag = true;
+ in_buffer_size = *std::max_element(neighbor_offset.data_ptr<int>(), 
+ neighbor_offset.data_ptr<int>() + kernel_volume/2);
+ in_buffer_size = std::max(in_buffer_size, 
+ *std::max_element(neighbor_offset.data_ptr<int>() + kernel_volume/2+1, 
+ neighbor_offset.data_ptr<int>() + kernel_volume));
+ in_buffer_size = std::max(in_buffer_size, 1);
+ 
+ torch::mm_out(out_feat, in_feat, kernel[kernel_volume / 2]);
+ }
+ else{
+ in_buffer_size = *std::max_element(neighbor_offset.data_ptr<int>(), 
+ neighbor_offset.data_ptr<int>() + kernel_volume);
+ }
 
+ auto options =
+ torch::TensorOptions().dtype(in_feat.dtype()).device(in_feat.device());
+ auto in_buffer = torch::zeros({in_buffer_size, in_feat.size(1)}, options);
+ auto out_buffer = torch::zeros({in_buffer_size, kernel.size(2)}, options);
+ int cur_offset = 0;
+ for(int i = 0; i < kernel_volume; i++){
+ if(flag && (i == kernel_volume / 2)){
+ cur_offset += 2 * neighbor_offset.data_ptr<int>()[i];
+ continue;
+ }
+ 
+ if(neighbor_offset.data_ptr<int>()[i]==0){
+ continue;
+ }
+ 
+ auto out_buffer_activated =
+ torch::from_blob(out_buffer.data_ptr<float>(), 
+ {neighbor_offset.data_ptr<int>()[i], kernel.size(2)}, options);
+ auto in_buffer_activated =
+ torch::from_blob(in_buffer.data_ptr<float>(), 
+ {neighbor_offset.data_ptr<int>()[i], in_feat.size(1)}, options);
+ // gather
+ gather_launch(in_buffer_activated.size(0), in_feat.size(0), kernel.size(1),
+ in_feat.data_ptr<float>(), in_buffer_activated.data_ptr<float>(), 
+ neighbor_map.data_ptr<int>() + cur_offset, transpose);
+ // GEMM
+ torch::mm_out(out_buffer_activated, in_buffer_activated, kernel[i]);
+ // scatter
+ scatter_launch(neighbor_offset.data_ptr<int>()[i], out_nrows, kernel.size(2), out_buffer_activated.data_ptr<float>(), 
+ out_feat.data_ptr<float>(), neighbor_map.data_ptr<int>() + cur_offset, transpose);
+ cur_offset += 2 * neighbor_offset.data_ptr<int>()[i];
+ }
+ 
+ 
+ /*
  cublasHandle_t handle =
  //THCState_getCurrentBlasHandle(at::globalContext().getTHCState());
  at::cuda::getCurrentCUDABlasHandle();
@@ -35,7 +88,7 @@ void ConvolutionForwardGPU(at::Tensor in_feat, at::Tensor out_feat,
  neighbor_offset.data_ptr<int>(), in_feat.size(0), out_feat.size(0), 
  kernel.size(0), transpose, handle, 
  at::cuda::getCurrentCUDAStream());
- 
+ */
 
 
 }
@@ -52,7 +105,63 @@ void ConvolutionBackwardGPU(
 
  int kernel_volume = kernel.size(0);
  bool flag = false;
+ int in_buffer_size;
+ in_buffer_size = *std::max_element(neighbor_offset.data_ptr<int>(), 
+ neighbor_offset.data_ptr<int>() + kernel_volume);
 
+ auto options =
+ torch::TensorOptions().dtype(in_feat.dtype()).device(in_feat.device());
+ auto in_buffer = torch::zeros({in_buffer_size, in_feat.size(1)}, options);
+ auto in_grad_buffer = torch::zeros({in_buffer_size, in_feat.size(1)}, options);
+ auto out_grad_buffer = torch::zeros({in_buffer_size, kernel.size(2)}, options);
+ 
+ 
+ int cur_offset = 0;
+ for(int i = 0; i < kernel_volume; i++){
+ auto kernel_grad_buffer = grad_kernel[i];
+ if(flag && (i == kernel_volume / 2)){
+ cur_offset += 2 * neighbor_offset.data_ptr<int>()[i];
+ continue;
+ }
+ 
+ if(neighbor_offset.data_ptr<int>()[i]==0){
+ continue;
+ }
+ 
+ auto out_grad_buffer_activated =
+ torch::from_blob(out_grad_buffer.data_ptr<float>(), 
+ {neighbor_offset.data_ptr<int>()[i], kernel.size(2)}, options);
+ auto in_grad_buffer_activated =
+ torch::from_blob(in_grad_buffer.data_ptr<float>(), 
+ {neighbor_offset.data_ptr<int>()[i], in_feat.size(1)}, options);
+ auto in_buffer_activated =
+ torch::from_blob(in_buffer.data_ptr<float>(), 
+ {neighbor_offset.data_ptr<int>()[i], in_feat.size(1)}, options);
+ // gather
+ 
+ gather_launch(out_grad_buffer_activated.size(0), grad_out_feat.size(0), kernel.size(2),
+ grad_out_feat.data_ptr<float>(), out_grad_buffer_activated.data_ptr<float>(), 
+ neighbor_map.data_ptr<int>() + cur_offset, !transpose);
+ 
+ gather_launch(in_buffer_activated.size(0), in_feat.size(0), kernel.size(1),
+ in_feat.data_ptr<float>(), in_buffer_activated.data_ptr<float>(), 
+ neighbor_map.data_ptr<int>() + cur_offset, transpose);
+ 
+ // GEMM
+ //torch::mm_out(out_buffer_activated, in_buffer_activated, kernel[i]);
+ torch::mm_out(in_grad_buffer_activated, out_grad_buffer_activated, torch::transpose(kernel[i], 0, 1));
+ torch::mm_out(kernel_grad_buffer, torch::transpose(in_buffer_activated, 0, 1), out_grad_buffer_activated);
+ // scatter
+ //grad_kernel[i] = kernel_grad_buffer;
+ 
+ scatter_launch(neighbor_offset.data_ptr<int>()[i], in_feat.size(0), kernel.size(1), in_grad_buffer_activated.data_ptr<float>(), 
+ grad_in_feat.data_ptr<float>(), neighbor_map.data_ptr<int>() + cur_offset, !transpose);
+ 
+ cur_offset += 2 * neighbor_offset.data_ptr<int>()[i];
+ 
+ }
+ 
+ /* 
  cublasHandle_t handle =
  //THCState_getCurrentBlasHandle(at::globalContext().getTHCState());
  at::cuda::getCurrentCUDABlasHandle();
@@ -62,7 +171,7 @@ void ConvolutionBackwardGPU(
  grad_kernel.data_ptr<float>(), neighbor_map.data_ptr<int>(), neighbor_offset.data_ptr<int>(), 
  in_feat.size(0), grad_out_feat.size(0), kernel.size(0), 
  transpose, handle, at::cuda::getCurrentCUDAStream());
- 
+ */
 }
 
 
@@ -72,3 +181,4 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("sparseconv_backward", &ConvolutionBackwardGPU, "point cloud convolution backward (CUDA)");
 }
 */
+