wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/s_tir/meta_schedule/testing/__init__.py
+19
View File
@@ -0,0 +1,19 @@
# 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 utilities in meta schedule"""
# NOTE: Do not import any module here by default
python/tvm/s_tir/meta_schedule/testing/custom_builder_runner.py
+54
View File
@@ -0,0 +1,54 @@
# 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.
"""Customized builder and runner methods"""
# pylint: disable=import-outside-toplevel
from collections.abc import Callable
import numpy as np # type: ignore
from tvm.runtime import Executable, Module
from tvm.s_tir.meta_schedule.runner import RPCConfig
def run_module_via_rpc(
rpc_config: RPCConfig,
lib: Module | Executable,
dev_type: str,
args: dict[int, np.ndarray] | dict[str, np.ndarray],
continuation: Callable,
):
"""Execute a tvm.runtime.Module on RPC remote"""
# pylint: disable=import-outside-toplevel
import os
import tempfile
from tvm.runtime import ndarray
from tvm.support.tar import tar
# pylint: enable=import-outside-toplevel
with tempfile.TemporaryDirectory() as tmp_dir:
filename = os.path.join(tmp_dir, "tvm_tmp_mod." + tar.output_format)
lib.export_library(filename, fcompile=tar)
session = rpc_config.connect_server()
session.upload(filename)
_, filename = os.path.split(filename)
rt_mod = session.load_module(filename)
dev = session.device(dev_type, 0)
nd_args = {k: ndarray.array(v, dev) for k, v in args.items()}
return continuation(rt_mod, dev, nd_args)
python/tvm/s_tir/meta_schedule/testing/distributed_measure_candidates.py
+199
View File
@@ -0,0 +1,199 @@
# 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=missing-docstring
import argparse
import glob
import os
from tqdm import tqdm # type: ignore
from tvm.s_tir import meta_schedule as ms
from tvm.target import Target
def _parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--candidate_cache_dir", type=str, help="Please provide the full path to the candidates."
)
parser.add_argument(
"--result_cache_dir", type=str, help="Please provide the full path to the result database."
)
parser.add_argument(
"--target",
type=str,
default="nvidia/nvidia-v100",
help="Please specify the target hardware for tuning context.",
)
parser.add_argument(
"--rpc_host", type=str, help="Please provide the private IPv4 address for the tracker."
)
parser.add_argument(
"--rpc_port", type=int, default=4445, help="Please provide the port for the tracker."
)
parser.add_argument(
"--rpc_key",
type=str,
default="p3.2xlarge",
help="Please provide the key for the rpc servers.",
)
parser.add_argument(
"--builder_timeout_sec",
type=int,
default=10,
help="The time for the builder session to time out.",
)
parser.add_argument(
"--min_repeat_ms", type=int, default=100, help="The time for preheating the gpu."
)
parser.add_argument(
"--runner_timeout_sec",
type=int,
default=100,
help="The time for the runner session to time out.",
)
parser.add_argument(
"--cpu_flush", type=bool, default=False, help="Whether to enable cpu cache flush or not."
)
parser.add_argument(
"--batch_size",
type=int,
default=128,
help="The batch size of candidates sent to builder and runner each time.",
)
return parser.parse_args()
# pylint: disable=too-many-locals
def measure_candidates(database, builder, runner):
"""Send the candidates to builder and runner for distributed measurement,
and save the results in a new json database.
Parameters
----------
database : JSONDatabase
The database for candidates to be measured.
builder : Builder
The builder for building the candidates.
runner : Runner
The runner for measuring the candidates.
Returns
-------
None
"""
candidates, runner_results, build_fail_indices, run_fail_indices = [], [], [], []
context = ms.TuneContext(target=Target(args.target))
tuning_records = database.get_all_tuning_records()
for record in tuning_records:
candidates.append(record.as_measure_candidate())
with ms.Profiler() as profiler:
for idx in range(0, len(candidates), args.batch_size):
batch_candidates = candidates[idx : idx + args.batch_size]
context._set_measure_candidates(batch_candidates) # pylint: disable=protected-access
with ms.Profiler.timeit("build"):
context._send_to_builder(builder) # pylint: disable=protected-access
with ms.Profiler.timeit("run"):
context._send_to_runner(runner) # pylint: disable=protected-access
batch_runner_results = context._join() # pylint: disable=protected-access
runner_results.extend(batch_runner_results)
for i, result in enumerate(context.builder_results):
if result.error_msg is None:
ms.utils.remove_build_dir(result.artifact_path)
else:
build_fail_indices.append(i + idx)
context._clear_measure_state() # pylint: disable=protected-access
model_name, workload_name = database.path_workload.split("/")[-2:]
record_name = database.path_tuning_record.split("/")[-1]
new_database = ms.database.JSONDatabase(
path_workload=os.path.join(args.result_cache_dir, model_name, workload_name),
path_tuning_record=os.path.join(args.result_cache_dir, model_name, record_name),
)
workload = tuning_records[0].workload
new_database.commit_workload(workload.mod)
for i, (record, result) in enumerate(zip(tuning_records, runner_results)):
if result.error_msg is None:
new_database.commit_tuning_record(
ms.database.TuningRecord(
trace=record.trace,
workload=workload,
run_secs=[v.value for v in result.run_secs],
target=Target(args.target),
)
)
else:
run_fail_indices.append(i)
fail_indices_name = workload_name.replace("_workload.json", "_failed_indices.txt")
with open(
os.path.join(args.result_cache_dir, model_name, fail_indices_name), "w", encoding="utf8"
) as file:
file.write(" ".join([str(n) for n in run_fail_indices]))
print(
f"Builder time: {profiler.get()['build']}, Runner time: {profiler.get()['run']}\n\
Failed number of builds: {len(build_fail_indices)},\
Failed number of runs: {len(run_fail_indices)}"
)
args = _parse_args() # pylint: disable=invalid-name
def main():
builder = ms.builder.LocalBuilder(timeout_sec=args.builder_timeout_sec)
runner = ms.runner.RPCRunner(
rpc_config=ms.runner.RPCConfig(
tracker_host=args.rpc_host,
tracker_port=args.rpc_port,
tracker_key=args.rpc_key,
session_timeout_sec=args.runner_timeout_sec,
),
evaluator_config=ms.runner.EvaluatorConfig(
number=3,
repeat=1,
min_repeat_ms=args.min_repeat_ms,
enable_cpu_cache_flush=args.cpu_flush,
),
max_workers=os.cpu_count(),
)
if not os.path.isdir(args.candidate_cache_dir):
raise Exception("Please provide a correct candidate cache dir.")
try:
os.makedirs(args.result_cache_dir, exist_ok=True)
except OSError:
print(f"Directory {args.result_cache_dir} cannot be created successfully.")
model_dirs = glob.glob(os.path.join(args.candidate_cache_dir, "*"))
for model_dir in model_dirs:
model_name = model_dir.split("/")[-1]
os.makedirs(os.path.join(args.result_cache_dir, model_name), exist_ok=True)
all_tasks = glob.glob(os.path.join(model_dir, "*.json"))
workload_paths = []
for path in all_tasks:
if path.endswith("_workload.json"):
workload_paths.append(path)
for workload_path in tqdm(workload_paths):
candidate_path = workload_path.replace("_workload.json", "_candidates.json")
database = ms.database.JSONDatabase(
path_workload=workload_path,
path_tuning_record=candidate_path,
)
measure_candidates(database, builder, runner)
if __name__ == "__main__":
main()
python/tvm/s_tir/meta_schedule/testing/dummy_object.py
+63
View File
@@ -0,0 +1,63 @@
# 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.
"""Dummy objects for testing."""
import random
from tvm.ir.utils import derived_object
from tvm.s_tir.schedule import Trace
from ..builder import BuilderInput, BuilderResult, PyBuilder
from ..mutator import PyMutator
from ..runner import PyRunner, PyRunnerFuture, RunnerFuture, RunnerInput, RunnerResult
from ..tune_context import TuneContext # pylint: disable=unused-import
@derived_object
class DummyRunnerFuture(PyRunnerFuture):
def done(self) -> bool:
return True
def result(self) -> RunnerResult:
run_secs = [random.uniform(5, 30) for _ in range(random.randint(1, 10))]
return RunnerResult(run_secs, None)
@derived_object
class DummyBuilder(PyBuilder):
def build(self, build_inputs: list[BuilderInput]) -> list[BuilderResult]:
return [BuilderResult("test_path", None) for _ in build_inputs]
@derived_object
class DummyRunner(PyRunner):
def run(self, runner_inputs: list[RunnerInput]) -> list[RunnerFuture]:
return [DummyRunnerFuture() for _ in runner_inputs] # type: ignore
@derived_object
class DummyMutator(PyMutator):
"""Dummy Mutator for testing"""
def _initialize_with_tune_context(self, context: "TuneContext") -> None:
pass
def apply(self, trace: Trace, _) -> Trace | None:
return Trace(trace.insts, {})
def clone(self):
return DummyMutator()
python/tvm/s_tir/meta_schedule/testing/local_rpc.py
+72
View File
@@ -0,0 +1,72 @@
# 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.
"""RPC tracker and server running locally"""
from tvm.rpc.server import Server
from tvm.rpc.tracker import Tracker
class LocalRPC:
"""A pair of RPC tracker/server running locally
Parameters
----------
tracker_host : str
The host URL of the tracker
tracker_port : int
The port of the tracker
tracker_key: str
The key used in the tracker to refer to a worker
"""
tracker_host: str
tracker_port: int
tracker_key: str
def __init__(
self,
tracker_key: str = "key",
silent: bool = False,
no_fork: bool = False,
) -> None:
self.tracker = Tracker(
silent=silent,
port=9190,
port_end=12345,
)
self.server = Server(
host="0.0.0.0",
is_proxy=False,
tracker_addr=(self.tracker.host, self.tracker.port),
key=tracker_key,
silent=silent,
no_fork=no_fork,
port=9190,
port_end=12345,
)
self.tracker_host = self.tracker.host
self.tracker_port = self.tracker.port
self.tracker_key = tracker_key
def __enter__(self):
return self
def __exit__(self, _type, _value, _traceback):
if hasattr(self, "server"):
del self.server
if hasattr(self, "tracker"):
del self.tracker
python/tvm/s_tir/meta_schedule/testing/space_generation.py
+150
View File
@@ -0,0 +1,150 @@
# 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=missing-module-docstring,missing-function-docstring,missing-class-docstring
# isort: off
from typing import Literal
# isort: on
import tvm_ffi
from tvm.ir import IRModule
from tvm.s_tir import Schedule
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.schedule import Trace
from tvm.s_tir.schedule.testing import verify_trace_roundtrip
from tvm.target import Target
def get_rules(
kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
types: type | tuple[type, ...],
) -> list[ms.ScheduleRule]:
"""Get default schedule rules"""
rules = ms.ScheduleRule.create(kind)
return [rule for rule in rules if isinstance(rule, types)]
def structural_equal_no_gs(mod1: IRModule, mod2: IRModule) -> bool:
"""
Checks structural equality but ignores global symbols
"""
# for every function in the modules, remove global symbols from the attrs and then compare
def remove_global_symbols(mod: IRModule) -> IRModule:
stripped_mod = IRModule()
for global_var in mod.get_global_vars():
func = mod[global_var]
stripped_mod[global_var] = func.without_attr("global_symbol")
return stripped_mod
return tvm_ffi.structural_equal(remove_global_symbols(mod1), remove_global_symbols(mod2))
def generate_design_space(
kind: Literal["llvm", "cuda", "cuda-tensorcore", "hexagon"],
mod: IRModule,
target: Target,
types: type | tuple[type, ...],
sch_rules: list[ms.ScheduleRule] | None = None,
) -> list[Schedule]:
if sch_rules is None:
sch_rules = get_rules(kind, types)
else:
assert types is None
return ms.TuneContext(
mod=mod,
target=target,
space_generator=ms.space_generator.PostOrderApply(
sch_rules=sch_rules,
postprocs=[],
mutator_probs={},
),
task_name="test",
).generate_design_space()
def _find_match_sketch_id(
mod: IRModule,
sketches: list[Schedule],
expected_mod: IRModule,
expected_decision: list[tuple[str, list[int]]],
*,
debug_mask="all",
) -> int | None:
for sketch_id, sketch in enumerate(sketches):
i = 0
new_decisions = {}
for inst in sketch.trace.insts:
if not inst.kind.name.startswith("Sample"):
continue
assert i < len(expected_decision)
if inst.kind.name == expected_decision[i][0]:
new_decisions[inst] = expected_decision[i][1]
i += 1
if len(new_decisions) != len(expected_decision):
continue
sch = Schedule(mod, debug_mask=debug_mask)
Trace(
insts=sketch.trace.insts,
decisions=new_decisions,
).apply_to_schedule(sch, remove_postproc=True)
if structural_equal_no_gs(sch.mod, expected_mod):
verify_trace_roundtrip(sch=sch, mod=mod, debug_mask=debug_mask, text_format="json")
return sketch_id
return None
def check_sketches(
mod: IRModule,
sketches: list[Schedule],
expected_mods: list[IRModule],
expected_decisions: list[list[tuple[str, list[int]]]],
*,
debug_mask="all",
):
assert len(expected_mods) == len(expected_decisions)
assert len(sketches) == len(expected_mods)
expected_mods = [
IRModule({"main": m}) if not isinstance(m, IRModule) else m for m in expected_mods
]
sketches = list(sketches)
for expected_id, (expected_mod, expected_decision) in enumerate(
zip(expected_mods, expected_decisions)
):
sketch_id = _find_match_sketch_id(
mod,
sketches,
expected_mod,
expected_decision,
debug_mask=debug_mask,
)
if sketch_id is None:
raise AssertionError(
f"Expected sketch #{expected_id} doesn't exist in the generated sketches."
)
sketches.pop(sketch_id)
def print_sketches(sketches: list[Schedule]):
for i, sch in enumerate(sketches):
print(f"###### {i}")
sch.mod.show(black_format=False)
for inst in sch.trace.insts:
if inst in sch.trace.decisions:
print(f'("{inst.kind.name}", {sch.trace.decisions[inst]}),')
python/tvm/s_tir/meta_schedule/testing/te_workload.py
+837
View File
@@ -0,0 +1,837 @@
# 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: E741
"""Workloads in TE"""
# pylint: disable=missing-docstring
from tvm import te, tirx, topi
from tvm.target import Target
def batch_matmul_nkkm( # pylint: disable=invalid-name,missing-docstring
B: int,
N: int,
M: int,
K: int,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
x = te.placeholder((B, N, K), name="X", dtype=in_dtype)
y = te.placeholder((B, K, M), name="Y", dtype=in_dtype)
k = te.reduce_axis((0, K), name="k")
z = te.compute( # pylint: disable=invalid-name
(B, N, M),
lambda b, i, j: te.sum(
x[b][i][k].astype(out_dtype) * y[b][k][j].astype(out_dtype),
axis=[k],
),
name="Z",
)
return (x, y, z)
def conv1d_nlc( # pylint: disable=invalid-name,missing-docstring
N: int,
L: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder((N, L, CI), name="inputs", dtype=in_dtype)
weight = te.placeholder((kernel_size, CI // groups, CO), name="weight", dtype=in_dtype)
batch_size, in_len, _ = inputs.shape
k_len, channel_per_group, out_channel = weight.shape
out_channel_per_group = out_channel // groups
out_len = (in_len + 2 * padding - dilation * (k_len - 1) - 1) // stride + 1
rc = te.reduce_axis((0, channel_per_group), name="rc")
rl = te.reduce_axis((0, k_len), name="rl")
padded = topi.nn.pad(inputs, [0, padding, 0])
output = te.compute(
(batch_size, out_len, out_channel),
lambda n, l, co: te.sum(
(
padded[
n,
l * stride + rl * dilation,
co // out_channel_per_group * channel_per_group + rc,
].astype(out_dtype)
* weight[rl, rc, co].astype(out_dtype)
),
axis=[rl, rc],
),
name="conv1d_nlc",
)
return (inputs, weight, output)
def conv2d_nhwc( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder((N, H, W, CI), name="inputs", dtype=in_dtype)
weight = te.placeholder(
(kernel_size, kernel_size, CI // groups, CO), name="weight", dtype=in_dtype
)
batch_size, in_h, in_w, _ = inputs.shape
k_h, k_w, channel_per_group, out_channel = weight.shape
out_channel_per_group = out_channel // groups
out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
rh = te.reduce_axis((0, k_h), name="rh")
rw = te.reduce_axis((0, k_w), name="rw")
rc = te.reduce_axis((0, channel_per_group), name="rc")
padded = topi.nn.pad(inputs, [0, padding, padding, 0])
output = te.compute(
(batch_size, out_h, out_w, out_channel),
lambda n, h, w, co: te.sum(
(
padded[
n,
h * stride + rh * dilation,
w * stride + rw * dilation,
co // out_channel_per_group * channel_per_group + rc,
].astype(out_dtype)
* weight[rh, rw, rc, co].astype(out_dtype)
),
axis=[rh, rw, rc],
),
name="conv2d_nhwc",
)
return (inputs, weight, output)
def conv3d_ndhwc( # pylint: disable=invalid-name,missing-docstring
N: int,
D: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
groups: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder((N, D, H, W, CI), name="inputs", dtype=in_dtype)
weight = te.placeholder(
(kernel_size, kernel_size, kernel_size, CI // groups, CO), name="weight", dtype=in_dtype
)
batch_size, in_d, in_h, in_w, _ = inputs.shape
k_d, k_h, k_w, channel_per_group, out_channel = weight.shape
out_channel_per_group = out_channel // groups
out_d = (in_d + 2 * padding - dilation * (k_d - 1) - 1) // stride + 1
out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
rd = te.reduce_axis((0, k_d), name="rd")
rh = te.reduce_axis((0, k_h), name="rh")
rw = te.reduce_axis((0, k_w), name="rw")
rc = te.reduce_axis((0, channel_per_group), name="rc")
padded = topi.nn.pad(inputs, [0, padding, padding, padding, 0])
output = te.compute(
(batch_size, out_d, out_h, out_w, out_channel),
lambda n, d, h, w, co: te.sum(
(
padded[
n,
d * stride + rd * dilation,
h * stride + rh * dilation,
w * stride + rw * dilation,
co // out_channel_per_group * channel_per_group + rc,
].astype(out_dtype)
* weight[rd, rh, rw, rc, co].astype(out_dtype)
),
axis=[rd, rh, rw, rc],
),
name="conv3d_ndhwc",
)
return (inputs, weight, output)
def depthwise_conv2d_nhwc( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
C: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
factor: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder((N, H, W, C), dtype=in_dtype)
weight = te.placeholder((factor, kernel_size, kernel_size, C), dtype=in_dtype)
batch_size, in_h, in_w, in_channel = inputs.shape
factor, k_h, k_w, in_channel = weight.shape
out_channel = in_channel * factor
assert int(factor) == 1, "Not optimized for factor != 1"
out_h = (in_h + 2 * padding - dilation * (k_h - 1) - 1) // stride + 1
out_w = (in_w + 2 * padding - dilation * (k_w - 1) - 1) // stride + 1
rh = te.reduce_axis((0, k_h), name="rh")
rw = te.reduce_axis((0, k_w), name="rw")
padded = topi.nn.pad(inputs, [0, padding, padding, 0])
output = te.compute(
(batch_size, out_h, out_w, out_channel),
lambda n, h, w, c: te.sum(
(
padded[
n,
h * stride + rh * dilation,
w * stride + rw * dilation,
c // factor,
].astype(out_dtype)
* weight[c % factor, rh, rw, c // factor].astype(out_dtype)
),
axis=[rh, rw],
),
name="depth_conv2d_nhwc",
)
return (inputs, weight, output)
def conv2d_transpose_nhwc( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder((N, H, W, CI), name="inputs", dtype=in_dtype)
weight = te.placeholder((kernel_size, kernel_size, CI, CO), name="weight", dtype=in_dtype)
batch, in_h, in_w, in_c = inputs.shape
filter_h, filter_w, in_c, out_c = weight.shape
stride_h, stride_w = (stride, stride)
# compute padding
fpad_top, fpad_left, fpad_bottom, fpad_right = topi.nn.get_pad_tuple(
padding, (filter_h, filter_w)
)
bpad_top = filter_h - 1 - fpad_top
bpad_bottom = filter_h - 1 - fpad_bottom
bpad_left = filter_w - 1 - fpad_left
bpad_right = filter_w - 1 - fpad_right
# padding stage
padded = topi.nn.pad(
inputs,
[
0,
(bpad_top + stride_h - 1) // stride_h,
(bpad_left + stride_w - 1) // stride_w,
0,
],
[
0,
(bpad_bottom + stride_h - 1) // stride_h,
(bpad_right + stride_w - 1) // stride_w,
0,
],
)
# remove extra padding introduced by dilatation
idx_div = te.indexdiv
idx_mod = te.indexmod
border_h = idx_mod(stride_h - idx_mod(bpad_top, stride_h), stride_h)
border_w = idx_mod(stride_w - idx_mod(bpad_left, stride_w), stride_w)
# dilation stage
strides = [1, stride_h, stride_w, 1]
n = len(padded.shape)
# We should embed this dilation directly into te.compute rather than creating a new te.compute.
# Only in this way can we use unroll to eliminate the multiplication of zeros.
def _dilate(*indices):
not_zero = []
index_tuple = []
for i in range(n):
if not strides[i] == 1:
index_tuple.append(idx_div(indices[i], strides[i]))
not_zero.append(idx_mod(indices[i], strides[i]).equal(0))
else:
index_tuple.append(indices[i])
if not_zero:
not_zero = te.all(*not_zero)
return te.if_then_else(not_zero, padded(*index_tuple), tirx.const(0.0, padded.dtype))
return padded(*index_tuple)
# convolution stage
out_h = (in_h - 1) * stride_h - fpad_top - fpad_bottom + filter_h
out_w = (in_w - 1) * stride_w - fpad_left - fpad_right + filter_w
rc = te.reduce_axis((0, in_c), name="rc")
rh = te.reduce_axis((0, filter_h), name="rh")
rw = te.reduce_axis((0, filter_w), name="rw")
output = te.compute(
(batch, out_h, out_w, out_c),
lambda n, h, w, co: te.sum(
_dilate(n, h + rh + border_h, w + rw + border_w, rc).astype(out_dtype)
* weight[filter_h - 1 - rh, filter_w - 1 - rw, rc, co].astype(out_dtype),
axis=[rh, rw, rc],
),
name="conv2d_transpose_nhwc",
)
return (inputs, weight, output)
def conv2d_capsule_nhwijc( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
capsule_size: int = 4,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
inputs = te.placeholder(
(N, H, W, capsule_size, capsule_size, CI), name="inputs", dtype=in_dtype
)
weight = te.placeholder(
(kernel_size, kernel_size, capsule_size, capsule_size, CI, CO),
name="weight",
dtype=in_dtype,
)
batch_size, in_h, in_w, _, _, in_channel = inputs.shape
k_h, k_w, _, _, _, out_channel = weight.shape
out_h = (in_h + 2 * padding - kernel_size) // stride + 1
out_w = (in_w + 2 * padding - kernel_size) // stride + 1
rh = te.reduce_axis((0, k_h), name="rh")
rw = te.reduce_axis((0, k_w), name="rw")
cap_k = te.reduce_axis((0, capsule_size), name="cap_k")
rc = te.reduce_axis((0, in_channel), name="rc")
padded = topi.nn.pad(inputs, [0, padding, padding, 0, 0, 0])
output = te.compute(
(batch_size, out_h, out_w, capsule_size, capsule_size, out_channel),
lambda n, h, w, cap_i, cap_j, co: te.sum(
(
padded[n, h * stride + rh, w * stride + rw, cap_i, cap_k, rc].astype(out_dtype)
* weight[rh, rw, cap_k, cap_j, rc, co].astype(out_dtype)
),
axis=[rh, rw, cap_k, rc],
),
name="conv2d_capsule_nhwijc",
)
return (inputs, weight, output)
def norm_bmn( # pylint: disable=invalid-name,missing-docstring
B: int,
M: int,
N: int,
) -> tuple[te.Tensor, te.Tensor]:
a = te.placeholder((B, M, N), name="A")
i = te.reduce_axis((0, M), name="i")
j = te.reduce_axis((0, N), name="j")
c = te.compute(
(B,),
lambda b: te.sum(a[b][i][j] * a[b][i][j], axis=[i, j]),
name="C",
)
d = te.compute((B,), lambda b: te.sqrt(c[b]), name="D")
return (a, d)
def conv2d_nhwc_without_layout_rewrite( # pylint: disable=invalid-name
Input: te.Tensor,
Filter: te.Tensor,
stride: int,
padding: int,
dilation: int,
out_dtype="float32",
):
"""A copy of `topi.nn.conv2d_nhwc` but without the 'layout_free` attribute.
We use this in single op and subgraph evaluation
because we don't want to introduce graph level optimization.
"""
assert isinstance(stride, int) or len(stride) == 2
assert isinstance(dilation, int) or len(dilation) == 2
if isinstance(stride, int):
stride_h = stride_w = stride
else:
stride_h, stride_w = stride
if isinstance(dilation, int):
dilation_h = dilation_w = dilation
else:
dilation_h, dilation_w = dilation
batch, in_height, in_width, in_channel = Input.shape # type: ignore
kernel_h, kernel_w, _channel, num_filter = Filter.shape # type: ignore
# compute the output shape
dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
pad_top, pad_left, pad_down, pad_right = topi.nn.get_pad_tuple(
padding, (dilated_kernel_h, dilated_kernel_w)
)
out_channel = num_filter
out_height = topi.utils.simplify(
(in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1
)
out_width = topi.utils.simplify(
(in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1
)
pad_before = [0, pad_top, pad_left, 0]
pad_after = [0, pad_down, pad_right, 0]
PaddedInput = topi.nn.pad(Input, pad_before, pad_after, name="PaddedInput")
rc = te.reduce_axis((0, in_channel), name="rc")
ry = te.reduce_axis((0, kernel_h), name="ry")
rx = te.reduce_axis((0, kernel_w), name="rx")
Output = te.compute(
(batch, out_height, out_width, out_channel),
lambda nn, yy, xx, ff: te.sum(
PaddedInput[
nn, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rc
].astype(out_dtype)
* Filter[ry, rx, rc, ff].astype(out_dtype), # type: ignore
axis=[ry, rx, rc],
),
name="Conv2dOutput",
tag="conv2d_nhwc",
)
return Output
def conv2d_nhwc_bn_relu( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
strides: int,
padding: int,
dilation: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor]:
data = te.placeholder((N, H, W, CI), name="data", dtype=in_dtype)
kernel = te.placeholder((kernel_size, kernel_size, CI, CO), name="kernel", dtype=in_dtype)
bias = te.placeholder((CO,), name="bias")
bn_scale = te.placeholder((CO,), name="bn_scale")
bn_offset = te.placeholder((CO,), name="bn_offset")
OH = (H + 2 * padding - (kernel_size - 1) * dilation - 1) // strides + 1
OW = (W + 2 * padding - (kernel_size - 1) * dilation - 1) // strides + 1
conv = conv2d_nhwc_without_layout_rewrite(data, kernel, strides, padding, dilation, out_dtype)
conv = te.compute(
(N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] + bias[l], name="bias_add"
)
conv = te.compute(
(N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] * bn_scale[l], name="bn_mul"
)
conv = te.compute(
(N, OH, OW, CO), lambda i, j, k, l: conv[i, j, k, l] + bn_offset[l], name="bn_add"
)
out = topi.nn.relu(conv)
return (data, kernel, bias, bn_offset, bn_scale, out)
def transpose_batch_matmul( # pylint: disable=invalid-name,missing-docstring
batch: int,
seq_len: int,
n_head: int,
n_dim: int,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
query = te.placeholder((batch, seq_len, n_head, n_dim), name="query", dtype=in_dtype)
value = te.placeholder((batch, seq_len, n_head, n_dim), name="value", dtype=in_dtype)
query_T = te.compute(
(batch, n_head, seq_len, n_dim),
lambda b, h, l, d: query[b, l, h, d],
name="query_T",
)
value_T = te.compute(
(batch, n_head, n_dim, seq_len),
lambda b, h, d, l: value[b, l, h, d],
name="value_T",
)
k = te.reduce_axis((0, n_dim), name="k")
out = te.compute(
(batch, n_head, seq_len, seq_len),
lambda b, h, i, j: te.sum(
query_T[b, h, i, k].astype(out_dtype) * value_T[b, h, k, j].astype(out_dtype), axis=[k]
),
name="C",
)
return (query, value, out)
def conv2d_winograd_nhwc( # pylint: disable=invalid-name,missing-docstring
N: int,
H: int,
W: int,
CI: int,
CO: int,
kernel_size: int,
stride: int = 1,
padding: int = 0,
dilation: int = 1,
tile_size: int = 4,
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
from tvm.topi.nn.conv2d import ( # pylint: disable=import-outside-toplevel
_conv2d_winograd_nhwc_impl,
)
target = Target.current(allow_none=True)
if target is not None and target.kind.name == "cuda":
write_cache_level = 3
else:
write_cache_level = 2
data = te.placeholder((N, H, W, CI), "float32", name="data")
weight = te.placeholder((kernel_size, kernel_size, CO, CI), "float32", name="weight")
out = _conv2d_winograd_nhwc_impl(
data,
weight,
stride,
padding,
dilation,
"float32",
pre_computed=True,
auto_scheduler_rewritten_layout="",
meta_schedule_original_shape=None,
tile_size=tile_size,
write_cache_level=write_cache_level,
)
return (data, weight, out)
def matmul(
n: int, m: int, k: int, in_dtype: str = "float32", out_dtype: str = "float32"
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
a = te.placeholder((n, k), name="A", dtype=in_dtype)
b = te.placeholder((k, m), name="B", dtype=in_dtype)
k = te.reduce_axis((0, k), name="k")
c = te.compute(
(n, m),
lambda i, j: te.sum(a[i, k].astype(out_dtype) * b[k, j].astype(out_dtype), axis=[k]),
name="C",
)
return (a, b, c)
def matmul_relu(
n: int, m: int, k: int, in_dtype: str = "float32", out_dtype: str = "float32"
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
a = te.placeholder((n, k), name="A", dtype=in_dtype)
b = te.placeholder((k, m), name="B", dtype=in_dtype)
k = te.reduce_axis((0, k), name="k")
c = te.compute(
(n, m),
lambda i, j: te.sum(a[i, k].astype(out_dtype) * b[k, j].astype(out_dtype), axis=[k]),
name="C",
)
d = topi.nn.relu(c) # pylint: disable=invalid-name
return (a, b, d)
def conv2d_nchw( # pylint: disable=invalid-name
n: int,
h: int,
w: int,
ci: int,
co: int,
kh: int,
kw: int,
stride: int,
padding: int,
dilation: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor]:
x = te.placeholder((n, ci, h, w), name="X", dtype=in_dtype)
w = te.placeholder((co, ci, kh, kw), name="W", dtype=in_dtype)
y = topi.nn.conv2d_nchw(
Input=x, Filter=w, stride=stride, padding=padding, dilation=dilation, out_dtype=out_dtype
)
return (x, w, y)
def conv2d_nchw_bias_bn_relu( # pylint: disable=invalid-name
n: int,
h: int,
w: int,
ci: int,
co: int,
kh: int,
kw: int,
stride: int,
padding: int,
dilation: int = 1,
in_dtype: str = "float32",
out_dtype: str = "float32",
) -> tuple[te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor, te.Tensor]:
oh = (h + 2 * padding - (kh - 1) * dilation - 1) // stride + 1 # pylint: disable=invalid-name
ow = (w + 2 * padding - (kw - 1) * dilation - 1) // stride + 1 # pylint: disable=invalid-name
x = te.placeholder((n, ci, h, w), name="X", dtype=in_dtype)
w = te.placeholder((co, ci, kh, kw), name="W", dtype=in_dtype)
b = te.placeholder((co, 1, 1), name="B", dtype=out_dtype)
bn_scale = te.placeholder((co, 1, 1), name="bn_scale", dtype=out_dtype)
bn_offset = te.placeholder((co, 1, 1), name="bn_offset", dtype=out_dtype)
y = topi.nn.conv2d_nchw(
Input=x, Filter=w, stride=stride, padding=padding, dilation=dilation, out_dtype=out_dtype
)
y = te.compute((n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] + b[j, 0, 0], name="bias_add")
y = te.compute(
(n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] * bn_scale[j, 0, 0], name="bn_mul"
)
y = te.compute(
(n, co, oh, ow), lambda i, j, k, l: y[i, j, k, l] + bn_offset[j, 0, 0], name="bn_add"
)
y = topi.nn.relu(y)
return (x, w, b, bn_scale, bn_offset, y)
def max_pool2d_nchw( # pylint: disable=invalid-name
n: int,
h: int,
w: int,
ci: int,
padding: int,
) -> tuple[te.Tensor, te.Tensor]: # pylint: disable=invalid-name
x = te.placeholder((n, ci, h, w), name="X")
y = topi.nn.pool2d(x, [2, 2], [1, 1], [1, 1], [padding, padding, padding, padding], "max")
return (x, y)
def softmax_mn(m, n) -> tuple[te.Tensor, te.Tensor]: # pylint: disable=invalid-name
a = te.placeholder((m, n), name="A")
b = topi.nn.softmax(a, axis=1)
return (a, b)
def create_te_workload(name: str, idx: int) -> tirx.PrimFunc:
workload_func, params = CONFIGS[name]
return te.create_prim_func(workload_func(*params[idx])) # type: ignore
CONFIGS = {
"C1D": (
conv1d_nlc,
[
# derived from conv2d_shapes
(1, 256, 64, 128, 3, 2, 1),
# (1, 256, 64, 128, 1, 2, 0),
# (1, 256, 64, 64, 1, 1, 0),
# (1, 128, 128, 256, 3, 2, 1),
(1, 128, 128, 256, 1, 2, 0),
# (1, 128, 128, 128, 3, 1, 1),
# (1, 64, 256, 512, 3, 2, 1),
# (1, 64, 256, 512, 1, 2, 0),
(1, 64, 256, 256, 5, 1, 2),
(1, 32, 512, 512, 3, 1, 1),
],
),
"C2D": (
conv2d_nhwc,
[
# all conv2d layers in resnet-18
(1, 224, 224, 3, 64, 7, 2, 3),
# (1, 56, 56, 64, 128, 3, 2, 1),
# (1, 56, 56, 64, 128, 1, 2, 0),
# (1, 56, 56, 64, 64, 3, 1, 1),
(1, 56, 56, 64, 64, 1, 1, 0),
# (1, 28, 28, 128, 256, 3, 2, 1),
# (1, 28, 28, 128, 256, 1, 2, 0),
# (1, 28, 28, 128, 128, 3, 1, 1),
# (1, 14, 14, 256, 512, 3, 2, 1),
# (1, 14, 14, 256, 512, 1, 2, 0),
(1, 14, 14, 256, 256, 3, 1, 1),
(1, 7, 7, 512, 512, 3, 1, 1),
],
),
"C3D": (
conv3d_ndhwc,
[
# Derived from conv2d_shapes. Use depth=16 for all configurations
(1, 16, 224, 224, 3, 64, 7, 2, 3),
# (1, 16, 56, 56, 64, 128, 3, 2, 1),
# (1, 16, 56, 56, 64, 128, 1, 2, 0),
# (1, 16, 56, 56, 64, 64, 3, 1, 1),
(1, 16, 56, 56, 64, 64, 1, 1, 0),
# (1, 16, 28, 28, 128, 256, 3, 2, 1),
# (1, 16, 28, 28, 128, 256, 1, 2, 0),
# (1, 16, 28, 28, 128, 128, 3, 1, 1),
# (1, 16, 14, 14, 256, 512, 3, 2, 1),
# (1, 16, 14, 14, 256, 512, 1, 2, 0),
(1, 16, 14, 14, 256, 256, 3, 1, 1),
(1, 16, 7, 7, 512, 512, 3, 1, 1),
],
),
"GMM": (
batch_matmul_nkkm,
[
(1, 128, 128, 128),
(1, 512, 32, 512),
(1, 512, 512, 512),
(1, 1024, 1024, 1024),
],
),
"GRP": (
conv2d_nhwc,
[
# Derived from conv2d_shapes. Use group=4 for all configurations
(1, 56, 56, 64, 128, 3, 2, 1, 1, 4),
# (1, 56, 56, 64, 128, 1, 2, 0 , 1, 4),
# (1, 56, 56, 64, 64, 3, 1, 1 , 1, 4),
(1, 56, 56, 64, 64, 1, 1, 0, 1, 4),
# (1, 28, 28, 128, 256, 3, 2, 1, 1, 4),
# (1, 28, 28, 128, 256, 1, 2, 0, 1, 4),
# (1, 28, 28, 128, 128, 3, 1, 1, 1, 4),
# (1, 14, 14, 256, 512, 3, 2, 1, 1, 4),
# (1, 14, 14, 256, 512, 1, 2, 0, 1, 4),
(1, 14, 14, 256, 256, 3, 1, 1, 1, 4),
(1, 7, 7, 512, 512, 3, 1, 1, 1, 4),
],
),
"DIL": (
conv2d_nhwc,
[
# Derived from conv2d_shapes. Use dilation=2 for all configurations
(1, 224, 224, 3, 64, 7, 2, 3, 2),
# (1, 56, 56, 64, 128, 3, 2, 1 , 2),
# (1, 56, 56, 64, 128, 1, 2, 0 , 2),
# (1, 56, 56, 64, 64, 3, 1, 1 , 2),
(1, 56, 56, 64, 64, 1, 1, 0, 2),
# (1, 28, 28, 128, 256, 3, 2, 1, 2),
# (1, 28, 28, 128, 256, 1, 2, 0, 2),
# (1, 28, 28, 128, 128, 3, 1, 1, 2),
# (1, 14, 14, 256, 512, 3, 2, 1, 2),
# (1, 14, 14, 256, 512, 1, 2, 0, 2),
(1, 14, 14, 256, 256, 3, 1, 1, 2),
(1, 7, 7, 512, 512, 3, 1, 1, 2),
],
),
"DEP": (
depthwise_conv2d_nhwc,
[
# all depthwise conv2d layers in mobilenet
(1, 112, 112, 32, 3, 1, 1),
(1, 112, 112, 64, 3, 2, 1),
# (1, 56, 56, 128, 3, 1, 1),
# (1, 56, 56, 128, 3, 2, 1),
# (1, 28, 28, 256, 3, 1, 1),
# (1, 28, 28, 256, 3, 2, 1),
# (1, 14, 14, 512, 3, 1, 1),
(1, 14, 14, 512, 3, 2, 1),
(1, 7, 7, 1024, 3, 1, 1),
],
),
"T2D": (
conv2d_transpose_nhwc,
[
# all conv2d transpose layers in DCGAN
(1, 4, 4, 512, 256, 4, 2, 1),
(1, 8, 8, 256, 128, 4, 2, 1),
(1, 16, 16, 128, 64, 4, 2, 1),
(1, 32, 32, 64, 3, 4, 2, 1),
],
),
"CAP": (
conv2d_capsule_nhwijc,
[
# all conv2d capsule layers in matrix capsules withemrouting (ICLR 2018)
(1, 16, 16, 32, 32, 3, 2, 1),
(1, 8, 8, 32, 32, 3, 1, 1),
(1, 16, 16, 8, 16, 3, 2, 1),
(1, 8, 8, 16, 16, 3, 1, 1),
],
),
"NRM": (
norm_bmn,
[
(1, 256, 256),
(1, 512, 512),
(1, 1024, 1024),
(1, 4096, 1024),
],
),
"SFM": (
softmax_mn,
[
(256, 256),
(512, 512),
(1024, 1024),
(2048, 2048),
],
),
"CBR": (
conv2d_nhwc_bn_relu,
[
(1, 224, 224, 3, 64, 7, 2, 3),
(1, 56, 56, 64, 128, 3, 2, 1),
(1, 28, 28, 128, 256, 1, 2, 0),
(1, 7, 7, 512, 512, 3, 1, 1),
],
),
"TBG": (
transpose_batch_matmul,
[
(1, 128, 12, 64),
(1, 128, 16, 64),
(1, 64, 12, 128),
(1, 128, 12, 128),
],
),
"C2D_WIN_NHWC": (
conv2d_winograd_nhwc,
[
(1, 14, 14, 128, 128, 6),
],
),
}
python/tvm/s_tir/meta_schedule/testing/tune_te.py
+150
View File
@@ -0,0 +1,150 @@
# 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=missing-docstring
# ruff: noqa: F821
import argparse
import logging
import tvm
from tvm.s_tir import meta_schedule as ms
from tvm.s_tir.meta_schedule.testing.te_workload import create_te_workload
from tvm.support import describe
from tvm.testing.utils import strtobool
def _parse_args():
args = argparse.ArgumentParser()
args.add_argument(
"--workload",
type=str,
required=True,
)
args.add_argument(
"--target",
type=str,
required=True,
)
args.add_argument(
"--num-trials",
type=int,
required=True,
)
args.add_argument(
"--rpc-host",
type=str,
required=True,
)
args.add_argument(
"--rpc-port",
type=int,
required=True,
)
args.add_argument(
"--rpc-key",
type=str,
required=True,
)
args.add_argument(
"--work-dir",
type=str,
required=True,
)
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(
"--adaptive-training",
type=lambda x: bool(strtobool(x)),
required=False,
help="example: True / False",
default=True,
)
args.add_argument(
"--cpu-flush",
type=lambda x: bool(strtobool(x)),
help="example: True / False",
required=True,
)
parsed = args.parse_args()
parsed.target = tvm.target.Target(parsed.target)
parsed.rpc_config = ms.runner.RPCConfig(
tracker_host=parsed.rpc_host,
tracker_port=parsed.rpc_port,
tracker_key=parsed.rpc_key,
session_timeout_sec=60,
)
return parsed
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)
ARGS = _parse_args()
def main():
describe()
print(f"Workload: {ARGS.workload}")
with ms.Profiler() as profiler:
sch: s_tir.Schedule | None = ms.tir_integration.tune_tir(
mod=create_te_workload(ARGS.workload, 0),
target=ARGS.target,
work_dir=ARGS.work_dir,
max_trials_global=ARGS.num_trials,
num_trials_per_iter=64,
runner=ms.runner.RPCRunner( # type: ignore
rpc_config=ARGS.rpc_config,
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,
),
alloc_repeat=1,
),
cost_model=ms.cost_model.XGBModel( # type: ignore
extractor=ms.feature_extractor.PerStoreFeature(),
adaptive_training=ARGS.adaptive_training,
),
strategy=ms.search_strategy.EvolutionarySearch(),
)
print("Tuning Time:")
print(profiler.table())
if sch is None:
print("No valid schedule found!")
else:
print(sch.mod.script())
print(sch.trace)
if __name__ == "__main__":
main()
python/tvm/s_tir/meta_schedule/testing/tune_utils.py
+113
View File
@@ -0,0 +1,113 @@
# 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
python/tvm/s_tir/meta_schedule/testing/validate_database.py
+784
View File
@@ -0,0 +1,784 @@
# 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()