()
| 32 | |
| 33 | |
| 34 | def test_rpc(): |
| 35 | if not tvm.runtime.enabled("rpc"): |
| 36 | return |
| 37 | # generate the wasm library |
| 38 | target = tvm.target.Target( |
| 39 | "webgpu", host={"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"} |
| 40 | ) |
| 41 | |
| 42 | n = te.var("n") |
| 43 | A = te.placeholder((n,), name="A") |
| 44 | B = te.compute(A.shape, lambda *i: te.log(te.abs(A(*i) + 1)), name="B") |
| 45 | mod = tvm.IRModule.from_expr(te.create_prim_func([A, B])) |
| 46 | sch = tvm.s_tir.Schedule(mod) |
| 47 | (i,) = sch.get_loops(block=sch.get_sblock("B")) |
| 48 | i0, i1 = sch.split(i, [None, 32]) |
| 49 | sch.bind(i0, "blockIdx.x") |
| 50 | sch.bind(i1, "threadIdx.x") |
| 51 | |
| 52 | fadd = tvm.build(sch.mod.with_attr("system_lib_prefix", ""), target=target) |
| 53 | temp = utils.tempdir() |
| 54 | |
| 55 | wasm_path = temp.relpath("addone_gpu.wasm") |
| 56 | fadd.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm) |
| 57 | |
| 58 | wasm_binary = open(wasm_path, "rb").read() |
| 59 | remote = rpc.connect( |
| 60 | proxy_host, |
| 61 | proxy_port, |
| 62 | key="wasm", |
| 63 | session_constructor_args=["rpc.WasmSession", wasm_binary], |
| 64 | ) |
| 65 | |
| 66 | def check(remote, size): |
| 67 | # basic function checks. |
| 68 | dev = remote.webgpu(0) |
| 69 | adata = np.random.uniform(size=size).astype(A.dtype) |
| 70 | a = tvm.runtime.tensor(adata, dev) |
| 71 | b = tvm.runtime.tensor(np.zeros(size, dtype=A.dtype), dev) |
| 72 | |
| 73 | np.testing.assert_equal(a.numpy(), adata) |
| 74 | f1 = remote.system_lib() |
| 75 | addone = f1.get_function("main") |
| 76 | addone(a, b) |
| 77 | tvm.testing.assert_allclose(b.numpy(), np.log(np.abs(a.numpy()) + 1), atol=1e-5, rtol=1e-5) |
| 78 | print("Test pass..") |
| 79 | |
| 80 | check(remote, 71821 * 32) |
| 81 | |
| 82 | |
| 83 | test_rpc() |
no test coverage detected
searching dependent graphs…