test_sort.py

def test_sort():
    dshape = (1, 500)
    data = tvm.placeholder(dshape, name="data")
    np_data = np.random.rand(dshape[0], dshape[1]).astype(data.dtype)
    print(np_data)
    np_result = np.argsort(-np_data)
    def check_device(device):
        ctx = tvm.context(device, 0)
        if not ctx.exist:
            print("Skip because %s is not enabled" % device)
            return
        print("Running on target: %s" % device)
        with tvm.target.create(device):
            out = topi.cuda.sort(data)
            s = topi.generic.schedule_sort(out)

        tvm_data = tvm.nd.array(np_data, ctx)
        tvm_out = tvm.nd.array(np.zeros(dshape, dtype="int32"), ctx)
        f = tvm.build(s, [data, out], device)
        f(tvm_data, tvm_out)
        tvm.testing.assert_allclose(tvm_out.asnumpy(), np_result, rtol=1e-4)

    for device in ['cuda']:
        check_device(device)