| # 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. |
| |
| import ctypes |
| from typing import Tuple, Callable |
| |
| import numpy as np |
| import pytest |
| |
| import tvm |
| import tvm.script |
| import tvm.testing |
| from tvm import relax, rpc, te, tir, topi |
| from tvm.contrib import utils, cc, popen_pool |
| from tvm.relax.testing import nn |
| from tvm.script import relax as R, tir as T, ir as I |
| from tvm.relax.testing.vm import check_saved_func |
| from tvm.runtime import ShapeTuple |
| |
| 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, sinfo_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, sinfo_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.Object) -> 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 |
| 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.block("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 |
| 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({"tir.noalias": True}) |
| for i0, i1 in T.grid(T.int64(2), T.int64(3)): |
| with T.block("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 |
| 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({"tir.noalias": True}) |
| for i0, i1 in T.grid(T.int64(2), T.int64(3)): |
| with T.block("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 = tir.Var("n", "int64"), tir.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 = tir.Var("n", "int64"), tir.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.tir.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 = tir.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) |
| np.testing.assert_allclose(res.numpy(), inp.numpy().astype("int16")) |
| |
| |
| def test_vm_emit_te_floor_symbolic_shape(exec_mode): |
| bb = relax.BlockBuilder() |
| n = tir.Var("n", "int64") |
| x = relax.Var("x", R.Tensor([n], "float32")) |
| |
| def te_func(A): |
| C = te.compute((tir.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) |
| |
| |
| @tvm.testing.requires_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.tir.Schedule(mod, debug_mask="all") |
| loops = sch.get_loops(sch.get_block(name="T_add", func_name="add")) |
| sch.bind(loops[0], "threadIdx.x") |
| |
| exec = relax.build(sch.mod, "cuda", exec_mode=exec_mode) |
| 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) |
| |
| |
| def test_vm_relax_symbolic_shape(exec_mode): |
| bb = relax.BlockBuilder() |
| n = tir.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 |
| 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(ShapeTuple([2, 3])) == (4, 9) |
| |
| with pytest.raises(ValueError): |
| func(ShapeTuple([2, 3, 4])) |
| |
| with pytest.raises(TypeError): |
| func(R.prim_value(2)) |
| |
| |
| def test_vm_relax_symbolic_prim_value(exec_mode): |
| @I.ir_module |
| class mod: |
| @R.function |
| def main(shape: R.Prim(value="n")): |
| n = T.int64() |
| return R.prim_value(n * 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(2) == 4 |
| |
| with pytest.raises(TypeError): |
| func(ShapeTuple([2])) |
| |
| |
| def test_vm_relax_multiple_symbolic_prim_value(exec_mode): |
| """Like test_vm_relax_symbolic_prim_value, but with multiple variables""" |
| |
| @I.ir_module |
| class mod: |
| @R.function |
| def main( |
| # Provides definition of "n" |
| _n: R.Prim(value="n"), |
| # Requires definitions of both "n" and "m", but cannot |
| # provide either. |
| _shape: R.Shape(["n*2", "m*2"]), |
| # Provides definition of "m" |
| _m: R.Prim(value="m"), |
| ): |
| n = T.int64() |
| m = T.int64() |
| return R.shape([n * n, m + 1]) |
| |
| 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(2, ShapeTuple([4, 12]), 6) == (4, 7) |
| |
| with pytest.raises(RuntimeError): |
| func(2, ShapeTuple([4, 12]), 1) |
| |
| with pytest.raises(tvm.TVMError): |
| func(ShapeTuple([2])) |
| |
| |
| @pytest.mark.xfail(reason="Current support for R.Prim with known value is primarily for int64") |
| @pytest.mark.parametrize("exec_mode", EXEC_MODE) |
| def test_vm_relax_prim_value_fp32(exec_mode): |
| """A PrimValue may be R.prim('float32') |
| |
| Unlike shape tuples, which must contain int64, a PrimValue may be |
| any type that can be represented as a single primitive value. |
| """ |
| |
| @I.ir_module |
| class mod: |
| @R.function |
| def main( |
| # First failure occurs during parsing. The syntactic |
| # sugar for symbolic variables assumes that all symbolic |
| # variables are int64, rather than using the type that is |
| # later declared. |
| _x: R.Prim(value="half_fill_value"), |
| ): |
| half_fill_value = T.float32() |
| # Second failure occurs when calling `relax.op.full`. The |
| # `fill_value` is expected to be a scalar constant |
| # (R.Tensor with 0-dim shape), not a primitive value, even |
| # though these are semantically the same. |
| return R.full(shape=[16, 16], fill_value=R.prim_value(2 * half_fill_value)) |
| |
| target = tvm.target.Target("llvm", host="llvm") |
| # Third failure occurs here. The current codegen assumes that all |
| # symbolic variables are int64. |
| ex = relax.build(mod, target, exec_mode=exec_mode) |
| vm = relax.VirtualMachine(ex, tvm.cpu()) |
| |
| func = vm["main"] |
| |
| res = func(16.0).numpy() |
| assert np.all(res == 32.0) |
| |
| |
| def test_vm_relax_dyn_tir_shape(exec_mode): |
| # case where TIR variables are unbound in generated PrimFunc |
| bb = relax.BlockBuilder() |
| n = tir.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 = tir.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, sinfo_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 |
| def copy(A: T.Buffer((2, 3), "float32"), B: T.Buffer((2, 3), "float32")): |
| for i0, i1 in T.grid(2, 3): |
| with T.block("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 |
| 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.block("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.Object: |
| gv0 = R.call_pure_packed( |
| "test.vm.mul", x, w, sinfo_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.Object: |
| 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, sinfo_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, sinfo_args=(R.Tensor(ndim=1, dtype="float32")) |
| ) |
| tmp = TestVMRecursion.recursion(gv0) |
| res = R.call_pure_packed( |
| "test.vm.add", tmp, tmp, sinfo_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) |
| |
| |
| @tvm.testing.requires_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) |
| 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()) |
| |
| |
| 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, sinfo_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,), sinfo_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, sinfo_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 |
| 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.block("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): |
| @tvm.script.ir_module |
| class ModA: |
| I.module_attrs({"system_lib_prefix": "libA_"}) |
| |
| @T.prim_func |
| 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 |
| 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.runtime.ShapeTuple([1])) |
| retB = vmB["main"](tvm.runtime.ShapeTuple([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): |
| 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): |
| 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): |
| 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): |
| 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): |
| with pytest.raises(RuntimeError): |
| run_on_rpc(TestVMSetInput, set_input_attempt_get, exec_mode) |
| |
| |
| @tvm.testing.requires_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 |
| 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.dlight.ApplyDefaultSchedule(tvm.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), |
| ) |
| |
| 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), |
| ) |
| |
| np_C = np_A + np_B |
| |
| tvm.testing.assert_allclose(llvm_output.numpy(), np_C) |
| tvm.testing.assert_allclose(cuda_output.numpy(), np_C) |
| |
| |
| if __name__ == "__main__": |
| tvm.testing.main() |
