wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
apache--tvm/python/tvm/s_tir/meta_schedule/testing/validate_database.py
T
wehub-resource-sync 26446540fa
Lint / lint (push) Waiting to run
Details
CI / MacOS (push) Waiting to run
Details
CI / Windows (push) Waiting to run
Details
chore: import upstream snapshot with attribution
2026-07-13 13:36:25 +08:00

785 lines
25 KiB
Python
Raw Permalink Blame History

# 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.
# ruff: noqa: F403
"""JSON Database validation script"""
import argparse
import itertools
import logging
import warnings
from collections.abc import Callable
from statistics import mean
from typing import Any
import numpy as np # type: ignore
import tvm_ffi
from tvm_ffi import get_global_func, register_global_func
import tvm
from tvm.ir import IRModule
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.tune_utils import generate_input_data
from tvm.s_tir.meta_schedule.utils import remove_build_dir
from tvm.s_tir.schedule import Trace
from tvm.s_tir.tensor_intrin import * # type: ignore # pylint: disable=wildcard-import,unused-wildcard-import
from tvm.support import describe
from tvm.target import Target
from tvm.testing.utils import strtobool
DELIMITOR = "\n" + "-" * 30 + "\n"
def _parse_args():
args = argparse.ArgumentParser()
args.add_argument(
"--work-dir",
type=str,
required=True,
help="The path to the work directory containing database files.",
)
args.add_argument(
"--target",
type=Target,
required=True,
)
args.add_argument(
"--baseline-target",
type=Target,
default='{"kind": "llvm", "num-cores": 1}',
required=False,
help="The baseline target to compile the original module.",
)
args.add_argument(
"--top-k",
type=int,
default=10**9,
required=False,
help="The number of top-k tuning records to validate for each unique original workload.",
)
args.add_argument(
"--rpc-host",
type=str,
)
args.add_argument(
"--rpc-port",
type=int,
)
args.add_argument(
"--rpc-key",
type=str,
)
args.add_argument(
"--number",
type=int,
default=3,
)
args.add_argument(
"--repeat",
type=int,
default=1,
)
args.add_argument(
"--min-repeat-ms",
type=int,
default=100,
)
args.add_argument(
"--cpu-flush",
type=lambda x: bool(strtobool(x)),
help="example: True / False",
required=True,
)
args.add_argument(
"--input-generator-func",
type=str,
default="tvm.s_tir.meta_schedule.testing.default_input_generator",
)
args.add_argument(
"--check-metric-func",
type=str,
default="tvm.s_tir.meta_schedule.testing.default_check_metric",
)
parsed = args.parse_args()
parsed.target = tvm.target.Target(parsed.target)
if parsed.rpc_host is not None and parsed.rpc_port is not None and parsed.rpc_key is not None:
parsed.rpc_config = ms.runner.RPCConfig(
tracker_host=parsed.rpc_host,
tracker_port=parsed.rpc_port,
tracker_key=parsed.rpc_key,
session_timeout_sec=600,
)
else:
parsed.rpc_config = None
warnings.warn("RPC config is not provided, will use local runner.")
if parsed.cpu_flush and parsed.target.kind.name != "llvm":
warnings.warn("cpu_flush is only supported on llvm target")
return parsed
# arg parser
ARGS = _parse_args()
# logging
logging.basicConfig(
format="%(asctime)s.%(msecs)03d %(levelname)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
)
logging.getLogger("tvm.s_tir.meta_schedule").setLevel(logging.DEBUG)
logging.getLogger("tvm.s_tir.meta_schedule.runner").setLevel(logging.WARN)
def get_device_type(target: Target) -> str:
"""Get the device type string from a target.
Parameters
----------
target : Target
The target to get the device type from.
Returns
-------
device_type : str
The device type string.
"""
if target.kind.name == "llvm":
return "cpu"
elif target.kind.name == "cuda":
return "cuda"
else:
raise RuntimeError(f"Unsupported target kind for device type: {target.kind.name}")
def get_runtime_device(target: Target) -> tvm.runtime.Device:
"""Get the runtime device from a target.
Parameters
----------
target : Target
The target to get the runtime device from.
Returns
-------
device : tvm.runtime.Device
The runtime device.
"""
if target.kind.name == "llvm":
return tvm.cpu()
elif target.kind.name == "cuda":
return tvm.cuda()
else:
raise RuntimeError(f"Unsupported target kind for runtime device: {target.kind.name}")
def check_and_run(func: str | Callable, *args, **kwargs) -> Any:
"""Check if the function is a string or a callable, and run it."""
if isinstance(func, str):
func = get_global_func(func)
return func(*args, **kwargs) # type: ignore
class OriginalModule:
"""Original module class for deduplication."""
def __init__(self, mod: IRModule):
self.mod = mod
def __eq__(self, __o: "OriginalModule") -> bool: # type: ignore
return tvm_ffi.structural_equal(self.mod, __o.mod)
def __hash__(self) -> int:
return tvm_ffi.structural_hash(self.mod)
def initializer() -> None:
"""Initializer function to register the functions on PopenWorker."""
@register_global_func("tvm.s_tir.meta_schedule.testing.default_check_metric")
def default_check_metric( # pylint: disable=unused-variable,unreachable-code
lhs: list[tvm.runtime.Tensor], rhs: list[tvm.runtime.Tensor]
) -> bool:
"""Check if the outputs are equal
Parameters
----------
lhs : List[tvm.runtime.Tensor]
The first list of Tensors to compare.
rhs : List[tvm.runtime.Tensor]
The second list of Tensors to compare.
Returns
-------
is_equal : bool
Whether the two lists of Tensors are equal.
"""
assert len(lhs) == len(rhs), "Different number of outputs from two modules"
for i in range(len(lhs)): # pylint: disable=consider-using-enumerate
if not np.allclose(lhs[i].numpy(), rhs[i].numpy(), rtol=1e-3, atol=2e-3):
return False
return True
@register_global_func("tvm.s_tir.meta_schedule.testing.default_input_generator")
def default_input_generator( # pylint: disable=unused-variable
mod: IRModule,
) -> list[tvm.runtime.Tensor]:
"""Default input generator function
Parameters
----------
mod : IRModule
The IRModule to generate the input data for.
Returns
-------
inputs : List[tvm.runtime.Tensor]
The generated input data.
"""
args_info = ms.arg_info.TensorInfo.from_prim_func(mod["main"])
inputs = [
tvm.runtime.tensor(
generate_input_data(input_shape=arg_info.shape, input_dtype=arg_info.dtype)
)
for arg_info in args_info
]
return inputs
def to_numpy(a: list[tvm.runtime.Tensor]) -> list[np.ndarray]:
"""Convert a list of TVM Tensor to a list of numpy array
Parameters
----------
a : List[tvm.runtime.Tensor]
The list of TVM Tensor to be converted
Returns
-------
b : List[np.ndarray]
The list of numpy array
"""
assert a is not None, "Empty result cannot be converted to numpy"
return [x.numpy() for x in a]
def to_tvm_tensor(a: list[np.ndarray]) -> list[tvm.runtime.Tensor]:
"""Convert a list of numpy array to a list of TVM Tensor
Parameters
----------
a : List[np.ndarray]
The list of numpy array to be converted.
Returns
-------
b : List[tvm.runtime.Tensor]
The list of TVM Tensor.
"""
assert a is not None, "Empty result cannot be converted to TVM Tensor"
return [tvm.runtime.tensor(x) for x in a]
def is_failed_record(record: ms.database.TuningRecord) -> bool:
"""Check if a tuning record is failed.
Parameters
----------
record : TuningRecord
The tuning record to check.
Returns
-------
is_failed : bool
"""
return len(record.run_secs) == 1 and record.run_secs[0] == 1e9
def print_with_counter_func(counter: int, total: int) -> Callable:
"""Print with counter
Parameters
----------
counter : int
The counter to print with.
total : int
The total number of items to print with.
Returns
-------
print_result : Callable
The print result function.
"""
def print_result(
result: str,
*,
original_mod: IRModule = None,
scheduled_mod: IRModule = None,
inputs: list[np.ndarray] | None = None,
original_res: list[np.ndarray] | None = None,
scheduled_res: list[np.ndarray] | None = None,
original_run_secs: list[float] | None = None,
scheduled_run_secs: list[float] | None = None,
exception: Exception | None = None,
trace: str | None = None,
) -> None:
"""Print the validation result."""
status = f"Progress {counter: 6d} / {total: 6d} (estimated) checked, result: {result:>10}, "
if result in ["pass", "wrong answer"]:
status += (
f"original: {mean(original_run_secs) * 1e3: 10.3f} ms, "
f"scheduled: {mean(scheduled_run_secs) * 1e3: 10.3f} ms"
)
output = [status]
if result not in ["pass", "skip"]:
output.extend(
[
"Original IRModule:" + DELIMITOR + original_mod.script(),
"Scheduled IRModule:" + DELIMITOR + scheduled_mod.script(),
"Trace" + DELIMITOR + str(trace),
]
)
if result == "wrong answer":
output.extend(
[
"Input:" + DELIMITOR + str(inputs),
"Original Result:" + DELIMITOR + str(original_res),
"Scheduled Result:" + DELIMITOR + str(scheduled_res),
"Max Diff:"
+ DELIMITOR
+ str(
[
np.max(np.abs(original_res[i] - scheduled_res[i]))
for i in range(len(original_res))
]
)
+ "\n",
]
)
elif result == "exception":
output.extend(["Exception:" + DELIMITOR + str(exception) + "\n"])
else:
raise ValueError(f"Unknown result: {result}")
print("\n\n".join(output))
return print_result
def make_alloc_arg_and_check(
inputs: list[np.ndarray],
original_mod: IRModule,
scheduled_mod: IRModule,
trace: str,
original_res: list[np.ndarray],
original_run_secs: list[float],
print_result: Callable,
) -> tuple[Callable, Callable]:
"""Make alloc_arg and check functions for the given inputs and collect results.
Parameters
----------
inputs : List[np.ndarray]
The inputs to the two modules.
original_mod : IRModule
The original IRModule.
scheduled_mod : IRModule
The scheduled IRModule.
trace : str
The trace of the scheduled IRModule.
original_res : List[np.ndarray]
The original results.
original_run_secs : List[float]
The original run times.
print_result : Callable
The print result function.
Returns
-------
f_with_args_alloc_argument : Callable
The function to allocate arguments.
f_with_args_run_evaluator : Callable
The function to run evaluator.
"""
def f_with_args_alloc_argument_common(
device: tvm.runtime.Device,
args_info: ms.runner.rpc_runner.T_ARG_INFO_JSON_OBJ_LIST, # pylint: disable=unused-argument
alloc_repeat: int,
) -> list[ms.runner.rpc_runner.T_ARGUMENT_LIST]:
"""Allocate arguments using the given inputs.
Parameters
----------
session : RPCSession
The RPC session.
device : Device
The device.
args_info : T_ARG_INFO_JSON_OBJ_LIST
argument information.
alloc_repeat : int
The number of times to repeat the allocation.
Returns
-------
args_list : List[T_ARGUMENT_LIST]
The list of argument lists.
"""
return [
[tvm.runtime.tensor(arg, device=device) for arg in inputs] for _ in range(alloc_repeat)
]
def f_with_args_run_evaluator_common(
rt_mod: tvm.runtime.Module,
device: tvm.runtime.Device,
evaluator_config: ms.runner.EvaluatorConfig,
repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST],
) -> list[float]:
"""With args function to run the evaluator
Parameters
----------
session : tvm.rpc.RPCSession
The RPC session
rt_mod: Module
The runtime module
device: Device
The device to run the evaluator
evaluator_config: EvaluatorConfig
The evaluator config
repeated_args: List[T_ARGUMENT_LIST]
The repeated arguments
Returns
-------
costs: List[float]
The evaluator results
"""
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)
costs = [float(cost) for cost in itertools.chain.from_iterable(repeated_costs)]
assert len(repeated_args) == 1, "Only support one set of arguments"
scheduled_res = [arg.numpy() for arg in repeated_args[0]] # type: ignore
# fetch comparison function
passed = check_and_run(
ARGS.check_metric_func,
to_tvm_tensor(original_res),
to_tvm_tensor(scheduled_res),
)
print_result(
result="pass" if passed else "wrong answer",
original_mod=original_mod,
scheduled_mod=scheduled_mod,
trace=trace,
inputs=inputs,
original_res=original_res,
scheduled_res=scheduled_res,
original_run_secs=original_run_secs,
scheduled_run_secs=costs,
)
return costs
def f_with_args_alloc_argument_rpc(
rpc_session: ms.runner.rpc_runner.RPCSession, # pylint: disable=unused-argument
device: tvm.runtime.Device,
args_info: ms.runner.rpc_runner.T_ARG_INFO_JSON_OBJ_LIST,
alloc_repeat: int,
) -> list[ms.runner.rpc_runner.T_ARGUMENT_LIST]:
return f_with_args_alloc_argument_common(device, args_info, alloc_repeat)
def f_with_args_run_evaluator_rpc(
rpc_session: ms.runner.rpc_runner.RPCSession, # pylint: disable=unused-argument
rt_mod: tvm.runtime.Module,
device: tvm.runtime.Device,
evaluator_config: ms.runner.EvaluatorConfig,
repeated_args: list[ms.runner.rpc_runner.T_ARGUMENT_LIST],
) -> list[float]:
return f_with_args_run_evaluator_common(rt_mod, device, evaluator_config, repeated_args)
if ARGS.rpc_config is None:
return f_with_args_alloc_argument_common, f_with_args_run_evaluator_common
else:
return f_with_args_alloc_argument_rpc, f_with_args_run_evaluator_rpc
def local_build_and_run(
mod: IRModule,
target: Target,
device: tvm.runtime.Device,
inputs: list[np.ndarray],
) -> tuple[list[np.ndarray], list[float]]:
"""Build and run the module locally.
Parameters
----------
mod: IRModule
The module to build and run
target: Target
The target to build the module
device: Device
The device to run the module
inputs: List[np.ndarray]
The inputs to run the module
Returns
-------
res: List[np.ndarray]
The results of running the module
run_secs: List[float]
The running time of running the module
"""
# potential memory leak https://github.com/apache/tvm/issues/11096
lib = tvm.compile(mod, target=target)
tvm_inputs = [tvm.runtime.tensor(inp, device=device) for inp in inputs]
device.sync()
func = lib.time_evaluator(lib.entry_name, dev=device, number=ARGS.number, repeat=ARGS.repeat)
benchmark_res = func(*tvm_inputs)
device.sync()
return [arg.numpy() for arg in tvm_inputs], list(benchmark_res.results)
def _check_builder_result(builder_result: ms.builder.BuilderResult) -> None:
"""Check if the builder result is defined.
Parameters
----------
builder_result: BuilderResult
The builder result
"""
assert builder_result.error_msg is None, "Builder failed: " + str(
builder_result.error_msg if builder_result.error_msg else "Empty error message"
)
def _apply_trace(mod: IRModule, trace: Trace) -> IRModule:
"""Apply the trace to the module.
Parameters
----------
mod: IRModule
The module to apply the trace to
trace: Trace
The trace to apply
Returns
-------
mod: IRModule
The module with the trace applied
"""
sch = Schedule(mod)
trace.apply_to_schedule(sch, remove_postproc=False)
return sch.mod
def _build_all_mods(
mods: list[IRModule], builder: ms.builder.Builder, target: Target
) -> list[ms.builder.BuilderResult]:
"""Build all the modules.
Parameters
----------
mods: List[IRModule]
The modules to build
builder: Builder
The builder to build the modules
target: Target
The target to build the modules
Returns
-------
builder_results: List[BuilderResult]
The builder results
"""
builder_results = builder.build([ms.builder.BuilderInput(mod, target) for mod in mods])
assert len(builder_results) == len(mods), (
f"Unexpected number of build results, expected {len(mods)} got {len(builder_results)}"
)
return builder_results
def _run_single_mod(
builder_result: ms.builder.BuilderResult,
runner: ms.runner.Runner,
dev_type: str,
) -> None:
"""Run a single module.
Parameters
----------
builder_result: BuilderResult
The builder result
runner: Runner
The runner to run the module
dev_type: str
The device type
"""
runner_futures = runner.run(
# arginfo is not used in this case so we can pass an empty list
[ms.runner.RunnerInput(builder_result.artifact_path, device_type=dev_type, args_info=[])]
)
assert len(runner_futures) == 1, (
f"Unexpected number of runner futures, expected 1 got {len(runner_futures)}"
)
(runner_future,) = runner_futures # pylint: disable=unbalanced-tuple-unpacking
runner_res = runner_future.result()
assert runner_res.error_msg is None, "Runner failed: " + (
runner_res.error_msg if runner_res.error_msg else "Empty error message"
)
def main():
"""Main function"""
describe()
with ms.Profiler() as profiler:
# initialize
target = ARGS.target
dev_type = get_device_type(target)
builder = ms.builder.LocalBuilder()
database = ms.database.create(work_dir=ARGS.work_dir)
# collect records
with profiler.timeit("collect records"):
records = database.get_all_tuning_records()
total = len(records)
print(
f"Total {total} records to be validated. "
f"Collected in {float(profiler.get()['collect records']): 3.3f} sec."
)
# collect unique original TIR
with profiler.timeit("deduplicate records"):
workloads = set()
for record in records:
workloads.add(OriginalModule(record.workload.mod))
print(
f"Total {len(workloads)} unique original TIR to validate. "
f"Deduplicated in {float(profiler.get()['deduplicate records']): 3.3f} sec."
)
if ARGS.top_k < 10**9:
print(f"Top {ARGS.top_k} records for each original TIR will be validated.")
total = len(workloads) * ARGS.top_k
print()
# validate correctness
counter = 0
for item in workloads:
original_mod = item.mod
records = database.get_top_k(
workload=database.commit_workload(original_mod), top_k=ARGS.top_k
)
if len(records) < ARGS.top_k:
total -= ARGS.top_k - len(records)
inputs = to_numpy(check_and_run(ARGS.input_generator_func, original_mod))
original_res, original_run_secs = local_build_and_run(
original_mod,
target=ARGS.baseline_target,
inputs=inputs,
device=get_runtime_device(ARGS.baseline_target),
)
scheduled_mods = [_apply_trace(original_mod, record.trace) for record in records]
builder_results = _build_all_mods(scheduled_mods, builder, target) # type: ignore
for i, record in enumerate(records):
counter += 1
print_result = print_with_counter_func(counter=counter, total=total)
if is_failed_record(record):
# skip failed records where run_secs is 1e9
# these records are only negative samples for cost model
print_result(result="skip")
continue
try:
# prepare scheduled module
scheduled_mod = scheduled_mods[i]
# check build result
builder_result = builder_results[i]
_check_builder_result(builder_result)
# fetch functions
(
f_with_args_alloc_argument,
f_with_args_run_evaluator,
) = make_alloc_arg_and_check(
inputs,
original_mod,
scheduled_mod,
str(record.trace),
original_res=original_res,
original_run_secs=original_run_secs,
print_result=print_result,
)
# create runner
evaluator_config = ms.runner.EvaluatorConfig(
number=ARGS.number,
repeat=ARGS.repeat,
min_repeat_ms=ARGS.min_repeat_ms,
enable_cpu_cache_flush=ARGS.cpu_flush,
)
if ARGS.rpc_config is not None:
runner: ms.Runner = ms.runner.RPCRunner( # type: ignore
ARGS.rpc_config,
evaluator_config=evaluator_config,
alloc_repeat=1,
f_alloc_argument=f_with_args_alloc_argument,
f_run_evaluator=f_with_args_run_evaluator,
initializer=initializer,
)
else:
runner: ms.Runner = ms.runner.LocalRunner( # type: ignore
evaluator_config=evaluator_config,
alloc_repeat=1,
f_alloc_argument=f_with_args_alloc_argument,
f_run_evaluator=f_with_args_run_evaluator,
initializer=initializer,
)
# run and validate
_run_single_mod(builder_result, runner, dev_type) # type: ignore
except Exception as e: # pylint: disable=broad-except, invalid-name
# validation failed with exception
print_result(
result="exception",
original_mod=original_mod,
scheduled_mod=scheduled_mod,
trace=str(record.trace),
exception=e,
)
# clean up
remove_build_dir(builder_result.artifact_path)
print(f"Validation finished! Total time spent: {float(profiler.get()['Total']): 3.3f} sec.")
if __name__ == "__main__":
main()
Reference in New Issue View Git Blame Copy Permalink