test_runtime_rpc.py 11.6 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# 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.
17
import tvm
18
import os
19 20
import logging
import time
21 22 23
import multiprocessing

import numpy as np
24 25
from tvm import rpc
from tvm.contrib import util
26
from tvm.rpc.tracker import Tracker
27

28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49

def test_bigendian_rpc():
    """Test big endian rpc when there is a PowerPC RPC server available"""
    host = os.environ.get("TVM_POWERPC_TEST_HOST", None)
    port = os.environ.get("TVM_POWERPC_TEST_PORT", 9090)
    if host is None:
        return
    def verify_rpc(remote, target, shape, dtype):
        A = tvm.placeholder(shape, dtype=dtype)
        B = tvm.compute(A.shape, lambda i: A[i]+tvm.const(1, A.dtype))
        s = tvm.create_schedule(B.op)
        f = tvm.build(s, [A, B], target, name="myadd")

        ctx = remote.cpu(0)
        a = tvm.nd.array(np.random.randint(0, 256, size=shape).astype(A.dtype), ctx=ctx)
        b = tvm.nd.array(np.zeros(shape).astype(A.dtype), ctx=ctx)
        temp = util.tempdir()
        path_dso = temp.relpath("dev_lib.o")
        f.save(path_dso)
        remote.upload(path_dso)
        f = remote.load_module("dev_lib.o")
        f(a, b)
50
        tvm.testing.assert_allclose(a.asnumpy() + 1, b.asnumpy())
51 52 53 54 55 56 57 58

    print("Test RPC connection to PowerPC...")
    remote = rpc.connect(host, port)
    target = "llvm -mtriple=powerpc-linux-gnu"
    for dtype in ["float32", "float64", "int32", "int8"]:
        verify_rpc(remote, target, (10,), dtype)


59 60 61 62 63 64 65
def test_rpc_simple():
    if not tvm.module.enabled("rpc"):
        return
    @tvm.register_func("rpc.test.addone")
    def addone(x):
        return x + 1
    @tvm.register_func("rpc.test.strcat")
66
    def strcat(name, x):
67 68
        return "%s:%d" % (name, x)

69 70 71 72
    @tvm.register_func("rpc.test.except")
    def remotethrow(name):
        raise ValueError("%s" % name)

73
    server = rpc.Server("localhost", key="x1")
74
    client = rpc.connect(server.host, server.port, key="x1")
75 76
    f1 = client.get_function("rpc.test.addone")
    assert f1(10) == 11
77 78 79 80 81 82 83
    f3 = client.get_function("rpc.test.except")
    try:
        f3("abc")
        assert False
    except tvm.TVMError as e:
        assert "abc" in str(e)

84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
    f2 = client.get_function("rpc.test.strcat")
    assert f2("abc", 11) == "abc:11"

def test_rpc_array():
    if not tvm.module.enabled("rpc"):
        return
    x = np.random.randint(0, 10, size=(3, 4))
    @tvm.register_func("rpc.test.remote_array_func")
    def remote_array_func(y):
        np.testing.assert_equal(y.asnumpy(), x)
    server = rpc.Server("localhost")
    remote = rpc.connect(server.host, server.port)
    r_cpu = tvm.nd.array(x, remote.cpu(0))
    assert str(r_cpu.context).startswith("remote")
    np.testing.assert_equal(r_cpu.asnumpy(), x)
    fremote = remote.get_function("rpc.test.remote_array_func")
    fremote(r_cpu)

def test_rpc_file_exchange():
    if not tvm.module.enabled("rpc"):
        return
    server = rpc.Server("localhost")
    remote = rpc.connect(server.host, server.port)
107
    blob = bytearray(np.random.randint(0, 10, size=(10)))
108 109
    remote.upload(blob, "dat.bin")
    rev = remote.download("dat.bin")
110
    assert(rev == blob)
111 112 113 114 115

def test_rpc_remote_module():
    if not tvm.module.enabled("rpc"):
        return
    server = rpc.Server("localhost")
116
    client = rpc.connect(server.host, server.port)
117 118 119 120 121 122
    # graph
    n = tvm.convert(1024)
    A = tvm.placeholder((n,), name='A')
    B = tvm.compute(A.shape, lambda *i: A(*i) + 1.0, name='B')
    s = tvm.create_schedule(B.op)

123
    def check_remote(remote):
124 125 126 127 128 129 130 131 132 133 134 135
        if not tvm.module.enabled("llvm"):
            print("Skip because llvm is not enabled")
            return
        temp = util.tempdir()
        ctx = remote.cpu(0)
        f = tvm.build(s, [A, B], "llvm", name="myadd")
        path_dso = temp.relpath("dev_lib.so")
        f.export_library(path_dso)
        remote.upload(path_dso)
        f1 = remote.load_module("dev_lib.so")
        a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
        b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
136
        time_f = f1.time_evaluator(f1.entry_name, remote.cpu(0), number=10)
137
        cost = time_f(a, b).mean
138
        print('%g secs/op' % cost)
139
        np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)
140

141
    def check_remote_link_cl(remote):
142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188
        """Test function to run remote code such as cl

        This is not enabled because there is forking issue
        of TVM runtime when server launches after OpenCL
        runtime initializes. We leave it as an example
        on how to do rpc when we want to do linking on remote.
        """
        if not tvm.module.enabled("llvm"):
            print("Skip because llvm is not enabled")
            return
        if not tvm.module.enabled("opencl"):
            print("Skip because opencl is not enabled")
            return
        temp = util.tempdir()
        ctx = remote.cl(0)
        s = tvm.create_schedule(B.op)
        xo, xi = s[B].split(B.op.axis[0], factor=32)
        s[B].bind(xo, tvm.thread_axis("blockIdx.x"))
        s[B].bind(xi, tvm.thread_axis("threadIdx.x"))
        f = tvm.build(s, [A, B], "opencl", target_host="llvm", name="myadd")
        # Option 1: save modules separately and rely on remote compiler
        path_o = temp.relpath("myadd.o")
        path_cl = temp.relpath("myadd.cl")
        path_json = temp.relpath("myadd.tvm_meta.json")
        f.save(path_o)
        f.imported_modules[0].save(path_cl)
        remote.upload(path_o)
        remote.upload(path_cl)
        # upload meta data
        remote.upload(path_json)
        fhost = remote.load_module("myadd.o")
        fdev = remote.load_module("myadd.cl")
        fhost.import_module(fdev)
        a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
        b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
        fhost(a, b)
        np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)
        # Option 2: export library as a tar ball then handled by remote compiler
        path_tar = temp.relpath("myadd.tar")
        f.export_library(path_tar)
        remote.upload(path_tar)
        fhost = remote.load_module("myadd.tar")
        a = tvm.nd.array(np.random.uniform(size=1024).astype(A.dtype), ctx)
        b = tvm.nd.array(np.zeros(1024, dtype=A.dtype), ctx)
        fhost(a, b)
        np.testing.assert_equal(b.asnumpy(), a.asnumpy() + 1)

189 190 191
    check_remote(client)
    check_remote(rpc.LocalSession())

192

193 194 195 196
def test_rpc_return_func():
    @tvm.register_func("rpc.test.remote_func")
    def addone(x):
        return lambda y: x+y
197

198
    server = rpc.Server("localhost", key="x1")
199 200 201 202 203
    client = rpc.connect(server.host, server.port, key="x1")
    f1 = client.get_function("rpc.test.remote_func")
    fadd = f1(10)
    assert fadd(12) == 22

204

205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244
def test_rpc_return_ndarray():
    # Use closure to check the ref counter correctness
    nd = tvm.nd.array(np.zeros(10).astype("float32"))
    @tvm.register_func("rpc.test.remote_return_nd")
    def my_module(name):
        if name == "get_arr":
            return lambda : nd
        elif name == "ref_count":
            return lambda : tvm._api_internal._ndarray_use_count(nd)
        elif name == "get_elem":
            return lambda idx: nd.asnumpy()[idx]
        elif name == "get_arr_elem":
            return lambda arr, idx: arr.asnumpy()[idx]

    # start server
    server = rpc.Server("localhost", key="x1")
    client = rpc.connect(server.host, server.port, key="x1")
    m = client.get_function("rpc.test.remote_return_nd")
    get_arr = m("get_arr")
    ref_count = m("ref_count")
    get_elem = m("get_elem")
    get_arr_elem = m("get_arr_elem")
    # array test
    def run_arr_test():
        arr = get_arr()
        assert ref_count() == 2
        arr2 = get_arr()
        assert ref_count() == 3
        assert arr.context == client.cpu(0)
        arr.copyfrom(np.ones(10).astype(arr.dtype))
        assert arr2.asnumpy()[0] == 1.0
        assert get_elem(0) == 1.0
        assert get_arr_elem(arr2, 0) == 1.0

    assert ref_count() == 1
    run_arr_test()
    # check recycle correctness
    assert ref_count() == 1


245 246 247 248 249 250 251 252 253 254 255 256 257 258
def test_local_func():
    @tvm.register_func("rpc.test.remote_func2")
    def addone(x):
        return lambda y: x+y
    client = rpc.LocalSession()
    f1 = client.get_function("rpc.test.remote_func2")
    fadd = f1(10)
    assert fadd(12) == 22

    blob = bytearray(np.random.randint(0, 10, size=(10)))
    client.upload(blob, "dat.bin")
    rev = client.download("dat.bin")
    assert rev == blob

259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331
def test_rpc_tracker_register():
    # test registration
    tracker = Tracker('localhost', port=9000, port_end=10000)
    device_key = 'test_device'
    server = rpc.Server('localhost', port=9000, port_end=10000,
                        key=device_key,
                        tracker_addr=(tracker.host, tracker.port))
    time.sleep(1)
    client = rpc.connect_tracker(tracker.host, tracker.port)

    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 1

    remote = client.request(device_key)
    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 0

    del remote
    time.sleep(1)

    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 1

    server.terminate()
    time.sleep(1)

    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 0

    tracker.terminate()

def test_rpc_tracker_request():
    # test concurrent request
    tracker = Tracker('localhost', port=9000, port_end=10000)
    device_key = 'test_device'
    server = rpc.Server('localhost', port=9000, port_end=10000,
                        key=device_key,
                        tracker_addr=(tracker.host, tracker.port))
    client = rpc.connect_tracker(tracker.host, tracker.port)

    def target(host, port, device_key, timeout):
        client = rpc.connect_tracker(host, port)
        remote = client.request(device_key, session_timeout=timeout)
        while True:
            pass
        remote.cpu()

    proc1 = multiprocessing.Process(target=target,
                                    args=(tracker.host, tracker.port, device_key, 4))
    proc2 = multiprocessing.Process(target=target,
                                    args=(tracker.host, tracker.port, device_key, 200))
    proc1.start()
    time.sleep(0.5)
    proc2.start()
    time.sleep(0.5)

    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 0
    assert summary['queue_info'][device_key]['pending'] == 1

    proc1.terminate()
    proc1.join()
    time.sleep(0.5)

    summary = client.summary()
    assert summary['queue_info'][device_key]['free'] == 0
    assert summary['queue_info'][device_key]['pending'] == 0

    proc2.terminate()
    proc2.join()
    server.terminate()
    tracker.terminate()

332

333 334
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
335 336
    test_rpc_return_ndarray()
    test_rpc_return_func()
337
    test_bigendian_rpc()
338
    test_rpc_remote_module()
339
    test_rpc_file_exchange()
340 341
    test_rpc_array()
    test_rpc_simple()
342
    test_local_func()
343 344
    test_rpc_tracker_register()
    test_rpc_tracker_request()