chore: import upstream snapshot with attribution

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}")