profile_argsort.py

import tvm 
from tvm import relay
from tvm.contrib import graph_runtime

import numpy as np

shapes = [ 
    [(2000, 2, 2), 0], 
    [(2, 2000, 2), 1], 
    [(2, 2, 2000), 2], 
    [(4000, 2, 2), 0], 
    [(2, 4000, 2), 1], 
    [(2, 2, 4000), 2], 
    [(2, 12000, 2), 1], 
    [(2, 2, 12000), 2], 
    [(12000, 2, 2), 0], 
    [(2000, 8, 8), 0], 
    [(8, 2000, 8), 1], 
    [(8, 8, 2000), 2], 
    [(4000, 8, 8), 0], 
    [(8, 4000, 8), 1], 
    [(8, 8, 4000), 2], 
    [(12000, 8, 8), 0], 
    [(8, 12000, 8), 1], 
    [(8, 8, 12000), 2], 
]
ctx = tvm.gpu(0)
target="cuda"
for shape, axis in shapes:
    x = relay.var("x", relay.TensorType(shape, "float32"))
    z = relay.argsort(x, axis=axis, is_ascend=True, dtype="int32")
    func = relay.Function([x], z)
    mod = tvm.ir.IRModule.from_expr(func)
    with tvm.transform.PassContext(opt_level=3):
        lib = relay.build(mod, target)
    m = graph_runtime.GraphModule(lib['default'](ctx))
    ftimer = m.module.time_evaluator("run", ctx, number=1, repeat=10)
    prof_res = np.array(ftimer().results)
    print(shape, axis, "Mean inference time (std dev): %.2f ms (%.2f ms)" %
            (np.mean(prof_res) * 1000, np.std(prof_res) * 1000))