Spaces:
Running on Zero
Running on Zero
| import torch | |
| import torch.nn.functional as F | |
| from typing import List, Tuple, Optional | |
| # CUDA is mandatory for this module; there is intentionally no CPU fallback. | |
| assert torch.cuda.is_available(), "cuda_knn requires a CUDA-capable device." | |
| def compute_nearest_distance_to_list_cuda( | |
| tensor_a: torch.Tensor, | |
| tensor_list: List[torch.Tensor], | |
| batch_mode: bool = False, | |
| return_vectors: bool = False, | |
| return_tensor_indices: bool = True | |
| ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: | |
| """ | |
| CUDA-accelerated parallel computation of nearest distances from query | |
| points to a list of reference point sets. All tensors must reside on the | |
| same device. | |
| Args: | |
| tensor_a: Query tensor | |
| - If batch_mode=False: shape [N, D] | |
| - If batch_mode=True: shape [B, N, D] | |
| tensor_list: List of reference tensors | |
| - If batch_mode=False: [M_i, D] each | |
| - If batch_mode=True: [B, M_i, D] each | |
| batch_mode: Whether to use batch processing | |
| return_vectors: If True, also return vectors from query to nearest points | |
| return_tensor_indices: If False, skip the per-list-entry loop that maps the nearest point to | |
| the index of its list entry (returned as None); distances and vectors are unaffected | |
| Returns: | |
| min_distances: Minimum distances [N] or [B, N] | |
| min_indices: Tensor list indices [N] or [B, N] | |
| vectors: (optional) Vectors from query to nearest points [N, D] or [B, N, D] | |
| vectors = nearest_point - query_point | |
| """ | |
| # All tensors must reside on the same device | |
| device = tensor_a.device | |
| for i, t in enumerate(tensor_list): | |
| if t.device != device: | |
| raise RuntimeError(f"All tensors must be on the same device: tensor_a is on {device}, tensor_list[{i}] is on {t.device}") | |
| if len(tensor_list) == 0: | |
| raise ValueError("tensor_list cannot be empty") | |
| device = tensor_a.device | |
| original_batch_mode = batch_mode | |
| if not batch_mode: | |
| # Add batch dimension for unified processing | |
| tensor_a = tensor_a.unsqueeze(0) # [1, N, D] | |
| tensor_list = [t.unsqueeze(0) for t in tensor_list] # Each becomes [1, M_i, D] | |
| B, N, D = tensor_a.shape | |
| # Strategy 1: Concatenate all references for single parallel computation | |
| ref_sizes = [ref.shape[1] for ref in tensor_list] | |
| all_refs = torch.cat(tensor_list, dim=1) # [B, total_M, D] | |
| # Single CUDA-parallel distance computation | |
| dist_matrix = torch.cdist(tensor_a, all_refs, p=2.0) # [B, N, total_M] | |
| # Find minimum distances and indices in parallel | |
| min_dists, min_idx = dist_matrix.min(dim=2) # Both [B, N] | |
| # Convert concatenated indices to tensor list indices | |
| min_tensor_indices = None | |
| if return_tensor_indices: | |
| min_tensor_indices = torch.zeros_like(min_idx, dtype=torch.long) | |
| cumsum = 0 | |
| for i, size in enumerate(ref_sizes): | |
| mask = (min_idx >= cumsum) & (min_idx < cumsum + size) | |
| min_tensor_indices[mask] = i | |
| cumsum += size | |
| # Compute vectors if requested | |
| vectors = None | |
| if return_vectors: | |
| # Get the nearest points using advanced indexing | |
| batch_indices = torch.arange(B, device=device).view(B, 1).expand(B, N) | |
| nearest_points = all_refs[batch_indices, min_idx] # [B, N, D] | |
| # Vector from query to nearest point | |
| vectors = nearest_points - tensor_a # [B, N, D] | |
| if not original_batch_mode: | |
| min_dists = min_dists.squeeze(0) | |
| if min_tensor_indices is not None: | |
| min_tensor_indices = min_tensor_indices.squeeze(0) | |
| if vectors is not None: | |
| vectors = vectors.squeeze(0) | |
| if return_vectors: | |
| return min_dists, min_tensor_indices, vectors | |
| else: | |
| return min_dists, min_tensor_indices | |
| def compute_knn_all_objects_vectorized( | |
| positions: torch.Tensor, | |
| valid_vtx_mask: torch.Tensor, | |
| FAR: float = 1e6, | |
| valid_obj_mask: torch.Tensor = None, | |
| num_groups: int = 1, | |
| active_slots: Optional[List[int]] = None | |
| ) -> torch.Tensor: | |
| """ | |
| Efficient batch-mode KNN for multi-object collision detection. | |
| Args: | |
| positions: (B, Nobj, Nv, 3) - vertex positions for all objects | |
| valid_vtx_mask: (B, Nobj, Nv) - True for valid vertices | |
| FAR: Distance offset for invalid vertices (default 1e6) | |
| valid_obj_mask: (B, Nobj) - True for valid objects | |
| num_groups: Number of groups to split objects into for GPU-parallel processing. | |
| num_groups=1: per-object loop (Nobj iterations) | |
| num_groups>1: GPU-parallel within each group (num_groups iterations) | |
| Example: Nobj=564, num_groups=6 -> 94 objects processed in parallel per group | |
| active_slots: Optional list of object slots that hold a valid object in at least one batch element | |
| (per-object path only). Slots outside it are padding in every sample and their output | |
| is zeroed either way, so they are not searched; valid objects still search every other slot. | |
| Returns: | |
| dist_vectors: (B, Nobj, Nv, 3) - vector from each vertex to nearest point on OTHER objects | |
| """ | |
| B, Nobj, Nv, D = positions.shape | |
| device = positions.device | |
| if Nobj < 2: | |
| raise RuntimeError("compute_knn_all_objects_vectorized requires at least 2 objects") | |
| if valid_obj_mask is None: | |
| valid_obj_mask = valid_vtx_mask.any(dim=2) # (B, Nobj) | |
| dist_vectors = torch.zeros(B, Nobj, Nv, D, device=device) | |
| # Pre-compute masked positions (FAR offset for invalid vertices) | |
| masked_positions = positions.clone() | |
| vtx_mask_3d = valid_vtx_mask.unsqueeze(-1) # (B, Nobj, Nv, 1) | |
| masked_positions = torch.where(vtx_mask_3d, masked_positions, masked_positions + FAR) | |
| # ========== GPU-PARALLEL GROUPED PROCESSING ========== | |
| if num_groups > 1 and Nobj > num_groups: | |
| group_size = (Nobj + num_groups - 1) // num_groups | |
| # All reference points (all objects, masked): (B, Nobj*Nv, 3) | |
| all_refs_flat = masked_positions.reshape(B, Nobj * Nv, D) | |
| for group_idx in range(num_groups): | |
| g_start = group_idx * group_size | |
| g_end = min(g_start + group_size, Nobj) | |
| G = g_end - g_start # Objects in this group | |
| # Query points for this group: (B, G, Nv, 3) -> (B, G*Nv, 3) | |
| group_queries = positions[:, g_start:g_end].reshape(B, G * Nv, D) | |
| # Compute full distance matrix: (B, G*Nv, Nobj*Nv) | |
| dist_matrix = torch.cdist(group_queries, all_refs_flat, p=2.0) | |
| # Create self-exclusion mask: for each query vertex, mask out vertices from same object | |
| # Query vertex q belongs to object (g_start + q // Nv) | |
| # Ref vertex r belongs to object (r // Nv) | |
| query_obj_idx = torch.arange(G, device=device).unsqueeze(1).expand(G, Nv).reshape(-1) + g_start # (G*Nv,) | |
| ref_obj_idx = torch.arange(Nobj, device=device).unsqueeze(1).expand(Nobj, Nv).reshape(-1) # (Nobj*Nv,) | |
| # Self mask: (G*Nv, Nobj*Nv), True where query and ref belong to same object | |
| self_mask = query_obj_idx.unsqueeze(1) == ref_obj_idx.unsqueeze(0) # (G*Nv, Nobj*Nv) | |
| # Apply FAR to self-distances so they're not selected as minimum | |
| dist_matrix = dist_matrix + self_mask.float() * FAR # (B, G*Nv, Nobj*Nv) | |
| # Find minimum distance and corresponding ref index | |
| min_dists, min_indices = dist_matrix.min(dim=2) # (B, G*Nv) | |
| # Get nearest points using advanced indexing | |
| batch_idx = torch.arange(B, device=device).view(B, 1).expand(B, G * Nv) | |
| nearest_points = all_refs_flat[batch_idx, min_indices] # (B, G*Nv, 3) | |
| # Compute vectors from query to nearest | |
| vectors = nearest_points - group_queries # (B, G*Nv, 3) | |
| # Reshape back and store: (B, G*Nv, 3) -> (B, G, Nv, 3) | |
| dist_vectors[:, g_start:g_end] = vectors.reshape(B, G, Nv, D) | |
| else: | |
| # ========== PER-OBJECT PROCESSING ========== | |
| # The references of object i are all other slots in slot order: the flattened (B, Nobj*Nv, D) | |
| # tensor without object i's block, the same tensor as concatenating the other slots one by one. | |
| refs_flat = masked_positions.reshape(B, Nobj * Nv, D) | |
| for i in (range(Nobj) if active_slots is None else active_slots): | |
| body_pos = positions[:, i] # (B, Nv, 3) | |
| ref_list = [refs_flat[:, :i * Nv], refs_flat[:, (i + 1) * Nv:]] | |
| _, _, body_dist_vec = compute_nearest_distance_to_list_cuda( | |
| body_pos, ref_list, batch_mode=True, return_vectors=True, return_tensor_indices=False | |
| ) | |
| dist_vectors[:, i] = body_dist_vec | |
| # Zero out invalid query vertices | |
| vtx_mask_expanded = valid_vtx_mask.unsqueeze(-1) | |
| dist_vectors = torch.where(vtx_mask_expanded, dist_vectors, torch.zeros_like(dist_vectors)) | |
| # Zero out invalid objects | |
| obj_mask_expanded = valid_obj_mask.unsqueeze(-1).unsqueeze(-1) | |
| dist_vectors = torch.where(obj_mask_expanded, dist_vectors, torch.zeros_like(dist_vectors)) | |
| return dist_vectors | |
| def knn_pack(valid_vtx_mask: torch.Tensor, valid_obj_mask: Optional[torch.Tensor] = None) -> dict: | |
| """Packed layout of every sample's valid vertices for compute_knn_packed. | |
| The masks do not change over a batch or a rollout, so this is computed once (one host sync for the packed | |
| width) and reused at every step. valid_vtx_mask (B, Nobj, Nv); valid_obj_mask (B, Nobj) or None. | |
| Packed order = slot order, then vertex order: the order in which the per-object path lists candidates.""" | |
| B, Nobj, Nv = valid_vtx_mask.shape | |
| vm = valid_vtx_mask.bool() | |
| if valid_obj_mask is not None: | |
| vm = vm & valid_obj_mask.bool().unsqueeze(-1) | |
| n_obj = vm.any(dim=2).sum(dim=1) | |
| if bool((n_obj < 2).any()): | |
| raise ValueError("compute_knn_packed: every sample needs at least two objects with valid vertices") | |
| flat = vm.reshape(B, Nobj * Nv) | |
| counts = flat.sum(dim=1) | |
| n = int(counts.max()) | |
| order = torch.argsort((~flat).to(torch.int8), dim=1, stable=True)[:, :n] # (B, n) valid first | |
| valid = torch.arange(n, device=flat.device).unsqueeze(0) < counts.unsqueeze(1) # (B, n) | |
| obj = torch.where(valid, torch.div(order, Nv, rounding_mode='floor'), torch.full_like(order, -1)) | |
| return {'order': order, 'valid': valid, 'obj': obj, 'shape': (B, Nobj, Nv)} | |
| def compute_knn_packed(positions: torch.Tensor, pack: dict, block_elems: int = 1 << 27) -> torch.Tensor: | |
| """Vector from every valid vertex to the nearest valid vertex of any OTHER object, over the packed valid | |
| vertices only (padded vertices and padded object slots never enter the distance matrix). | |
| Same candidates in the same order and the same first-minimum tie-breaking as compute_knn_all_objects_vectorized, | |
| and the squared distance is cdist's own construction ([-2x, |x|^2, 1] . [y, 1, |y|^2]) without the clamp and | |
| square root over the whole matrix (monotone, so the argmin is unchanged). Results agree with the per-object | |
| path up to the floating-point rounding of the distance values (a near-tie may pick a different, equally near | |
| vertex). positions (B, Nobj, Nv, D); pack from knn_pack; returns (B, Nobj, Nv, D), zero at invalid entries.""" | |
| B, Nobj, Nv = pack['shape'] | |
| D = positions.shape[-1] | |
| order, valid, obj = pack['order'], pack['valid'], pack['obj'] | |
| n = order.shape[1] | |
| with torch.autocast(device_type='cuda', enabled=False): | |
| flat = positions.float().reshape(B, Nobj * Nv, D) | |
| P = torch.gather(flat, 1, order.unsqueeze(-1).expand(B, n, D)) # (B, n, D) | |
| sq = P.pow(2).sum(dim=-1, keepdim=True) | |
| ones = torch.ones_like(sq) | |
| Xq = torch.cat([P.mul(-2), sq, ones], dim=-1) # (B, n, D+2) | |
| Yr = torch.cat([P, ones, sq], dim=-1).transpose(1, 2) # (B, D+2, n) | |
| rows = max(1, min(n, block_elems // max(1, B * n))) | |
| idx = torch.empty(B, n, dtype=torch.long, device=P.device) | |
| for s0 in range(0, n, rows): | |
| s1 = min(n, s0 + rows) | |
| d2 = torch.bmm(Xq[:, s0:s1], Yr) # (B, rows, n) | |
| d2.masked_fill_((obj[:, s0:s1, None] == obj[:, None, :]) | ~valid[:, None, :], float('inf')) | |
| idx[:, s0:s1] = d2.argmin(dim=2) | |
| nearest = torch.gather(P, 1, idx.unsqueeze(-1).expand(B, n, D)) | |
| vec = torch.where(valid.unsqueeze(-1), nearest - P, torch.zeros((), device=P.device)) | |
| out = torch.zeros(B, Nobj * Nv, D, device=P.device, dtype=torch.float32) | |
| out.scatter_(1, order.unsqueeze(-1).expand(B, n, D), vec) | |
| return out.reshape(B, Nobj, Nv, D).to(positions.dtype) | |
| def compute_nearest_distance_to_list_cuda_chunked( | |
| tensor_a: torch.Tensor, | |
| tensor_list: List[torch.Tensor], | |
| chunk_size: int = 10000 | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Memory-efficient CUDA version using chunked processing for large tensors. | |
| Processes distance computation in chunks to avoid OOM on large datasets. | |
| Args: | |
| tensor_a: Query tensor on CUDA, shape [N, D] | |
| tensor_list: List of reference tensors on CUDA, each [M_i, D] | |
| chunk_size: Maximum number of reference points to process at once | |
| Returns: | |
| min_distances: shape [N] | |
| min_tensor_indices: shape [N] | |
| """ | |
| if not tensor_a.is_cuda: | |
| raise RuntimeError("tensor_a must be on CUDA; CPU execution is not supported.") | |
| for i, t in enumerate(tensor_list): | |
| if not t.is_cuda: | |
| raise RuntimeError(f"tensor_list[{i}] must be on CUDA; CPU execution is not supported.") | |
| N, D = tensor_a.shape | |
| device = tensor_a.device | |
| # Initialize results | |
| global_min_dists = torch.full((N,), float('inf'), device=device) | |
| global_tensor_idx = torch.zeros(N, dtype=torch.long, device=device) | |
| for tensor_idx, ref_tensor in enumerate(tensor_list): | |
| M = ref_tensor.shape[0] | |
| # Process this reference tensor in chunks | |
| chunk_min_dists = torch.full((N,), float('inf'), device=device) | |
| for start_idx in range(0, M, chunk_size): | |
| end_idx = min(start_idx + chunk_size, M) | |
| ref_chunk = ref_tensor[start_idx:end_idx] | |
| # Compute distances for this chunk | |
| dist_matrix = torch.cdist(tensor_a, ref_chunk, p=2.0) # [N, chunk_size] | |
| chunk_dists, _ = dist_matrix.min(dim=1) # [N] | |
| # Update chunk minimums | |
| chunk_min_dists = torch.minimum(chunk_min_dists, chunk_dists) | |
| # Update global minimums | |
| update_mask = chunk_min_dists < global_min_dists | |
| global_min_dists[update_mask] = chunk_min_dists[update_mask] | |
| global_tensor_idx[update_mask] = tensor_idx | |
| return global_min_dists, global_tensor_idx | |
| def compute_nearest_distance_to_list_cuda_optimized( | |
| tensor_a: torch.Tensor, | |
| tensor_list: List[torch.Tensor], | |
| max_memory_mb: float = 1000 | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Adaptive CUDA implementation that chooses strategy based on memory constraints. | |
| Args: | |
| tensor_a: Query tensor on CUDA, shape [N, D] | |
| tensor_list: List of reference tensors on CUDA | |
| max_memory_mb: Maximum memory to use in MB | |
| Returns: | |
| min_distances: shape [N] | |
| min_tensor_indices: shape [N] | |
| """ | |
| if not tensor_a.is_cuda: | |
| raise RuntimeError("All tensors must be on CUDA; CPU execution is not supported.") | |
| N, D = tensor_a.shape | |
| # Estimate memory requirement | |
| total_ref_points = sum(t.shape[0] for t in tensor_list) | |
| memory_needed_mb = (N * total_ref_points * 4) / (1024 * 1024) # float32 = 4 bytes | |
| if memory_needed_mb <= max_memory_mb: | |
| # Use fast concatenated version | |
| return compute_nearest_distance_to_list_cuda(tensor_a, tensor_list, batch_mode=False) | |
| else: | |
| # Use memory-efficient chunked version | |
| chunk_size = int(max_memory_mb * 1024 * 1024 / (N * 4)) | |
| return compute_nearest_distance_to_list_cuda_chunked(tensor_a, tensor_list, chunk_size) | |
| # Self-tests | |
| if __name__ == "__main__": | |
| print("=" * 60) | |
| print("Testing CUDA-accelerated nearest distance computation") | |
| print("=" * 60) | |
| # Create test data ON CUDA | |
| tensor_a = torch.rand(1000, 3).cuda() # 1000 query points | |
| tensor_list = [ | |
| torch.rand(500, 3).cuda(), # 500 reference points | |
| torch.rand(800, 3).cuda(), # 800 reference points | |
| torch.rand(1200, 3).cuda(), # 1200 reference points | |
| ] | |
| print("\nTest 1: Basic CUDA version") | |
| try: | |
| min_dists, min_indices = compute_nearest_distance_to_list_cuda( | |
| tensor_a, tensor_list, batch_mode=False | |
| ) | |
| print(f"Success. Min distances shape: {min_dists.shape}") | |
| print(f" First 5 distances: {min_dists[:5]}") | |
| print(f" Tensor indices: {min_indices[:5]}") | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| print("\nTest 1b: With vector output") | |
| try: | |
| min_dists, min_indices, vectors = compute_nearest_distance_to_list_cuda( | |
| tensor_a, tensor_list, batch_mode=False, return_vectors=True | |
| ) | |
| print(f"Success. Vectors shape: {vectors.shape}") | |
| print(f" First vector: {vectors[0]}") | |
| # Verify vector length matches distance | |
| vec_norm = vectors[0].norm() | |
| print(f" Vector norm: {vec_norm:.6f}, Distance: {min_dists[0]:.6f}") | |
| print(f" Match: {torch.allclose(vec_norm, min_dists[0], rtol=1e-5)}") | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| print("\nTest 2: Batch mode") | |
| tensor_a_batch = torch.rand(4, 100, 3).cuda() # Batch of 4 | |
| tensor_list_batch = [ | |
| torch.rand(4, 50, 3).cuda(), | |
| torch.rand(4, 80, 3).cuda(), | |
| ] | |
| try: | |
| min_dists_batch, min_indices_batch = compute_nearest_distance_to_list_cuda( | |
| tensor_a_batch, tensor_list_batch, batch_mode=True | |
| ) | |
| print(f"Success. Batch distances shape: {min_dists_batch.shape}") | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| print("\nTest 3: Chunked version for large tensors") | |
| try: | |
| min_dists_chunked, min_indices_chunked = compute_nearest_distance_to_list_cuda_chunked( | |
| tensor_a, tensor_list, chunk_size=100 | |
| ) | |
| print(f"Success. Chunked distances shape: {min_dists_chunked.shape}") | |
| # Verify chunked gives same result as full | |
| diff = (min_dists - min_dists_chunked).abs().max() | |
| print(f" Difference from full version: {diff:.6f}") | |
| except Exception as e: | |
| print(f"Error: {e}") | |
| print("\nTest 4: Testing with CPU tensor (should fail)") | |
| tensor_cpu = torch.rand(10, 3) # CPU tensor | |
| try: | |
| min_dists, min_indices = compute_nearest_distance_to_list_cuda( | |
| tensor_cpu, tensor_list, batch_mode=False | |
| ) | |
| print("FAILED: expected an error for the CPU tensor") | |
| except RuntimeError as e: | |
| print(f"Correctly rejected CPU tensor: {e}") | |
| print("\nTest 5: Performance comparison") | |
| import time | |
| # Large test case | |
| large_query = torch.rand(5000, 3).cuda() | |
| large_refs = [torch.rand(2000, 3).cuda() for _ in range(5)] | |
| # Time the optimized version | |
| start = time.time() | |
| result = compute_nearest_distance_to_list_cuda_optimized(large_query, large_refs) | |
| cuda_time = time.time() - start | |
| print(f"Processed 5000 queries against 10000 references in {cuda_time:.3f}s") | |
| print(f" Throughput: {5000 * 10000 / cuda_time / 1e6:.1f} million distance computations/second") | |
| print("\n" + "=" * 60) | |
| print("Test 6: Vectorized multi-object KNN (correctness verification)") | |
| print("=" * 60) | |
| # Create test data: B batches, Nobj objects, Nv vertices each | |
| B, Nobj, Nv = 4, 5, 100 | |
| FAR = 1e6 | |
| positions = torch.rand(B, Nobj, Nv, 3).cuda() | |
| # Create a valid mask (80% of vertices are valid) | |
| valid_mask = torch.rand(B, Nobj, Nv).cuda() > 0.2 | |
| print(f" Test setup: B={B}, Nobj={Nobj}, Nv={Nv}") | |
| print(f" Valid vertices per object: {valid_mask.float().mean(dim=(0,2)) * Nv}") | |
| # Method 1: loop-based reference implementation | |
| print("\n Method 1: Loop-based KNN...") | |
| dist_vectors_loop = torch.zeros(B, Nobj, Nv, 3, device=positions.device) | |
| start = time.time() | |
| for i in range(Nobj): | |
| body_pos = positions[:, i] # (B, Nv, 3) | |
| ref_list = [] | |
| for j in range(Nobj): | |
| if j == i: | |
| continue | |
| other_pos = positions[:, j].clone() | |
| other_vmask = valid_mask[:, j].unsqueeze(-1) # (B, Nv, 1) | |
| other_pos = torch.where(other_vmask, other_pos, other_pos + FAR) | |
| ref_list.append(other_pos) | |
| _, _, body_dist_vec = compute_nearest_distance_to_list_cuda( | |
| body_pos, ref_list, batch_mode=True, return_vectors=True | |
| ) | |
| dist_vectors_loop[:, i] = body_dist_vec | |
| loop_time = time.time() - start | |
| print(f" Loop-based completed in {loop_time*1000:.2f}ms") | |
| # Method 2: Vectorized approach | |
| print("\n Method 2: Vectorized KNN...") | |
| start = time.time() | |
| dist_vectors_vec = compute_knn_all_objects_vectorized(positions, valid_mask, FAR) | |
| vec_time = time.time() - start | |
| print(f" Vectorized completed in {vec_time*1000:.2f}ms") | |
| # Compare results | |
| print("\n Comparing results...") | |
| # Only compare valid query vertices (invalid queries may have different results) | |
| valid_query_mask = valid_mask.unsqueeze(-1).expand_as(dist_vectors_loop) | |
| diff = (dist_vectors_loop - dist_vectors_vec).abs() | |
| diff_valid = diff[valid_query_mask] | |
| max_diff = diff_valid.max().item() if diff_valid.numel() > 0 else 0 | |
| mean_diff = diff_valid.mean().item() if diff_valid.numel() > 0 else 0 | |
| print(f" Max difference (valid vertices): {max_diff:.6e}") | |
| print(f" Mean difference (valid vertices): {mean_diff:.6e}") | |
| if max_diff < 1e-5: | |
| print(f" PASSED: Results match within tolerance (1e-5)") | |
| else: | |
| print(f" FAILED: Results differ by more than tolerance") | |
| # Debug: find where they differ | |
| diff_mask = diff > 1e-5 | |
| num_diff = diff_mask.sum().item() | |
| print(f" Number of differing elements: {num_diff}") | |
| # Speedup | |
| speedup = loop_time / vec_time if vec_time > 0 else float('inf') | |
| print(f"\n Speedup: {speedup:.2f}x") | |
| # Test with larger data | |
| print("\n Testing with larger data (B=8, Nobj=10, Nv=500)...") | |
| B2, Nobj2, Nv2 = 8, 10, 500 | |
| positions2 = torch.rand(B2, Nobj2, Nv2, 3).cuda() | |
| valid_mask2 = torch.rand(B2, Nobj2, Nv2).cuda() > 0.1 | |
| # Warmup | |
| _ = compute_knn_all_objects_vectorized(positions2, valid_mask2, FAR) | |
| torch.cuda.synchronize() | |
| start = time.time() | |
| _ = compute_knn_all_objects_vectorized(positions2, valid_mask2, FAR) | |
| torch.cuda.synchronize() | |
| vec_time2 = time.time() - start | |
| print(f" Vectorized (large): {vec_time2*1000:.2f}ms") | |
| print(f" Throughput: {B2 * Nobj2 * Nv2 / vec_time2 / 1e6:.2f}M vertices/sec") | |
| print("\n" + "=" * 60) | |
| print("Test 7: Invalid vertex/object masking verification") | |
| print("=" * 60) | |
| B, Nobj, Nv = 2, 4, 50 | |
| positions = torch.rand(B, Nobj, Nv, 3).cuda() | |
| # Create mask with some fully invalid objects and some invalid vertices | |
| valid_vtx_mask = torch.ones(B, Nobj, Nv, dtype=torch.bool, device='cuda') | |
| valid_obj_mask = torch.ones(B, Nobj, dtype=torch.bool, device='cuda') | |
| # Make object 2 fully invalid (no valid vertices) | |
| valid_vtx_mask[:, 2, :] = False | |
| valid_obj_mask[:, 2] = False | |
| # Make some vertices invalid in valid objects | |
| valid_vtx_mask[0, 0, 10:20] = False # batch 0, object 0, vertices 10-19 invalid | |
| valid_vtx_mask[1, 1, 30:40] = False # batch 1, object 1, vertices 30-39 invalid | |
| print(f" Setup: B={B}, Nobj={Nobj}, Nv={Nv}") | |
| print(f" Object 2 is fully invalid (all vertices masked)") | |
| print(f" Some vertices in valid objects are also masked") | |
| dist_vectors = compute_knn_all_objects_vectorized( | |
| positions, valid_vtx_mask, FAR, valid_obj_mask | |
| ) | |
| # Check 1: Invalid object (object 2) should have ALL zero distance vectors | |
| obj2_dist = dist_vectors[:, 2, :, :] # (B, Nv, 3) | |
| obj2_is_zero = (obj2_dist.abs() < 1e-10).all() | |
| print(f"\n Check 1: Invalid object (obj 2) has zero dist_vec: {obj2_is_zero.item()}") | |
| if obj2_is_zero: | |
| print(" PASSED: Invalid object has all zeros") | |
| else: | |
| print(f" FAILED: Non-zero values found: max={obj2_dist.abs().max().item():.6e}") | |
| # Check 2: Invalid vertices in valid objects should have zero distance vectors | |
| invalid_vtx_dist_0 = dist_vectors[0, 0, 10:20, :] # batch 0, obj 0, vtx 10-19 | |
| invalid_vtx_dist_1 = dist_vectors[1, 1, 30:40, :] # batch 1, obj 1, vtx 30-39 | |
| vtx0_is_zero = (invalid_vtx_dist_0.abs() < 1e-10).all() | |
| vtx1_is_zero = (invalid_vtx_dist_1.abs() < 1e-10).all() | |
| print(f"\n Check 2: Invalid vertices have zero dist_vec:") | |
| print(f" batch 0, obj 0, vtx 10-19: {vtx0_is_zero.item()}") | |
| print(f" batch 1, obj 1, vtx 30-39: {vtx1_is_zero.item()}") | |
| if vtx0_is_zero and vtx1_is_zero: | |
| print(" PASSED: Invalid vertices have all zeros") | |
| else: | |
| print(f" FAILED: Non-zero values found") | |
| # Check 3: Valid vertices in valid objects should have NON-zero distance vectors | |
| # (at least most of them, since they should find nearest neighbors on other objects) | |
| valid_vtx_dist_0 = dist_vectors[0, 0, :10, :] # batch 0, obj 0, vtx 0-9 (valid) | |
| valid_vtx_dist_1 = dist_vectors[0, 1, :30, :] # batch 0, obj 1, vtx 0-29 (valid) | |
| valid_has_nonzero_0 = (valid_vtx_dist_0.abs() > 1e-10).any() | |
| valid_has_nonzero_1 = (valid_vtx_dist_1.abs() > 1e-10).any() | |
| print(f"\n Check 3: Valid vertices have NON-zero dist_vec (found neighbors):") | |
| print(f" batch 0, obj 0, vtx 0-9: {valid_has_nonzero_0.item()}") | |
| print(f" batch 0, obj 1, vtx 0-29: {valid_has_nonzero_1.item()}") | |
| if valid_has_nonzero_0 and valid_has_nonzero_1: | |
| print(" PASSED: Valid vertices have non-zero distances") | |
| else: | |
| print(" FAILED: Valid vertices should find neighbors on other objects") | |
| print("\n" + "=" * 60) | |
| print("All tests completed! CUDA acceleration is working.") |