()
| 243 | |
| 244 | |
| 245 | def test_vm_invoke_closure(): |
| 246 | ib = relax.ExecBuilder() |
| 247 | with ib.function("lifted_func_1", num_inputs=4): |
| 248 | ib.emit_call("test.vm.add", args=[ib.r(0), ib.r(1)], dst=ib.r(4)) |
| 249 | ib.emit_call("test.vm.add", args=[ib.r(2), ib.r(4)], dst=ib.r(5)) |
| 250 | ib.emit_call("test.vm.add", args=[ib.r(3), ib.r(5)], dst=ib.r(6)) |
| 251 | ib.emit_ret(ib.r(6)) |
| 252 | with ib.function("main", num_inputs=2): |
| 253 | ib.emit_call( |
| 254 | "vm.builtin.make_closure", args=[ib.f("lifted_func_1"), ib.r(0), ib.r(1)], dst=ib.r(2) |
| 255 | ) |
| 256 | ib.emit_ret(ib.r(2)) |
| 257 | |
| 258 | ex = ib.get() |
| 259 | vm = relax.VirtualMachine(ex, tvm.cpu()) |
| 260 | w_inp = tvm.runtime.tensor(np.random.rand(2, 3)) |
| 261 | x_inp = tvm.runtime.tensor(np.random.rand(2, 3)) |
| 262 | y_inp = tvm.runtime.tensor([[3.1, 4.0, 5.0], [6.0, 7.1, 9.0]]) |
| 263 | z_inp = tvm.runtime.tensor(np.random.rand(2, 3)) |
| 264 | clo = vm["main"](w_inp, x_inp) |
| 265 | res = vm.invoke_closure(clo, y_inp, z_inp) |
| 266 | tvm.testing.assert_allclose( |
| 267 | res.numpy(), w_inp.numpy() + x_inp.numpy() + y_inp.numpy() + z_inp.numpy() |
| 268 | ) |
| 269 | |
| 270 | |
| 271 | def test_vm_stack_restore_after_failure(): |
nothing calls this directly
no test coverage detected
searching dependent graphs…