wehub-resource-sync
114 lines
3.4 KiB
Python
114 lines
3.4 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.
|
|
"""Testing utility functions in meta schedule"""
|
|
|
|
from collections.abc import Callable
|
|
|
|
import numpy as np # type: ignore
|
|
|
|
import tvm
|
|
from tvm.runtime import Tensor
|
|
|
|
|
|
def generate_input_data(
|
|
input_shape: list[int],
|
|
input_dtype: str,
|
|
*,
|
|
low: int | None = None,
|
|
high: int | None = None,
|
|
) -> np.ndarray:
|
|
"""Generate input date with given shape and data type.
|
|
|
|
Parameters
|
|
----------
|
|
input_shape : List[int]
|
|
The shape of the input data.
|
|
input_dtype : str
|
|
The data type of the input date.
|
|
|
|
Returns
|
|
-------
|
|
input_data : np.ndarray
|
|
The generated input data with given shape and data type in numpy ndarray.
|
|
"""
|
|
if input_dtype.startswith("float"):
|
|
return np.random.uniform(size=input_shape).astype(input_dtype)
|
|
range_map = {
|
|
"uint8": (0, 255),
|
|
"int8": (-128, 127),
|
|
"int32": (0, 10000),
|
|
"uint32": (0, 10000),
|
|
"int64": (0, 10000),
|
|
"uint64": (0, 10000),
|
|
}
|
|
if input_dtype in range_map:
|
|
_low, _high = range_map[input_dtype]
|
|
return np.random.randint(
|
|
low=_low if low is None else low,
|
|
high=_high if high is None else high,
|
|
size=input_shape,
|
|
dtype=input_dtype,
|
|
)
|
|
raise ValueError("Unsupported input datatype!")
|
|
|
|
|
|
def create_calculator(backend: str) -> Callable:
|
|
"""Create a function to fetch the computing result of running the given runtime module.
|
|
|
|
Parameters
|
|
----------
|
|
backend : str
|
|
The backend to use, only tirx is supported for now.
|
|
|
|
Returns
|
|
-------
|
|
func : Callable
|
|
The function to fetch the computing result.
|
|
"""
|
|
|
|
def f_calculator(
|
|
rt_mod: tvm.runtime.Module,
|
|
dev: tvm.runtime.Device, # pylint: disable=unused-argument
|
|
input_data: dict[str, Tensor],
|
|
) -> list[Tensor]:
|
|
"""Fetch the result of running the given runtime module.
|
|
|
|
Parameters
|
|
----------
|
|
rt_mod : tvm.runtime.Module
|
|
The runtime module.
|
|
dev : tvm.device
|
|
The device type to run workload.
|
|
input_data : Dict[str, np.ndarray]
|
|
The input data as a dictionary.
|
|
"""
|
|
try:
|
|
if backend == "tirx":
|
|
data = [v for _, v in sorted(input_data.items(), key=lambda x: x[0])]
|
|
rt_mod(*data)
|
|
return data
|
|
else:
|
|
raise ValueError(f"Backend {backend} not supported in f_calculator!")
|
|
|
|
except Exception as exc: # pylint: disable=broad-except
|
|
print(
|
|
f"Run module f_calculator via RPC failed, exception: {exc}",
|
|
)
|
|
return None
|
|
|
|
return f_calculator
|