From a2f7561cf73825030649acdb411728fadcf458b3 Mon Sep 17 00:00:00 2001 From: AI Lab <129358391+1038lab@users.noreply.github.com> Date: Thu, 1 Jan 2026 19:04:31 -0800 Subject: [PATCH] Add files via upload --- models/sam3/perflib/__init__.py | 8 + models/sam3/perflib/associate_det_trk.py | 137 +++++ models/sam3/perflib/compile.py | 99 ++++ models/sam3/perflib/connected_components.py | 84 ++++ models/sam3/perflib/fa3.py | 27 + models/sam3/perflib/masks_ops.py | 69 +++ models/sam3/perflib/nms.py | 91 ++++ models/sam3/perflib/tests/assets/masks.tiff | Bin 0 -> 352484 bytes models/sam3/perflib/tests/tests.py | 59 +++ .../perflib/triton/connected_components.py | 468 ++++++++++++++++++ models/sam3/perflib/triton/nms.py | 124 +++++ 11 files changed, 1166 insertions(+) create mode 100644 models/sam3/perflib/__init__.py create mode 100644 models/sam3/perflib/associate_det_trk.py create mode 100644 models/sam3/perflib/compile.py create mode 100644 models/sam3/perflib/connected_components.py create mode 100644 models/sam3/perflib/fa3.py create mode 100644 models/sam3/perflib/masks_ops.py create mode 100644 models/sam3/perflib/nms.py create mode 100644 models/sam3/perflib/tests/assets/masks.tiff create mode 100644 models/sam3/perflib/tests/tests.py create mode 100644 models/sam3/perflib/triton/connected_components.py create mode 100644 models/sam3/perflib/triton/nms.py diff --git a/models/sam3/perflib/__init__.py b/models/sam3/perflib/__init__.py new file mode 100644 index 0000000..5c3823b --- /dev/null +++ b/models/sam3/perflib/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +is_enabled = False +if os.getenv("USE_PERFLIB", "1") == "1": + # print("Enabled the use of perflib.\n", end="") + is_enabled = True diff --git a/models/sam3/perflib/associate_det_trk.py b/models/sam3/perflib/associate_det_trk.py new file mode 100644 index 0000000..508ae81 --- /dev/null +++ b/models/sam3/perflib/associate_det_trk.py @@ -0,0 +1,137 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +from collections import defaultdict + +import torch +import torch.nn.functional as F +from sam3.perflib.masks_ops import mask_iou +from scipy.optimize import linear_sum_assignment + + +def associate_det_trk( + det_masks, + track_masks, + iou_threshold=0.5, + iou_threshold_trk=0.5, + det_scores=None, + new_det_thresh=0.0, +): + """ + Optimized implementation of detection <-> track association that minimizes DtoH syncs. + + Args: + det_masks: (N, H, W) tensor of predicted masks + track_masks: (M, H, W) tensor of track masks + + Returns: + new_det_indices: list of indices in det_masks considered 'new' + unmatched_trk_indices: list of indices in track_masks considered 'unmatched' + """ + with torch.autograd.profiler.record_function("perflib: associate_det_trk"): + assert isinstance(det_masks, torch.Tensor), "det_masks should be a tensor" + assert isinstance(track_masks, torch.Tensor), "track_masks should be a tensor" + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {}, {} # all detections are new + + if list(det_masks.shape[-2:]) != list(track_masks.shape[-2:]): + # resize to the smaller size to save GPU memory + if torch.numel(det_masks[-2:]) < torch.numel(track_masks[-2:]): + track_masks = ( + F.interpolate( + track_masks.unsqueeze(1).float(), + size=det_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + else: + # resize detections to track size + det_masks = ( + F.interpolate( + det_masks.unsqueeze(1).float(), + size=track_masks.shape[-2:], + mode="bilinear", + align_corners=False, + ).squeeze(1) + > 0 + ) + + det_masks = det_masks > 0 + track_masks = track_masks > 0 + + iou = mask_iou(det_masks, track_masks) # (N, M) + igeit = iou >= iou_threshold + igeit_any_dim_1 = igeit.any(dim=1) + igeit_trk = iou >= iou_threshold_trk + + iou_list = iou.cpu().numpy().tolist() + igeit_list = igeit.cpu().numpy().tolist() + igeit_any_dim_1_list = igeit_any_dim_1.cpu().numpy().tolist() + igeit_trk_list = igeit_trk.cpu().numpy().tolist() + + det_scores_list = ( + det_scores + if det_scores is None + else det_scores.cpu().float().numpy().tolist() + ) + + # Hungarian matching for tracks (one-to-one: each track matches at most one detection) + # For detections: allow many tracks to match to the same detection (many-to-one) + + # If either is empty, return all detections as new + if det_masks.size(0) == 0 or track_masks.size(0) == 0: + return list(range(det_masks.size(0))), [], {} + + # Hungarian matching: maximize IoU for tracks + cost_matrix = 1 - iou.cpu().numpy() # Hungarian solves for minimum cost + row_ind, col_ind = linear_sum_assignment(cost_matrix) + + def branchy_hungarian_better_uses_the_cpu( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ): + matched_trk = set() + matched_det = set() + matched_det_scores = {} # track index -> [det_score, det_score * iou] det score of matched detection mask + for d, t in zip(row_ind, col_ind): + matched_det_scores[t] = [ + det_scores_list[d], + det_scores_list[d] * iou_list[d][t], + ] + if igeit_trk_list[d][t]: + matched_trk.add(t) + matched_det.add(d) + + # Tracks not matched by Hungarian assignment above threshold are unmatched + unmatched_trk_indices = [ + t for t in range(track_masks.size(0)) if t not in matched_trk + ] + + # For detections: allow many tracks to match to the same detection (many-to-one) + # So, a detection is 'new' if it does not match any track above threshold + assert track_masks.size(0) == igeit.size( + 1 + ) # Needed for loop optimizaiton below + new_det_indices = [] + for d in range(det_masks.size(0)): + if not igeit_any_dim_1_list[d]: + if det_scores is not None and det_scores[d] >= new_det_thresh: + new_det_indices.append(d) + + # for each detection, which tracks it matched to (above threshold) + det_to_matched_trk = defaultdict(list) + for d in range(det_masks.size(0)): + for t in range(track_masks.size(0)): + if igeit_list[d][t]: + det_to_matched_trk[d].append(t) + + return ( + new_det_indices, + unmatched_trk_indices, + det_to_matched_trk, + matched_det_scores, + ) + + return (branchy_hungarian_better_uses_the_cpu)( + cost_matrix, row_ind, col_ind, iou_list, det_masks, track_masks + ) diff --git a/models/sam3/perflib/compile.py b/models/sam3/perflib/compile.py new file mode 100644 index 0000000..f427aa7 --- /dev/null +++ b/models/sam3/perflib/compile.py @@ -0,0 +1,99 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def recursive_fn_factory(fn): + def recursive_fn(b): + if isinstance(b, dict): + return {k: recursive_fn(b[k]) for k in b} + if isinstance(b, list): + return [recursive_fn(t) for t in b] + if isinstance(b, tuple): + return tuple(recursive_fn(t) for t in b) + if isinstance(b, torch.Tensor): + return fn(b) + # Yes, writing out an explicit white list of + # trivial types is tedious, but so are bugs that + # come from not applying fn, when expected to have + # applied it. + if b is None: + return b + trivial_types = [bool, int] + for t in trivial_types: + if isinstance(b, t): + return b + raise TypeError(f"Unexpected type {type(b)}") + + return recursive_fn + + +recursive_contiguous = recursive_fn_factory(lambda x: x.contiguous()) +recursive_clone = recursive_fn_factory(torch.clone) + + +def compile_wrapper( + fn, *, mode="max-autotune", fullgraph=True, dynamic=False, name=None +): + compiled_fn = torch.compile(fn, mode=mode, fullgraph=fullgraph, dynamic=dynamic) + + def compiled_fn_wrapper(*args, **kwargs): + with torch.autograd.profiler.record_function( + f"compiled {fn}" if name is None else name + ): + cont_args = recursive_contiguous(args) + cont_kwargs = recursive_contiguous(kwargs) + result = compiled_fn(*cont_args, **cont_kwargs) + cloned_result = recursive_clone(result) + return cloned_result + + return compiled_fn_wrapper + + +def shape_logging_wrapper(fn, keep_kwargs, enable_logging=False): + """ + Wraps a function and prints the shapes of all tensor inputs. + Only prints when a new combination of shapes is seen. + Thread-safe. + + Args: + fn: Function to wrap + enable_logging: Boolean flag to enable/disable logging + """ + seen_shapes = set() + + def get_shape(obj): + if isinstance(obj, torch.Tensor): + return obj.shape + elif isinstance(obj, (list, tuple)): + if len(obj) > 1: + return tuple(get_shape(x) for x in obj) + return get_shape(obj[0]) + elif isinstance(obj, dict): + return tuple(sorted((k, get_shape(v)) for k, v in obj.items())) + else: + return type(obj).__name__ + + def wrapper(*args, **kwargs): + shapes = tuple(get_shape(arg) for arg in args) + tuple( + (k, get_shape(v)) + for k, v in kwargs.items() + if isinstance(v, (torch.Tensor, list)) + and (len(keep_kwargs) > 0 and k in keep_kwargs) + ) + if shapes not in seen_shapes: + seen_shapes.add(shapes) + if enable_logging: + print(f"[ShapeLogger] New input shapes for {fn.__qualname__}: {shapes}") + return fn(*args, **kwargs) + + # Allow toggling the flag at runtime + wrapper.enable_logging = enable_logging + + def set_logging(enabled=False): + nonlocal enable_logging + enable_logging = enabled + wrapper.enable_logging = enable_logging + + wrapper.set_logging = set_logging + return wrapper diff --git a/models/sam3/perflib/connected_components.py b/models/sam3/perflib/connected_components.py new file mode 100644 index 0000000..c96932a --- /dev/null +++ b/models/sam3/perflib/connected_components.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import logging + +import torch + +try: + from cc_torch import get_connected_components + + HAS_CC_TORCH = True +except ImportError: + logging.debug( + "cc_torch not found. Consider installing for better performance. Command line:" + " pip install git+https://github.com/ronghanghu/cc_torch.git" + ) + HAS_CC_TORCH = False + + +def connected_components_cpu_single(values: torch.Tensor): + assert values.dim() == 2 + from skimage.measure import label + + labels, num = label(values.cpu().numpy(), return_num=True) + labels = torch.from_numpy(labels) + counts = torch.zeros_like(labels) + for i in range(1, num + 1): + cur_mask = labels == i + cur_count = cur_mask.sum() + counts[cur_mask] = cur_count + return labels, counts + + +def connected_components_cpu(input_tensor: torch.Tensor): + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + batch_size = input_tensor.shape[0] + labels_list = [] + counts_list = [] + for b in range(batch_size): + labels, counts = connected_components_cpu_single(input_tensor[b]) + labels_list.append(labels) + counts_list.append(counts) + labels_tensor = torch.stack(labels_list, dim=0).to(input_tensor.device) + counts_tensor = torch.stack(counts_list, dim=0).to(input_tensor.device) + return labels_tensor.view(out_shape), counts_tensor.view(out_shape) + + +def connected_components(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D tensors, using the best available backend. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, torch.Tensor]: Both tensors have the same shape as input_tensor. + - A tensor with dense labels. Background is 0. + - A tensor with the size of the connected component for each pixel. + """ + if input_tensor.dim() == 3: + input_tensor = input_tensor.unsqueeze(1) + + assert ( + input_tensor.dim() == 4 and input_tensor.shape[1] == 1 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + if input_tensor.is_cuda: + if HAS_CC_TORCH: + return get_connected_components(input_tensor.to(torch.uint8)) + else: + # triton fallback + from sam3.perflib.triton.connected_components import ( + connected_components_triton, + ) + + return connected_components_triton(input_tensor) + + # CPU fallback + return connected_components_cpu(input_tensor) diff --git a/models/sam3/perflib/fa3.py b/models/sam3/perflib/fa3.py new file mode 100644 index 0000000..8f8c9bd --- /dev/null +++ b/models/sam3/perflib/fa3.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +@torch.library.custom_op("flash::flash_attn_func", mutates_args=()) +def flash_attn_func_op( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor +) -> torch.Tensor: + from flash_attn_interface import flash_attn_func as fa3 + + return fa3(q, k, v) + + +def flash_attn_func(q, k, v): + dtype = torch.float8_e4m3fn + return flash_attn_func_op(q.to(dtype), k.to(dtype), v.to(dtype)).to(q.dtype) + + +@flash_attn_func_op.register_fake +def _(q, k, v, **kwargs): + # two outputs: + # 1. output: (batch, seq_len, num_heads, head_dim) + # 2. softmax_lse: (batch, num_heads, seq_len) with dtype=torch.float32 + # output needs to be bfloat16, not float8! + meta_q = torch.empty_like(q, dtype=torch.bfloat16).contiguous() + return meta_q diff --git a/models/sam3/perflib/masks_ops.py b/models/sam3/perflib/masks_ops.py new file mode 100644 index 0000000..48299d5 --- /dev/null +++ b/models/sam3/perflib/masks_ops.py @@ -0,0 +1,69 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import torch + + +def masks_to_boxes(masks: torch.Tensor, obj_ids: list[int]): + with torch.autograd.profiler.record_function("perflib: masks_to_boxes"): + # Sanity check based on callsite for replacement + assert masks.shape[0] == len(obj_ids) + assert masks.dim() == 3 + + # Based on torchvision masks_to_boxes + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device, dtype=torch.float) + + N, H, W = masks.shape + device = masks.device + y = torch.arange(H, device=device).view(1, H) + x = torch.arange(W, device=device).view(1, W) + + masks_with_obj = masks != 0 # N, H, W + masks_with_obj_x = masks_with_obj.amax( + dim=1 + ) # N, H (which columns have objects) + masks_with_obj_y = masks_with_obj.amax(dim=2) # N, W (which rows have objects) + masks_without_obj_x = ~masks_with_obj_x + masks_without_obj_y = ~masks_with_obj_y + + bounding_boxes_0 = torch.amin( + (masks_without_obj_x * W) + (masks_with_obj_x * x), dim=1 + ) + bounding_boxes_1 = torch.amin( + (masks_without_obj_y * H) + (masks_with_obj_y * y), dim=1 + ) + bounding_boxes_2 = torch.amax(masks_with_obj_x * x, dim=1) + bounding_boxes_3 = torch.amax(masks_with_obj_y * y, dim=1) + + bounding_boxes = torch.stack( + [bounding_boxes_0, bounding_boxes_1, bounding_boxes_2, bounding_boxes_3], + dim=1, + ).to(dtype=torch.float) + assert bounding_boxes.shape == (N, 4) + assert bounding_boxes.device == masks.device + assert bounding_boxes.dtype == torch.float + return bounding_boxes + + +def mask_iou(pred_masks: torch.Tensor, gt_masks: torch.Tensor) -> torch.Tensor: + """ + Compute the IoU (Intersection over Union) between predicted masks and ground truth masks. + Args: + - pred_masks: (N, H, W) bool Tensor, containing binary predicted segmentation masks + - gt_masks: (M, H, W) bool Tensor, containing binary ground truth segmentation masks + Returns: + - ious: (N, M) float Tensor, containing IoUs for each pair of predicted and ground truth masks + """ + assert pred_masks.dtype == gt_masks.dtype == torch.bool + N, H, W = pred_masks.shape + M, _, _ = gt_masks.shape + + # Flatten masks: (N, 1, H*W) and (1, M, H*W) + pred_flat = pred_masks.view(N, 1, H * W) + gt_flat = gt_masks.view(1, M, H * W) + + # Compute intersection and union: (N, M) + intersection = (pred_flat & gt_flat).sum(dim=2).float() + union = (pred_flat | gt_flat).sum(dim=2).float() + ious = intersection / union.clamp(min=1) + return ious # shape: (N, M) diff --git a/models/sam3/perflib/nms.py b/models/sam3/perflib/nms.py new file mode 100644 index 0000000..b3efc59 --- /dev/null +++ b/models/sam3/perflib/nms.py @@ -0,0 +1,91 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import logging + +import numpy as np +import torch + +from sam3.perflib.masks_ops import mask_iou + + +try: + from torch_generic_nms import generic_nms as generic_nms_cuda + + GENERIC_NMS_AVAILABLE = True +except ImportError: + logging.debug( + "Falling back to triton or CPU mask NMS implementation -- please install `torch_generic_nms` via\n\t" + 'pip uninstall -y torch_generic_nms; TORCH_CUDA_ARCH_LIST="8.0 9.0" pip install git+https://github.com/ronghanghu/torch_generic_nms' + ) + GENERIC_NMS_AVAILABLE = False + + +def nms_masks( + pred_probs: torch.Tensor, + pred_masks: torch.Tensor, + prob_threshold: float, + iou_threshold: float, +) -> torch.Tensor: + """ + Args: + - pred_probs: (num_det,) float Tensor, containing the score (probability) of each detection + - pred_masks: (num_det, H_mask, W_mask) float Tensor, containing the binary segmentation mask of each detection + - prob_threshold: float, score threshold to prefilter detections (NMS is performed on detections above threshold) + - iou_threshold: float, mask IoU threshold for NMS + + Returns: + - keep: (num_det,) bool Tensor, indicating whether each detection is kept after score thresholding + NMS + """ + # prefilter the detections with prob_threshold ("valid" are those above prob_threshold) + is_valid = pred_probs > prob_threshold # (num_det,) + probs = pred_probs[is_valid] # (num_valid,) + masks_binary = pred_masks[is_valid] > 0 # (num_valid, H_mask, W_mask) + if probs.numel() == 0: + return is_valid # no valid detection, return empty keep mask + + ious = mask_iou(masks_binary, masks_binary) # (num_valid, num_valid) + kept_inds = generic_nms(ious, probs, iou_threshold) + + # valid_inds are the indices among `probs` of valid detections before NMS (or -1 for invalid) + valid_inds = torch.where(is_valid, is_valid.cumsum(dim=0) - 1, -1) # (num_det,) + keep = torch.isin(valid_inds, kept_inds) # (num_det,) + return keep + + +def generic_nms( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix.""" + + assert ious.dim() == 2 and ious.size(0) == ious.size(1) + assert scores.dim() == 1 and scores.size(0) == ious.size(0) + + if ious.is_cuda: + if GENERIC_NMS_AVAILABLE: + return generic_nms_cuda(ious, scores, iou_threshold, use_iou_matrix=True) + else: + from sam3.perflib.triton.nms import nms_triton + + return nms_triton(ious, scores, iou_threshold) + + return generic_nms_cpu(ious, scores, iou_threshold) + + +def generic_nms_cpu( + ious: torch.Tensor, scores: torch.Tensor, iou_threshold=0.5 +) -> torch.Tensor: + """ + A generic version of `torchvision.ops.nms` that takes a pairwise IoU matrix. (CPU implementation + based on https://github.com/jwyang/faster-rcnn.pytorch/blob/master/lib/model/nms/nms_cpu.py) + """ + ious_np = ious.float().detach().cpu().numpy() + scores_np = scores.float().detach().cpu().numpy() + order = scores_np.argsort()[::-1] + kept_inds = [] + while order.size > 0: + i = order.item(0) + kept_inds.append(i) + inds = np.where(ious_np[i, order[1:]] <= iou_threshold)[0] + order = order[inds + 1] + + return torch.tensor(kept_inds, dtype=torch.int64, device=scores.device) diff --git a/models/sam3/perflib/tests/assets/masks.tiff b/models/sam3/perflib/tests/assets/masks.tiff new file mode 100644 index 0000000000000000000000000000000000000000..7a8efc6dd0ee958b3ebf2733a0bfb32d39449478 GIT binary patch literal 352484 zcmeI&&#Gis6$ju`_x8QXAOw1l1d~a~D2NCKL2%$apidy}pFtV}4gxw6@d11R(TPJJ z!6%51;5yZd~*>UKZguDfr>?{02}?~I@Qb9_JF@7->CuUYecbLI!H)_gq- zA8zgr_kOe=|7Ses(`N4v#`E79Kl|_a9$wrHfBfWkzxnl7Uwrc8kAME7?|uB)XFuHi z^5Re1FZZu)rS1FoU%q_#^2^_T@%^v<&@VJ2FjHXDzGj|vecRh`yYcmvIHRHm>9!xt zIPc1~9*tKf;E8lqT6iX2g%+NTS0G^Vd_lbI4GiNY4_G|kgm|eNSUlgPc!?V*-%jAl zez!cm9R1V3=}*hkuk?I+KA--nn|`C`)AKp~X?gmU)t}GtZ+a%Yp5A@28z-bo{J#$) z(}aMFVXkpd^?Y#y*HB;b&Y*)!UbV_u)$^TyLRTSQ_|h{ST==H7jysWXt#X#X_h`W7 zuUYZL2MJfK<#Dd=5;WXFX;m+-$7uN=WFhzhw9_O+^bgEAICbZSf#$XmmaE*2Xn7l zWq%y&uws?^=3aWJJ|4`yYL)$Qtiy^`>YID%q560*_o`L)$FUA8R;h39rHAU{!Q87> z*&oL`tXQSK*_R&5j|Q`ETIYZ$*5T{xlsEg4tt^ewm|{Zdz-Y2i{=tb5MMfa{jE-0dd?cx<~VDWtWcvTyCh*vaV@qA6Zk_{}LuZvf(fyMK^7Oz?ZAH*vb zuz0>B@k%wYc)p|Y3N^5Jz7z4vG_ZKSlksXa@JzfC0gLB58?QhEi{}gCrEg&Id||xY z4V(}!b-?2JCdEtKz~cEP#!K74;`t`WOWVL~|N23_1PBlyK!5-N0t5&UAV7cs0RjXF z5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk z1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs z0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZ zfB*pk1PBlya0LR-+Se7-vup`%xZQZ!<5xcILAvb+E7!E(^+)5?33wu2m4GMX)d+Yd zUWI^XqHz8i?2Ff=IEchG2r#?#m^zW&kxl8Wze0n}Fd0w7+ zt>@G8`P9$cCHGZ-K8FbbFY&s%wbk>*4ZOORTvkxnz{`%izSTj%>nr%uip~bS^ssAO ze2_K>@~+5wB_8{?I3VDWsrcy$}Nk5~35%&UAw10LcPeL{=p zYvPq`VDWriym}3MEnc}Fp|9sF7Vts5VozxCd`IGyYGCnvN8{CL;E8x;euSRNS0v!c zctxJj;`z?RE78E>`Oe16-@qVV`X8ZDzT^SJc*&p8;`t`TOWnZY`6k6n-N38;i&ljZ zAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C7 z2oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N z0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+ z009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBng zVuAZBu5E$R?(b^{ES_(SSGs}4^X=jlZea0z`*>9wc!*atVDWrSypjzpp0A5nuz|(% zy%w)p10TdI7O;4}Bk@W#uz0?s@d`Dtc)kBctVTkI~%V+ z1B>Si;-zn3@qA&t+zp%%FLl7;`6k6n+`!`bCdNzKz~cEP$IIEk3*w~=SUlfF@e($$ zc)kncrE6gEd>6+{*T4n-v6>?Y5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs z0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZ zfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&U zAV7cs0RjXF5FkK+009C72oNAZfB*pk1PBlyK!5-N0t5&UAV7cs0RjXF5FpSCymiT| z-#jL459;Qgt7a50YR`gp_%$n-s zVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0IsXiX2NYob@ zhiX>zTT{%M>f>RGM17HQsAfgKHN~u{J|3n>)E60tYF6}HQ_Py`<6(+KeUWjfW<|d> z#jL459;Qgt7a50YR`gp_%$n-sVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_W zibQ>paj0fRzct0IsXiX2NYob@hiX>zTT{%M>f>RGM17HQsAfgKHN~u{J|3n>)E60t zYF6}HQ_Py`<6(+KeUWjfW<|d>#jL459;Qgt7a50YR`gp_%$n-sVTwe3k#VSIMZY!0 ztf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0IsXiX2NYob@hiX>zTT{%M>f>RG zM17HQsAfgKHN~u{J|3n>)E60tYF6}HQ_Py`<6(+KeUWjfW<|d>#jL459;Qgt7a50Y zR`gp_%$n-sVTwe3k#VSIMZY!0tf@X8rbyHm8HZ|C^jlNRn(E_WibQ>paj0fRzct0I zsXiX2NYob@hiX>zTT{%M>f>RGM17HQsAfgK)rFOn z(~NI6x4X~Z9zVa`kGJdYoALAg`r%*0@Lu!z`^`BYyjt`1FnqYV+nxEB7k9%)`|;n$ z@u$t+e~k0(J?(p;D*wD$tjeGM;?b)7+09~Ae*1V;KEGY8%D?{d(JcSHT$Mlh)uUDU Rw`Yr4j^8-lch2^^=zrO2=;8nX literal 0 HcmV?d00001 diff --git a/models/sam3/perflib/tests/tests.py b/models/sam3/perflib/tests/tests.py new file mode 100644 index 0000000..0fb88ad --- /dev/null +++ b/models/sam3/perflib/tests/tests.py @@ -0,0 +1,59 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +import os + +import numpy as np +import pytest +import torch +from PIL import Image +from sam3.perflib.masks_ops import masks_to_boxes + + +class TestMasksToBoxes: + def test_masks_box(self): + def masks_box_check(masks, expected, atol=1e-4): + out = masks_to_boxes(masks, [1 for _ in range(masks.shape[0])]) + assert out.dtype == torch.float + print("out: ", out) + print("expected: ", expected) + torch.testing.assert_close( + out, expected, rtol=0.0, check_dtype=True, atol=atol + ) + + # Check for int type boxes. + def _get_image(): + assets_directory = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "assets" + ) + mask_path = os.path.join(assets_directory, "masks.tiff") + image = Image.open(mask_path) + return image + + def _create_masks(image, masks): + for index in range(image.n_frames): + image.seek(index) + frame = np.array(image) + masks[index] = torch.tensor(frame) + + return masks + + expected = torch.tensor( + [ + [127, 2, 165, 40], + [2, 50, 44, 92], + [56, 63, 98, 100], + [139, 68, 175, 104], + [160, 112, 198, 145], + [49, 138, 99, 182], + [108, 148, 152, 213], + ], + dtype=torch.float, + ) + + image = _get_image() + for dtype in [torch.float16, torch.float32, torch.float64]: + masks = torch.zeros( + (image.n_frames, image.height, image.width), dtype=dtype + ) + masks = _create_masks(image, masks) + masks_box_check(masks, expected) diff --git a/models/sam3/perflib/triton/connected_components.py b/models/sam3/perflib/triton/connected_components.py new file mode 100644 index 0000000..253ca9d --- /dev/null +++ b/models/sam3/perflib/triton/connected_components.py @@ -0,0 +1,468 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved +import math + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _any_combine(a, b): + return a | b + + +@triton.jit +def tl_any(a, dim=0): + return tl.reduce(a, dim, _any_combine) + + +# ============================================================================== +# ## Phase 1: Initialization Kernel +# ============================================================================== +# Each foreground pixel (value > 0) gets a unique label equal to its +# linear index. Background pixels (value == 0) get a sentinel label of -1. +# Note that the indexing is done across batch boundaries for simplicity +# (i.e., the first pixel of image 1 gets label H*W, etc.) + + +@triton.jit +def _init_labels_kernel( + input_ptr, labels_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + input_values = tl.load(input_ptr + offsets, mask=mask, other=0) + + indices = tl.where((input_values != 0), offsets, -1) + tl.store(labels_ptr + offsets, indices, mask=mask) + + +# ============================================================================== +# ## Phase 2: Local merging +# ============================================================================== +# Each pixel tries to merge with its 8-connected neighbors (up, down, left, right) +# if they have the same value. This is done using a disjoint-set union operation. + + +@triton.jit +def find(labels_ptr, indices, mask): + current_pids = indices + + # 'is_done' tracks lanes that have finished their work. + # A lane is initially "done" if it's not active (mask is False). + is_done = ~mask + + # Loop as long as there is at least one lane that is NOT done. + while tl_any(~is_done): + # The work_mask is for lanes that are still active and seeking their root. + work_mask = ~is_done + parents = tl.load(labels_ptr + current_pids, mask=work_mask, other=-1) + # A lane is now done if its parent is itself (it's a root) + # or if it hits a -1 sentinel (a safe exit condition). + is_root = parents == current_pids + is_sentinel = parents == -1 + is_done |= is_root | is_sentinel + + # For lanes that are not yet done, update their pid to their parent to continue traversal. + current_pids = tl.where(is_done, current_pids, parents) + # We could add the following line to do path compression, but experimentally it's slower + # tl.atomic_min(labels_ptr + indices, current_pids, mask=mask) + return current_pids + + +@triton.jit +def union(labels_ptr, a, b, process_mask): + # This function implements a disjoint-set union + # As an invariant, we use the fact that the roots have the lower id. That helps parallelization + # However, that is not sufficient by itself. Suppose two threads want to do union(0,2) and union(1,2) at the same time + # Then if we do a naive atomic_min, 0 and 1 will compete to be the new parent of 2 and min(0, 1) will win. + # However, 1 still needs to be merged with the new {0, 2} component. + # To ensure that merge is also done, we need to detect whether the merge was successful, and if not retry until it is + + current_a = a + current_b = b + + final_root = a + # A mask to track which lanes have successfully completed their union. + done_mask = ~process_mask # tl.zeros_like(a) == 1 # Init with all False + + while tl_any(~done_mask): + # Define the mask for lanes that still need work in this iteration + work_mask = process_mask & ~done_mask + + # Find the roots for the current a and b values in the active lanes + root_a = find(labels_ptr, current_a, work_mask) + tl.debug_barrier() + root_b = find(labels_ptr, current_b, work_mask) + + # 7. Merge logic + # If roots are already the same, the sets are already merged. Mark as done. + are_equal = root_a == root_b + final_root = tl.where(are_equal & work_mask & ~done_mask, root_a, final_root) + done_mask |= are_equal & work_mask + + # Define masks for the two merge cases (a < b or b < a) + a_is_smaller = root_a < root_b + + # Case 1: root_a < root_b. Attempt to set parent[root_b] = root_a + merge_mask_a_smaller = work_mask & a_is_smaller & ~are_equal + ptr_b = labels_ptr + root_b + old_val_b = tl.atomic_min(ptr_b, root_a, mask=merge_mask_a_smaller) + + # A lane is done if its atomic op was successful (old value was what we expected) + success_b = old_val_b == root_b + final_root = tl.where(success_b & work_mask & ~done_mask, root_a, final_root) + done_mask |= success_b & merge_mask_a_smaller + + # *** Crucial Retry Logic *** + # If the update failed (old_val_b != root_b), another thread interfered. + # We update `current_b` to this new root (`old_val_b`) and will retry in the next loop iteration. + current_b = tl.where(success_b | ~merge_mask_a_smaller, current_b, old_val_b) + + # Case 2: root_b < root_a. Attempt to set parent[root_a] = root_b + merge_mask_b_smaller = work_mask & ~a_is_smaller & ~are_equal + ptr_a = labels_ptr + root_a + old_val_a = tl.atomic_min(ptr_a, root_b, mask=merge_mask_b_smaller) + + success_a = old_val_a == root_a + final_root = tl.where(success_a & work_mask & ~done_mask, root_b, final_root) + done_mask |= success_a & merge_mask_b_smaller + + # *** Crucial Retry Logic *** + # Similarly, update `current_a` if the atomic operation failed. + current_a = tl.where(success_a | ~merge_mask_b_smaller, current_a, old_val_a) + + return final_root + + +@triton.jit +def _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + dx: tl.constexpr, + dy: tl.constexpr, +): + # Helper functions to compute merge with a specific neighbor offset (dx, dy) + + neighbor_h = offsets_h + dy + neighbor_w = offsets_w + dx + # Proper bounds checking: all four bounds must be satisfied + mask_n = ( + mask_2d + & (neighbor_h[:, None] >= 0) + & (neighbor_h[:, None] < H) + & (neighbor_w[None, :] >= 0) + & (neighbor_w[None, :] < W) + ) + + offsets_neighbor = neighbor_h[:, None] * W + neighbor_w[None, :] + neighbor_values = tl.load( + input_ptr + base_offset + offsets_neighbor, mask=mask_n, other=-1 + ) + + mask_n = tl.ravel(mask_n) + neighbor_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_neighbor), mask=mask_n, other=-1 + ) + + to_merge = ( + mask_n & (neighbor_labels != -1) & tl.ravel(current_values == neighbor_values) + ) + valid_write = valid_current & to_merge + + # returns new parents for the pixels that were merged (otherwise keeps current labels) + return tl.where( + valid_write, + union(labels_ptr, current_labels, neighbor_labels, valid_write), + current_labels, + ) + + +@triton.autotune( + configs=[ + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 16}, num_stages=1, num_warps=2 + ), + triton.Config( + {"BLOCK_SIZE_H": 4, "BLOCK_SIZE_W": 32}, num_stages=2, num_warps=4 + ), + ], + key=["H", "W"], + restore_value=["labels_ptr"], +) +@triton.jit +def _local_prop_kernel( + labels_ptr, + input_ptr, + H: tl.constexpr, + W: tl.constexpr, + BLOCK_SIZE_H: tl.constexpr, + BLOCK_SIZE_W: tl.constexpr, +): + # This is the meat of the Phase 2 to do local merging + # It will be launched with a 2D grid: + # - dim 0: batch index + # - dim 1: block index over HxW image (2D tiling) + pid_b = tl.program_id(0) + pid_hw = tl.program_id(1) + + # Calculate offsets for the core block + offsets_h = (pid_hw // tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_H + tl.arange( + 0, BLOCK_SIZE_H + ) + offsets_w = (pid_hw % tl.cdiv(W, BLOCK_SIZE_W)) * BLOCK_SIZE_W + tl.arange( + 0, BLOCK_SIZE_W + ) + + base_offset = pid_b * H * W + offsets_2d = offsets_h[:, None] * W + offsets_w[None, :] + mask_2d = (offsets_h[:, None] < H) & (offsets_w[None, :] < W) + mask_1d = tl.ravel(mask_2d) + + # Load the current labels for the block - these are parent pointers + current_labels = tl.load( + labels_ptr + tl.ravel(base_offset + offsets_2d), mask=mask_1d, other=-1 + ) + current_values = tl.load( + input_ptr + base_offset + offsets_2d, mask=mask_2d, other=-1 + ) + valid_current = mask_1d & (current_labels != -1) + + # Horizontal merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 0, + ) + # Vertical merge + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + 0, + -1, + ) + # Diagonal merges + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + -1, + ) + current_labels = _merge_helper( + input_ptr, + labels_ptr, + base_offset, + offsets_h, + offsets_w, + mask_2d, + valid_current, + current_values, + current_labels, + H, + W, + -1, + 1, + ) + + # This actually does some path compression, in a lightweight but beneficial way + tl.atomic_min( + labels_ptr + tl.ravel(base_offset + offsets_2d), current_labels, mask=mask_1d + ) + + +# ============================================================================== +# ## Phase 3: Pointer Jumping Kernel +# ============================================================================== +# This kernel performs pointer jumping to ensure that all pixels point directly to their root labels. +# This is done in a loop until convergence. + + +@triton.jit +def _pointer_jump_kernel( + labels_in_ptr, labels_out_ptr, numel: tl.constexpr, BLOCK_SIZE: tl.constexpr +): + """ + Pointer jumping kernel with double buffering to avoid race conditions. + Reads from labels_in_ptr and writes to labels_out_ptr. + """ + # This kernel is launched with a 1D grid, and does not care about batching explicitly. + # By construction, the labels are global indices across the batch, and we never perform + # cross-batch merges, so this is safe. + + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load current labels from input buffer + current_labels = tl.load(labels_in_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (current_labels != -1) + + # A mask to track which lanes have successfully completed their union. + done_mask = ~valid_mask + while tl_any(~(done_mask | ~valid_mask)): + parent_labels = tl.load( + labels_in_ptr + current_labels, mask=valid_mask, other=-1 + ) + + are_equal = current_labels == parent_labels + done_mask |= are_equal & valid_mask + + current_labels = tl.where( + ~done_mask, tl.minimum(current_labels, parent_labels), current_labels + ) + + # Write to output buffer (safe because we're not reading from it) + tl.store(labels_out_ptr + offsets, current_labels, mask=mask) + + +# ============================================================================== +# ## Phase 4: Kernels for Computing Component Sizes +# ============================================================================== + + +# Step 4.1: Count occurrences of each root label using atomic adds. +@triton.jit +def _count_labels_kernel(labels_ptr, sizes_ptr, numel, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final, converged labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Atomically increment the counter for each label. This builds a histogram. + tl.atomic_add(sizes_ptr + labels, 1, mask=valid_mask) + + +# Step 4.2: Broadcast the computed sizes back to the output tensor. +@triton.jit +def _broadcast_sizes_kernel( + labels_ptr, sizes_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < numel + + # Load the final labels + labels = tl.load(labels_ptr + offsets, mask=mask, other=-1) + valid_mask = mask & (labels != -1) + + # Look up the size for each label from the histogram + component_sizes = tl.load(sizes_ptr + labels, mask=valid_mask, other=0) + + # Write the size to the final output tensor. Background pixels get size 0. + tl.store(out_ptr + offsets, component_sizes, mask=mask) + + +def connected_components_triton(input_tensor: torch.Tensor): + """ + Computes connected components labeling on a batch of 2D integer tensors using Triton. + + Args: + input_tensor (torch.Tensor): A BxHxW integer tensor or Bx1xHxW. Non-zero values are considered foreground. Bool tensor also accepted + + Returns: + Tuple[torch.Tensor, int]: A tuple containing: + - A BxHxW output tensor with dense labels. Background is 0. + - A BxHxW tensor with the size of the connected component for each pixel. + """ + assert ( + input_tensor.is_cuda and input_tensor.is_contiguous() + ), "Input tensor must be a contiguous CUDA tensor." + out_shape = input_tensor.shape + if input_tensor.dim() == 4 and input_tensor.shape[1] == 1: + input_tensor = input_tensor.squeeze(1) + else: + assert ( + input_tensor.dim() == 3 + ), "Input tensor must be (B, H, W) or (B, 1, H, W)." + + B, H, W = input_tensor.shape + numel = B * H * W + device = input_tensor.device + + # --- Allocate Tensors --- + labels = torch.empty_like(input_tensor, dtype=torch.int32) + output = torch.empty_like(input_tensor, dtype=torch.int32) + + # --- Phase 1 --- + BLOCK_SIZE = 256 + grid_init = (triton.cdiv(numel, BLOCK_SIZE),) + _init_labels_kernel[grid_init]( + input_tensor, + labels, + numel, + BLOCK_SIZE=BLOCK_SIZE, + ) + + # --- Phase 2 --- + grid_local_prop = lambda meta: ( + B, + triton.cdiv(H, meta["BLOCK_SIZE_H"]) * triton.cdiv(W, meta["BLOCK_SIZE_W"]), + ) + _local_prop_kernel[grid_local_prop](labels, input_tensor, H, W) + + # --- Phase 3 --- + BLOCK_SIZE = 256 + grid_jump = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),) + _pointer_jump_kernel[grid_jump](labels, output, numel, BLOCK_SIZE=BLOCK_SIZE) + + # --- Phase 4 --- + # Allocate tensor to store the final output sizes + component_sizes_out = torch.empty_like(input_tensor, dtype=torch.int32) + + # Allocate a temporary 1D tensor to act as the histogram + # Size is numel because labels can be up to numel-1 + sizes_histogram = torch.zeros(numel, dtype=torch.int32, device=device) + + # 4.1: Count the occurrences of each label + grid_count = (triton.cdiv(numel, BLOCK_SIZE),) + _count_labels_kernel[grid_count]( + output, sizes_histogram, numel, BLOCK_SIZE=BLOCK_SIZE + ) + + # 2.2: Broadcast the counts to the final output tensor + grid_broadcast = (triton.cdiv(numel, BLOCK_SIZE),) + _broadcast_sizes_kernel[grid_broadcast]( + output, sizes_histogram, component_sizes_out, numel, BLOCK_SIZE=BLOCK_SIZE + ) + return output.view(out_shape) + 1, component_sizes_out.view(out_shape) diff --git a/models/sam3/perflib/triton/nms.py b/models/sam3/perflib/triton/nms.py new file mode 100644 index 0000000..ed800a1 --- /dev/null +++ b/models/sam3/perflib/triton/nms.py @@ -0,0 +1,124 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. All Rights Reserved + +# Adapted from https://github.com/stackav-oss/conch/blob/main/conch/kernels/vision/nms.py + +import torch +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({"cxpr_block_size": 128}), + triton.Config({"cxpr_block_size": 256}), + triton.Config({"cxpr_block_size": 512}), + triton.Config({"cxpr_block_size": 1024}), + triton.Config({"cxpr_block_size": 2048}), + triton.Config({"cxpr_block_size": 4096}), + triton.Config({"cxpr_block_size": 8192}), + ], + key=["num_boxes"], +) +@triton.jit +def _nms_suppression_kernel( + # Tensors + iou_mask_ptr: tl.tensor, # [N, N] + keep_mask_ptr: tl.tensor, # [N] + # Scalars + num_boxes: tl.int32, + # Strides + iou_mask_stride: tl.int32, + # Constexprs + cxpr_block_size: tl.constexpr, +) -> None: + """NMS suppression kernel. + + Args: + iou_mask_ptr: Pointer to precomputed IoU mask, shape: (N, N). + keep_mask_ptr: Pointer to keep mask tensor, shape: (N,). + num_boxes: Number of boxes. + iou_mask_stride: Stride for IoU mask tensor. + cxpr_block_size: Block size for processing. + """ + # Sequential NMS: for each box in sorted order, suppress later boxes + for current_box_idx in range(num_boxes - 1): + # Check if current box is still kept + is_kept = tl.load(keep_mask_ptr + current_box_idx) + if is_kept: + # IoU mask row offset for the current box + # Because the IoU mask is sorted by score, we will only consider boxes that come after the current box. + # This means we only need to read the upper triangular part of the IoU mask. + iou_row_offset = current_box_idx * iou_mask_stride + + # Only process boxes that come after the current box + next_box_idx = current_box_idx + 1 + remaining_boxes = num_boxes - next_box_idx + + # Iterate blockwise through the columns + for block_idx in range(tl.cdiv(remaining_boxes, cxpr_block_size)): + # Masked load of indices for the target boxes in the current block + block_start = next_box_idx + block_idx * cxpr_block_size + target_box_offsets = block_start + tl.arange(0, cxpr_block_size) + target_box_mask = target_box_offsets < num_boxes + + # Suppress boxes with lower scores that have high IoU + suppression_mask = tl.load( + iou_mask_ptr + iou_row_offset + target_box_offsets, + mask=target_box_mask, + other=False, + ) + suppression_mask = tl.cast(suppression_mask, tl.int1) + + # Conditionally store suppression result for high-IoU boxes + tl.store( + keep_mask_ptr + target_box_offsets, False, mask=suppression_mask + ) + + # Potential race condition: we need to ensure all threads complete the store before the next + # iteration otherwise we may load stale data for whether or not a box has been suppressed. + tl.debug_barrier() + + +def nms_triton( + ious: torch.Tensor, + scores: torch.Tensor, + iou_threshold: float, +) -> torch.Tensor: + """Perform NMS given the iou matrix, the scores and the iou threshold + + Args: + ious: Pairwise IoU tensor of shape (N, N). + scores: Scores tensor of shape (N,). + iou_threshold: IoU threshold for suppression. + + Returns: + Tensor: Indices of kept boxes, sorted by decreasing score. + """ + assert scores.dim() == 1, "Scores must be 1D" + iou_mask = ious > iou_threshold + assert iou_mask.dim() == 2 + assert iou_mask.shape[0] == iou_mask.shape[1] == scores.shape[0] + assert iou_mask.device == scores.device + assert iou_mask.dtype == torch.bool + + num_boxes = scores.size(0) + keep_mask = torch.ones(len(scores), device=scores.device, dtype=torch.bool) + + # Sort boxes by scores in descending order + _, sorted_indices = torch.sort(scores, dim=0, stable=True, descending=True) + iou_mask = iou_mask[sorted_indices][:, sorted_indices].contiguous() + + # For the suppression stage, we need to process sequentially, but we'll still take + # advantage of parallelism by processing in blocks in one program. + stage2_grid = (1,) + _nms_suppression_kernel[stage2_grid]( + # Tensors + iou_mask_ptr=iou_mask, + keep_mask_ptr=keep_mask, + # Scalars + num_boxes=num_boxes, + # Strides + iou_mask_stride=iou_mask.stride(0), + ) + # Extract indices of kept boxes + return sorted_indices[keep_mask]