chore: import upstream snapshot with attribution

python/tvm/topi/vision/__init__.py

@@ -0,0 +1,23 @@
+# isort: skip_file
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Vision operators."""
+
+from .multibox_transform_loc import *
+from .nms import *
+from .roi_align import *
+from .roi_pool import *

python/tvm/topi/vision/multibox_transform_loc.py

@@ -0,0 +1,121 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""Multibox location transform (SSD / TFLite DetectionPostProcess decode)."""
+
+import tvm
+from tvm import te, topi
+
+
+def multibox_transform_loc(
+    cls_pred,
+    loc_pred,
+    anchor,
+    variances,
+    clip=False,
+    threshold=0.0,
+    keep_background=True,
+):
+    """TFLite ``DecodeCenterSizeBoxes``-style decode + softmax score post-process.
+
+    Inputs must match Relax op contracts: ``cls_pred [B,C,N]``, ``loc_pred [B,4*N]``,
+    ``anchor [1,N,4]`` ltrb; per-anchor loc order ``(x,y,w,h)`` after yxhw→xywh reorder.
+
+    Parameters
+    ----------
+    cls_pred : te.Tensor
+        ``[B, C, N]`` logits.
+    loc_pred : te.Tensor
+        ``[B, 4*N]`` encodings ``(x,y,w,h)`` per anchor.
+    anchor : te.Tensor
+        ``[1, N, 4]`` ``(left, top, right, bottom)``.
+    variances : tuple of 4 float
+        ``(x,y,w,h)`` = ``1/x_scale, 1/y_scale, 1/w_scale, 1/h_scale`` (TFLite).
+    clip : bool
+        Clip ``ymin,xmin,ymax,xmax`` to ``[0,1]``.
+    threshold : float
+        After softmax: ``scores *= (scores >= threshold)``.
+    keep_background : bool
+        If False: ``scores[:,0,:] = 0``.
+
+    Returns
+    -------
+    boxes : te.Tensor
+        ``[B, N, 4]`` as ``(ymin,xmin,ymax,xmax)``.
+    scores : te.Tensor
+        ``[B, C, N]`` softmax, then threshold mask and optional background zero.
+    """
+    dtype = cls_pred.dtype
+    B = cls_pred.shape[0]
+    num_anchors = cls_pred.shape[2]
+    loc_reshaped = topi.reshape(loc_pred, [B, num_anchors, 4])
+
+    vx = tvm.tirx.const(float(variances[0]), dtype)
+    vy = tvm.tirx.const(float(variances[1]), dtype)
+    vw = tvm.tirx.const(float(variances[2]), dtype)
+    vh = tvm.tirx.const(float(variances[3]), dtype)
+    half = tvm.tirx.const(0.5, dtype)
+    zero = tvm.tirx.const(0.0, dtype)
+    one = tvm.tirx.const(1.0, dtype)
+    th = tvm.tirx.const(float(threshold), dtype)
+
+    def decode_bbox(b, a, k):
+        left = anchor[0, a, 0]
+        top = anchor[0, a, 1]
+        right = anchor[0, a, 2]
+        bottom = anchor[0, a, 3]
+        ay = (top + bottom) * half
+        ax = (left + right) * half
+        ah = bottom - top
+        aw = right - left
+        ex = loc_reshaped[b, a, 0]
+        ey = loc_reshaped[b, a, 1]
+        ew = loc_reshaped[b, a, 2]
+        eh = loc_reshaped[b, a, 3]
+        ycenter = ey * vy * ah + ay
+        xcenter = ex * vx * aw + ax
+        half_h = half * te.exp(eh * vh) * ah
+        half_w = half * te.exp(ew * vw) * aw
+        ymin = ycenter - half_h
+        xmin = xcenter - half_w
+        ymax = ycenter + half_h
+        xmax = xcenter + half_w
+        if clip:
+            ymin = te.max(zero, te.min(one, ymin))
+            xmin = te.max(zero, te.min(one, xmin))
+            ymax = te.max(zero, te.min(one, ymax))
+            xmax = te.max(zero, te.min(one, xmax))
+        return tvm.tirx.Select(
+            k == 0,
+            ymin,
+            tvm.tirx.Select(k == 1, xmin, tvm.tirx.Select(k == 2, ymax, xmax)),
+        )
+
+    boxes = te.compute((B, num_anchors, 4), decode_bbox, name="multibox_boxes")
+
+    scores = topi.nn.softmax(cls_pred, axis=1)
+    mask = topi.cast(topi.greater_equal(scores, th), dtype)
+    scores = scores * mask
+    if not keep_background:
+
+        def zero_bg(b, c, n):
+            s = scores[b, c, n]
+            return te.if_then_else(c == 0, zero, s)
+
+        scores = te.compute(scores.shape, zero_bg, name="multibox_scores_bg")
+
+    return [boxes, scores]

python/tvm/topi/vision/nms_util.py

@@ -0,0 +1,487 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+# ruff: noqa: E741
+"""Common utilities used in Non-maximum suppression operators"""
+
+import tvm
+from tvm import te
+from tvm.script.ir_builder import IRBuilder
+from tvm.script.ir_builder import tirx as T
+
+
+def _get_boundaries(output, box_idx):
+    l = tvm.te.min(
+        output[box_idx],
+        output[box_idx + 2],
+    )
+    t = tvm.te.min(
+        output[box_idx + 1],
+        output[box_idx + 3],
+    )
+    r = tvm.te.max(
+        output[box_idx],
+        output[box_idx + 2],
+    )
+    b = tvm.te.max(
+        output[box_idx + 1],
+        output[box_idx + 3],
+    )
+    return l, t, r, b
+
+
+def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
+    """Calculate overlap of two boxes."""
+    a_l, a_t, a_r, a_b = _get_boundaries(out_tensor, box_a_idx)
+    b_l, b_t, b_r, b_b = _get_boundaries(out_tensor, box_b_idx)
+
+    # Overlapping width and height
+    w = tvm.te.max(0.0, tvm.te.min(a_r, b_r) - tvm.te.max(a_l, b_l))
+    h = tvm.te.max(0.0, tvm.te.min(a_b, b_b) - tvm.te.max(a_t, b_t))
+
+    # Overlapping area
+    area = h * w
+
+    # total area of the figure formed by box a and box b
+    # except for overlapping area
+    u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area
+    return tvm.tirx.Select(u <= 0.0, 0.0, area / u)
+
+
+def binary_search(y, num_boxes, scores, score_threshold, out):
+    """Binary search for score_threshold on scores sorted in descending order.
+
+    Must be called within an IRBuilder context.
+    """
+    out = T.buffer_proxy(out)
+    lo_buf = T.decl_buffer([1], "int32", scope="local")
+    hi_buf = T.decl_buffer([1], "int32", scope="local")
+    lo = T.buffer_proxy(lo_buf)
+    hi = T.buffer_proxy(hi_buf)
+    lo[0] = T.int32(0)
+    hi[0] = tvm.tirx.Cast("int32", num_boxes)
+    with T.While(lo[0] < hi[0]):
+        mid = (hi[0] + lo[0]) >> 1
+        with T.If(scores[y, mid] > score_threshold):
+            with T.Then():
+                lo[0] = mid + 1
+            with T.Else():
+                hi[0] = mid
+    out[y] = lo[0]
+
+
+def _estimate_max_detections(batch_class, input_image_size=None):
+    """Estimate maximum detections based on input image size and number of classes.
+
+    This provides a more intelligent default for production environments.
+    """
+    if input_image_size is not None:
+        # Estimate based on image size: larger images typically have more objects
+        if len(input_image_size) >= 2:
+            height, width = input_image_size[-2], input_image_size[-1]
+            total_pixels = height * width
+
+            # Base estimation per class based on image size
+            if total_pixels < 300000:  # Small images (< 300k pixels)
+                base_detections_per_class = min(50, max(10, total_pixels // 2000))
+            elif total_pixels < 1000000:  # Medium images (< 1M pixels)
+                base_detections_per_class = min(100, max(25, total_pixels // 3000))
+            else:  # Large images (>= 1M pixels)
+                base_detections_per_class = min(200, max(50, total_pixels // 4000))
+
+            # Scale down for many classes (more realistic for multi-class scenarios)
+            if batch_class > 20:
+                # For many classes, reduce per-class detections to avoid explosion
+                detections_per_class = min(base_detections_per_class, 50)
+            else:
+                detections_per_class = base_detections_per_class
+        else:
+            detections_per_class = 50  # fallback
+    else:
+        # Fallback to class-based estimation
+        if batch_class == 1:
+            detections_per_class = 100  # Single class detection
+        elif batch_class <= 10:
+            detections_per_class = 50  # Small multi-class
+        else:
+            detections_per_class = 25  # Large multi-class (COCO-like)
+
+    return batch_class * detections_per_class
+
+
+def collect_selected_indices(
+    num_class,
+    selected_indices,
+    num_detections,
+    row_offsets,
+    ir,
+    max_output_boxes_per_class=None,
+    output_shape=None,
+    num_total_detections=None,
+    input_image_size=None,
+):
+    """Collect selected indices from the core NMS loop into one linear output.
+
+    Parameters
+    ----------
+    num_class : int
+    selected_indices : tvm.te.Tensor
+        2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
+        of selected boxes by the core NMS loop.
+    num_detections : tvm.te.Tensor
+        1-D tensor with shape (batch_size * num_classes,), representing
+        the number of boxes selected by the core NMS loop, per batch and class.
+    row_offsets : tvm.te.Tensor
+        1-D tensor with shape (batch_size * num_classes,), this should be the exclusive scan
+        of num_detections.
+    ir : function
+        A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py.
+
+    Returns
+    -------
+    out : tvm.te.Tensor
+        The output is indices of size (batch_size * num_class* num_boxes , 3).
+        Rows of indices are ordered such that selected boxes from batch 0, class 0 come
+        first, in descending of scores, followed by boxes from batch 0, class 1 etc.
+    """
+    batch_class, num_boxes = selected_indices.shape
+
+    if output_shape is not None:
+        return te.extern(
+            [output_shape],
+            [selected_indices, num_detections, row_offsets],
+            lambda ins, outs: ir(
+                num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
+            ),
+            dtype=["int64"],
+            name="collect_indices",
+            tag="collect_indices",
+        )
+
+    # TODO: Implement dynamic trimming based on num_total_detections
+    if num_total_detections is not None:
+        if isinstance(max_output_boxes_per_class, int):
+            out_rows = batch_class * max_output_boxes_per_class
+        else:
+            # Smart fallback based on input image size and typical production scenarios
+            out_rows = _estimate_max_detections(batch_class, input_image_size)
+
+        return te.extern(
+            [(out_rows, 3)],
+            [selected_indices, num_detections, row_offsets],
+            lambda ins, outs: ir(
+                num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
+            ),
+            dtype=["int64"],
+            name="collect_indices",
+            tag="collect_indices",
+        )
+
+    if isinstance(max_output_boxes_per_class, int):
+        out_rows = batch_class * max_output_boxes_per_class
+        return te.extern(
+            [(out_rows, 3)],
+            [selected_indices, num_detections, row_offsets],
+            lambda ins, outs: ir(
+                num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
+            ),
+            dtype=["int64"],
+            name="collect_indices",
+            tag="collect_indices",
+        )
+
+    if isinstance(max_output_boxes_per_class, te.Tensor):
+        try:
+            if len(max_output_boxes_per_class.shape) == 0:
+                max_boxes_val = int(max_output_boxes_per_class.data.numpy())
+            elif (
+                len(max_output_boxes_per_class.shape) == 1
+                and max_output_boxes_per_class.shape[0] == 1
+            ):
+                max_boxes_val = int(max_output_boxes_per_class.data.numpy()[0])
+            else:
+                max_boxes_val = num_boxes
+        except (ValueError, IndexError, AttributeError):
+            max_boxes_val = num_boxes
+
+        out_rows = batch_class * max_boxes_val
+        return te.extern(
+            [(out_rows, 3)],
+            [selected_indices, num_detections, row_offsets],
+            lambda ins, outs: ir(
+                num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
+            ),
+            dtype=["int64"],
+            name="collect_indices",
+            tag="collect_indices",
+        )
+
+    return te.extern(
+        [(batch_class * num_boxes, 3)],
+        [selected_indices, num_detections, row_offsets],
+        lambda ins, outs: ir(
+            num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
+        ),
+        dtype=["int64"],
+        name="collect_indices",
+        tag="collect_indices",
+    )
+
+
+def collect_selected_indices_and_scores(
+    selected_indices, selected_scores, num_detections, row_offsets, num_total_detections, ir
+):
+    """Collect selected indices and scores from the core NMS loop into one linear output.
+
+    Parameters
+    ----------
+    selected_indices : tvm.te.Tensor
+        2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
+        of selected boxes by the core NMS loop.
+    selected_scores : tvm.te.Tensor
+        2-D tensor with shape (batch_size * num_classes, num_boxes), representing the scores
+        of selected boxes by the core NMS loop.
+    num_detections : tvm.te.Tensor
+        2-D tensor with shape (batch_size, num_classes), representing
+        the number of boxes selected by the core NMS loop, per batch and class.
+    row_offsets : tvm.te.Tensor
+        2-D tensor with shape (batch_size, num_classes), this should be the exclusive scan
+        of num_detections along axis 1.
+    num_total_detections : tvm.te.Tensor
+        Total number of detections.
+    ir : function
+        A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py.
+
+    Returns
+    -------
+    out : [tvm.te.Tensor, tvm.te.Tensor]
+        The output is two tensors. The first is indices of size
+        (batch_size, num_class* num_boxes, 2), and the second is scores of size
+        (batch_size, num_class* num_boxes).
+    """
+    batch_size, num_class = row_offsets.shape
+    num_boxes = selected_indices.shape[1]
+    return te.extern(
+        [(batch_size, num_class * num_boxes, 2), (batch_size, num_class * num_boxes)],
+        [selected_indices, selected_scores, num_detections, row_offsets, num_total_detections],
+        lambda ins, outs: ir(ins[0], ins[1], ins[2], ins[3], ins[4], outs[0], outs[1]),
+        dtype=["int64", "float32"],
+        name="collect_indices_and_scores",
+        tag="collect_indices_and_scores",
+    )
+
+
+def _all_class_nms_ir(
+    boxes,
+    sorted_scores,
+    sorted_indices,
+    valid_count,
+    batch_class,
+    num_class,
+    num_anchors,
+    iou_threshold,
+    max_output_size_per_class,
+    box_indices,
+    selected_scores,
+    num_valid_boxes,
+    nms_loop,
+    score_threshold=None,
+):
+    with IRBuilder() as ib:
+        # Wrap buffers with T.buffer_proxy for flat indexing support
+        boxes = T.buffer_proxy(boxes)
+        box_indices = T.buffer_proxy(box_indices)
+        if selected_scores is not None:
+            selected_scores = T.buffer_proxy(selected_scores)
+
+        if isinstance(iou_threshold, float | int):
+            iou_threshold = tvm.tirx.FloatImm("float32", float(iou_threshold))
+        elif isinstance(iou_threshold, te.Tensor):
+            if len(iou_threshold.shape) == 0:
+                iou_threshold = iou_threshold()
+            elif len(iou_threshold.shape) == 1 and iou_threshold.shape[0] == 1:
+                iou_threshold = iou_threshold[0]
+            else:
+                iou_threshold = tvm.tirx.FloatImm("float32", 0.5)
+
+        if isinstance(max_output_size_per_class, int):
+            max_output_size_per_class = tvm.tirx.const(max_output_size_per_class)
+        elif isinstance(max_output_size_per_class, te.Tensor):
+            if len(max_output_size_per_class.shape) == 0:
+                max_output_size_per_class = max_output_size_per_class()
+            elif (
+                len(max_output_size_per_class.shape) == 1
+                and max_output_size_per_class.shape[0] == 1
+            ):
+                max_output_size_per_class = max_output_size_per_class[0]
+            else:
+                max_output_size_per_class = tvm.tirx.const(1000)
+
+        def calc_overlap(i, j, k):
+            offset_j = sorted_indices[i, j] * 4
+            offset_k = sorted_indices[i, k] * 4
+            batch_id = i // num_class
+            base_bbox_idx = batch_id * num_anchors * 4
+            return calculate_overlap(
+                boxes,
+                base_bbox_idx + offset_j,
+                base_bbox_idx + offset_k,
+            )
+
+        def on_new_valid_box(tid, num_current_valid_box, i, j):
+            with T.If(tid + 0 == 0):
+                with T.Then():
+                    box_indices[i, num_current_valid_box] = sorted_indices[i, j]
+                    if selected_scores is not None:
+                        selected_scores[i, num_current_valid_box] = sorted_scores[i, j]
+
+        def on_new_invalidated_box(*_):
+            pass
+
+        def needs_bbox_check(*_):
+            return tvm.tirx.const(True)
+
+        nms_loop(
+            batch_class,
+            tvm.tirx.IntImm("int32", -1),  # top_k
+            iou_threshold,
+            max_output_size_per_class,
+            valid_count,
+            on_new_valid_box,
+            on_new_invalidated_box,
+            needs_bbox_check,
+            calc_overlap,
+            sorted_scores,
+            num_valid_boxes,
+            score_threshold,
+        )
+
+        return ib.get()
+
+
+def run_all_class_nms(
+    boxes,
+    sorted_scores,
+    sorted_indices,
+    valid_count,
+    max_output_size_per_class,
+    iou_threshold,
+    nms_loop,
+    return_scores=False,
+    score_threshold=None,
+):
+    """The core all class NMS routine.
+
+    Parameters
+    ----------
+    boxes : tvm.te.Tensor
+        3-D tensor with shape (batch_size, num_boxes, 4)
+    sorted_scores : tvm.te.Tensor
+        2-D tensor with shape (batch_size * num_classes, num_boxes).
+        One of the outputs from argsort.
+    sorted_indices : tvm.te.Tensor
+        2-D tensor with shape (batch_size * num_classes, num_boxes).
+        The other output from argsort.
+    valid_count : tvm.te.Tensor
+        1-D tensor with shape (batch_size * num_classes,), representing
+        the number of boxes whose score is above score_threshold, per batch and class.
+    max_output_size_per_class : int or tvm.te.Tensor, optional
+        The maxinum number of output selected boxes per class.
+    iou_threshold : float or tvm.te.Tensor, optional
+        IoU test threshold.
+    nms_loop : function
+        A core NMS loop, see its usage in vision/nms.py and cuda/nms.py.
+    return_scores : bool, optional
+        Whether or not to return selected scores, needed by the tensorflow output format.
+
+    Returns
+    -------
+    out : a list of tvm.te.Tensor
+        The output is three tensors, the first and second are indices and scores of size
+        (batch_size * num_class, num_boxes), and the third is a tensor
+        num_selected_boxes of shape (batch_size * num_class,) representing the total number of
+        selected boxes per batch and class. If return_scores is False, the second output is
+        None.
+    """
+    batch, num_boxes, _ = boxes.shape
+    batch_class = sorted_scores.shape[0]
+    num_class = batch_class // batch
+
+    if return_scores is False:
+        all_class_num0_buf = tvm.tirx.decl_buffer(
+            (batch_class, num_boxes), "int32", "all_class_nms0", data_alignment=8, layout=None
+        )
+        all_class_num1_buf = tvm.tirx.decl_buffer(
+            (batch_class,), "int32", "all_class_nms1", data_alignment=8, layout=None
+        )
+        extern_inputs = [boxes, sorted_scores, sorted_indices, valid_count]
+        if score_threshold is not None:
+            extern_inputs.append(score_threshold)
+
+        selected_indices, num_detections = te.extern(
+            [(batch_class, num_boxes), (batch_class,)],
+            extern_inputs,
+            lambda ins, outs: _all_class_nms_ir(
+                ins[0],  # boxes
+                ins[1],  # sorted_scores
+                ins[2],  # sorted_indices
+                ins[3],  # valid_count
+                batch_class,
+                num_class,
+                num_boxes,
+                iou_threshold,
+                max_output_size_per_class,
+                outs[0],  # box_indices
+                None,  # scores
+                outs[1],  # num_selected_boxes
+                nms_loop,
+                ins[4] if score_threshold is not None else None,  # score_threshold
+            ),
+            out_buffers=[all_class_num0_buf, all_class_num1_buf],
+            dtype=["int32", "int32"],
+            name="all_class_nms",
+            tag="all_class_nms",
+        )
+        return selected_indices, None, num_detections
+
+    extern_inputs = [boxes, sorted_scores, sorted_indices, valid_count]
+    if score_threshold is not None:
+        extern_inputs.append(score_threshold)
+
+    return te.extern(
+        [(batch_class, num_boxes), (batch_class, num_boxes), (batch_class,)],
+        extern_inputs,
+        lambda ins, outs: _all_class_nms_ir(
+            ins[0],  # boxes
+            ins[1],  # sorted_scores
+            ins[2],  # sorted_indices
+            ins[3],  # valid_count
+            batch_class,
+            num_class,
+            num_boxes,
+            iou_threshold,
+            max_output_size_per_class,
+            outs[0],  # box_indices
+            outs[1],  # selected scores
+            outs[2],  # num_selected_boxes
+            nms_loop,
+            ins[4] if score_threshold is not None else None,  # score_threshold
+        ),
+        dtype=["int32", "float32", "int32"],
+        name="all_class_nms",
+        tag="all_class_nms",
+    )

python/tvm/topi/vision/roi_align.py

@@ -0,0 +1,204 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""ROI Align operator"""
+
+import tvm
+from tvm import te
+
+from ..cpp.utils import bilinear_sample_nchw, bilinear_sample_nhwc
+
+
+def _sample_common(
+    i,
+    c,
+    ph,
+    pw,
+    rois,
+    pooled_size_h,
+    pooled_size_w,
+    spatial_scale,
+    sample_ratio,
+    aligned,
+    dtype,
+    avg_mode,
+    bilinear_func,
+):
+    roi = rois[i]
+    batch_index = roi[0].astype("int32")
+    roi_start_w = roi[1] * spatial_scale
+    roi_start_h = roi[2] * spatial_scale
+    roi_end_w = roi[3] * spatial_scale
+    roi_end_h = roi[4] * spatial_scale
+
+    if aligned:
+        roi_h = roi_end_h - roi_start_h
+        roi_w = roi_end_w - roi_start_w
+    else:
+        roi_h = te.max(roi_end_h - roi_start_h, tvm.tirx.const(1.0, dtype))
+        roi_w = te.max(roi_end_w - roi_start_w, tvm.tirx.const(1.0, dtype))
+
+    pooled_size_h_const = tvm.tirx.const(pooled_size_h, dtype)
+    pooled_size_w_const = tvm.tirx.const(pooled_size_w, dtype)
+    bin_h = roi_h / pooled_size_h_const
+    bin_w = roi_w / pooled_size_w_const
+
+    if sample_ratio > 0:
+        roi_bin_grid_h = tvm.tirx.const(sample_ratio, "int32")
+        roi_bin_grid_w = tvm.tirx.const(sample_ratio, "int32")
+    else:
+        roi_bin_grid_h = te.ceil(roi_h / pooled_size_h_const).astype("int32")
+        roi_bin_grid_w = te.ceil(roi_w / pooled_size_w_const).astype("int32")
+
+    count = roi_bin_grid_h * roi_bin_grid_w
+    rh = te.reduce_axis((0, roi_bin_grid_h), name="rh")
+    rw = te.reduce_axis((0, roi_bin_grid_w), name="rw")
+    roi_start_h = roi_start_h + tvm.tirx.Cast(dtype, ph) * bin_h
+    roi_start_w = roi_start_w + tvm.tirx.Cast(dtype, pw) * bin_w
+
+    def sample_value(rh_idx, rw_idx):
+        return bilinear_func(
+            batch_index,
+            c,
+            roi_start_h
+            + (tvm.tirx.Cast(dtype, rh_idx) + tvm.tirx.const(0.5, dtype))
+            * bin_h
+            / tvm.tirx.Cast(dtype, roi_bin_grid_h),
+            roi_start_w
+            + (tvm.tirx.Cast(dtype, rw_idx) + tvm.tirx.const(0.5, dtype))
+            * bin_w
+            / tvm.tirx.Cast(dtype, roi_bin_grid_w),
+        )
+
+    if avg_mode:
+        return te.sum(
+            sample_value(rh, rw) / tvm.tirx.Cast(dtype, count),
+            axis=[rh, rw],
+        )
+    return te.max(sample_value(rh, rw), axis=[rh, rw])
+
+
+def roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
+    """ROI align operator in NCHW layout."""
+    avg_mode = mode in (b"avg", "avg", 0)
+    max_mode = mode in (b"max", "max", 1)
+    assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."
+
+    _, channel, height, width = data.shape
+    num_roi, _ = rois.shape
+    dtype = rois.dtype
+
+    if isinstance(pooled_size, int):
+        pooled_size_h = pooled_size_w = pooled_size
+    else:
+        pooled_size_h, pooled_size_w = pooled_size
+
+    height_f = tvm.tirx.Cast(dtype, height)
+    width_f = tvm.tirx.Cast(dtype, width)
+    zero = tvm.tirx.const(0.0, data.dtype)
+
+    def _bilinear(n, c, y, x):
+        outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
+        y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
+        x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
+        val = bilinear_sample_nchw(data, (n, c, y, x), height - 1, width - 1)
+        return tvm.tirx.if_then_else(outside, zero, val)
+
+    return te.compute(
+        (num_roi, channel, pooled_size_h, pooled_size_w),
+        lambda i, c, ph, pw: _sample_common(
+            i,
+            c,
+            ph,
+            pw,
+            rois,
+            pooled_size_h,
+            pooled_size_w,
+            spatial_scale,
+            sample_ratio,
+            aligned,
+            dtype,
+            avg_mode,
+            _bilinear,
+        ),
+        tag="pool,roi_align_nchw",
+    )
+
+
+def roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio=-1, aligned=False):
+    """ROI align operator in NHWC layout."""
+    avg_mode = mode in (b"avg", "avg", 0)
+    max_mode = mode in (b"max", "max", 1)
+    assert avg_mode or max_mode, "Mode must be avg or max. Please pass in a valid mode."
+
+    _, height, width, channel = data.shape
+    num_roi, _ = rois.shape
+    dtype = rois.dtype
+
+    if isinstance(pooled_size, int):
+        pooled_size_h = pooled_size_w = pooled_size
+    else:
+        pooled_size_h, pooled_size_w = pooled_size
+
+    height_f = tvm.tirx.Cast(dtype, height)
+    width_f = tvm.tirx.Cast(dtype, width)
+    zero = tvm.tirx.const(0.0, data.dtype)
+
+    def _bilinear(n, c, y, x):
+        outside = tvm.tirx.any(y < -1.0, x < -1.0, y > height_f, x > width_f)
+        y = te.min(te.max(y, 0.0), tvm.tirx.Cast(dtype, height - 1))
+        x = te.min(te.max(x, 0.0), tvm.tirx.Cast(dtype, width - 1))
+        val = bilinear_sample_nhwc(data, (n, y, x, c), height - 1, width - 1)
+        return tvm.tirx.if_then_else(outside, zero, val)
+
+    return te.compute(
+        (num_roi, pooled_size_h, pooled_size_w, channel),
+        lambda i, ph, pw, c: _sample_common(
+            i,
+            c,
+            ph,
+            pw,
+            rois,
+            pooled_size_h,
+            pooled_size_w,
+            spatial_scale,
+            sample_ratio,
+            aligned,
+            dtype,
+            avg_mode,
+            _bilinear,
+        ),
+        tag="pool,roi_align_nhwc",
+    )
+
+
+def roi_align(
+    data,
+    rois,
+    pooled_size,
+    spatial_scale,
+    mode="avg",
+    sample_ratio=-1,
+    aligned=False,
+    layout="NCHW",
+):
+    """ROI align operator."""
+    if layout == "NCHW":
+        return roi_align_nchw(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
+    if layout == "NHWC":
+        return roi_align_nhwc(data, rois, pooled_size, spatial_scale, mode, sample_ratio, aligned)
+    raise ValueError(f"Unsupported layout for roi_align: {layout}")

python/tvm/topi/vision/roi_pool.py

@@ -0,0 +1,101 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=invalid-name
+"""ROI Pool operator"""
+
+import tvm
+from tvm import te
+
+
+def roi_pool_nchw(data, rois, pooled_size, spatial_scale):
+    """ROI pool operator in NCHW layout."""
+    _, channel, height, width = data.shape
+    num_roi, _ = rois.shape
+
+    if isinstance(pooled_size, int):
+        pooled_size_h = pooled_size_w = pooled_size
+    else:
+        pooled_size_h, pooled_size_w = pooled_size
+
+    zero = tvm.tirx.const(0.0, data.dtype)
+    roi_dtype = rois.dtype
+
+    neg_inf = tvm.tirx.const(float("-inf"), data.dtype)
+
+    def _round_away(x):
+        # ONNX MaxRoiPool spec uses ties-away-from-zero rounding for coordinate
+        # mapping (matching std::round semantics in the reference implementation).
+        # Use floor(x + 0.5) to be explicit and independent of tir.round semantics.
+        half = tvm.tirx.const(0.5, roi_dtype)
+        return te.floor(x + half)
+
+    def _bin_bounds(i, ph, pw):
+        roi = rois[i]
+        roi_start_w = _round_away(roi[1] * spatial_scale).astype("int32")
+        roi_start_h = _round_away(roi[2] * spatial_scale).astype("int32")
+        roi_end_w = _round_away(roi[3] * spatial_scale).astype("int32")
+        roi_end_h = _round_away(roi[4] * spatial_scale).astype("int32")
+
+        roi_h = te.max(roi_end_h - roi_start_h + 1, tvm.tirx.const(1, "int32"))
+        roi_w = te.max(roi_end_w - roi_start_w + 1, tvm.tirx.const(1, "int32"))
+
+        bin_h = tvm.tirx.Cast(roi_dtype, roi_h) / tvm.tirx.const(float(pooled_size_h), roi_dtype)
+        bin_w = tvm.tirx.Cast(roi_dtype, roi_w) / tvm.tirx.const(float(pooled_size_w), roi_dtype)
+
+        hstart = te.floor(tvm.tirx.Cast(roi_dtype, ph) * bin_h).astype("int32")
+        wstart = te.floor(tvm.tirx.Cast(roi_dtype, pw) * bin_w).astype("int32")
+        hend = te.ceil(tvm.tirx.Cast(roi_dtype, ph + 1) * bin_h).astype("int32")
+        wend = te.ceil(tvm.tirx.Cast(roi_dtype, pw + 1) * bin_w).astype("int32")
+
+        hstart = te.min(te.max(hstart + roi_start_h, 0), height)
+        hend = te.min(te.max(hend + roi_start_h, 0), height)
+        wstart = te.min(te.max(wstart + roi_start_w, 0), width)
+        wend = te.min(te.max(wend + roi_start_w, 0), width)
+        return hstart, hend, wstart, wend
+
+    def _sample(i, c, ph, pw):
+        roi = rois[i]
+        batch_index = roi[0].astype("int32")
+        hstart, hend, wstart, wend = _bin_bounds(i, ph, pw)
+        valid = tvm.tirx.all(hstart <= rh, rh < hend, wstart <= rw, rw < wend)
+        return tvm.tirx.if_then_else(valid, data[batch_index, c, rh, rw], neg_inf)
+
+    def _is_empty(i, ph, pw):
+        hstart, hend, wstart, wend = _bin_bounds(i, ph, pw)
+        return tvm.tirx.any(hend <= hstart, wend <= wstart)
+
+    rh = te.reduce_axis((0, height), name="rh")
+    rw = te.reduce_axis((0, width), name="rw")
+    pooled = te.compute(
+        (num_roi, channel, pooled_size_h, pooled_size_w),
+        lambda i, c, ph, pw: te.max(_sample(i, c, ph, pw), axis=[rh, rw]),
+        tag="pool,roi_pool_nchw",
+    )
+
+    return te.compute(
+        (num_roi, channel, pooled_size_h, pooled_size_w),
+        lambda i, c, ph, pw: tvm.tirx.if_then_else(
+            _is_empty(i, ph, pw), zero, pooled[i, c, ph, pw]
+        ),
+    )
+
+
+def roi_pool(data, rois, pooled_size, spatial_scale, layout="NCHW"):
+    """ROI pool operator."""
+    if layout == "NCHW":
+        return roi_pool_nchw(data, rois, pooled_size, spatial_scale)
+    raise ValueError(f"Unsupported layout for roi_pool: {layout}")