wehub-resource-sync
98 lines
3.6 KiB
Python
98 lines
3.6 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.
|
|
"""Testcode for iOS RPC.
|
|
|
|
To use it, start a rpc proxy with "python -m tvm.exec.rpc_proxy".
|
|
And configure the proxy host field as commented.
|
|
"""
|
|
|
|
import argparse
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
from tvm import rpc, te
|
|
from tvm.support import utils, xcode
|
|
|
|
# Change target configuration, this is setting for iphone6s
|
|
arch = "arm64"
|
|
sdk = "iphoneos"
|
|
target = {"kind": "llvm", "mtriple": f"{arch}-apple-darwin"}
|
|
|
|
MODES = {"proxy": rpc.connect, "tracker": rpc.connect_tracker, "standalone": rpc.connect}
|
|
|
|
|
|
# override metal compiler to compile to iphone
|
|
@tvm.register_global_func("tvm_callback_metal_compile")
|
|
def compile_metal(src, target):
|
|
return xcode.compile_metal(src, sdk=sdk)
|
|
|
|
|
|
def test_rpc_module(host, port, key, mode):
|
|
# graph
|
|
n = tvm.runtime.convert(1024)
|
|
A = te.placeholder((n,), name="A")
|
|
B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
|
|
temp = utils.tempdir()
|
|
mod = tvm.IRModule.from_expr(te.create_prim_func([A, B]).with_attr("global_symbol", "myadd"))
|
|
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")
|
|
|
|
# Build the dynamic lib.
|
|
# If we don't want to do metal and only use cpu, just set target to be target
|
|
f = tvm.compile(sch.mod, target=tvm.target.Target("metal", host=target))
|
|
path_dso1 = temp.relpath("dev_lib.dylib")
|
|
f.export_library(path_dso1, fcompile=xcode.create_dylib, arch=arch, sdk=sdk)
|
|
|
|
# connect to the proxy
|
|
if mode == "tracker":
|
|
remote = MODES[mode](host, port).request(key)
|
|
else:
|
|
remote = MODES[mode](host, port, key=key)
|
|
remote.upload(path_dso1)
|
|
dev = remote.metal(0)
|
|
f1 = remote.load_module("dev_lib.dylib")
|
|
a_np = np.random.uniform(size=1024).astype(A.dtype)
|
|
a = tvm.runtime.tensor(a_np, dev)
|
|
b = tvm.runtime.tensor(np.zeros(1024, dtype=A.dtype), dev)
|
|
time_f = f1.time_evaluator(f1.entry_name, dev, number=10)
|
|
cost = time_f(a, b).mean
|
|
print(f"Metal: {cost:g} secs/op")
|
|
np.testing.assert_equal(b.numpy(), a.numpy() + 1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Demo app demonstrates how ios_rpc works.")
|
|
parser.add_argument("--host", required=True, type=str, help="Address of rpc server")
|
|
parser.add_argument("--port", type=int, default=9090, help="rpc port (default: 9090)")
|
|
parser.add_argument("--key", type=str, default="iphone", help="device key (default: iphone)")
|
|
parser.add_argument(
|
|
"--mode",
|
|
type=str,
|
|
default="tracker",
|
|
help="type of RPC connection (default: tracker), possible values: {}".format(
|
|
", ".join(MODES.keys())
|
|
),
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
assert args.mode in MODES.keys()
|
|
test_rpc_module(args.host, args.port, args.key, args.mode)
|