Marching Cubes cuda extension

Summary:
Torch CUDA extension for Marching Cubes
- MC involving 3 steps:
  - 1st forward pass to collect vertices and occupied state for each voxel
  - Compute compactVoxelArray to skip non-empty voxels
  - 2nd pass to genereate interpolated vertex positions and faces by marching through the grid
- In contrast to existing MC:
   - Bind each interpolated vertex with a global edge_id to address floating-point precision
   - Added deduplication process to remove redundant vertices and faces

Benchmarks (ms):

| N / V(^3)      | python          | C++             |   CUDA   | Speedup |
| 2 / 20          |    12176873  |       24338     |     4363   | 2790x/5x|
| 1 / 100          |     -             |    3070511     |   27126   |    113x    |
| 2 / 100          |     -             |    5968934     |   53129   |    112x    |
| 1 / 256          |     -             |  61278092     | 430900   |    142x    |
| 2 / 256          |     -             |125687930     | 856941   |    146x   |

Reviewed By: kjchalup

Differential Revision: D39644248

fbshipit-source-id: d679c0c79d67b98b235d12296f383d760a00042a

pytorch3d/ops/marching_cubes.py

@@ -230,18 +230,19 @@ def marching_cubes_naive(
 
 
 ########################################
-# Marching Cubes Implementation in C++
+# Marching Cubes Implementation in C++/Cuda
 ########################################
 class _marching_cubes(Function):
     """
-    Torch Function wrapper for marching_cubes C++ implementation
-    Backward is not supported.
+    Torch Function wrapper for marching_cubes implementation.
+    This function is not differentiable. An autograd wrapper is used
+    to ensure an error if user tries to get gradients.
     """
 
     @staticmethod
     def forward(ctx, vol, isolevel):
-        verts, faces = _C.marching_cubes(vol, isolevel)
-        return verts, faces
+        verts, faces, ids = _C.marching_cubes(vol, isolevel)
+        return verts, faces, ids
 
     @staticmethod
     def backward(ctx, grad_verts, grad_faces):
@@ -268,7 +269,6 @@ def marching_cubes(
             the range [-1, 1] x [-1, 1] x [-1, 1]. If False they will be in the range
             [0, W-1] x [0, H-1] x [0, D-1]
 
-
     Returns:
         verts: [{V_0}, {V_1}, ...] List of N sets of vertices of shape (|V_i|, 3) in FloatTensor
         faces: [{F_0}, {F_1}, ...] List of N sets of faces of shape (|F_i|, 3) in LongTensors
@@ -279,7 +279,7 @@ def marching_cubes(
         vol = vol_batch[i]
         thresh = ((vol.max() + vol.min()) / 2).item() if isolevel is None else isolevel
         # pyre-fixme[16]: `_marching_cubes` has no attribute `apply`.
-        verts, faces = _marching_cubes.apply(vol, thresh)
+        verts, faces, ids = _marching_cubes.apply(vol, thresh)
         if len(faces) > 0 and len(verts) > 0:
             # Convert from world coordinates ([0, D-1], [0, H-1], [0, W-1]) to
             # local coordinates in the range [-1, 1]
@@ -289,6 +289,13 @@ def marching_cubes(
                     .scale((vol.new_tensor([W, H, D])[None] - 1) * 0.5)
                     .inverse()
                 ).transform_points(verts[None])[0]
+            # deduplication for cuda
+            if vol.is_cuda:
+                unique_ids, inverse_idx = torch.unique(ids, return_inverse=True)
+                verts_ = verts.new_zeros(unique_ids.shape[0], 3)
+                verts_[inverse_idx] = verts
+                verts = verts_
+                faces = inverse_idx[faces]
             batched_verts.append(verts)
             batched_faces.append(faces)
         else: