From d6830a2d9c1ab9a3df3c8b4ceb0231abda159716 Mon Sep 17 00:00:00 2001 From: Leyuan Wang Date: Mon, 18 Mar 2019 04:29:48 +0000 Subject: [PATCH] error fixed --- topi/python/topi/cuda/nms.py | 10 +++------- topi/python/topi/cuda/ssd/multibox.py | 4 ++-- topi/python/topi/cuda/vision.py | 8 ++++++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/topi/python/topi/cuda/nms.py b/topi/python/topi/cuda/nms.py index 6118ff11f9ce6..53f413f837d6e 100644 --- a/topi/python/topi/cuda/nms.py +++ b/topi/python/topi/cuda/nms.py @@ -44,13 +44,13 @@ def get_valid_counts_pre(data, flag, idx, score_threshold): idx = ib.buffer_ptr(idx) score_threshold = tvm.make.node("FloatImm", dtype="float32", value=score_threshold) - max_threads = int(math.sqrt(tvm.target.current_target(allow_none=False).max_num_threads)) + max_threads = int(tvm.target.current_target(allow_none=False).max_num_threads) nthread_tx = max_threads nthread_bx = batch_size * num_anchors // max_threads + 1 tx = tvm.thread_axis("threadIdx.x") - bx = tvm.thread_axis("blockIdx.x") + bx = tvm.thread_axis("vthread") ib.scope_attr(tx, "thread_extent", nthread_tx) - ib.scope_attr(bx, "thread_extent", nthread_bx) + ib.scope_attr(bx, "virtual_thread", nthread_bx) tid = bx * max_threads + tx with ib.if_scope(tid < batch_size * num_anchors): @@ -64,10 +64,6 @@ def get_valid_counts_pre(data, flag, idx, score_threshold): flag[tid] = 0 idx[tid] = 0 - ib.emit(tvm.make.Call(None, 'tvm_storage_sync', - tvm.convert(['shared']), - tvm.expr.Call.Intrinsic, None, 0)) - with ib.if_scope(tid < batch_size): with ib.for_range(0, num_anchors) as k: with ib.if_scope(k > 0): diff --git a/topi/python/topi/cuda/ssd/multibox.py b/topi/python/topi/cuda/ssd/multibox.py index 599c729d95cae..0cb6f70d9ee6e 100644 --- a/topi/python/topi/cuda/ssd/multibox.py +++ b/topi/python/topi/cuda/ssd/multibox.py @@ -181,7 +181,7 @@ def transform_loc_pre(cls_prob, valid_count, temp_valid_count, temp_cls_id, temp threshold = tvm.make.node("FloatImm", dtype="float32", value=threshold) max_threads = int( - tvm.target.current_target(allow_none=False).max_num_threads) + math.sqrt(tvm.target.current_target(allow_none=False).max_num_threads)) nthread_tx = max_threads nthread_bx = (batch_size * num_anchors) // max_threads + 1 tx = tvm.thread_axis("threadIdx.x") @@ -293,7 +293,7 @@ def transform_loc(loc, loc_base_idx, anchor, anchor_base_idx, clip, vx, vy, vw, out_loc = ib.buffer_ptr(out) max_threads = int( - tvm.target.current_target(allow_none=False).max_num_threads) + math.sqrt(tvm.target.current_target(allow_none=False).max_num_threads)) nthread_tx = max_threads nthread_bx = (batch_size * num_anchors) // max_threads + 1 tx = tvm.thread_axis("threadIdx.x") diff --git a/topi/python/topi/cuda/vision.py b/topi/python/topi/cuda/vision.py index f32e21183f037..1e6aa87843275 100644 --- a/topi/python/topi/cuda/vision.py +++ b/topi/python/topi/cuda/vision.py @@ -17,10 +17,14 @@ def _default_schedule(outs): def traverse(op): """inline all one-to-one-mapping operators except the last stage (output)""" if op.tag in ["nms", "invalid_to_bottom"]: - sort = op.input_tensors[1] + if op.tag == "nms": + sort = op.input_tensors[1] + else: + out = op.input_tensors[0] + sort = s[out].op.input_tensors[1] score = s[sort].op.input_tensors[0] fused = s[score].fuse(*s[score].op.axis) - num_thread = tvm.target.current_target(allow_none=False).max_num_threads + num_thread = int(tvm.target.current_target(allow_none=False).max_num_threads) bx, tx = s[score].split(fused, factor=num_thread) s[score].bind(bx, tvm.thread_axis("blockIdx.x")) s[score].bind(tx, tvm.thread_axis("threadIdx.x"))