| 313 | class TestCollectPointsOnRay(unittest.TestCase): |
| 314 | |
| 315 | def test(self, b=1, n=100, m=2, max_size=40, seed=0): |
| 316 | if seed is not None: |
| 317 | torch.manual_seed(seed) |
| 318 | random.seed(seed) |
| 319 | np.random.seed(seed) |
| 320 | |
| 321 | |
| 322 | points = torch.randn(b, n, 3) |
| 323 | ray_origins = torch.randn(b, m, 3) # (b, m, 3) |
| 324 | ray_directions = torch.nn.functional.normalize(torch.randn(b, m, 3), dim=-1) # (b, m, 3) |
| 325 | ray_radius = torch.rand(b,)*0.8 + 1e-3 # (b,) |
| 326 | grid_size = (torch.rand(b, 3) * (max_size-1) + 1).floor().long() # (b, 3) |
| 327 | total_cells = torch.prod(grid_size, dim=-1) |
| 328 | grid_center = torch.randn(b, 3) |
| 329 | grid_width = torch.rand(b, 3) * 20 + 10 |
| 330 | |
| 331 | points = points.cuda() |
| 332 | ray_origins = ray_origins.cuda() |
| 333 | ray_directions = ray_directions.cuda() |
| 334 | ray_radius = ray_radius.cuda() |
| 335 | grid_size = grid_size.cuda() |
| 336 | grid_center = grid_center.cuda() |
| 337 | grid_width = grid_width.cuda() |
| 338 | total_cells = total_cells.cuda() |
| 339 | |
| 340 | stime = timer() |
| 341 | grid_idxs, valid_mask, cell_counts = pr_cuda.get_grid_idx( |
| 342 | points, |
| 343 | grid_size, |
| 344 | grid_center, |
| 345 | grid_width, |
| 346 | 'ind', |
| 347 | ) |
| 348 | time_cuda = timer() - stime |
| 349 | print('CollectPointsOnRay::get_grid_idx:') |
| 350 | print(f'cuda: {time_cuda * 1000.:.3f} ms ') |
| 351 | |
| 352 | stime = timer() |
| 353 | gidx2pidx_bank, gidx_start_idx = pr_cuda.gather_points_v2( |
| 354 | grid_idxs, |
| 355 | total_cells, |
| 356 | valid_mask, |
| 357 | cell_counts, |
| 358 | ) # membank: (b, n), cell_start_idxs: (b, n_cell+1) |
| 359 | time_cuda = timer() - stime |
| 360 | print('CollectPointsOnRay::gather_points_v2:') |
| 361 | print(f'cuda: {time_cuda * 1000.:.3f} ms ') |
| 362 | |
| 363 | stime = timer() |
| 364 | ray2gidx, n_ray2gidx = pr_cuda.grid_ray_intersection( |
| 365 | ray_origins, |
| 366 | ray_directions, |
| 367 | ray_radius, |
| 368 | grid_size, |
| 369 | grid_center, |
| 370 | grid_width, |
| 371 | ) |
| 372 | time_cuda = timer() - stime |