Modify topi tests (apache#9)

* Add pooling, reorg, softmax and vision

* Add lrn

Author: kevinthesun

topi/tests/python/common.py

@@ -64,6 +64,14 @@ def get_schedule_reduce(target):
             return _reduce_schedule[key]
     return _reduce_schedule["generic"]
 
+def get_schedule(target, schedule):
+    if isinstance(target, str):
+        target = tvm.target.create(target)
+    for key in target.keys:
+        if key in schedule:
+            return schedule[key]
+    return schedule["generic"]
+
 get_schedule_broadcast = get_schedule_injective
 get_schedule_elemwise = get_schedule_injective
 

topi/tests/python/test_topi_lrn.py

@@ -21,6 +21,18 @@
 from topi.util import get_const_tuple
 import topi.testing
 
+from common import get_schedule
+
+_lrn_schedule = {
+    "generic": topi.generic.schedule_lrn,
+    "gpu": topi.cuda.schedule_lrn,
+    "opencl": topi.cuda.schedule_lrn,
+    "metal": topi.cuda.schedule_lrn,
+    "rocm": topi.cuda.schedule_lrn,
+    "vulkan": topi.cuda.schedule_lrn,
+    "nvptx": topi.cuda.schedule_lrn,
+}
+
 def verify_lrn(shape, size, axis, bias, alpha, beta):
     A = tvm.placeholder(shape, name='A')
     B = topi.nn.lrn(A, size, axis, alpha, beta, bias)
@@ -35,10 +47,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            if device == 'llvm':
-                s = topi.generic.schedule_lrn([B])
-            else:
-                s = topi.cuda.schedule_lrn([B])
+            s_func = get_schedule(device, _lrn_schedule)
+            s = s_func([B])
         ctx = tvm.context(device, 0)
         a = tvm.nd.array(a_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), ctx)

topi/tests/python/test_topi_pooling.py

@@ -21,7 +21,26 @@
 import topi
 import topi.testing
 from topi.util import get_const_tuple
-from common import get_all_backend
+from common import get_all_backend, get_schedule
+
+_pool_schedule = {
+    "generic": topi.generic.schedule_pool,
+    "cpu": topi.x86.schedule_pool,
+    "gpu": topi.cuda.schedule_pool,
+    "hls": topi.hls.schedule_pool,
+}
+
+_adaptive_pool_schedule = {
+    "generic": topi.generic.schedule_adaptive_pool,
+    "cpu": topi.x86.schedule_adaptive_pool,
+    "gpu": topi.cuda.schedule_adaptive_pool,
+    "hls": topi.hls.schedule_adaptive_pool,
+}
+
+_pool_grad_schedule = {
+    "generic": topi.generic.schedule_pool_grad,
+    "gpu": topi.cuda.schedule_pool_grad,
+}
 
 def verify_pool(n, ic, ih, kh, sh, padding, pool_type, ceil_mode, count_include_pad=True):
     iw = ih
@@ -74,7 +93,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_pool(B, layout)
+            s_func = get_schedule(device, _pool_schedule)
+            s = s_func(B, layout)
 
         a = tvm.nd.array(a_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), ctx)
@@ -129,7 +149,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_pool_grad(PoolGrad)
+            s_func = get_schedule(device, _pool_grad_schedule)
+            s = s_func(PoolGrad)
 
         a = tvm.nd.array(a_np, ctx)
         out_grad = tvm.nd.array(out_grad_np, ctx)
@@ -201,7 +222,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_adaptive_pool(B)
+            s_func = get_schedule(device, _adaptive_pool_schedule)
+            s = s_func(B)
         a = tvm.nd.array(a_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), ctx)
         f = tvm.build(s, [A, B], device)
@@ -255,7 +277,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_adaptive_pool(out)
+            s_func = get_schedule(device, _adaptive_pool_schedule)
+            s = s_func(out)
         a = tvm.nd.array(np_data, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(oshape), dtype=out.dtype), ctx)
         f = tvm.build(s, [data, out], device)
@@ -298,7 +321,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_pool(B, layout)
+            s_func = get_schedule(device, _pool_schedule)
+            s = s_func(B, layout)
 
         a = tvm.nd.array(input_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), ctx)
@@ -350,7 +374,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            s = topi.generic.schedule_pool(B, layout)
+            s_func = _pool_schedule[device] if device in _pool_schedule else _pool_schedule["generic"]
+            s = s_func(B, layout)
 
         a = tvm.nd.array(input_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=dtype), ctx)

topi/tests/python/test_topi_reorg.py

@@ -20,6 +20,12 @@
 from topi.util import get_const_tuple
 import tvm
 import topi.testing
+from common import get_schedule
+
+_reorg_schedule = {
+    "generic": topi.generic.schedule_reorg,
+    "gpu": topi.cuda.schedule_reorg,
+}
 
 def verify_reorg(batch, in_size, in_channel, stride):
     '''Verify reorg operator by comparing outputs from tvm and numpy implementation'''
@@ -46,10 +52,8 @@ def check_device(device):
             return
         print("Running on target: %s" % device)
         with tvm.target.create(device):
-            if device == 'llvm':
-                s = topi.generic.schedule_reorg([B])
-            else:
-                s = topi.cuda.schedule_reorg([B])
+            s_func = get_schedule(device, _reorg_schedule)
+            s = s_func([B])
         a = tvm.nd.array(a_np, ctx)
         b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), ctx)
         func = tvm.build(s, [A, B], device)

topi/tests/python/test_topi_softmax.py

@@ -23,7 +23,15 @@
 import logging
 from topi.util import get_const_tuple
 
-from common import get_all_backend
+from common import get_all_backend, get_schedule
+
+_softmax_schedule = {
+    "generic": topi.generic.schedule_softmax,
+    "cpu": topi.x86.schedule_softmax,
+    "gpu": topi.cuda.schedule_softmax,
+    "hls": topi.hls.schedule_softmax,
+    "opengl": topi.opengl.schedule_softmax,
+}
 
 def check_device(A, B, a_np, b_np, device, name):
     ctx = tvm.context(device, 0)
@@ -32,11 +40,12 @@ def check_device(A, B, a_np, b_np, device, name):
         return
     print("Running on target: %s" % device)
     with tvm.target.create(device):
-        s = topi.generic.schedule_softmax(B)
+        s_func = get_schedule(device, _softmax_schedule)
+        s = s_func(B)
 
     a = tvm.nd.array(a_np, ctx)
     b = tvm.nd.array(np.zeros(get_const_tuple(B.shape), dtype=B.dtype), ctx)
-    f = tvm.build(s, [A, B], device, name="softmax")
+    f = tvm.build(s, [A, B], device, name=name)
     f(a, b)
     tvm.testing.assert_allclose(b.asnumpy(), b_np, rtol=1e-5)
 
@@ -50,7 +59,7 @@ def verify_softmax(m, n, dtype="float32"):
     a_np = np.random.uniform(size=get_const_tuple(A.shape)).astype(A.dtype)
     b_np = topi.testing.softmax_python(a_np)
 
-    for device in ['cuda', 'opencl', 'metal', 'rocm', 'vulkan', 'nvptx']:
+    for device in get_all_backend():
         check_device(A, B, a_np, b_np, device, "softmax")
 
 def verify_softmax_4d(shape, dtype="float32"):
@@ -62,7 +71,7 @@ def verify_softmax_4d(shape, dtype="float32"):
     b_np = topi.testing.softmax_python(a_np.transpose(0, 2, 3, 1).reshape(h*w, c))
     b_np = b_np.reshape(1, h, w, c).transpose(0, 3, 1, 2)
 
-    for device in ['cuda', 'opencl', 'metal', 'rocm', 'vulkan', 'nvptx']:
+    for device in get_all_backend():
         check_device(A, B, a_np, b_np, device, "softmax")
 
 def test_softmax():

topi/tests/python/test_topi_vision.py

@@ -26,6 +26,42 @@
 from topi.util import get_const_tuple
 from topi.vision import ssd, non_max_suppression, get_valid_counts
 
+from common import get_schedule
+
+_get_valid_counts_schedule = {
+    "generic": topi.generic.schedule_get_valid_counts,
+    "gpu": topi.cuda.schedule_get_valid_counts,
+}
+
+_nms_schedule = {
+    "generic": topi.generic.schedule_nms,
+    "gpu": topi.cuda.schedule_nms,
+}
+
+_multibox_prior_schedule = {
+    "generic": topi.generic.schedule_multibox_prior,
+    "gpu": topi.cuda.schedule_multibox_prior,
+}
+
+_multibox_detection_schedule = {
+    "generic": topi.generic.schedule_multibox_detection,
+    "gpu": topi.cuda.schedule_multibox_detection,
+}
+
+_roi_align_schedule = {
+    "generic": topi.generic.schedule_roi_align,
+    "gpu": topi.cuda.schedule_roi_align,
+}
+
+_roi_pool_schedule = {
+    "generic": topi.generic.schedule_roi_pool,
+    "gpu": topi.cuda.schedule_roi_pool,
+}
+
+_proposal_schedule = {
+    "generic": topi.generic.schedule_proposal,
+    "gpu": topi.cuda.schedule_proposal,
+}
 
 def verify_get_valid_counts(dshape, score_threshold, id_index, score_index):
     dtype = "float32"
@@ -56,7 +92,8 @@ def check_device(device):
         with tvm.target.create(device):
             data = tvm.placeholder(dshape, name="data", dtype=dtype)
             outs = get_valid_counts(data, score_threshold, id_index, score_index)
-            s = topi.generic.schedule_get_valid_counts(outs)
+            s_func = get_schedule(device, _get_valid_counts_schedule)
+            s = s_func(outs)
 
         tvm_input_data = tvm.nd.array(np_data, ctx)
         tvm_out1 = tvm.nd.array(np.zeros(np_out1.shape, dtype="int32"), ctx)
@@ -68,8 +105,6 @@ def check_device(device):
 
     for device in ['llvm', 'cuda', 'opencl']:
         # Disable gpu test for now
-        if device != "llvm":
-            continue
         check_device(device)
 
 
@@ -107,7 +142,8 @@ def check_device(device):
                                                     return_indices=False)
                 indices_out = topi.cuda.non_max_suppression(data, valid_count, -1, iou_threshold, force_suppress, top_k,
                                                             coord_start=coord_start, score_index=score_index, id_index=id_index)
-            s = topi.generic.schedule_nms(out)
+            s_func = get_schedule(device, _nms_schedule)
+            s = s_func(out)
             indices_s = topi.generic.schedule_nms(indices_out)
 
         tvm_data = tvm.nd.array(np_data, ctx)
@@ -198,7 +234,8 @@ def check_device(device):
                 out = ssd.multibox_prior(data, sizes, ratios, steps, offsets, clip)
             else:
                 out = topi.cuda.ssd.multibox_prior(data, sizes, ratios, steps, offsets, clip)
-            s = topi.generic.schedule_multibox_prior(out)
+            s_func = get_schedule(device, _multibox_prior_schedule)
+            s = s_func(out)
 
         tvm_input_data = tvm.nd.array(input_data, ctx)
         tvm_out = tvm.nd.array(np.zeros(oshape, dtype=dtype), ctx)
@@ -244,7 +281,8 @@ def check_device(device):
                 out = ssd.multibox_detection(cls_prob, loc_preds, anchors)
             else:
                 out = topi.cuda.ssd.multibox_detection(cls_prob, loc_preds, anchors)
-            s = topi.generic.schedule_multibox_detection(out)
+            s_func = get_schedule(device, _multibox_detection_schedule)
+            s = s_func(out)
 
         tvm_cls_prob = tvm.nd.array(np_cls_prob.astype(cls_prob.dtype), ctx)
         tvm_loc_preds = tvm.nd.array(np_loc_preds.astype(loc_preds.dtype), ctx)
@@ -289,7 +327,8 @@ def check_device(device):
             b = topi.vision.rcnn.roi_align_nchw(a, rois, pooled_size=pooled_size,
                                                 spatial_scale=spatial_scale,
                                                 sample_ratio=sample_ratio)
-            s = topi.generic.schedule_roi_align(b)
+            s_func = get_schedule(device, _roi_align_schedule)
+            s = s_func(b)
 
         tvm_a = tvm.nd.array(a_np, ctx)
         tvm_rois = tvm.nd.array(rois_np, ctx)
@@ -338,7 +377,8 @@ def check_device(device):
         with tvm.target.create(device):
             b = topi.vision.rcnn.roi_pool_nchw(a, rois, pooled_size=pooled_size,
                                                 spatial_scale=spatial_scale)
-            s = topi.generic.schedule_roi_pool(b)
+            s_func = get_schedule(device, _roi_pool_schedule)
+            s = s_func(b)
 
         tvm_a = tvm.nd.array(a_np, ctx)
         tvm_rois = tvm.nd.array(rois_np, ctx)
@@ -369,7 +409,8 @@ def check_device(device):
         print("Running on target: %s" % device)
         with tvm.target.create(device):
             out = topi.vision.proposal(cls_prob, bbox_pred, im_info, **attrs)
-            s = topi.generic.schedule_proposal(out)
+            s_func = get_schedule(device, _proposal_schedule)
+            s = s_func(out)
             f = tvm.build(s, [cls_prob, bbox_pred, im_info, out], device)
             tvm_cls_prob = tvm.nd.array(np_cls_prob, ctx=ctx)
             tvm_bbox_pred = tvm.nd.array(np_bbox_pred, ctx=ctx)