diff --git a/embodichain/utils/nms.py b/embodichain/utils/nms.py index ca1047405..9ebac9480 100644 --- a/embodichain/utils/nms.py +++ b/embodichain/utils/nms.py @@ -20,8 +20,8 @@ import numpy as np import torch -import warp as wp +from embodichain.utils import logger from embodichain.utils.math import quat_from_matrix __all__ = ["pose_nms", "pose_nms_indices"] @@ -29,114 +29,10 @@ _POSE_NMS_CHUNK_SIZE = 2048 -@wp.func -def _poses_are_close( - positions: wp.array(dtype=wp.float32, ndim=2), - quaternions: wp.array(dtype=wp.float32, ndim=2), - reference_idx: int, - target_idx: int, - rotation_cosine_threshold: float, - distance_threshold_squared: float, - rotation_always_close: bool, -) -> bool: - """Compare poses through their relative rotation and translation.""" - # For unit xyzw quaternions, the real component of - # inverse(q_reference) * q_target is their dot product. Its absolute value - # gives the shortest relative rotation while treating q and -q equally. - relative_rotation_w = ( - quaternions[reference_idx, 0] * quaternions[target_idx, 0] - + quaternions[reference_idx, 1] * quaternions[target_idx, 1] - + quaternions[reference_idx, 2] * quaternions[target_idx, 2] - + quaternions[reference_idx, 3] * quaternions[target_idx, 3] - ) - relative_translation_x = positions[target_idx, 0] - positions[reference_idx, 0] - relative_translation_y = positions[target_idx, 1] - positions[reference_idx, 1] - relative_translation_z = positions[target_idx, 2] - positions[reference_idx, 2] - - rotation_close = rotation_always_close or ( - wp.abs(relative_rotation_w) > rotation_cosine_threshold - ) - translation_close = ( - relative_translation_x * relative_translation_x - + relative_translation_y * relative_translation_y - + relative_translation_z * relative_translation_z - < distance_threshold_squared - ) - return rotation_close and translation_close - - -@wp.kernel(enable_backward=False) -def _pose_pair_close_kernel( - positions: wp.array(dtype=wp.float32, ndim=2), - quaternions: wp.array(dtype=wp.float32, ndim=2), - reference_offset: int, - target_offset: int, - num_targets: int, - rotation_cosine_threshold: float, - distance_threshold_squared: float, - rotation_always_close: bool, - close: wp.array(dtype=wp.uint8), -) -> None: - """Compute a tile of the pairwise pose-closeness matrix.""" - pair_idx = wp.tid() - reference_local_idx = pair_idx // num_targets - target_local_idx = pair_idx - reference_local_idx * num_targets - reference_idx = reference_offset + reference_local_idx - target_idx = target_offset + target_local_idx - - if reference_idx == target_idx: - close[pair_idx] = wp.uint8(0) - return - - close[pair_idx] = wp.uint8( - _poses_are_close( - positions, - quaternions, - reference_idx, - target_idx, - rotation_cosine_threshold, - distance_threshold_squared, - rotation_always_close, - ) - ) - - -@wp.kernel(enable_backward=False) -def _count_close_poses_kernel( - positions: wp.array(dtype=wp.float32, ndim=2), - quaternions: wp.array(dtype=wp.float32, ndim=2), - reference_offset: int, - target_offset: int, - num_targets: int, - rotation_cosine_threshold: float, - distance_threshold_squared: float, - rotation_always_close: bool, - close_counts: wp.array(dtype=wp.int32), -) -> None: - """Accumulate close-neighbor counts for a pairwise tile.""" - pair_idx = wp.tid() - reference_local_idx = pair_idx // num_targets - target_local_idx = pair_idx - reference_local_idx * num_targets - reference_idx = reference_offset + reference_local_idx - target_idx = target_offset + target_local_idx - - if reference_idx != target_idx and _poses_are_close( - positions, - quaternions, - reference_idx, - target_idx, - rotation_cosine_threshold, - distance_threshold_squared, - rotation_always_close, - ): - wp.atomic_add(close_counts, reference_idx, 1) - - def _poses_to_components(poses: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """Convert pose matrices to positions and normalized xyzw quaternions.""" - # Warp composite arrays currently use float32 storage. NMS only uses these - # values for threshold decisions; the returned poses retain their original - # dtype and autograd relationship. + # NMS only uses these float32 values for threshold decisions; the returned + # poses retain their original dtype and autograd relationship. poses_f32 = poses.detach().to(dtype=torch.float32).contiguous() positions = poses_f32[:, :3, 3].contiguous() quaternions_wxyz = quat_from_matrix(poses_f32[:, :3, :3]) @@ -147,6 +43,57 @@ def _poses_to_components(poses: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso return positions, quaternions +def _reference_block_rows(num_poses: int, chunk_size: int) -> int: + """Bound reference-block rows so a block holds ~``chunk_size**2`` entries.""" + return max(1, min(chunk_size, chunk_size * chunk_size // max(1, num_poses))) + + +def _close_tile( + ref_positions: torch.Tensor, + ref_quaternions: torch.Tensor, + tile_positions: torch.Tensor, + tile_quaternions: torch.Tensor, + rotation_cosine_threshold: float, + distance_threshold_squared: float, + rotation_always_close: bool, +) -> torch.Tensor: + """Compute one bounded ``(R, T)`` tile of the pairwise closeness matrix. + + For unit xyzw quaternions, the real component of + ``inverse(q_reference) * q_target`` is their dot product. Its absolute + value gives the shortest relative rotation while treating ``q`` and ``-q`` + equally. Translation closeness compares the squared Euclidean distance. + The arithmetic uses explicit per-component multiplies and left-to-right + additions — the exact association of the previous scalar implementation — + so decisions are identical across CPU and CUDA backends. + + Args: + ref_positions: Reference positions with shape ``(R, 3)``. + ref_quaternions: Reference quaternions with shape ``(R, 4)``. + tile_positions: Target-tile positions with shape ``(T, 3)``. + tile_quaternions: Target-tile quaternions with shape ``(T, 4)``. + rotation_cosine_threshold: ``cos(angle_th / 2)`` decision value. + distance_threshold_squared: Squared translation threshold. + rotation_always_close: Whether rotation is trivially satisfied. + + Returns: + Boolean closeness tile with shape ``(R, T)``. + """ + dx = tile_positions[None, :, 0] - ref_positions[:, None, 0] + dy = tile_positions[None, :, 1] - ref_positions[:, None, 1] + dz = tile_positions[None, :, 2] - ref_positions[:, None, 2] + tile_close = dx * dx + dy * dy + dz * dz < distance_threshold_squared + if not rotation_always_close: + dots = ( + ref_quaternions[:, None, 0] * tile_quaternions[None, :, 0] + + ref_quaternions[:, None, 1] * tile_quaternions[None, :, 1] + + ref_quaternions[:, None, 2] * tile_quaternions[None, :, 2] + + ref_quaternions[:, None, 3] * tile_quaternions[None, :, 3] + ) + tile_close &= dots.abs() > rotation_cosine_threshold + return tile_close + + def _count_close_poses( positions: torch.Tensor, quaternions: torch.Tensor, @@ -155,40 +102,44 @@ def _count_close_poses( rotation_always_close: bool, chunk_size: int, ) -> torch.Tensor: - """Count close neighbors using bounded Warp pairwise tiles.""" - num_poses = positions.shape[0] - positions_wp = wp.from_torch(positions, dtype=wp.float32) - quaternions_wp = wp.from_torch(quaternions, dtype=wp.float32) - if positions_wp.device.is_cuda: - # The components were produced by Torch immediately before this call. - # Make them visible to Warp before launching on its stream. - torch.cuda.synchronize(positions.device) - close_counts_wp = wp.zeros(num_poses, dtype=wp.int32, device=positions_wp.device) + """Count close neighbors, accumulating per bounded pairwise tile. + Counts are summed tile by tile, so no ``(rows, num_poses)`` matrix is ever + materialized and peak memory stays bounded by ``chunk_size ** 2`` entries. + """ + num_poses = positions.shape[0] + close_counts = torch.zeros(num_poses, dtype=torch.int64, device=positions.device) for reference_offset in range(0, num_poses, chunk_size): - num_references = min(chunk_size, num_poses - reference_offset) + reference_end = min(reference_offset + chunk_size, num_poses) + ref_positions = positions[reference_offset:reference_end] + ref_quaternions = quaternions[reference_offset:reference_end] + block_counts = torch.zeros( + reference_end - reference_offset, + dtype=torch.int64, + device=positions.device, + ) for target_offset in range(0, num_poses, chunk_size): - num_targets = min(chunk_size, num_poses - target_offset) - wp.launch( - kernel=_count_close_poses_kernel, - dim=num_references * num_targets, - inputs=[ - positions_wp, - quaternions_wp, - reference_offset, - target_offset, - num_targets, - rotation_cosine_threshold, - distance_threshold_squared, - rotation_always_close, - close_counts_wp, - ], - device=positions_wp.device, + target_end = min(target_offset + chunk_size, num_poses) + tile = _close_tile( + ref_positions, + ref_quaternions, + positions[target_offset:target_end], + quaternions[target_offset:target_end], + rotation_cosine_threshold, + distance_threshold_squared, + rotation_always_close, ) - if positions_wp.device.is_cuda: - wp.synchronize_device(positions_wp.device) - torch.cuda.synchronize(positions.device) - return wp.to_torch(close_counts_wp).clone() + # A pose is not its own neighbor: clear the diagonal overlap. + overlap_start = max(reference_offset, target_offset) + overlap_end = min(reference_end, target_end) + if overlap_start < overlap_end: + diagonal = torch.arange( + overlap_start, overlap_end, device=positions.device + ) + tile[diagonal - reference_offset, diagonal - target_offset] = False + block_counts += tile.sum(dim=1) + close_counts[reference_offset:reference_end] = block_counts + return close_counts def _greedy_keep_indices( @@ -200,66 +151,62 @@ def _greedy_keep_indices( rotation_always_close: bool, chunk_size: int, ) -> torch.Tensor: - """Apply greedy suppression while computing pairwise tiles with Warp.""" + """Apply greedy suppression over batched pairwise closeness rows. + + References are visited strictly in ``visit_order``. Closeness rows are + computed in bounded blocks, and only for references that are still alive + when their block starts; a reference suppressed earlier can never be kept, + so skipping its row is semantics-preserving and makes the run time scale + with the number of survivors rather than with the raw candidate count. + """ num_poses = positions.shape[0] ordered_positions = positions[visit_order].contiguous() ordered_quaternions = quaternions[visit_order].contiguous() - positions_wp = wp.from_torch(ordered_positions, dtype=wp.float32) - quaternions_wp = wp.from_torch(ordered_quaternions, dtype=wp.float32) - if positions_wp.device.is_cuda: - # The indexing operations above run on Torch's stream. - torch.cuda.synchronize(ordered_positions.device) - - # Keep the host-side closeness block bounded to roughly chunk_size**2 - # entries even when there are far more poses than one target tile. - reference_chunk_size = max(1, min(chunk_size, chunk_size**2 // num_poses)) + + # Keep every closeness block bounded to roughly chunk_size**2 entries even + # when there are far more poses than one target tile. + reference_block_rows = _reference_block_rows(num_poses, chunk_size) suppressed = np.zeros(num_poses, dtype=np.bool_) keep_ordered_indices: list[int] = [] - for reference_offset in range(0, num_poses, reference_chunk_size): - num_references = min(reference_chunk_size, num_poses - reference_offset) - max_num_targets = min(chunk_size, num_poses) - close_buffer_wp = wp.empty( - num_references * max_num_targets, - dtype=wp.uint8, - device=positions_wp.device, + for reference_offset in range(0, num_poses, reference_block_rows): + reference_end = min(reference_offset + reference_block_rows, num_poses) + alive_local = np.flatnonzero(~suppressed[reference_offset:reference_end]) + if alive_local.size == 0: + continue + alive_ordered = alive_local + reference_offset + alive_torch = torch.as_tensor( + alive_ordered, dtype=torch.long, device=positions.device + ) + ref_positions = ordered_positions[alive_torch] + ref_quaternions = ordered_quaternions[alive_torch] + block = torch.empty( + alive_torch.numel(), + num_poses, + dtype=torch.bool, + device=positions.device, ) - close_block = np.empty((num_references, num_poses), dtype=np.bool_) - for target_offset in range(0, num_poses, chunk_size): - num_targets = min(chunk_size, num_poses - target_offset) - num_pairs = num_references * num_targets - wp.launch( - kernel=_pose_pair_close_kernel, - dim=num_pairs, - inputs=[ - positions_wp, - quaternions_wp, - reference_offset, - target_offset, - num_targets, - rotation_cosine_threshold, - distance_threshold_squared, - rotation_always_close, - close_buffer_wp, - ], - device=positions_wp.device, - ) - # The tile is consumed by NumPy immediately, so make the Warp - # launch complete before copying it to host memory. - if close_buffer_wp.device.is_cuda: - wp.synchronize_device(close_buffer_wp.device) - close_block[:, target_offset : target_offset + num_targets] = ( - close_buffer_wp.numpy()[:num_pairs].reshape(num_references, num_targets) - != 0 + target_end = min(target_offset + chunk_size, num_poses) + block[:, target_offset:target_end] = _close_tile( + ref_positions, + ref_quaternions, + ordered_positions[target_offset:target_end], + ordered_quaternions[target_offset:target_end], + rotation_cosine_threshold, + distance_threshold_squared, + rotation_always_close, ) + # A pose never suppresses itself. + rows = torch.arange(alive_torch.numel(), device=positions.device) + block[rows, alive_torch] = False + block_np = block.cpu().numpy() - for reference_local_idx in range(num_references): - reference_idx = reference_offset + reference_local_idx + for row_idx, reference_idx in enumerate(alive_ordered): if suppressed[reference_idx]: continue - keep_ordered_indices.append(reference_idx) - suppressed |= close_block[reference_local_idx] + keep_ordered_indices.append(int(reference_idx)) + suppressed |= block_np[row_idx] suppressed[reference_idx] = True ordered_keep = torch.tensor( @@ -278,8 +225,8 @@ def pose_nms_indices( """Return pose indices after removing poses that are too close. Pose matrices are first converted into ``(N, 3)`` positions and unit - ``(N, 4)`` xyzw quaternions. Warp kernels compare their relative rotation - and Euclidean relative translation in bounded pairwise tiles. + ``(N, 4)`` xyzw quaternions. Relative rotation and Euclidean relative + translation are compared in bounded batched pairwise tiles. Args: poses: Input pose matrices. Shape is ``(N, 4, 4)``. @@ -290,7 +237,7 @@ def pose_nms_indices( preserve_order: Whether to greedily select poses in input order. If ``False``, poses with fewer close neighbors are selected first. Defaults to ``False``. - chunk_size: Maximum size of either dimension of a Warp pairwise tile. + chunk_size: Maximum size of either dimension of a pairwise tile. Defaults to 2048. Returns: @@ -298,7 +245,7 @@ def pose_nms_indices( Raises: ValueError: If ``poses`` is not shaped as ``(N, 4, 4)``, is not on a - Warp-supported device, or ``chunk_size`` is not positive. + CPU or CUDA device, or ``chunk_size`` is not positive. """ if poses.ndim != 3 or poses.shape[-2:] != (4, 4): raise ValueError(f"Invalid input shape {poses.shape}, expected (N, 4, 4).") @@ -317,40 +264,53 @@ def pose_nms_indices( if angle_th <= 0.0 or dist_th <= 0.0: return torch.arange(num_poses, dtype=torch.long, device=poses.device) - # ``pose_nms`` may be called without a SimulationManager. ``wp.init`` is - # idempotent when the simulation has already initialized Warp. - wp.init() - positions, quaternions = _poses_to_components(poses) + base_positions, base_quaternions = _poses_to_components(poses) rotation_always_close = angle_th > math.pi rotation_cosine_threshold = ( math.cos(0.5 * float(angle_th)) if not rotation_always_close else 0.0 ) distance_threshold_squared = float(dist_th * dist_th) - if preserve_order: - visit_order = torch.arange(num_poses, dtype=torch.long, device=poses.device) - else: - close_counts = _count_close_poses( + def _select_indices( + positions: torch.Tensor, quaternions: torch.Tensor + ) -> torch.Tensor: + if preserve_order: + visit_order = torch.arange(num_poses, dtype=torch.long, device=poses.device) + else: + close_counts = _count_close_poses( + positions, + quaternions, + rotation_cosine_threshold, + distance_threshold_squared, + rotation_always_close, + chunk_size, + ).to(poses.device) + tie_breaker = torch.arange(num_poses, dtype=torch.long, device=poses.device) + visit_priority = close_counts * (num_poses + 1) + tie_breaker + visit_order = torch.argsort(visit_priority) + return _greedy_keep_indices( positions, quaternions, + visit_order.to(positions.device), rotation_cosine_threshold, distance_threshold_squared, rotation_always_close, chunk_size, - ).to(dtype=torch.long) - tie_breaker = torch.arange(num_poses, dtype=torch.long, device=poses.device) - visit_priority = close_counts * (num_poses + 1) + tie_breaker - visit_order = torch.argsort(visit_priority) - - return _greedy_keep_indices( - positions, - quaternions, - visit_order, - rotation_cosine_threshold, - distance_threshold_squared, - rotation_always_close, - chunk_size, - ) + ).to(poses.device) + + # The pairwise threshold math is elementwise float32 with explicit + # component association, so offloading it to CUDA produces the same + # decisions as the CPU path. Indices return on the input device, and any + # CUDA failure (initialization, memory pressure) falls back to the CPU + # path that previously served these requests. + if poses.device.type == "cpu" and torch.cuda.is_available(): + try: + return _select_indices(base_positions.cuda(), base_quaternions.cuda()) + except RuntimeError as error: + logger.log_warning( + f"pose_nms CUDA offload failed ({error}); falling back to CPU." + ) + return _select_indices(base_positions, base_quaternions) def pose_nms( @@ -365,7 +325,7 @@ def pose_nms( poses: Input pose matrices. Shape is ``(N, 4, 4)``. angle_th: Rotation threshold in radians. Defaults to pi / 36. dist_th: Translation threshold. Defaults to 0.003. - chunk_size: Maximum size of either dimension of a Warp pairwise tile. + chunk_size: Maximum size of either dimension of a pairwise tile. Defaults to 2048. Returns: diff --git a/tests/utils/test_nms.py b/tests/utils/test_nms.py index ffc7f0b8b..8c32f057a 100644 --- a/tests/utils/test_nms.py +++ b/tests/utils/test_nms.py @@ -18,6 +18,7 @@ import math +import numpy as np import pytest import torch import warp as wp @@ -135,3 +136,148 @@ def test_pose_nms_cuda_input_keeps_indices_on_cuda() -> None: assert indices.tolist() == [2, 0] assert indices.device.type == "cuda" + + +def _reference_pose_nms_indices( + poses: torch.Tensor, + angle_th: float, + dist_th: float, + preserve_order: bool, +) -> list[int]: + """Literal O(N^2) implementation of the documented NMS semantics.""" + from embodichain.utils.nms import _poses_to_components + + num_poses = poses.shape[0] + positions, quaternions = _poses_to_components(poses) + rotation_always_close = angle_th > np.pi + cos_th = 0.0 if rotation_always_close else float(np.cos(0.5 * angle_th)) + dist_sq = float(dist_th * dist_th) + + diff = positions[:, None, :] - positions[None, :, :] + close = diff.pow(2).sum(-1) < dist_sq + if not rotation_always_close: + dots = (quaternions[:, None, :] * quaternions[None, :, :]).sum(-1) + close &= dots.abs() > cos_th + close.fill_diagonal_(False) + + if preserve_order: + order = list(range(num_poses)) + else: + counts = close.sum(dim=1).cpu() + priority = counts * (num_poses + 1) + torch.arange(num_poses) + order = torch.argsort(priority).tolist() + + close_np = close.cpu().numpy() + suppressed = np.zeros(num_poses, dtype=bool) + keep: list[int] = [] + for idx in order: + if suppressed[idx]: + continue + keep.append(idx) + suppressed |= close_np[idx] + suppressed[idx] = True + return keep + + +def _clustered_poses(n: int, seed: int, device: str = "cpu") -> torch.Tensor: + """Random pose set with deliberate near-duplicate clusters.""" + g = torch.Generator().manual_seed(seed) + n_centers = max(1, n // 8) + centers_pos = torch.rand(n_centers, 3, generator=g) * 0.5 + centers_aa = torch.rand(n_centers, 3, generator=g) * 2.0 - 1.0 + assign = torch.randint(0, n_centers, (n,), generator=g) + # Jitter spans both sides of the thresholds (3 mm / 5 deg defaults). + pos = centers_pos[assign] + (torch.rand(n, 3, generator=g) - 0.5) * 0.01 + aa = centers_aa[assign] + (torch.rand(n, 3, generator=g) - 0.5) * 0.3 + + angle = aa.norm(dim=-1, keepdim=True).clamp_min(1e-8) + axis = aa / angle + k = torch.zeros(n, 3, 3) + k[:, 0, 1], k[:, 0, 2] = -axis[:, 2], axis[:, 1] + k[:, 1, 0], k[:, 1, 2] = axis[:, 2], -axis[:, 0] + k[:, 2, 0], k[:, 2, 1] = -axis[:, 1], axis[:, 0] + eye = torch.eye(3).expand(n, 3, 3) + sin, cos = torch.sin(angle)[..., None], torch.cos(angle)[..., None] + rot = eye + sin * k + (1 - cos) * (k @ k) + + poses = torch.eye(4).repeat(n, 1, 1) + poses[:, :3, :3] = rot + poses[:, :3, 3] = pos + return poses.to(device) + + +@pytest.mark.parametrize("preserve_order", [True, False]) +@pytest.mark.parametrize("chunk_size", [1, 7, 128, 2048]) +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_pose_nms_matches_reference_semantics(preserve_order, chunk_size, seed): + """Blocked batched implementation must match the literal reference.""" + poses = _clustered_poses(400, seed) + expected = _reference_pose_nms_indices( + poses, angle_th=np.pi / 36, dist_th=0.003, preserve_order=preserve_order + ) + got = pose_nms_indices( + poses, + angle_th=np.pi / 36, + dist_th=0.003, + preserve_order=preserve_order, + chunk_size=chunk_size, + ) + assert got.tolist() == expected + + +@pytest.mark.parametrize("preserve_order", [True, False]) +def test_pose_nms_rotation_always_close_branch(preserve_order): + """angle_th > pi treats every rotation as close (translation only).""" + poses = _clustered_poses(200, seed=3) + expected = _reference_pose_nms_indices( + poses, angle_th=4.0, dist_th=0.02, preserve_order=preserve_order + ) + got = pose_nms_indices( + poses, angle_th=4.0, dist_th=0.02, preserve_order=preserve_order + ) + assert got.tolist() == expected + + +def test_pose_nms_all_duplicates_keep_single(): + poses = _pose().repeat(500, 1, 1) + indices = pose_nms_indices(poses, preserve_order=True) + assert indices.tolist() == [0] + + +@pytest.mark.gpu +def test_pose_nms_cuda_matches_cpu_reference(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is not available") + poses = _clustered_poses(400, seed=4) + expected = _reference_pose_nms_indices( + poses, angle_th=np.pi / 36, dist_th=0.003, preserve_order=False + ) + got = pose_nms_indices(poses.cuda(), preserve_order=False) + assert got.cpu().tolist() == expected + + +def _tight_clustered_poses(n: int, seed: int) -> torch.Tensor: + """Heavily duplicated pose set: few survivors, exercising alive filtering.""" + g = torch.Generator().manual_seed(seed) + n_centers = max(1, n // 100) + base = _clustered_poses(n_centers, seed) + assign = torch.randint(0, n_centers, (n,), generator=g) + poses = base[assign].clone() + # Jitter well below the 3 mm / 5 deg thresholds. + poses[:, :3, 3] += (torch.rand(n, 3, generator=g) - 0.5) * 0.001 + return poses + + +@pytest.mark.parametrize("preserve_order", [True, False]) +def test_pose_nms_heavy_suppression_matches_reference(preserve_order): + """Large tight-cluster input (the realistic grasp-candidate profile): + the CUDA-offloaded pairwise math must reproduce the literal CPU + reference exactly.""" + poses = _tight_clustered_poses(6000, seed=5) + expected = _reference_pose_nms_indices( + poses, angle_th=np.pi / 36, dist_th=0.003, preserve_order=preserve_order + ) + got = pose_nms_indices(poses, preserve_order=preserve_order) + assert got.tolist() == expected + # Suppression must actually be heavy for this test to mean anything. + assert len(expected) < 6000 // 10