Function test_vm_goto

tests/python/relax/test_vm_execbuilder.py:196–216

()

def test_vm_goto():
    ib = relax.ExecBuilder()
    with ib.function("main", num_inputs=2):
        ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(2))
        ib.emit_goto(2)
        ib.emit_call("test.vm.mul", args=[ib.r(2), ib.r(1)], dst=ib.r(2))
        ib.emit_ret(ib.r(2))
    ex = ib.get()
    vm = relax.VirtualMachine(ex, tvm.cpu())
    a = tvm.runtime.tensor(
        np.random.rand(
            4,
        )
    )
    b = tvm.runtime.tensor(
        np.random.rand(
            4,
        )
    )
    res = check_saved_func(vm, "main", a, b)
    tvm.testing.assert_allclose(res.numpy(), a.numpy() + b.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_if():