Function test_vm_storage

tests/python/relax/test_vm_execbuilder.py:169–193

()

def test_vm_storage():
    dtype = tvm.DataType("float32")
    shape = (4, 6)
    ib = relax.ExecBuilder()
    with ib.function("main", num_inputs=0):
        ib.emit_call(
            "vm.builtin.alloc_storage",
            args=[
                ib.vm_state(),
                (24,),
                ib.convert_constant(0),
                dtype,
                ib.convert_constant("global"),
            ],
            dst=ib.r(1),
        )
        ib.emit_call(
            "vm.builtin.alloc_tensor", args=[ib.r(1), ib.imm(0), shape, dtype], dst=ib.r(2)
        )
        ib.emit_ret(ib.r(2))
    ex = ib.get()
    vm = relax.VirtualMachine(ex, tvm.cpu())
    res = vm["main"]()
    assert res.device == tvm.cpu()
    assert res.shape == shape


def test_vm_goto():