wehub-resource-sync
77 lines
2.3 KiB
Python
77 lines
2.3 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.
|
|
|
|
# pylint: disable=invalid-name,unnecessary-comprehension
|
|
"""Testing functions for the RPC server."""
|
|
|
|
import numpy as np
|
|
|
|
import tvm
|
|
|
|
|
|
# RPC test functions to be registered for unit-tests purposes
|
|
@tvm.register_global_func("rpc.test.addone")
|
|
def _addone(x):
|
|
return x + 1
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.strcat")
|
|
def _strcat(name, x):
|
|
return f"{name}:{x}"
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.except")
|
|
def _remotethrow(name):
|
|
raise ValueError(f"{name}")
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.runtime_str_concat")
|
|
def _strcat(x, y):
|
|
return x + y
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.remote_tensor_func")
|
|
def _remote_tensor_func(y):
|
|
x = np.ones((3, 4))
|
|
np.testing.assert_equal(y.numpy(), x)
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.add_to_lhs")
|
|
def _add_to_lhs(x):
|
|
return lambda y: x + y
|
|
|
|
|
|
@tvm.register_global_func("rpc.test.remote_return_nd")
|
|
def _my_module(name):
|
|
# Use closure to check the ref counter correctness
|
|
nd = tvm.runtime.tensor(np.zeros(10).astype("float32"))
|
|
|
|
if name == "get_arr":
|
|
return lambda: nd
|
|
if name == "ref_count":
|
|
# Imported lazily: rpc.server imports this module to register the
|
|
# rpc.test.* helpers, so a top-level ``import tvm.testing`` would drag
|
|
# tvm.testing (which imports pytest) into the plain ``import tvm`` path.
|
|
from tvm.testing import object_use_count
|
|
|
|
return lambda: object_use_count(nd)
|
|
if name == "get_elem":
|
|
return lambda idx: nd.numpy()[idx]
|
|
if name == "get_arr_elem":
|
|
return lambda arr, idx: arr.numpy()[idx]
|
|
raise RuntimeError("unknown name")
|