# 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.device("llvm") 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.device("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()