(remote, tracker, ex, inputs)
| 89 | |
| 90 | |
| 91 | def _run(remote, tracker, ex, inputs): |
| 92 | tmp = utils.tempdir() |
| 93 | so_name = "test_mod.so" |
| 94 | so_path = tmp / so_name |
| 95 | ex.export_library(str(so_path), fcompile=ndk.create_shared, options=["-shared", "-fPIC", "-lm"]) |
| 96 | |
| 97 | remote.upload(so_path) |
| 98 | dev = remote.cpu(0) |
| 99 | |
| 100 | try: |
| 101 | # Execute the model on the remote. |
| 102 | remote_ex = remote.load_module(so_name) |
| 103 | vm = tvm.relax.VirtualMachine(remote_ex, device=dev) |
| 104 | |
| 105 | inputs = [x.copyto(dev) for x in inputs] |
| 106 | |
| 107 | vm.set_input("main", *inputs) |
| 108 | vm.invoke_stateful("main") |
| 109 | output = vm.get_outputs("main") |
| 110 | output = output.numpy() |
| 111 | except Exception as e: |
| 112 | # Re-raise all exceptions |
| 113 | raise e |
| 114 | finally: |
| 115 | # Manually close the connection. |
| 116 | # See https://discuss.tvm.apache.org/t/trouble-with-rpc-session/14008/. |
| 117 | # |
| 118 | # TODO: Remove if it does not happen on Python 3.11. |
| 119 | remote._sess.get_function("CloseRPCConnection")() |
| 120 | tracker.close() |
| 121 | pass |
| 122 | |
| 123 | return output |
| 124 | |
| 125 | |
| 126 | def build_and_run( |
no test coverage detected
searching dependent graphs…