def test_meta_schedule_rpc_runner_time_out():
"""Test meta schedule RPC Runner time out"""
def initializer():
@register_func("meta_schedule.runner.test_time_out")
def timeout_session_creator( # pylint: disable=unused-variable
rpc_config: RPCConfig, # pylint: disable=unused-argument
) -> RPCSession:
time.sleep(2)
runner_input = RunnerInput(
"test",
"llvm",
[
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
],
)
with LocalRPC() as rpc:
rpc_config = RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=1,
)
evaluator_config = EvaluatorConfig(
number=1,
repeat=1,
min_repeat_ms=0,
enable_cpu_cache_flush=False,
)
runner = RPCRunner(
rpc_config,
evaluator_config,
initializer=initializer,
f_create_session="meta_schedule.runner.test_time_out",
)
# Run the module
(runner_future,) = runner.run([runner_input])
runner_result = runner_future.result()
assert runner_result.error_msg is not None and runner_result.error_msg.startswith(
"RPCRunner: Timeout, killed after"
)
assert runner_result.run_secs is None
def test_meta_schedule_rpc_single_run():
"""Test meta schedule rpc runner for a single run"""
# Build the module
mod = MatmulModule
builder = LocalBuilder()
(builder_result, ) = builder.build([BuilderInput(mod, Target("llvm"))])
assert builder_result.artifact_path is not None
assert builder_result.error_msg is None
runner_input = RunnerInput(
builder_result.artifact_path,
"llvm",
[
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
],
)
with LocalRPC() as rpc:
rpc_config = RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
evaluator_config = EvaluatorConfig(
number=1,
repeat=1,
min_repeat_ms=0,
enable_cpu_cache_flush=False,
)
runner = RPCRunner(rpc_config, evaluator_config)
# Run the module
(runner_future, ) = runner.run([runner_input])
runner_result = runner_future.result()
assert runner_result.error_msg is None
for result in runner_result.run_secs:
if isinstance(result, FloatImm):
result = result.value
assert isinstance(result, float)
assert result >= 0.0
_clean_build(builder_result.artifact_path)
def test_tune_run_module_via_rpc():
target = tvm.target.Target("llvm")
rt_mod = tvm.build(matmul, target)
# construct the input
input_data = {}
input_shape = (128, 128)
input_dtype = "float32"
a_np = np.random.uniform(size=input_shape).astype(input_dtype)
b_np = np.random.uniform(size=input_shape).astype(input_dtype)
c_np = np.zeros(input_shape).astype(input_dtype)
for i in range(128):
for j in range(128):
for k in range(128):
c_np[i, j] = c_np[i, j] + a_np[i, k] * b_np[j, k]
input_data["a"] = a_np
input_data["b"] = b_np
input_data["c"] = np.zeros(input_shape).astype(input_dtype)
with LocalRPC() as rpc:
rpc_config = ms.runner.RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
def f_timer(rt_mod, dev, input_data):
rt_mod(input_data["a"], input_data["b"], input_data["c"])
return input_data["c"]
result = run_module_via_rpc(
rpc_config=rpc_config,
lib=rt_mod,
dev_type=target.kind.name,
args=input_data,
continuation=f_timer,
)
tvm.testing.assert_allclose(result.numpy(), c_np, rtol=1e-3)
def test_meta_schedule_runner_add_test():
"""Test meta schedule runner with add module"""
def _check_correct_add(args_before: List[np.ndarray],
args_after: List[np.ndarray]) -> None:
a_before, b_before, c_before = args_before
a_after, b_after, c_after = args_after
c_before = a_before + b_before
assert (a_before == a_after).all()
assert (b_before == b_after).all()
assert (c_before == c_after).all()
def test_alloc_argument(
session: RPCSession,
device: Device,
args_info: Any,
alloc_repeat: int,
) -> List[Any]:
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
repeated_args_before = [] # type: ignore
repeated_args = rpc_default_alloc_argument(
session,
device,
args_info,
alloc_repeat,
)
for args in repeated_args:
repeated_args_before.append([arg.numpy()
for arg in args]) # type: ignore
return repeated_args
def test_run_evaluator(
session: RPCSession, # pylint: disable=unused-argument
rt_mod: Module,
device: Device,
evaluator_config: EvaluatorConfig,
repeated_args: List[Any],
) -> List[float]:
global repeated_args_before # pylint: disable=global-variable-undefined, invalid-name
repeated_args_after = []
evaluator = rt_mod.time_evaluator(
func_name=rt_mod.entry_name,
dev=device,
number=evaluator_config.number,
repeat=evaluator_config.repeat,
min_repeat_ms=evaluator_config.min_repeat_ms,
f_preproc="cache_flush_cpu_non_first_arg"
if evaluator_config.enable_cpu_cache_flush else "",
)
repeated_costs: List[List[float]] = []
for args in repeated_args:
device.sync()
profile_result = evaluator(*args)
repeated_costs.append(profile_result.results)
repeated_args_after.append([arg.numpy() for arg in args])
costs = [
float(cost)
for cost in itertools.chain.from_iterable(repeated_costs)
]
for args_before, args_after in zip(
repeated_args_before, # type: ignore
repeated_args_after,
):
_check_correct_add(args_before, args_after)
del repeated_args_before # type: ignore
return costs
# Build the module
mod = AddModule
builder = LocalBuilder()
(builder_result, ) = builder.build([BuilderInput(mod, Target("llvm"))])
assert builder_result.artifact_path is not None
assert builder_result.error_msg is None
runner_input = RunnerInput(
builder_result.artifact_path,
"llvm",
[
TensorInfo("float32", [MATMUL_M]),
TensorInfo("float32", [MATMUL_M]),
TensorInfo("float32", [MATMUL_M]),
],
)
with LocalRPC() as rpc:
rpc_config = RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
evaluator_config = EvaluatorConfig(
number=1,
repeat=1,
min_repeat_ms=0,
enable_cpu_cache_flush=False,
)
runner = RPCRunner(
rpc_config,
evaluator_config,
f_alloc_argument=test_alloc_argument,
f_run_evaluator=test_run_evaluator,
)
# Run the module
(runner_future, ) = runner.run([runner_input])
runner_result = runner_future.result()
assert runner_result.error_msg is None
for result in runner_result.run_secs:
if isinstance(result, FloatImm):
result = result.value
assert isinstance(result, float)
assert result >= 0.0
_clean_build(builder_result.artifact_path)
def test_meta_schedule_rpc_multiple_runs():
"""Test meta schedule rpc runner for multiple runs"""
# Build the module
mods = [
MatmulModule,
MatmulReluModule,
BatchMatmulModule,
]
builder = LocalBuilder()
builder_inputs = [BuilderInput(mod, Target("llvm")) for mod in mods]
builder_results = builder.build(builder_inputs)
for builder_result in builder_results:
assert builder_result.artifact_path is not None
assert builder_result.error_msg is None
args_infos = [
[
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
],
[
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
TensorInfo("float32", (MATMUL_N, MATMUL_N)),
],
[
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
TensorInfo("float32", [16, MATMUL_M, MATMUL_M]),
],
]
runner_inputs = [
RunnerInput(builder_results[i].artifact_path, "llvm", args_infos[i])
for i in range(len(mods))
]
with LocalRPC() as rpc:
rpc_config = RPCConfig(
tracker_host=rpc.tracker_host,
tracker_port=rpc.tracker_port,
tracker_key=rpc.tracker_key,
session_priority=1,
session_timeout_sec=100,
)
evaluator_config = EvaluatorConfig(
number=1,
repeat=1,
min_repeat_ms=0,
enable_cpu_cache_flush=False,
)
runner = RPCRunner(rpc_config, evaluator_config)
# Run the module
runner_futures = runner.run(runner_inputs)
runner_results = [
runner_future.result() for runner_future in runner_futures
]
for runner_result in runner_results:
assert runner_result.error_msg is None
for result in runner_result.run_secs:
if isinstance(result, FloatImm):
result = result.value
assert isinstance(result, float)
assert result >= 0.0
for builder_result in builder_results:
_clean_build(builder_result.artifact_path)