wehub-resource-sync
92 lines
2.7 KiB
Python
92 lines
2.7 KiB
Python
# 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()
|