PassTimingInstrument

tvm.instrument.PassTimingInstrument 是用 C++ 实现的创建 pass 时间计时的包装器

import tvm
import tvm.testing
from tvm import relax, topi
from tvm.script import ir as I, relax as R, tir as T

def test():
    bb = relax.BlockBuilder()
    x = relax.Var("x", R.Tensor([10, 20], "float32"))
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit_te(topi.add, x, relax.const(1, "float32"))
            lv1 = bb.emit_te(topi.exp, lv0)
            gv = bb.emit_output(bb.call_te(topi.squeeze, lv1))
        bb.emit_func_output(gv)

    return bb.get()

timing_inst = tvm.instrument.PassTimingInstrument()
relax_mod = test()
with tvm.transform.PassContext(instruments=[timing_inst]):
    relax_mod = relax.transform.FuseOps()(relax_mod)
    # before exiting the context, get profile results.
    profiles = timing_inst.render()

print(profiles)

FuseOps: 343056us [343056us] (100.00%; 100.00%)