# 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, no-member, too-many-locals, too-many-arguments, too-many-statements, singleton-comparison
# pylint: disable=bad-continuation, unused-argument
"""Non-maximum suppression operator"""
import math
import tvm
from tvm import te

from tvm.tir import if_then_else
from .sort import argsort, argsort_thrust
from .. import tag


def cuda_atomic_add_rule(op):
    if op.dtype == "float32":
        return tvm.tir.call_pure_extern("float32", "atomicAdd", op.args[0], op.args[1])
    if op.dtype == "float64":
        return tvm.tir.call_pure_extern("float64", "atomicAdd", op.args[0], op.args[1])
    if op.dtype == "int32":
        return tvm.tir.call_pure_extern("int32", "atomicAdd", op.args[0], op.args[1])
    raise RuntimeError("only support int32, float32 and float64")


tvm.target.intrin.register_intrin_rule(
    "cuda", "atomic_add", cuda_atomic_add_rule, override=True)


def atomic_add(x, y):
    return tvm.tir.call_pure_intrin(y.dtype, "atomic_add", x, y)


def get_valid_counts_ir(data, valid_count, flag, score_threshold, id_index, score_index):
    """Low level IR to get valid count of bounding boxes
    given a score threshold. Also prepares to move valid boxes to the
    top of input data.

    Parameters
    ----------
    data : Buffer
        Input data. 3-D Buffer with shape [batch_size, num_anchors, elem_length].

    valid_count : Buffer
        1D buffer for valid number of boxes with shape [batch_size, ].

    flag : Buffer
        2D Buffer of flag indicating valid data with shape [batch_size, num_anchors].

    score_threshold : float32
        Lower limit of score for valid bounding boxes.

    id_index : optional, int
        index of the class categories, -1 to disable.

    score_index: optional, int
        Index of the scores/confidence of boxes.

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    batch_size = data.shape[0]
    num_anchors = data.shape[1]
    elem_length = data.shape[2]

    ib = tvm.tir.ir_builder.create()

    data = ib.buffer_ptr(data)

    valid_count = ib.buffer_ptr(valid_count)
    flag = ib.buffer_ptr(flag)
    atomic_add_return = ib.allocate(
        valid_count.dtype, (1,), name='atomic_add_return', scope='local')
    one_count = tvm.tir.const(1, dtype=valid_count.dtype)
    score_threshold = tvm.ir.make_node(
        "FloatImm", dtype="float32", value=score_threshold)
    id_index = tvm.ir.make_node("IntImm", dtype="int32", value=id_index)
    score_index = tvm.ir.make_node("IntImm", dtype="int32", value=score_index)

    max_threads = int(tvm.target.Target.current(
        allow_none=False).max_num_threads)
    nthread_tx = max_threads
    nthread_bx = batch_size * num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    tid = bx * max_threads + tx
    idxd = tvm.tir.indexdiv

    # initialize valid_count
    with ib.if_scope(tid < batch_size):
        valid_count[tid] = 0
    # initialize flag
    with ib.if_scope(tid < batch_size * num_anchors):
        flag[tid] = 0
    with ib.if_scope(tid < batch_size * num_anchors):
        i = idxd(tid, num_anchors)
        with ib.if_scope(
                tvm.tir.all(data[tid * elem_length + score_index] > score_threshold,
                            tvm.tir.any(id_index < 0, data[tid * elem_length + id_index] >= 0))):
            flag[tid] = 1
            atomic_add_return[0] = atomic_add(tvm.tir.call_pure_intrin("handle", "tvm_address_of",
                                                                       valid_count[i]), one_count)

    return ib.get()


def flag_scan(flag, prefix_sum):
    """Low level IR to calculate correct positions for valid boxes.

    Parameters
    ----------
    flag : Buffer
        2D Buffer of flag indicating valid data with shape [batch_size, num_anchors].

    prefix_sum : Buffer
        2D Buffer of prefix sum of flags indicating new locations of valid boxes
        with same shape as flag.

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    batch_size = flag.shape[0]
    num_anchors = flag.shape[1]

    ib = tvm.tir.ir_builder.create()

    flag = ib.buffer_ptr(flag)
    prefix_sum = ib.buffer_ptr(prefix_sum)

    max_threads = int(tvm.target.Target.current(
        allow_none=False).max_num_threads)
    nthread_tx = max_threads
    nthread_bx = batch_size * num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    tid = bx * max_threads + tx
    idxd = tvm.tir.indexdiv
    idxm = tvm.tir.indexmod

    # initialize prefix_sum
    with ib.if_scope(tid < batch_size * num_anchors):
        prefix_sum[tid] = 0
    with ib.if_scope(tid < batch_size * num_anchors):
        i = idxd(tid, num_anchors)
        j = idxm(tid, num_anchors)
        with ib.for_range(0, j) as r:
            prefix_sum[tid] += flag[i * num_anchors + r]

    return ib.get()


def out_rewrite(data, flag, prefix_sum, valid_count, out):
    """Low level IR to move valid boxes to the
    top of input data.

    Parameters
    ----------
    data : Buffer
        Input data. 3-D Buffer with shape [batch_size, num_anchors, elem_length].

    flag : Buffer
        2D Buffer of flag indicating valid data with shape [batch_size, num_anchors].

    prefix_sum : Buffer
        2D Buffer of prefix sum of flags indicating new locations of valid boxes
        with same shape as flag.

    valid_count : Buffer
        1D buffer for valid number of boxes with shape [batch_size, ].

    out : Buffer
        Rearranged data buffer.

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    batch_size = out.shape[0]
    num_anchors = out.shape[1]
    elem_length = out.shape[2]

    ib = tvm.tir.ir_builder.create()

    one = tvm.tir.const(1, dtype=out.dtype)
    data = ib.buffer_ptr(data)
    flag = ib.buffer_ptr(flag)
    valid_count = ib.buffer_ptr(valid_count)
    prefix_sum = ib.buffer_ptr(prefix_sum)
    out = ib.buffer_ptr(out)

    max_threads = int(tvm.target.Target.current(
        allow_none=False).max_num_threads)
    nthread_tx = max_threads
    nthread_bx = batch_size * num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    tid = bx * max_threads + tx
    idxd = tvm.tir.indexdiv
    idxm = tvm.tir.indexmod

    with ib.if_scope(tid < batch_size * num_anchors):
        i = idxd(tid, num_anchors)
        j = idxm(tid, num_anchors)
        base_idx = i * num_anchors * elem_length
        with ib.if_scope(tvm.tir.all(flag[tid] > 0, prefix_sum[tid] >= 0,
                                     prefix_sum[tid] < num_anchors)):
            with ib.for_range(0, elem_length) as k:
                out[base_idx + prefix_sum[tid] * elem_length +
                    k] = data[tid * elem_length + k]
        with ib.if_scope(j >= valid_count[i]):
            with ib.for_range(0, elem_length) as k:
                out[tid * elem_length + k] = -one

    return ib.get()


def get_valid_counts(data, score_threshold=0, id_index=0, score_index=1):
    """Get valid count of bounding boxes given a score threshold.
    Also moves valid boxes to the top of input data.

    Parameters
    ----------
    data : tvm.te.Tensor
        Input data. 3-D tensor with shape [batch_size, num_anchors, elem_length].

    score_threshold : optional, float
        Lower limit of score for valid bounding boxes.

    id_index : optional, int
        index of the class categories, -1 to disable.

    score_index: optional, int
        Index of the scores/confidence of boxes.

    Returns
    -------
    valid_count : tvm.te.Tensor
        1-D tensor for valid number of boxes.

    out_tensor : tvm.te.Tensor
        Rearranged data tensor.
    """
    batch_size = data.shape[0]
    num_anchors = data.shape[1]
    data_buf = tvm.tir.decl_buffer(
        data.shape, data.dtype, "data_buf", data_alignment=8)
    valid_count_buf = tvm.tir.decl_buffer(
        (batch_size,), "int32", "valid_count_buf", data_alignment=8)
    temp_flag_buf = tvm.tir.decl_buffer(
        (batch_size, num_anchors,), "int32", "temp_flag", data_alignment=8)
    temp_partial_buf = tvm.tir.decl_buffer(
        (batch_size, num_anchors), "int32", "temp_partial", data_alignment=8)
    out_buf = tvm.tir.decl_buffer(
        data.shape, data.dtype, "out_buf", data_alignment=8)

    valid_count, temp_flag = \
        te.extern([(batch_size,), (batch_size, num_anchors)], [data],
                  lambda ins, outs: get_valid_counts_ir(
            ins[0], outs[0], outs[1], score_threshold, id_index, score_index),
            dtype=["int32", "int32"],
            in_buffers=[data_buf],
            out_buffers=[valid_count_buf, temp_flag_buf],
            name="get_valid_counts",
            tag="get_valid_counts_gpu")

    temp_partial = \
        te.extern([(batch_size, num_anchors)], [temp_flag],
                  lambda ins, outs: flag_scan(
            ins[0], outs[0]),
            dtype=["int32"],
            in_buffers=[temp_flag_buf],
            out_buffers=[temp_partial_buf],
            name="flag_scan")

    out = \
        te.extern([data.shape], [data, temp_flag, temp_partial, valid_count],
                  lambda ins, outs: out_rewrite(
            ins[0], ins[1], ins[2], ins[3], outs[0]),
            dtype=[data.dtype],
            in_buffers=[data_buf, temp_flag_buf,
                        temp_partial_buf, valid_count_buf],
            out_buffers=[out_buf],
            name="out_rewrite")

    return [valid_count, out]


def nms_ir(data, sorted_index, valid_count, out, box_indices,
           max_output_size, iou_threshold, force_suppress,
           top_k, coord_start, id_index, score_index):
    """Low level IR routing for transform location in multibox_detection operator.

    Parameters
    ----------
    data : Buffer
        Buffer of output boxes with class and score.

    sort_index : Buffer
        Buffer of output box indexes sorted by score.

    valid_count : Buffer
        Buffer of number of valid output boxes.

    out : Buffer
        Output buffer.

    max_output_size : int
        Max number of output valid boxes for each instance.
        By default all valid boxes are returned.

    iou_threshold : float
        Overlapping(IoU) threshold to suppress object with smaller score.

    force_suppress : boolean
        Whether to suppress all detections regardless of class_id.

    top_k : int
        Keep maximum top k detections before nms, -1 for no limit.

    coord_start : int
        Start index of the consecutive 4 coordinates.

    id_index : int
        index of the class categories, -1 to disable.

    score_index : optional, int
        Index of the scores/confidence of boxes.

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
        """Calculate overlap of two boxes.
        """
        w = tvm.te.max(0.0, tvm.te.min(out_tensor[box_a_idx + 2], out_tensor[box_b_idx + 2])
                       - tvm.te.max(out_tensor[box_a_idx], out_tensor[box_b_idx]))
        h = tvm.te.max(0.0, tvm.te.min(out_tensor[box_a_idx + 3], out_tensor[box_b_idx + 3])
                       - tvm.te.max(out_tensor[box_a_idx + 1], out_tensor[box_b_idx + 1]))
        i = w * h
        u = (out_tensor[box_a_idx + 2] - out_tensor[box_a_idx]) * \
            (out_tensor[box_a_idx + 3] - out_tensor[box_a_idx + 1]) + \
            (out_tensor[box_b_idx + 2] - out_tensor[box_b_idx]) * \
            (out_tensor[box_b_idx + 3] - out_tensor[box_b_idx + 1]) - i
        return tvm.tir.Select(u <= 0.0, 0.0, i / u)

    batch_size = data.shape[0]
    num_anchors = data.shape[1]
    box_data_length = data.shape[2]

    ib = tvm.tir.ir_builder.create()

    data = ib.buffer_ptr(data)
    sorted_index = ib.buffer_ptr(sorted_index)
    valid_count = ib.buffer_ptr(valid_count)
    out = ib.buffer_ptr(out)
    box_indices = ib.buffer_ptr(box_indices)
    num_valid_boxes = ib.allocate(
        "int32", (1,), name="num_valid_boxes", scope="local")

    max_threads = int(
        tvm.target.Target.current(allow_none=False).max_num_threads)
    nthread_tx = max_threads
    nthread_bx = num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    j = bx * max_threads + tx

    iou_threshold = tvm.ir.make_node(
        "FloatImm", dtype="float32", value=iou_threshold)
    top_k = tvm.ir.make_node("IntImm", dtype="int32", value=top_k)
    coord_start = tvm.ir.make_node("IntImm", dtype="int32", value=coord_start)
    id_index = tvm.ir.make_node("IntImm", dtype="int32", value=id_index)
    score_index = tvm.ir.make_node("IntImm", dtype="int32", value=score_index)
    force_suppress = tvm.ir.make_node(
        "IntImm", dtype="int32", value=1 if force_suppress else 0)

    with ib.for_range(0, batch_size, for_type="unroll") as i:
        base_idx = i * num_anchors * box_data_length
        with ib.if_scope(tvm.tir.all(iou_threshold > 0, valid_count[i] > 0)):
            # Reorder output
            nkeep = if_then_else(
                tvm.tir.all(top_k > 0, top_k < valid_count[i]),
                top_k, valid_count[i])
            with ib.if_scope(j < nkeep):
                with ib.for_range(0, box_data_length) as k:
                    out[(base_idx + j * box_data_length + k)] = \
                        data[(base_idx + sorted_index[i * num_anchors + j]
                              * box_data_length + k)]
                box_indices[i * num_anchors +
                            j] = sorted_index[i * num_anchors + j]
            with ib.if_scope(tvm.tir.all(top_k > 0, top_k < valid_count[i])):
                with ib.if_scope(j < valid_count[i] - nkeep):
                    with ib.for_range(0, box_data_length) as k:
                        out[(base_idx + (j + nkeep) * box_data_length + k)] = -1.0
                    box_indices[i * num_anchors + (j + nkeep)] = -1
            # Apply nms
            with ib.for_range(0, valid_count[i]) as k:
                offset_k = k * box_data_length
                with ib.if_scope(
                        tvm.tir.all(out[base_idx + offset_k + score_index] > 0,
                                    tvm.tir.any(id_index < 0, out[base_idx +
                                                                  offset_k + id_index] >= 0))):
                    with ib.if_scope(j < valid_count[i]):
                        offset_j = j * box_data_length
                        with ib.if_scope(
                                tvm.tir.all(j > k,
                                            out[base_idx + offset_j +
                                                score_index] > 0,
                                            tvm.tir.any(id_index < 0,
                                                        out[base_idx + offset_j + id_index] >= 0),
                                            tvm.tir.any(force_suppress > 0, id_index < 0,
                                                    out[base_idx + offset_k + id_index] ==
                                                        out[base_idx + offset_j + id_index]))):
                            iou = calculate_overlap(out, base_idx + offset_j + coord_start,
                                                    base_idx + offset_k + coord_start)
                            with ib.if_scope(iou >= iou_threshold):
                                out[base_idx + offset_j + score_index] = -1.0
                                with ib.if_scope(id_index >= 0):
                                    out[base_idx + offset_j + id_index] = -1.0
                                box_indices[i * num_anchors + j] = -1
        with ib.else_scope():
            with ib.if_scope(j < valid_count[i]):
                offset_j = j * box_data_length
                with ib.for_range(0, box_data_length) as k:
                    out[(base_idx + offset_j + k)
                        ] = data[base_idx + offset_j + k]
                box_indices[i * num_anchors + j] = j
        # Set invalid entry to be -1
        with ib.if_scope(j < num_anchors - valid_count[i]):
            with ib.for_range(0, box_data_length) as k:
                out[base_idx + (j + valid_count[i]) *
                    box_data_length + k] = -1.0
            box_indices[i * num_anchors + j + valid_count[i]] = -1
        # Only return max_output_size number of valid boxes
        num_valid_boxes[0] = 0
        with ib.if_scope(max_output_size > 0):
            with ib.if_scope(j < valid_count[i]):
                offset_j = j * box_data_length
                with ib.if_scope(out[base_idx + offset_j] >= 0):
                    with ib.if_scope(num_valid_boxes[0] == max_output_size):
                        with ib.for_range(0, box_data_length) as k:
                            out[base_idx + offset_j + k] = -1.0
                        box_indices[i * num_anchors + j] = -1
                    with ib.else_scope():
                        num_valid_boxes[0] += 1

    return ib.get()


def invalid_to_bottom_pre(data, flag, idx):
    """Low level IR to rearrange nms output to move all valid entries to top.

    Parameters
    ----------
    data: Buffer
        3D Buffer with shape [batch_size, num_anchors, elem_length], output of nms.

    flag : Buffer
        1D Buffer of flag indicating valid data with [num_anchors].

    idx : Buffer
        1D Buffer of valid data indices with [num_anchors].

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    batch_size = data.shape[0]
    num_anchors = data.shape[1]
    elem_length = data.shape[2]

    ib = tvm.tir.ir_builder.create()

    data = ib.buffer_ptr(data)
    flag = ib.buffer_ptr(flag)
    idx = ib.buffer_ptr(idx)

    max_threads = int(math.sqrt(
        tvm.target.Target.current(allow_none=False).max_num_threads))
    nthread_tx = max_threads
    nthread_bx = num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    j = bx * max_threads + tx

    with ib.for_range(0, batch_size, for_type="unroll") as i:
        base_idx = i * num_anchors * elem_length
        with ib.if_scope(j < num_anchors):
            with ib.if_scope(data[base_idx + j * elem_length] >= 0):
                flag[i * num_anchors + j] = 1
                idx[i * num_anchors + j] = 1
            with ib.else_scope():
                flag[i * num_anchors + j] = 0
                idx[i * num_anchors + j] = 0

    with ib.if_scope(j < batch_size):
        with ib.for_range(0, num_anchors) as k:
            with ib.if_scope(k > 0):
                idx[j * num_anchors + k] += idx[j * num_anchors + k - 1]
    return ib.get()


def invalid_to_bottom_ir(data, flag, idx, out):
    """Low level IR to rearrange nms output to move all valid entries to top.

    Parameters
    ----------
    data: Buffer
        3D Buffer with shape [batch_size, num_anchors, elem_length], output of nms.

    flag : Buffer
        1D Buffer of flag indicating valid data with [num_anchors].

    idx : Buffer
        1D Buffer of valid data indices with [num_anchors].

    out : Buffer
        3D Buffer of rearranged nms output with shape [batch_size, num_anchors, elem_length].

    Returns
    -------
    stmt : Stmt
        The result IR statement.
    """
    batch_size = data.shape[0]
    num_anchors = data.shape[1]
    elem_length = data.shape[2]

    ib = tvm.tir.ir_builder.create()

    data = ib.buffer_ptr(data)
    flag = ib.buffer_ptr(flag)
    idx = ib.buffer_ptr(idx)
    out = ib.buffer_ptr(out)

    max_threads = int(math.sqrt(
        tvm.target.Target.current(allow_none=False).max_num_threads))
    nthread_tx = max_threads
    nthread_bx = num_anchors // max_threads + 1
    tx = te.thread_axis("threadIdx.x")
    bx = te.thread_axis("blockIdx.x")
    ib.scope_attr(tx, "thread_extent", nthread_tx)
    ib.scope_attr(bx, "thread_extent", nthread_bx)
    j = bx * max_threads + tx

    with ib.for_range(0, batch_size, for_type="unroll") as i:
        base_idx = i * num_anchors * elem_length
        with ib.if_scope(j < num_anchors):
            with ib.for_range(0, elem_length) as k:
                out[base_idx + j * elem_length + k] = -1.0
            with ib.if_scope(flag[i * num_anchors + j] > 0):
                with ib.for_range(0, elem_length) as k:
                    out[base_idx + (idx[i * num_anchors + j] - 1) * elem_length + k] \
                        = data[base_idx + j * elem_length + k]
    return ib.get()


def non_max_suppression(data, valid_count, max_output_size=-1,
                        iou_threshold=0.5, force_suppress=False, top_k=-1,
                        coord_start=2, score_index=1, id_index=0,
                        return_indices=True, invalid_to_bottom=False):
    """Non-maximum suppression operator for object detection.

    Parameters
    ----------
    data : tvm.te.Tensor
        3-D tensor with shape [batch_size, num_anchors, elem_length].
        The last dimension should be in format of
        [class_id, score, box_left, box_top, box_right, box_bottom].

    valid_count : tvm.te.Tensor
        1-D tensor for valid number of boxes.

    max_output_size : optional, int
        Max number of output valid boxes for each instance.
        By default all valid boxes are returned.

    iou_threshold : optional, float
        Non-maximum suppression threshold.

    force_suppress : optional, boolean
        Whether to suppress all detections regardless of class_id.

    top_k : optional, int
        Keep maximum top k detections before nms, -1 for no limit.

    coord_start : required, int
        Start index of the consecutive 4 coordinates.

    score_index : optional, int
        Index of the scores/confidence of boxes.

    id_index : optional, int
        index of the class categories, -1 to disable.

    return_indices : boolean
        Whether to return box indices in input data.

    invalid_to_bottom : optional, boolean
        Whether to move all valid bounding boxes to the top.

    Returns
    -------
    out : tvm.te.Tensor
        3-D tensor with shape [batch_size, num_anchors, elem_length].

    Example
    --------
    .. code-block:: python

        # An example to use nms
        dshape = (1, 5, 6)
        data = te.placeholder(dshape, name="data")
        valid_count = te.placeholder((dshape[0],), dtype="int32", name="valid_count")
        iou_threshold = 0.7
        force_suppress = True
        top_k = -1
        out = non_max_suppression(data=data, valid_count=valid_count, iou_threshold=iou_threshold,
                                 force_suppress=force_supress, top_k=top_k, return_indices=False)
        np_data = np.random.uniform(dshape)
        np_valid_count = np.array([4])
        s = topi.generic.schedule_nms(out)
        f = tvm.build(s, [data, valid_count, out], "cuda")
        ctx = tvm.gpu(0)
        tvm_data = tvm.nd.array(np_data, ctx)
        tvm_valid_count = tvm.nd.array(np_valid_count, ctx)
        tvm_out = tvm.nd.array(np.zeros(dshape, dtype=data.dtype), ctx)
        f(tvm_data, tvm_valid_count, tvm_out)
    """
    batch_size = data.shape[0]
    num_anchors = data.shape[1]

    valid_count_dtype = "int32"
    valid_count_buf = tvm.tir.decl_buffer(valid_count.shape, valid_count_dtype,
                                          "valid_count_buf", data_alignment=4)
    score_axis = score_index
    score_shape = (batch_size, num_anchors)
    score_tensor = te.compute(
        score_shape, lambda i, j: data[i, j, score_axis], tag=tag.ELEMWISE)
    if tvm.get_global_func("tvm.contrib.thrust.sort_nms", allow_missing=True):
        sort_tensor = argsort_thrust(
            score_tensor, valid_count=valid_count, axis=1, is_ascend=False)
    else:
        sort_tensor = argsort(
            score_tensor, valid_count=valid_count, axis=1, is_ascend=False)

    sort_tensor_buf = tvm.tir.decl_buffer(sort_tensor.shape, sort_tensor.dtype,
                                          "sort_tensor_buf", data_alignment=8)

    data_buf = tvm.tir.decl_buffer(
        data.shape, data.dtype, "data_buf", data_alignment=8)

    out_buf = tvm.tir.decl_buffer(
        data.shape, data.dtype, "out_buf", data_alignment=8)

    out, box_indices = \
        te.extern([data.shape, score_shape],
                  [data, sort_tensor, valid_count],
                  lambda ins, outs: nms_ir(
            ins[0], ins[1], ins[2], outs[0], outs[1],
            max_output_size, iou_threshold, force_suppress,
            top_k, coord_start, id_index, score_index),
            dtype=[data.dtype, "int32"],
            in_buffers=[data_buf, sort_tensor_buf, valid_count_buf],
            name="nms",
            tag="nms")

    if return_indices:
        return box_indices

    if invalid_to_bottom:
        output_buf = tvm.tir.decl_buffer(
            data.shape, data.dtype, "output_buf", data_alignment=8)
        temp_flag_buf = tvm.tir.decl_buffer(
            score_shape, valid_count_dtype, "temp_flag", data_alignment=8)
        temp_idx_buf = tvm.tir.decl_buffer(
            score_shape, valid_count_dtype, "temp_idx", data_alignment=8)
        temp_flag, temp_idx = te.extern([score_shape, score_shape], [out],
                                        lambda ins, outs: invalid_to_bottom_pre(
            ins[0], outs[0], outs[1]),
            dtype=["int32", "int32"],
            in_buffers=[out_buf],
            out_buffers=[
            temp_flag_buf, temp_idx_buf],
            name="invalid_to_bottom_phase_one")

        output = te.extern([data.shape], [out, temp_flag, temp_idx],
                           lambda ins, outs: invalid_to_bottom_ir(
            ins[0], ins[1], ins[2], outs[0]),
            dtype=[data.dtype],
            in_buffers=[out_buf, temp_flag_buf, temp_idx_buf],
            out_buffers=[output_buf],
            name="invalid_to_bottom",
            tag="invalid_to_bottom")
        return output

    return out
