Pytorch implementation of `gsplat.cuda._wrapper.isect_offset_encode()`. .. note:: This is a minimal implementation of the fully fused version, which has more arguments. Not all arguments are supported.
(
isect_ids: Tensor, C: int, tile_width: int, tile_height: int
)
| 401 | |
| 402 | @torch.no_grad() |
| 403 | def _isect_offset_encode( |
| 404 | isect_ids: Tensor, C: int, tile_width: int, tile_height: int |
| 405 | ) -> Tensor: |
| 406 | """Pytorch implementation of `gsplat.cuda._wrapper.isect_offset_encode()`. |
| 407 | |
| 408 | .. note:: |
| 409 | |
| 410 | This is a minimal implementation of the fully fused version, which has more |
| 411 | arguments. Not all arguments are supported. |
| 412 | """ |
| 413 | tile_n_bits = (tile_width * tile_height).bit_length() |
| 414 | tile_counts = torch.zeros( |
| 415 | (C, tile_height, tile_width), dtype=torch.int64, device=isect_ids.device |
| 416 | ) |
| 417 | |
| 418 | isect_ids_uq, counts = torch.unique_consecutive(isect_ids >> 32, return_counts=True) |
| 419 | |
| 420 | cam_ids_uq = isect_ids_uq >> tile_n_bits |
| 421 | tile_ids_uq = isect_ids_uq & ((1 << tile_n_bits) - 1) |
| 422 | tile_ids_x_uq = tile_ids_uq % tile_width |
| 423 | tile_ids_y_uq = tile_ids_uq // tile_width |
| 424 | |
| 425 | tile_counts[cam_ids_uq, tile_ids_y_uq, tile_ids_x_uq] = counts |
| 426 | |
| 427 | cum_tile_counts = torch.cumsum(tile_counts.flatten(), dim=0).reshape_as(tile_counts) |
| 428 | offsets = cum_tile_counts - tile_counts |
| 429 | return offsets.int() |
| 430 | |
| 431 | |
| 432 | def accumulate( |