Rename torchsparse_cuda as torchsparse_backend (#49)

* Rename `torchsparse_cuda` as `torchsparse_backend`

Co-authored-by: github-actions <[email protected]>

Author: Zhijian Liu

setup.py

@@ -65,7 +65,7 @@
     version=__version__,
     packages=find_packages(),
     ext_modules=[
-        extension_type('torchsparse_cuda',
+        extension_type('torchsparse_backend',
                        file_lis,
                        extra_compile_args=extra_compile_args)
     ],

torchsparse/__init__.py

@@ -4,13 +4,15 @@
 
 __version__ = '1.1.0'
 
+
 def cat(input_list, dim=1):
     assert len(input_list) > 0
     inputs = input_list[0]
     features = inputs.F
     coords = inputs.C
     cur_stride = inputs.s
-    output_tensor = SparseTensor(torch.cat([inputs.F for inputs in input_list], 1), coords, cur_stride)
+    output_tensor = SparseTensor(
+        torch.cat([inputs.F for inputs in input_list], 1), coords, cur_stride)
     output_tensor.coord_maps = inputs.coord_maps
     output_tensor.kernel_maps = inputs.kernel_maps
     return output_tensor

torchsparse/nn/functional/conv.py

@@ -1,7 +1,7 @@
 import copy
 
 import torch
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 from torchsparse import *
 from torchsparse.nn.functional.convert_neighbor_map import *
@@ -35,9 +35,9 @@ def forward(ctx,
                               device=features.device)
 
         if 'cuda' in str(features.device):
-            torchsparse_cuda.sparseconv_forward(features, out, kernel,
-                                                neighbor_map, neighbor_offset,
-                                                transpose)
+            torchsparse_backend.sparseconv_forward(features, out, kernel,
+                                                   neighbor_map,
+                                                   neighbor_offset, transpose)
         else:
             # use the native pytorch XLA APIs for the TPU.
             cur_st = 0
@@ -69,10 +69,11 @@ def backward(ctx, grad_out):
         grad_kernel = torch.zeros(K, c_in, c_out, device=kernel.device)
 
         if 'cuda' in str(features.device):
-            torchsparse_cuda.sparseconv_backward(features, grad_features,
-                                                 grad_out.contiguous(), kernel,
-                                                 grad_kernel, neighbor_map,
-                                                 neighbor_offset, transpose)
+            torchsparse_backend.sparseconv_backward(features, grad_features,
+                                                    grad_out.contiguous(),
+                                                    kernel, grad_kernel,
+                                                    neighbor_map,
+                                                    neighbor_offset, transpose)
         else:
             raise NotImplementedError
         return grad_features, grad_kernel, None, None, None, None

torchsparse/nn/functional/convert_neighbor_map.py

@@ -1,5 +1,5 @@
 import torch
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 
 
@@ -8,14 +8,14 @@ class ConvertNeighborMap(Function):
     def forward(ctx, neighbor_map):
         idx_batch, idx_point = torch.where(neighbor_map != -1)
         if 'cuda' in str(neighbor_map.device):
-            map_converted = torchsparse_cuda.convert_map_forward(
+            map_converted = torchsparse_backend.convert_map_forward(
                 neighbor_map.int(), idx_batch.int(), idx_point.int())
         elif 'cpu' in str(neighbor_map.device):
-            map_converted = torchsparse_cuda.cpu_convert_map_forward(
+            map_converted = torchsparse_backend.cpu_convert_map_forward(
                 neighbor_map.int(), idx_batch.int(), idx_point.int())
         else:
             device = neighbor_map.device
-            map_converted = torchsparse_cuda.cpu_convert_map_forward(
+            map_converted = torchsparse_backend.cpu_convert_map_forward(
                 neighbor_map.int().cpu(),
                 idx_batch.int().cpu(),
                 idx_point.int().cpu())

torchsparse/nn/functional/count.py

@@ -1,4 +1,4 @@
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 
 __all__ = ['spcount']
@@ -8,9 +8,9 @@ class CountGPU(Function):
     @staticmethod
     def forward(ctx, idx, num):
         if 'cuda' in str(idx.device):
-            outs = torchsparse_cuda.count_forward(idx.contiguous(), num)
+            outs = torchsparse_backend.count_forward(idx.contiguous(), num)
         else:
-            outs = torchsparse_cuda.cpu_count_forward(idx.contiguous(), num)
+            outs = torchsparse_backend.cpu_count_forward(idx.contiguous(), num)
         return outs
 
 

torchsparse/nn/functional/devox.py

@@ -1,5 +1,5 @@
 import torch
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 
 __all__ = ['spdevoxelize', 'calc_ti_weights']
@@ -63,11 +63,11 @@ class DevoxelizationGPU(Function):
     @staticmethod
     def forward(ctx, feat, indices, weights):
         if 'cuda' in str(feat.device):
-            out = torchsparse_cuda.devoxelize_forward(
+            out = torchsparse_backend.devoxelize_forward(
                 feat.contiguous(),
                 indices.contiguous().int(), weights.contiguous())
         else:
-            out = torchsparse_cuda.cpu_devoxelize_forward(
+            out = torchsparse_backend.cpu_devoxelize_forward(
                 feat.contiguous(),
                 indices.contiguous().int(), weights.contiguous())
 
@@ -81,10 +81,10 @@ def backward(ctx, grad_out):
         indices, weights, n = ctx.for_backwards
 
         if 'cuda' in str(grad_out.device):
-            grad_features = torchsparse_cuda.devoxelize_backward(
+            grad_features = torchsparse_backend.devoxelize_backward(
                 grad_out.contiguous(), indices, weights, n)
         else:
-            grad_features = torchsparse_cuda.cpu_devoxelize_backward(
+            grad_features = torchsparse_backend.cpu_devoxelize_backward(
                 grad_out.contiguous(), indices, weights, n)
 
         return grad_features, None, None

torchsparse/nn/functional/downsample.py

@@ -1,5 +1,5 @@
 import torch
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 from torchsparse.nn.functional.hash import *
 
@@ -24,16 +24,16 @@ def forward(ctx, coords, ratio):
         # gpu
         if 'cuda' in str(coords.device):
             uq_coords = torch.round(
-                torchsparse_cuda.insertion_forward(coords_new.float(), inv,
-                                                   cnt))
+                torchsparse_backend.insertion_forward(coords_new.float(), inv,
+                                                      cnt))
         elif 'cpu' in str(coords.device):
             uq_coords = torch.round(
-                torchsparse_cuda.cpu_insertion_forward(coords_new.float(), inv,
-                                                       cnt))
+                torchsparse_backend.cpu_insertion_forward(
+                    coords_new.float(), inv, cnt))
         else:
             device = coords.device
             uq_coords = torch.round(
-                torchsparse_cuda.cpu_insertion_forward(
+                torchsparse_backend.cpu_insertion_forward(
                     coords_new.float().cpu(), inv.cpu(), cnt.cpu()))
             uq_coords = uq_coords.to(device)
         uq_coords = uq_coords.int()

torchsparse/nn/functional/hash.py

@@ -1,4 +1,4 @@
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 
 __all__ = ['sphash']
@@ -8,28 +8,28 @@ class HashGPU(Function):
     @staticmethod
     def forward(ctx, idx):
         if 'cuda' in str(idx.device):
-            return torchsparse_cuda.hash_forward(idx.contiguous())
+            return torchsparse_backend.hash_forward(idx.contiguous())
         elif 'cpu' in str(idx.device):
-            return torchsparse_cuda.cpu_hash_forward(idx.int().contiguous())
+            return torchsparse_backend.cpu_hash_forward(idx.int().contiguous())
         else:
             device = idx.device
-            return torchsparse_cuda.cpu_hash_forward(
+            return torchsparse_backend.cpu_hash_forward(
                 idx.int().contiguous().cpu()).to(device)
 
 
 class KernelHashGPU(Function):
     @staticmethod
     def forward(ctx, idx, koffset):
         if 'cuda' in str(idx.device):
-            return torchsparse_cuda.kernel_hash_forward(
+            return torchsparse_backend.kernel_hash_forward(
                 idx.contiguous(), koffset.contiguous())
         elif 'cpu' in str(idx.device):
-            return torchsparse_cuda.cpu_kernel_hash_forward(
+            return torchsparse_backend.cpu_kernel_hash_forward(
                 idx.int().contiguous(),
                 koffset.int().contiguous())
         else:
             device = idx.device
-            return torchsparse_cuda.cpu_kernel_hash_forward(
+            return torchsparse_backend.cpu_kernel_hash_forward(
                 idx.int().contiguous().cpu(),
                 koffset.int().contiguous().cpu()).to(device)
 

torchsparse/nn/functional/query.py

@@ -1,5 +1,5 @@
 import torch
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 
 __all__ = ['sphashquery']
@@ -18,16 +18,16 @@ def forward(ctx, hash_query, hash_target):
                                   dtype=torch.long)
 
         if 'cuda' in str(hash_query.device):
-            out, key_buf, val_buf, key = torchsparse_cuda.query_forward(
+            out, key_buf, val_buf, key = torchsparse_backend.query_forward(
                 hash_query.view(-1).contiguous(), hash_target.contiguous(),
                 idx_target)
         elif 'cpu' in str(hash_query.device):
-            out = torchsparse_cuda.cpu_query_forward(
+            out = torchsparse_backend.cpu_query_forward(
                 hash_query.view(-1).contiguous(), hash_target.contiguous(),
                 idx_target)
         else:
             device = hash_query.device
-            out = torchsparse_cuda.cpu_query_forward(
+            out = torchsparse_backend.cpu_query_forward(
                 hash_query.view(-1).contiguous().cpu(),
                 hash_target.contiguous().cpu(), idx_target.cpu()).to(device)
 

torchsparse/nn/functional/voxelize.py

@@ -1,4 +1,4 @@
-import torchsparse_cuda
+import torchsparse_backend
 from torch.autograd import Function
 from torchsparse.nn.functional.hash import *
 
@@ -8,15 +8,16 @@
 class VoxelizeGPU(Function):
     @staticmethod
     def forward(ctx, feat, idx, cnt):
-        out = torchsparse_cuda.insertion_forward(feat.float().contiguous(),
-                                                 idx.int().contiguous(), cnt)
+        out = torchsparse_backend.insertion_forward(feat.float().contiguous(),
+                                                    idx.int().contiguous(),
+                                                    cnt)
         ctx.for_backwards = (idx.int().contiguous(), cnt, feat.shape[0])
         return out
 
     @staticmethod
     def backward(ctx, top_grad):
         idx, cnt, N = ctx.for_backwards
-        bottom_grad = torchsparse_cuda.insertion_backward(
+        bottom_grad = torchsparse_backend.insertion_backward(
             top_grad.float().contiguous(), idx, cnt, N)
         return bottom_grad, None, None