Laurawly commented on a change in pull request #6108:
URL: https://github.com/apache/incubator-tvm/pull/6108#discussion_r460479420
##########
File path: topi/python/topi/cuda/nms.py
##########
@@ -184,7 +155,84 @@ def get_valid_counts(data, score_threshold=0, id_index=0,
score_index=1):
return [valid_count, out, out_indices]
-def nms_ir(data, sorted_index, valid_count, out, box_indices,
+def rearrange_indices_out_ir(data, output, valid_box_count):
+ """Low level IR to get rearrange_indices_out.
+ Parameters
+ ----------
+ data : Buffer
+ Input data. 2-D Buffer with shape [batch_size, num_anchors].
+
+ output: Buffer
+ 2-D Buffer with shape [batch_size, num_anchors].
+
+ valid_box_count : Buffer
+ 2-D Buffer with shape [batch_size, 1].
+
+ Returns
+ -------
+ stmt : Stmt
+ The result IR statement.
+ """
+ batch_size = data.shape[0]
+ num_anchors = data.shape[1]
+ ib = tvm.tir.ir_builder.create()
+
+ data = ib.buffer_ptr(data)
+ output = ib.buffer_ptr(output)
+ valid_box_count = ib.buffer_ptr(valid_box_count)
+
+ nthread_tx = batch_size
+ nthread_bx = 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 = tx
+
+ valid_box_count[tid] = 0
+ with ib.for_range(0, num_anchors) as anchor_ind:
+ output[tid * num_anchors + anchor_ind] = data[tid * num_anchors +
anchor_ind]
+ return ib.get()
+
+
+def rearrange_indices_out(data):
Review comment:
This will regress performance a lot, don't recommend.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
