# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# ruff: noqa: F841

import ctypes
from collections.abc import Callable

import numpy as np
import pytest
import tvm_ffi
from tvm_ffi import Shape

import tvm
import tvm.script
import tvm.testing
from tvm import relax, rpc, te, tirx, topi
from tvm.relax.testing import nn
from tvm.relax.testing.vm import check_saved_func
from tvm.script import ir as I
from tvm.script import relax as R
from tvm.script import tirx as T
from tvm.support import cc, popen_pool, utils
from tvm.testing import env

EXEC_MODE = ["bytecode", "compiled"]


@pytest.fixture(params=EXEC_MODE)
def exec_mode(request):
    return request.param


def test_vm_compile_simple(exec_mode):
    @tvm.script.ir_module
    class TestVMCompileStage0:
        @R.function
        def foo(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
            z = R.call_pure_packed(
                "test.vm.identity", x, y, ty_args=(R.Tensor(ndim=2, dtype="float32"))
            )
            return y

    mod = TestVMCompileStage0
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    inp1 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
    inp2 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
    vm = relax.VirtualMachine(ex, tvm.cpu())
    vm["foo"](inp1, inp2)
    tvm.testing.assert_allclose(inp2.numpy(), inp1.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_compile_without_target_arg(exec_mode):
    """Like test_vm_compile_simple, but with a default target"""

    @tvm.script.ir_module
    class mod:
        @R.function
        def foo(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
            z = R.call_pure_packed(
                "test.vm.identity", x, y, ty_args=(R.Tensor(ndim=2, dtype="float32"))
            )
            return y

    ex = relax.build(mod, exec_mode=exec_mode)
    inp1 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
    inp2 = tvm.runtime.tensor(np.random.rand(3, 4).astype(np.float32))
    vm = relax.VirtualMachine(ex, tvm.cpu())
    vm["foo"](inp1, inp2)
    tvm.testing.assert_allclose(inp2.numpy(), inp1.numpy(), rtol=1e-7, atol=1e-7)


def test_match_check(exec_mode):
    @tvm.script.ir_module
    class TestMatchCheck:
        @R.function
        def foo(x: R.Tensor(["n", "m"], "int32"), y: R.Any) -> R.Tensor(["m", "n"], dtype=None):
            return y

    mod = TestMatchCheck
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x0 = tvm.runtime.tensor(np.zeros((1, 2)).astype("int32"))
    y0 = tvm.runtime.tensor(np.zeros((2, 1)).astype("float32"))
    y1 = tvm.runtime.tensor(np.zeros((1, 2)).astype("float32"))
    y2 = tvm.runtime.tensor(np.zeros((2, 1, 1)).astype("float32"))

    vm["foo"](x0, y0)

    with pytest.raises(RuntimeError, match=".*return.*"):
        vm["foo"](x0, y1)

    with pytest.raises(ValueError, match=".*return.*"):
        vm["foo"](x0, y2)


def test_vm_compile_stage2(exec_mode):
    @tvm.script.ir_module
    class TestVMCompileStage2:
        @R.function
        def foo(x: R.Tensor(dtype="float32")) -> R.Shape:
            n, m = T.int64(), T.int64()
            _ = R.match_cast(x, R.Tensor((n, m), "float32"))
            return R.shape([n * 2, m * 3])

    mod = TestVMCompileStage2
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    shape = (32, 16)
    arr = tvm.runtime.tensor(np.random.rand(*shape).astype("float32"))
    res = vm["foo"](arr)
    assert res[0] == shape[0] * 2
    assert res[1] == shape[1] * 3

    # dtype mismatch
    with pytest.raises(ValueError, match=".*dtype.*"):
        vm["foo"](tvm.runtime.tensor(np.zeros((1, 2)).astype("int32")))

    # ndim mismatch
    with pytest.raises(ValueError, match=".*match_cast.*ndim.*"):
        vm["foo"](tvm.runtime.tensor(np.zeros((1,)).astype("float32")))

    # type mismach
    with pytest.raises(TypeError):
        vm["foo"]([])


def test_vm_compile_stage3(exec_mode):
    @tvm.script.ir_module
    class TestVMCompileStage3:
        @R.function
        def foo(x: R.Tensor((32, 16), "float32")) -> R.Tensor:
            with R.dataflow():
                y = R.call_dps_packed("test.vm.identity", (x), R.Tensor((32, 16), dtype="float32"))
                R.output(y)
            return y

    mod = TestVMCompileStage3
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    shape = (32, 16)
    inp = tvm.runtime.tensor(np.random.rand(*shape).astype(np.float32))
    res = vm["foo"](inp)
    tvm.testing.assert_allclose(res.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_compile_e2e(exec_mode):
    @tvm.script.ir_module
    class TestVMCompileE2E:
        @R.function
        def foo(x: R.Tensor(dtype="float32")) -> R.Tensor:
            with R.dataflow():
                n, m = T.int64(), T.int64()
                _ = R.match_cast(x, R.Tensor((n, m), "float32"))
                y = R.call_dps_packed("test.vm.tile", (x), R.Tensor((n, m * 2), dtype="float32"))
                R.output(y)
            return y

    mod = TestVMCompileE2E

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    shape = (32, 16)
    inp = tvm.runtime.tensor(np.random.rand(*shape).astype(np.float32))
    res = check_saved_func(vm, "foo", inp)
    tvm.testing.assert_allclose(res.numpy(), np.tile(inp.numpy(), (1, 2)), rtol=1e-7, atol=1e-7)


def test_vm_compile_e2e_func_param_with_shape(exec_mode):
    @tvm.script.ir_module
    class TestVMCompileE2E2:
        @T.prim_func(s_tir=True)
        def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None:
            T.func_attr({"global_symbol": "tir_matmul"})
            m = T.int32()
            n = T.int32()
            k = T.int32()
            A = T.match_buffer(x, (m, n))
            B = T.match_buffer(y, (n, k))
            C = T.match_buffer(z, (m, k))

            for i, j, k in T.grid(m, k, n):
                with T.sblock("matmul"):
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    with T.init():
                        C[vi, vj] = T.float32(0)
                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]

        @R.function
        def func(
            x: R.Tensor(("m", "n"), "float32"), w: R.Tensor(("n", "k"), "float32")
        ) -> R.Tensor:
            m, k = T.int64(), T.int64()
            cls = TestVMCompileE2E2
            gv0 = R.call_tir(cls.tir_matmul, (x, w), R.Tensor((m, k), dtype="float32"))
            return gv0

    mod = TestVMCompileE2E2

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    data = tvm.runtime.tensor(np.random.rand(32, 16).astype(np.float32))
    weight = tvm.runtime.tensor(np.random.rand(16, 32).astype(np.float32))
    res = check_saved_func(vm, "func", data, weight)
    expected = np.dot(data.numpy(), weight.numpy())
    tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)


def test_call_tir_inplace_e2e_simple(exec_mode):
    @tvm.script.ir_module
    class TestCallTIRInplaceE2ESimple:
        @T.prim_func(s_tir=True)
        def copy(
            A: T.Buffer((2, 3), "int32"),
            B: T.Buffer((2, 3), "int32"),
            C: T.Buffer((2, 3), "int32"),
            out1: T.Buffer((2, 3), "int32"),
        ):
            # copies the contents of C into A, B, and out1
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_zeros"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(C[ax0, ax1])
                    T.writes(A[ax0, ax1], B[ax0, ax1], out1[ax0, ax1])
                    A[ax0, ax1] = C[ax0, ax1]
                    B[ax0, ax1] = C[ax0, ax1]
                    out1[ax0, ax1] = C[ax0, ax1]

        @R.function
        def main(
            x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32"), z: R.Tensor((2, 3), "int32")
        ) -> R.Tuple(
            R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")
        ):
            res = R.call_tir_inplace(
                TestCallTIRInplaceE2ESimple.copy,
                (x, y, z),
                [0, 1, -1],
                [R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32"), R.Tensor((2, 3), "int32")],
            )
            return res

    mod = TestCallTIRInplaceE2ESimple

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    x = tvm.runtime.tensor(np.zeros((2, 3)).astype(np.int32))
    y = tvm.runtime.tensor(np.zeros((2, 3)).astype(np.int32))
    z = tvm.runtime.tensor(np.ones((2, 3)).astype(np.int32))
    vm.set_input("main", x, y, z)
    vm.invoke_stateful("main")
    outs = vm.get_outputs("main")
    # check the expected aliasing (the last result is newly allocated)
    assert x == outs[0]
    assert y == outs[1]
    assert x != y
    assert x != outs[2]
    assert y != outs[2]
    tvm.testing.assert_allclose(x.numpy(), z.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(y.numpy(), z.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(outs[2].numpy(), z.numpy(), rtol=1e-7, atol=1e-7)


def test_call_tir_inplace_e2e_rw(exec_mode):
    # read and write from the same tensor
    @tvm.script.ir_module
    class TestCallTIRInplaceE2ERW:
        @T.prim_func(s_tir=True)
        def inplace_add(A: T.Buffer((2, 3), "int32"), B: T.Buffer((2, 3), "int32")):
            # sums A and B, storing the result in A
            T.func_attr({"tirx.noalias": True})
            for i0, i1 in T.grid(T.int64(2), T.int64(3)):
                with T.sblock("T_add"):
                    ax0, ax1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[ax0, ax1], B[ax0, ax1])
                    T.writes(A[ax0, ax1])
                    A[ax0, ax1] = A[ax0, ax1] + B[ax0, ax1]

        @R.function
        def main(x: R.Tensor((2, 3), "int32"), y: R.Tensor((2, 3), "int32")) -> R.Tensor(
            (2, 3), "int32"
        ):
            res = R.call_tir_inplace(
                TestCallTIRInplaceE2ERW.inplace_add, (x, y), [0], R.Tensor((2, 3), "int32")
            )
            return res

    mod = TestCallTIRInplaceE2ERW

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    x = tvm.runtime.tensor(np.ones((2, 3)).astype(np.int32))
    y = tvm.runtime.tensor(np.ones((2, 3)).astype(np.int32))
    vm.set_input("main", x, y)
    vm.invoke_stateful("main")
    out = vm.get_outputs("main")
    expected = tvm.runtime.tensor(np.full((2, 3), 2).astype(np.int32))

    assert x == out
    tvm.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_emit_te_extern(exec_mode):
    if not tvm.get_global_func("tvm.contrib.cblas.matmul", True):
        print("skip because extern function is not available")
        return
    bb = relax.BlockBuilder()
    n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
    x = relax.Var("x", R.Tensor([n, m], "float32"))
    y = relax.Var("y", R.Tensor([m, n], "float32"))

    with bb.function("rx_cblas_matmul", [x, y]):
        out = bb.emit_te(tvm.contrib.cblas.matmul, x, y, transa=False, transb=False)
        bb.emit_func_output(out)

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    data = tvm.runtime.tensor(np.random.rand(16, 32).astype(np.float32))
    weight = tvm.runtime.tensor(np.random.rand(32, 16).astype(np.float32))
    res = check_saved_func(vm, "rx_cblas_matmul", data, weight)
    expected = np.dot(data.numpy(), weight.numpy())
    tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)


def test_vm_emit_te_concat(exec_mode):
    # concatenate of two vectors of size (n,) and (m,)
    bb = relax.BlockBuilder()
    n, m = tirx.Var("n", "int64"), tirx.Var("m", "int64")
    x = relax.Var("x", R.Tensor([n], "float32"))
    y = relax.Var("y", R.Tensor([m], "float32"))

    def te_func(A, B):
        C = te.compute((n + m), lambda i: tvm.tirx.if_then_else(i < n, A[i], B[i - n]))
        return C

    with bb.function("rx_func", [x, y]):
        x1 = bb.emit_te(te_func, x, y)
        bb.emit_func_output(x1)

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    vm = relax.VirtualMachine(ex, tvm.cpu())
    inp = tvm.runtime.tensor(
        np.random.rand(
            1,
        ).astype(np.float32)
    )
    inp2 = tvm.runtime.tensor(
        np.random.rand(
            2,
        ).astype(np.float32)
    )
    res = check_saved_func(vm, "rx_func", inp, inp2)
    tvm.testing.assert_allclose(
        res.numpy(), np.append(inp.numpy(), inp2.numpy()), rtol=1e-7, atol=1e-7
    )


def test_vm_emit_te_dtype_change(exec_mode):
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")
    x = relax.Var("x", R.Tensor([n], "float32"))

    # convert a tensor with dtype of float32 to int16
    def te_func(A):
        B = te.compute((n,), lambda i: A[i].astype("int16"))
        return B

    with bb.function("rx_func", [x]):
        y = bb.emit_te(te_func, x)
        bb.emit_func_output(y)

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    vm = relax.VirtualMachine(ex, tvm.cpu())
    inp = tvm.runtime.tensor(
        np.random.rand(
            1,
        ).astype(np.float32)
    )
    res = check_saved_func(vm, "rx_func", inp)
    tvm.testing.assert_allclose(res.numpy(), inp.numpy().astype("int16"))


def test_vm_emit_te_floor_symbolic_shape(exec_mode):
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")
    x = relax.Var("x", R.Tensor([n], "float32"))

    def te_func(A):
        C = te.compute((tirx.floordiv(n, 2),), lambda i: A[i] + 1)
        return C

    with bb.function("rx_func", [x]):
        x1 = bb.emit_te(te_func, x)
        bb.emit_func_output(x1)

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    vm = relax.VirtualMachine(ex, tvm.cpu())
    shape = (9,)
    inp = tvm.runtime.tensor(np.random.rand(*shape).astype(np.float32))
    res = check_saved_func(vm, "rx_func", inp)

    def expected_output():
        output_shape = (shape[0] // 2,)
        return inp.numpy()[: output_shape[0]] + 1

    tvm.testing.assert_allclose(res.numpy(), expected_output(), rtol=1e-7, atol=1e-7)


def test_vm_emit_te_constant_param_cpu(exec_mode):
    x_np = np.random.rand(2, 2).astype("float32")
    c_np = np.random.rand(2, 2).astype("float32")

    bb = relax.BlockBuilder()
    x = relax.Var("x", R.Tensor((2, 2), "float32"))
    c = relax.const(c_np, "float32")
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit_te(topi.add, x, c)
            gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)

    mod = bb.get()
    exec = relax.build(mod, "llvm", exec_mode=exec_mode)
    dev = tvm.cpu()
    vm = relax.VirtualMachine(exec, dev)

    add_res = check_saved_func(vm, "main", tvm.runtime.tensor(x_np, dev))
    tvm.testing.assert_allclose(add_res.numpy(), x_np + c_np, rtol=1e-7, atol=1e-7)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_vm_emit_te_constant_param_gpu(exec_mode):
    x_np = np.random.rand(2, 2).astype("float32")
    c_np = np.random.rand(2, 2).astype("float32")

    bb = relax.BlockBuilder()
    x = relax.Var("x", R.Tensor((2, 2), "float32"))
    c = relax.const(c_np, "float32")
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit_te(topi.add, x, c)
            gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)

    mod = bb.get()
    sch = tvm.s_tir.Schedule(mod, debug_mask="all")
    loops = sch.get_loops(sch.get_sblock(name="T_add", func_name="add"))
    sch.bind(loops[0], "threadIdx.x")

    exec = relax.build(sch.mod, "cuda", exec_mode=exec_mode)

    def run_and_check():
        dev = tvm.cuda()
        vm = relax.VirtualMachine(exec, dev)
        add_res = check_saved_func(vm, "main", tvm.runtime.tensor(x_np, dev))
        tvm.testing.assert_allclose(add_res.numpy(), x_np + c_np, rtol=1e-7, atol=1e-7)

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_vm_relax_symbolic_shape(exec_mode):
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")
    x = relax.Var("x", R.Tensor([n], "float32"))
    y = relax.Var("y", R.Tensor([(n // 2) + 1], "float32"))

    def te_func(A, B):
        C = te.compute((n,), lambda i: A[i] + B[i // 2])
        return C

    with bb.function("rx_func", [x, y]):
        x1 = bb.emit_te(te_func, x, y)
        bb.emit_func_output(x1)

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    vm = relax.VirtualMachine(ex, tvm.cpu())
    shape1 = (5,)
    shape2 = (3,)
    inp = tvm.runtime.tensor(np.random.rand(*shape1).astype(np.float32))
    inp2 = tvm.runtime.tensor(np.random.rand(*shape2).astype(np.float32))
    res = check_saved_func(vm, "rx_func", inp, inp2)

    def expected_output():
        return inp.numpy() + np.repeat(inp2.numpy(), 2)[:5]

    tvm.testing.assert_allclose(res.numpy(), expected_output(), rtol=1e-7, atol=1e-7)


def test_vm_relax_symbolic_shape_tuple(exec_mode):
    @I.ir_module(s_tir=True)
    class mod:
        @R.function
        def main(shape: R.Shape(["m", "n"])):
            m = T.int64()
            n = T.int64()
            return R.shape([2 * m, 3 * n])

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    func = vm["main"]

    assert func(Shape([2, 3])) == (4, 9)

    with pytest.raises(ValueError):
        func(Shape([2, 3, 4]))

    with pytest.raises(TypeError):
        func(R.prim_value(2))


def test_vm_relax_dyn_tir_shape(exec_mode):
    # case where TIR variables are unbound in generated PrimFunc
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")

    def te_func(A):
        C = te.compute((n + 1), lambda i: A[i])
        return C

    with bb.function("rx_func"):
        x = nn.Placeholder((n,), dtype="float32", name="x")
        y = nn.Placeholder((n + 1,), dtype="float32", name="y")

        x1 = bb.emit_te(te_func, y)
        bb.emit_func_output(x1, params=[x, y])

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    with utils.tempdir() as temp:
        ex.export_library(temp.relpath("exec.so"))
        vm = relax.VirtualMachine(tvm.runtime.load_module(temp.relpath("exec.so")), tvm.cpu())

    inp = tvm.runtime.tensor(np.random.rand(2).astype(np.float32))
    inp2 = tvm.runtime.tensor(np.random.rand(3).astype(np.float32))

    res = check_saved_func(vm, "rx_func", inp, inp2)

    tvm.testing.assert_allclose(res.numpy(), inp2.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_tuple(exec_mode):
    bb = relax.BlockBuilder()
    n = tirx.Var("n", "int64")

    with bb.function("rx_func"):
        x = nn.Placeholder((n,), dtype="float32", name="x")
        y = nn.Placeholder((n,), dtype="float32", name="y")
        tup = relax.Tuple([x, y])
        item = tup[0]
        bb.emit_func_output([tup, item], params=[x, y])

    mod = bb.get()

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    vm = relax.VirtualMachine(ex, tvm.cpu())
    shape = (5,)
    inp = tvm.runtime.tensor(np.random.rand(*shape).astype(np.float32))
    inp2 = tvm.runtime.tensor(np.random.rand(*shape).astype(np.float32))
    (res1, res2), res3 = vm["rx_func"](inp, inp2)

    tvm.testing.assert_allclose(res1.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(res2.numpy(), inp2.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(res3.numpy(), inp.numpy(), rtol=1e-7, atol=1e-7)


def test_vm_tuplegetitem(exec_mode):
    @tvm.script.ir_module
    class TestVMTupleGetItem:
        @R.function
        def tuple_get_item(
            x: R.Tensor(ndim=2, dtype="float32"),
            y: R.Tensor(ndim=2, dtype="float32"),
        ):
            t = (x, y)
            a = t[0]
            b = t[1]
            c = R.call_pure_packed("test.vm.add", a, b, ty_args=(R.Tensor(ndim=2, dtype="float32")))
            return c

    mod = TestVMTupleGetItem
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    y_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    res = check_saved_func(vm, "tuple_get_item", x_inp, y_inp)
    tvm.testing.assert_allclose(res.numpy(), x_inp.numpy() + y_inp.numpy(), rtol=1e-7, atol=1e-7)


def test_lower_memory_alloc_storage_tensor(exec_mode):
    @tvm.script.ir_module
    class TestMemoryAllocStorageTensor:
        @R.function
        def main(x: R.Tensor((2, 3), dtype="float32")):
            R.func_attr({"relax.force_pure": True})
            cls = TestMemoryAllocStorageTensor
            storage = R.memory.alloc_storage(
                R.shape([24]), virtual_device_index=0, storage_scope="global", dtype="float32"
            )
            y = R.memory.alloc_tensor(storage, 0, R.shape([2, 3]), dtype="float32")
            # this is an impure operation, but the overall function is pure so we force purity
            _ = cls.copy(x, y)
            return y

        @T.prim_func(s_tir=True)
        def copy(A: T.Buffer((2, 3), "float32"), B: T.Buffer((2, 3), "float32")):
            for i0, i1 in T.grid(2, 3):
                with T.sblock("block"):
                    vi0, vi1 = T.axis.remap("SS", [i0, i1])
                    B[vi0, vi1] = A[vi0, vi1]

    mod = TestMemoryAllocStorageTensor
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    y = vm["main"](x)
    tvm.testing.assert_allclose(y.numpy(), x.numpy(), rtol=1e-7, atol=1e-7)


def test_sub_func_call(exec_mode):
    @tvm.script.ir_module
    class TestVMSubFunction:
        @T.prim_func(s_tir=True)
        def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None:
            T.func_attr({"global_symbol": "tir_matmul"})
            m = T.int32()
            n = T.int32()
            k = T.int32()
            A = T.match_buffer(x, (m, n))
            B = T.match_buffer(y, (n, k))
            C = T.match_buffer(z, (m, k))

            for i, j, k in T.grid(m, k, n):
                with T.sblock("matmul"):
                    vi, vj, vk = T.axis.remap("SSR", [i, j, k])
                    with T.init():
                        C[vi, vj] = T.float32(0)
                    C[vi, vj] = C[vi, vj] + A[vi, vk] * B[vk, vj]

        @R.function
        def relax_matmul_tir(
            x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")
        ) -> R.Tensor((32, 32), dtype="float32"):
            cls = TestVMSubFunction
            with R.dataflow():
                gv0 = R.call_tir(cls.tir_matmul, (x, w), R.Tensor((32, 32), dtype="float32"))
                R.output(gv0)
            return gv0

        @R.function
        def relax_matmul_packed(
            x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")
        ) -> R.Any:
            gv0 = R.call_pure_packed(
                "test.vm.mul", x, w, ty_args=(R.Tensor(ndim=2, dtype="float32"))
            )
            return gv0

        @R.function
        def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Any:
            cls = TestVMSubFunction
            gv0 = cls.relax_matmul_tir(x, w)
            gv1 = cls.relax_matmul_packed(gv0, gv0)
            return gv1

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(TestVMSubFunction, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x_inp = tvm.runtime.tensor(np.random.rand(32, 32).astype(np.float32))
    y_inp = tvm.runtime.tensor(np.random.rand(32, 32).astype(np.float32))
    res = check_saved_func(vm, "main", x_inp, y_inp)
    product = np.dot(x_inp.numpy(), y_inp.numpy())
    expected = product * product
    tvm.testing.assert_allclose(res.numpy(), expected, rtol=1e-6, atol=1e-6)


def test_recursion(exec_mode):
    @tvm.script.ir_module
    class TestVMRecursion:
        @R.function
        def recursion(n: R.Tensor((1,), "float32")) -> R.Tensor:
            cond = R.call_pure_packed(
                "test.vm.equal_zero", n, ty_args=(R.Tensor(ndim=1, dtype="float32"))
            )
            if cond:
                res = R.const(1.0)
            else:
                gv0 = R.call_pure_packed(
                    "test.vm.subtract_one", n, ty_args=(R.Tensor(ndim=1, dtype="float32"))
                )
                tmp = TestVMRecursion.recursion(gv0)
                res = R.call_pure_packed(
                    "test.vm.add", tmp, tmp, ty_args=(R.Tensor(ndim=1, dtype="float32"))
                )
            return res

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(TestVMRecursion, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())

    inp = np.empty(1).astype("float32")
    recursion_runs = np.random.randint(1, 10)
    inp.fill(recursion_runs)
    inp = tvm.runtime.tensor(inp)
    res = check_saved_func(vm, "recursion", inp)
    tvm.testing.assert_allclose(res.numpy(), np.power(2.0, recursion_runs), rtol=1e-7, atol=1e-7)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_vm_to_device(exec_mode):
    @tvm.script.ir_module
    class TestToVDevice:
        @R.function
        def foo1(
            x: R.Tensor((2, 3), "float32"),
        ) -> R.Tensor((2, 3), "float32"):
            copied = R.to_vdevice(x, tvm.ir.VDevice("cuda", 0, "global"))
            return copied

        @R.function
        def foo2(
            x: R.Tensor((2, 3), "float32"),
        ) -> R.Tensor((2, 3), "float32"):
            copied = R.to_vdevice(x, tvm.ir.VDevice("llvm", 0, "global"))
            return copied

    mod = TestToVDevice
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)

    def run_and_check():
        vm = relax.VirtualMachine(ex, tvm.cpu())
        x_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
        res_1 = check_saved_func(vm, "foo1", x_inp)
        res_2 = check_saved_func(vm, "foo2", x_inp)

        # check the copied tensor's device
        assert res_1.device == tvm.cuda(0)
        assert res_2.device == tvm.cpu(0)

        tvm.testing.assert_allclose(res_1.numpy(), x_inp.numpy())
        tvm.testing.assert_allclose(res_2.numpy(), x_inp.numpy())

    tvm.testing.run_with_gpu_lock(run_and_check)


def test_vm_closure(exec_mode):
    @tvm.script.ir_module
    class TestClosure:
        @R.function
        def lifted_func_1(x: R.Tensor((2, 3), "float32"), env: R.Tensor((2, 3), "float32")):
            return R.call_pure_packed("test.vm.add", x, env, ty_args=(R.Tensor()))

        @R.function
        def main(
            x: R.Tensor((2, 3), "float32"),
            y: R.Tensor((2, 3), "float32"),
        ):
            cls = TestClosure
            clo = R.make_closure(cls.lifted_func_1, (x,))
            res = R.invoke_pure_closure(clo, (y,), ty_args=(R.Tensor()))
            return res

    mod = TestClosure
    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(mod, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x_inp = tvm.runtime.tensor(np.random.rand(2, 3).astype("float32"))
    y_inp = tvm.runtime.tensor(np.array([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]], dtype="float32"))
    res = check_saved_func(vm, "main", x_inp, y_inp)
    tvm.testing.assert_allclose(res.numpy(), x_inp.numpy() + y_inp.numpy())


def test_time_evaluator(exec_mode):
    @tvm.script.ir_module
    class TestTimeEvaluator:
        @R.function
        def main(x: R.Tensor((1,), "float32"), y: R.Tensor((1,), "float32")):
            return R.call_pure_packed(
                "test.vm.add", x, y, ty_args=(R.Tensor(ndim=1, dtype="float32"))
            )

    target = tvm.target.Target("llvm", host="llvm")
    ex = relax.build(TestTimeEvaluator, target, exec_mode=exec_mode)
    vm = relax.VirtualMachine(ex, tvm.cpu())
    x = tvm.runtime.tensor(np.random.rand(1).astype("float32"))
    y = tvm.runtime.tensor(np.random.rand(1).astype("float32"))

    # ensure we can use time_evaluator with the stateful API
    vm.set_input("main", x, y)
    timing_res = vm.time_evaluator("invoke_stateful", tvm.cpu())("main")
    # just checking that it has some results at all
    assert timing_res.results

    # ensure we can use it with a closure
    vm.save_function("main", "saved_main", x, y)
    timing_res = vm.time_evaluator("saved_main", tvm.cpu())()
    assert timing_res.results


@tvm.script.ir_module
class TestVMSetInput:
    @T.prim_func(s_tir=True)
    def test_vm_mul(x: T.handle, y: T.handle, z: T.handle):
        T.func_attr({"global_symbol": "test_vm_mul"})
        m = T.int32()
        n = T.int32()
        A = T.match_buffer(x, (m, n))
        B = T.match_buffer(y, (m, n))
        C = T.match_buffer(z, (m, n))

        for i, j in T.grid(m, n):
            with T.sblock("mul"):
                vi = T.axis.spatial(m, i)
                vj = T.axis.spatial(n, j)
                with T.init():
                    C[vi, vj] = T.float32(0)
                C[vi, vj] = A[vi, vj] * B[vi, vj]

    # test returning a tuple
    @R.function
    def test_vm_tuple(
        x: R.Tensor((), "int32"),
    ) -> R.Tuple(R.Tensor((), "int32"), R.Tensor((), "int32")):
        return (x, x)

    # nested tuple too
    @R.function
    def test_vm_nested_tuple(x: R.Tensor((), "int32")) -> R.Tuple(
        R.Tuple(
            R.Tensor((), "int32"),
            R.Tuple(
                R.Tensor((), "int32"),
            ),
        ),
        R.Tensor((), "int32"),
    ):
        return ((x, (x,)), x)

    @R.function
    def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
        cls = TestVMSetInput
        gv0 = R.call_tir(cls.test_vm_mul, (x, w), R.Tensor((32, 32), dtype="float32"))
        return gv0


def test_multi_systemlib(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by popen_pool.PopenWorker

    @tvm.script.ir_module
    class ModA:
        I.module_attrs({"system_lib_prefix": "libA_"})

        @T.prim_func(s_tir=True)
        def tir_init(x_handle: T.handle):
            N = T.int64()
            x = T.match_buffer(x_handle, [N], "float32")
            for i in range(N):
                x[i] = T.float32(0)

        @R.function
        def main(s: R.Shape(["m"])) -> R.Tensor:
            m = T.int64()
            gv0 = R.call_tir(ModA.tir_init, (), R.Tensor((m + 1,), dtype="float32"))
            return gv0

    @tvm.script.ir_module
    class ModB:
        I.module_attrs({"system_lib_prefix": "libB_"})

        @T.prim_func(s_tir=True)
        def tir_init(x_handle: T.handle):
            N = T.int64()
            x = T.match_buffer(x_handle, [N], "float32")
            for i in range(N):
                x[i] = T.float32(1)

        @R.function
        def main(s: R.Shape(["m"])) -> R.Tensor:
            m = T.int64()
            gv0 = R.call_tir(ModB.tir_init, (), R.Tensor((m,), dtype="float32"))
            return gv0

    target = tvm.target.Target("llvm", host="llvm")
    libA = relax.build(ModA, target, exec_mode=exec_mode)
    libB = relax.build(ModB, target, exec_mode=exec_mode)

    temp = utils.tempdir()
    pathA = temp.relpath("libA.a")
    pathB = temp.relpath("libB.a")
    path_dso = temp.relpath("mylibAll.so")
    libA.export_library(pathA, fcompile=cc.create_staticlib)
    libB.export_library(pathB, fcompile=cc.create_staticlib)

    # package two static libs together
    # check that they do not interfere with each other
    # even though they have shared global var names
    # intentionally craft same gvar function with different behaviors
    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathA, pathB, "-Wl,--no-whole-archive"])

    def popen_check():
        # Load dll, will trigger system library registration
        ctypes.CDLL(path_dso)
        # Load the system wide library
        vmA = relax.VirtualMachine(tvm.runtime.system_lib("libA_"), tvm.cpu())
        vmB = relax.VirtualMachine(tvm.runtime.system_lib("libB_"), tvm.cpu())

        retA = vmA["main"](tvm_ffi.Shape([1]))
        retB = vmB["main"](tvm_ffi.Shape([2]))
        np.testing.assert_equal(retA.numpy(), np.array([0, 0]).astype("float32"))
        np.testing.assert_equal(retB.numpy(), np.array([1, 1]).astype("float32"))

    # system lib should be loaded in different process
    worker = popen_pool.PopenWorker()
    worker.send(popen_check)


def set_input_trial(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.set_input("main", a, b)
    vm.invoke_stateful("main")
    res0 = vm.get_outputs("main")

    data_dict = {"x": a, "w": b}
    vm.set_input("main", **data_dict)
    vm.invoke_stateful("main")
    res1 = vm.get_outputs("main")
    tvm.testing.assert_allclose(res0.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)
    tvm.testing.assert_allclose(res0.numpy(), res1.numpy(), rtol=1e-7, atol=1e-7)

    # bug! If you don't bind the Tensor to a var, the memory will get corrupted.
    # Possibly due to object lifecycles and other FFI issues
    a = tvm.runtime.tensor(np.array(2).astype("int32"), device)
    vm.set_input("test_vm_tuple", a)
    vm.invoke_stateful("test_vm_tuple")
    res2 = vm.get_outputs("test_vm_tuple")
    # the results are Tensors wrapped around scalars,
    # so we have to get the scalar out of the Tensor
    assert tuple(map(lambda a: int(a.numpy()), res2)) == (2, 2)

    b = tvm.runtime.tensor(np.array(1).astype("int32"), device)
    vm.set_input("test_vm_nested_tuple", b)
    vm.invoke_stateful("test_vm_nested_tuple")
    res3 = vm.get_outputs("test_vm_nested_tuple")
    assert len(res3) == 2 and len(res3[0]) == 2 and len(res3[0][1]) == 1
    result_cast = ((int(res3[0][0].numpy()), (int(res3[0][1][0].numpy()),)), int(res3[1].numpy()))
    assert result_cast == ((1, (1,)), 1)


def set_input_attempt_stateless(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    # this should fail: once you set inputs, you cannot run statelessly
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.set_input("main", a, b)
    # must use invoke stateful!
    vm["main"]()


def set_input_attempt_invoke(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    # this should fail: if the function needs inputs, you can't invoke directly
    vm.invoke_stateful("main")


def set_input_attempt_get(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    # this should fail: you can't get outputs without invoking the function first
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.set_input("main", a, b)
    _ = vm.get_outputs("main")


def make_vm(mod, exec_mode, temp) -> tuple[relax.VirtualMachine, tvm.runtime.Device]:
    """Returns a local VM for the given mod and the device"""
    target = tvm.target.Target("llvm", host="llvm")
    exec = relax.build(mod, target, exec_mode=exec_mode)
    libname = temp.relpath("exec.so")
    exec.export_library(libname)
    exec_loaded = tvm.runtime.load_module(libname)
    device = tvm.cpu()
    return relax.VirtualMachine(exec_loaded, device), device


def run_on_rpc(
    mod: tvm.IRModule,
    trial_func: Callable[[relax.VirtualMachine, tvm.runtime.Device], None],
    exec_mode: str,
):
    """
    Sets up a VM over localhost using the given mod and runs the given trial function.
    The trial function should take a VM and a device
    """
    target = tvm.target.Target("llvm", host="llvm")
    exec = relax.build(mod, target, exec_mode=exec_mode)
    temp = utils.tempdir()
    path = temp.relpath("vm_library.so")
    exec.export_library(path)

    # Use local rpc server for testing.
    # Server must use popen so it doesn't inherit the current process state. It
    # will crash otherwise.
    def check_remote(server):
        remote = rpc.connect(server.host, server.port, session_timeout=10)

        # Upload the serialized Executable.
        remote.upload(path)
        # Get a handle to remote Executable.
        rexec = remote.load_module("vm_library.so")

        device = remote.cpu()
        # Build a VM out of the executable and context.
        vm = relax.VirtualMachine(rexec, device=device)
        trial_func(vm, device)

    check_remote(rpc.Server("127.0.0.1"))


def test_set_input(exec_mode):
    temp = utils.tempdir()
    set_input_trial(*make_vm(TestVMSetInput, exec_mode, temp))


def test_set_input_tuple(exec_mode):
    @tvm.script.ir_module
    class MyMod:
        @R.function
        def main(x: R.Tuple([R.Tensor((32,), "float32"), R.Tensor((32,), "float32")])) -> R.Tensor:
            y = x[0]
            return y

    temp = utils.tempdir()
    vm, device = make_vm(MyMod, exec_mode, temp)
    device = tvm.cpu(0)
    a = tvm.runtime.empty((32,), "float32", device=device)
    b = tvm.runtime.empty((32,), "float32", device=device)
    vm.set_input("main", (a, b))
    vm.invoke_stateful("main")


def save_function_kwargs_trial(vm: relax.VirtualMachine, device: tvm.runtime.Device) -> None:
    # just checking that we can use kwargs for the args when saving a function
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.save_function("main", "saved_main", x=a, w=b)
    res0 = vm["saved_main"]()
    tvm.testing.assert_allclose(res0.numpy(), a.numpy() * b.numpy(), rtol=1e-7, atol=1e-7)


def test_save_function_kwargs(exec_mode):
    temp = utils.tempdir()
    save_function_kwargs_trial(*make_vm(TestVMSetInput, exec_mode, temp))


def test_save_function_kwargs_rpc(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by the popen RPC server
    run_on_rpc(TestVMSetInput, save_function_kwargs_trial, exec_mode)


def save_function_time_evaluator_trial(
    vm: relax.VirtualMachine, device: tvm.runtime.Device
) -> None:
    # just checking that the saved function can be called in the time evaluator
    a = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    b = tvm.runtime.tensor(np.random.rand(32, 32).astype("float32"), device)
    vm.save_function("main", "saved_main", a, b)
    vm.time_evaluator("saved_main", device)()


def test_save_function_time_evaluator(exec_mode):
    temp = utils.tempdir()
    save_function_time_evaluator_trial(*make_vm(TestVMSetInput, exec_mode, temp))


def test_save_function_time_evaluator_rpc(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by the popen RPC server
    run_on_rpc(TestVMSetInput, save_function_time_evaluator_trial, exec_mode)


# if you set an input, you should not be able to call statelessly


def test_set_input_stateless_failure(exec_mode):
    temp = utils.tempdir()
    args = make_vm(TestVMSetInput, exec_mode, temp)
    with pytest.raises(RuntimeError):
        set_input_attempt_stateless(*args)


def test_set_input_stateless_failure_rpc(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by the popen RPC server
    with pytest.raises(RuntimeError):
        run_on_rpc(TestVMSetInput, set_input_attempt_stateless, exec_mode)


def test_set_input_invoke_failure(exec_mode):
    temp = utils.tempdir()
    args = make_vm(TestVMSetInput, exec_mode, temp)
    with pytest.raises(ValueError):
        set_input_attempt_invoke(*args)


def test_set_input_invoke_failure_rpc(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by the popen RPC server
    with pytest.raises(RuntimeError):
        run_on_rpc(TestVMSetInput, set_input_attempt_invoke, exec_mode)


def test_set_input_get_failure(exec_mode):
    temp = utils.tempdir()
    args = make_vm(TestVMSetInput, exec_mode, temp)
    with pytest.raises(ValueError):
        set_input_attempt_get(*args)


def test_set_input_get_failure_rpc(exec_mode):
    pytest.importorskip("cloudpickle")  # needed by the popen RPC server
    with pytest.raises(RuntimeError):
        run_on_rpc(TestVMSetInput, set_input_attempt_get, exec_mode)


@pytest.mark.gpu
@pytest.mark.skipif(not env.has_gpu(), reason="need gpu")
def test_relax_module_with_multiple_targets(exec_mode):
    """Relax functions may contain kernels for multiple targets

    In this example, the module contains one function to execute on
    LLVM, and one function to execute on CUDA.

    """

    @I.ir_module(s_tir=True)
    class Module:
        I.module_global_infos({"vdevice": [I.vdevice("llvm")]})

        @R.function
        def func_cuda(A: R.Tensor([32, 32], "float32"), B: R.Tensor([32, 32], "float32")):
            C = R.add(A, B)
            return C

        @R.function
        def func_llvm(
            A: R.Tensor([32, 32], "float32", "llvm"), B: R.Tensor([32, 32], "float32", "llvm")
        ):
            C = R.add(A, B)
            return C

    seq = tvm.ir.transform.Sequential(
        [
            tvm.relax.transform.LegalizeOps(),
            tvm.s_tir.dlight.ApplyDefaultSchedule(tvm.s_tir.dlight.gpu.Fallback()),
        ],
        name="LegalizeAndSchedule",
    )
    with tvm.target.Target("cuda"):
        built = tvm.relax.build(seq(Module))

    np_A = np.random.random([32, 32]).astype("float32")
    np_B = np.random.random([32, 32]).astype("float32")

    dev_llvm = tvm.cpu()
    vm_llvm = tvm.relax.VirtualMachine(built, device=dev_llvm)
    llvm_output = vm_llvm["func_llvm"](
        tvm.runtime.tensor(np_A, dev_llvm),
        tvm.runtime.tensor(np_B, dev_llvm),
    )

    np_C = np_A + np_B

    tvm.testing.assert_allclose(llvm_output.numpy(), np_C)

    def run_and_check():
        dev_cuda = tvm.cuda()
        vm_cuda = tvm.relax.VirtualMachine(built, device=dev_cuda)
        cuda_output = vm_cuda["func_cuda"](
            tvm.runtime.tensor(np_A, dev_cuda),
            tvm.runtime.tensor(np_B, dev_cuda),
        )
        tvm.testing.assert_allclose(cuda_output.numpy(), np_C)

    tvm.testing.run_with_gpu_lock(run_and_check)


if __name__ == "__main__":
    tvm.testing.main()
