Merge pull request #40 from stackav-oss/feature/jmanning/voxelization-cleanup-v2

Change voxelization interface to match CUDA for 1:1 benchmarking

Author: jmanning-stackav

conch/ops/vision/voxelization.py

@@ -47,19 +47,19 @@ def _compute_grid_dim(self) -> tuple[int, int, int]:
 
 def generate_voxels(
     points: torch.Tensor, param: VoxelizationParameter
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Generates voxels from input points, output voxels and points are randomly ordered due to use of atomics.
 
     Args:
         points: input points; expected dimensions (num_points, 4), last dimension should have fields of x,y,z,_.
         param: parameters.
 
     Returns:
-        tuple of voxels:
-            capped_num_points_per_voxel, shape [num_filled_voxels], note per voxel point counters are capped with max_num_points_per_voxel.
-            point_features, shape [num_filled_voxels, max_num_points_per_voxel, num_features_per_point],
-            empty points are filled with 0.
+        tuple of tensors:
+            num_of_filled_voxels, shape [1, cpu]
+            point_features, shape [num_filled_voxels, max_num_points_per_voxel, num_features_per_point], empty points are filled with 0.
             voxel_indices, shape [num_filled_voxels, 4], only first 3 fields are used for x,y,z indices.
+            capped_num_points_per_voxel, shape [num_filled_voxels], note per voxel point counters are capped with max_num_points_per_voxel.
     """
     device = points.device
     num_points, num_features_per_point = points.shape
@@ -124,10 +124,6 @@ def generate_voxels(
         num_warps=block_size // num_threads_per_warp,
     )
 
-    total_filled_voxels = num_filled_voxels.cpu()[0]
+    total_filled_voxels = num_filled_voxels.cpu()
 
-    return (
-        num_points_per_voxel[:total_filled_voxels],
-        point_features[:total_filled_voxels, :, :],
-        voxel_indices[:total_filled_voxels, :],
-    )
+    return total_filled_voxels, point_features, voxel_indices, num_points_per_voxel

tests/voxelization_test.py

@@ -55,7 +55,12 @@ def test_voxelization(
     print(f"Max number of points per voxel: {param.max_num_points_per_voxel}")
 
     # pure triton version
-    num_points_per_voxel, point_features, voxel_indices = generate_voxels(points, param)
+    total_filled_voxels, point_features, voxel_indices, num_points_per_voxel = generate_voxels(points, param)
+
+    filled_voxels = total_filled_voxels[0]
+    num_points_per_voxel = num_points_per_voxel[:filled_voxels]
+    point_features = point_features[:filled_voxels, :, :]
+    voxel_indices = voxel_indices[:filled_voxels, :]
 
     # triton/torch hybrid, 2-step, stable voxelization first then generate a feature tensor the same format as above
     actual_num_points_per_voxel, point_raw_indices, flat_voxel_indices_stable = voxelization_stable(