chore: import upstream snapshot with attribution

web/tests/python/prepare_test_libs.py

@@ -0,0 +1,65 @@
+# 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.
+# Prepare test library for standalone wasm runtime test.
+
+import os
+
+import tvm
+from tvm import relax, te
+from tvm.contrib import tvmjs
+from tvm.script import relax as R
+
+
+def prepare_relax_lib(base_path):
+    pipeline = relax.get_pipeline()
+
+    @tvm.script.ir_module
+    class Mod:
+        @R.function
+        def main(x: R.Tensor(["n"], "float32"), y: R.Tensor(["n"], "float32")):
+            lv0 = R.add(x, y)
+            return lv0
+
+    target = tvm.target.Target({"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"})
+
+    mod = pipeline(Mod)
+    ex = relax.build(mod, target)
+    wasm_path = os.path.join(base_path, "test_relax.wasm")
+    ex.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
+
+
+def prepare_tir_lib(base_path):
+    target = {"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
+    if not tvm.runtime.enabled("llvm"):
+        raise RuntimeError(f"Target {target} is not enbaled")
+    n = te.var("n")
+    A = te.placeholder((n,), name="A")
+    B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
+    mod = tvm.IRModule.from_expr(
+        te.create_prim_func([A, B]).with_attr("global_symbol", "add_one")
+    ).with_attr("system_lib_prefix", "")
+
+    fadd = tvm.build(mod, target)
+    wasm_path = os.path.join(base_path, "test_addone.wasm")
+    fadd.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
+
+
+if __name__ == "__main__":
+    curr_path = os.path.dirname(os.path.abspath(os.path.expanduser(__file__)))
+    base_path = os.path.join(curr_path, "../../dist/wasm")
+    prepare_tir_lib(base_path)
+    prepare_relax_lib(base_path)

web/tests/python/relax_rpc_test.py

@@ -0,0 +1,91 @@
+# 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.
+"""Test relax vm through rpc."""
+
+import numpy as np
+
+import tvm
+from tvm import relax, rpc
+from tvm.contrib import tvmjs
+from tvm.script import relax as R
+from tvm.support import utils
+
+proxy_host = "127.0.0.1"
+proxy_port = 9090
+
+
+def get_model():
+    pipeline = relax.get_pipeline()
+
+    @tvm.script.ir_module
+    class Mod:
+        @R.function
+        def main(x: R.Tensor([1024], "float32"), y: R.Tensor([1024], "float32")):
+            lv0 = R.add(x, y)
+            return lv0
+
+    mod = pipeline(Mod)
+    sch = tvm.s_tir.Schedule(mod)
+    # manually transform loop
+    sch.work_on("add")
+    (i,) = sch.get_loops(block=sch.get_sblock("T_add"))
+    i0, i1 = sch.split(i, [None, 128])
+    sch.bind(i0, "blockIdx.x")
+    sch.bind(i1, "threadIdx.x")
+    return sch.mod
+
+
+def test_rpc():
+    if not tvm.runtime.enabled("rpc"):
+        return
+    n = 1024
+    dtype = "float32"
+    temp = utils.tempdir()
+    wasm_path = temp.relpath("relax.wasm")
+    target = tvm.target.Target(
+        "webgpu", host={"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
+    )
+
+    mod = get_model()
+    ex = relax.build(mod, target)
+    ex.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
+    wasm_binary = open(wasm_path, "rb").read()
+
+    remote = rpc.connect(
+        proxy_host,
+        proxy_port,
+        key="wasm",
+        session_constructor_args=["rpc.WasmSession", wasm_binary],
+    )
+
+    def check(remote):
+        dev = remote.webgpu(0)
+        # invoke the function
+        vm = relax.VirtualMachine(remote.system_lib(), device=dev)
+        adata = np.random.uniform(size=n).astype(dtype)
+        bdata = np.random.uniform(size=n).astype(dtype)
+        a = tvm.runtime.tensor(adata, dev)
+        b = tvm.runtime.tensor(bdata, dev)
+        vm.set_input("main", a, b)
+        vm.invoke_stateful("main")
+        c = vm.get_outputs("main")
+        np.testing.assert_equal(c.numpy(), a.numpy() + b.numpy())
+
+    check(remote)
+
+
+test_rpc()

web/tests/python/webgpu_rpc_test.py

@@ -0,0 +1,83 @@
+# 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.
+"""Simple testcode to test Javascript RPC
+
+To use it, start a rpc proxy with "python -m tvm.exec.rpc_proxy".
+Connect javascript end to the websocket port and connect to the RPC.
+"""
+
+import numpy as np
+
+import tvm
+from tvm import rpc, te
+from tvm.contrib import tvmjs
+from tvm.support import utils
+
+proxy_host = "127.0.0.1"
+proxy_port = 9090
+
+
+def test_rpc():
+    if not tvm.runtime.enabled("rpc"):
+        return
+    # generate the wasm library
+    target = tvm.target.Target(
+        "webgpu", host={"kind": "llvm", "mtriple": "wasm32-unknown-unknown-wasm"}
+    )
+
+    n = te.var("n")
+    A = te.placeholder((n,), name="A")
+    B = te.compute(A.shape, lambda *i: te.log(te.abs(A(*i) + 1)), name="B")
+    mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]))
+    sch = tvm.s_tir.Schedule(mod)
+    (i,) = sch.get_loops(block=sch.get_sblock("B"))
+    i0, i1 = sch.split(i, [None, 32])
+    sch.bind(i0, "blockIdx.x")
+    sch.bind(i1, "threadIdx.x")
+
+    fadd = tvm.build(sch.mod.with_attr("system_lib_prefix", ""), target=target)
+    temp = utils.tempdir()
+
+    wasm_path = temp.relpath("addone_gpu.wasm")
+    fadd.export_library(wasm_path, fcompile=tvmjs.create_tvmjs_wasm)
+
+    wasm_binary = open(wasm_path, "rb").read()
+    remote = rpc.connect(
+        proxy_host,
+        proxy_port,
+        key="wasm",
+        session_constructor_args=["rpc.WasmSession", wasm_binary],
+    )
+
+    def check(remote, size):
+        # basic function checks.
+        dev = remote.webgpu(0)
+        adata = np.random.uniform(size=size).astype(A.dtype)
+        a = tvm.runtime.tensor(adata, dev)
+        b = tvm.runtime.tensor(np.zeros(size, dtype=A.dtype), dev)
+
+        np.testing.assert_equal(a.numpy(), adata)
+        f1 = remote.system_lib()
+        addone = f1.get_function("main")
+        addone(a, b)
+        tvm.testing.assert_allclose(b.numpy(), np.log(np.abs(a.numpy()) + 1), atol=1e-5, rtol=1e-5)
+        print("Test pass..")
+
+    check(remote, 71821 * 32)
+
+
+test_rpc()