def main():
alloc_repeat = 1
runner = ms.runner.RPCRunner(
rpc_config=ARGS.rpc_config,
evaluator_config=ms.runner.EvaluatorConfig(
number=3,
repeat=1,
min_repeat_ms=100,
enable_cpu_cache_flush=False,
),
alloc_repeat=alloc_repeat,
max_workers=ARGS.rpc_workers,
)
sch: Optional[tir.Schedule] = ms.tune_tir(
mod=create_te_workload(ARGS.workload, 0),
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
),
runner=runner, # type: ignore
task_name=ARGS.workload,
work_dir=ARGS.work_dir,
num_threads=cpu_count(),
)
if sch is None:
print("No valid schedule found!")
else:
print(sch.mod.script())
print(sch.trace)
def main():
describe()
print(f"Workload: {ARGS.workload}")
runner = ms.runner.RPCRunner(
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,
)
with ms.Profiler() as profiler:
sch: Optional[tir.Schedule] = ms.tune_tir(
mod=create_te_workload(ARGS.workload, 0),
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
),
runner=runner, # type: ignore
task_name=ARGS.workload,
work_dir=ARGS.work_dir,
num_threads=cpu_count(),
)
print("Tuning Time:")
print(profiler.table())
if sch is None:
print("No valid schedule found!")
else:
print(sch.mod.script())
print(sch.trace)
def main():
describe()
print(f"Workload: {ARGS.workload}")
mod, params, (input_name, input_shape, input_dtype) = get_network(
ARGS.workload,
ARGS.input_shape,
cache_dir=ARGS.cache_dir,
)
input_info = {input_name: input_shape}
input_data = {
item["name"]: generate_input_data(item["shape"], item["dtype"])
for item in ARGS.input_shape
}
for input_name, input_shape in input_info.items():
print(f" input_name : {input_name}")
print(f" input_shape: {input_shape}")
print(f" input_dtype: {input_dtype}")
runner = ms.runner.RPCRunner(
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,
)
with ms.Profiler() as profiler:
lib = ms.tune_relay(
mod=mod,
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
adaptive_training=ARGS.adaptive_training,
),
runner=runner, # type: ignore
work_dir=ARGS.work_dir,
params=params,
backend=ARGS.backend,
)
print("Tuning Time:")
print(profiler.table())
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=lib,
dev_type=ARGS.target.kind.name,
args=input_data,
continuation=create_timer(ARGS.backend),
backend=ARGS.backend,
)
def test_cuda_tensor_core(model_name, input_shape):
"""Integration tests of auto tensorization with CUDA tensor core"""
target = tvm.target.Target("nvidia/geforce-rtx-3070")
dev = tvm.cuda()
if model_name.startswith("bert"):
data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape), dev) # embedding size
else:
data = tvm.nd.array(np.random.randn(*input_shape).astype("float32"), dev)
mod, params, (input_name, _, _) = relay_workload.get_network(model_name, input_shape)
seq = tvm.transform.Sequential(
[
relay.transform.ToMixedPrecision(),
]
)
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
def convert_layout(mod):
seq = tvm.transform.Sequential(
[relay.transform.ConvertLayout({"nn.conv2d": ["NHWC", "OHWI"]})]
)
with tvm.transform.PassContext(opt_level=3):
mod = seq(mod)
return mod
with tempfile.TemporaryDirectory() as work_dir:
with ms.Profiler() as profiler:
rt_mod1: tvm.runtime.Module = ms.tune_relay(
mod=convert_layout(mod),
params=params,
target=target,
config=ms.TuneConfig(
num_trials_per_iter=32,
max_trials_per_task=200,
max_trials_global=3000,
),
sch_rules=ms.default_config._DefaultCUDATensorCore.schedule_rules,
postprocs=ms.default_config._DefaultCUDATensorCore.postprocs,
work_dir=work_dir,
)
print(profiler.table())
# Compile without MetaSchedule for correctness check
with tvm.transform.PassContext(opt_level=0):
rt_mod2 = relay.build(mod, target=target, params=params)
def get_output(data, lib):
module = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
module.set_input(input_name, data)
module.run()
return module.get_output(0).numpy()
# Check correctness
actual_output = get_output(data, rt_mod1)
expected_output = get_output(data, rt_mod2)
assert np.allclose(actual_output, expected_output, rtol=1e-2, atol=2e-2)
def main():
describe()
print(f"Workload: {ARGS.model_name}")
onnx_model = onnx.load(ARGS.onnx_path)
shape_dict = {}
for item in ARGS.input_shape:
print(f" input_name : {item['name']}")
print(f" input_shape: {item['shape']}")
print(f" input_dtype: {item['dtype']}")
shape_dict[item["name"]] = item["shape"]
mod, params = from_onnx(onnx_model, shape_dict, freeze_params=True)
input_data = {
item["name"]: generate_input_data(item["shape"], item["dtype"]) for item in ARGS.input_shape
}
runner = ms.runner.RPCRunner(
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,
)
with ms.Profiler() as profiler:
lib = ms.tune_relay(
mod=mod,
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
adaptive_training=ARGS.adaptive_training,
),
runner=runner, # type: ignore
work_dir=ARGS.work_dir,
params=params,
backend=ARGS.backend,
)
print("Tuning Time:")
print(profiler.table())
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=lib,
dev_type=ARGS.target.kind.name,
args=input_data,
continuation=create_timer(ARGS.backend),
backend=ARGS.backend,
)
def test_meta_schedule_tune_relay(
model_name: str,
input_shape: List[int],
target: str,
):
dev = tvm.cpu() if str(target).startswith("llvm") else tvm.cuda()
if model_name.startswith("bert"):
data = tvm.nd.array(np.random.randint(0, 30521, size=input_shape),
dev) # embedding size
else:
data = tvm.nd.array(
np.random.randn(*input_shape).astype("float32"), dev)
mod, params, (input_name, _, _) = get_network(name=model_name,
input_shape=input_shape)
target = Target(target)
with tempfile.TemporaryDirectory() as work_dir:
with ms.Profiler() as profiler:
rt_mod1: tvm.runtime.Module = ms.tune_relay(
mod=mod,
params=params,
target=target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=32,
max_trials_per_task=20000,
max_trials_global=20000,
),
work_dir=work_dir,
)
print(profiler.table())
# Compile without meta-schedule for correctness check
with tvm.transform.PassContext(opt_level=0):
rt_mod2 = relay.build(mod, target=target, params=params)
def get_output(data, lib):
module = graph_executor.GraphModule(lib["default"](dev))
module.set_input(input_name, data)
module.run()
return module.get_output(0).numpy()
# Check correctness
actual_output = get_output(data, rt_mod1)
expected_output = get_output(data, rt_mod2)
assert np.allclose(actual_output,
expected_output,
rtol=1e-4,
atol=2e-4)
def main():
mod, params, (input_name, input_shape, input_dtype) = get_network(
ARGS.workload,
ARGS.input_shape,
cache_dir=ARGS.cache_dir,
)
print(f"Workload: {ARGS.workload}")
print(f" input_name: {input_name}")
print(f" input_shape: {input_shape}")
print(f" input_dtype: {input_dtype}")
alloc_repeat = 1
runner = ms.runner.RPCRunner(
rpc_config=ARGS.rpc_config,
evaluator_config=ms.runner.EvaluatorConfig(
number=3,
repeat=1,
min_repeat_ms=100,
enable_cpu_cache_flush=False,
),
alloc_repeat=alloc_repeat,
max_workers=ARGS.rpc_workers,
)
lib = ms.tune_relay(
mod=mod,
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
),
runner=runner, # type: ignore
work_dir=ARGS.work_dir,
params=params,
)
graph, rt_mod, params = lib.graph_json, lib.lib, lib.params
if input_dtype.startswith("float"):
input_data = np.random.uniform(size=input_shape).astype(input_dtype)
else:
input_data = np.random.randint(low=0, high=10000, size=input_shape, dtype=input_dtype)
def f_timer(rt_mod, dev, input_data):
# pylint: disable=import-outside-toplevel
from tvm.contrib.graph_executor import GraphModule
# pylint: enable=import-outside-toplevel
mod = GraphModule(rt_mod["default"](dev))
mod.set_input(input_name, input_data)
ftimer = mod.module.time_evaluator(
"run",
dev,
min_repeat_ms=500,
repeat=3,
)
results = list(np.array(ftimer().results) * 1000.0) # type: ignore
print("Running time in time_evaluator: ", results)
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=lib,
dev_type=ARGS.target.kind.name,
args=[input_data],
continuation=f_timer,
)
def f_per_layer(rt_mod, dev, input_data):
# pylint: disable=import-outside-toplevel
from tvm.contrib.debugger.debug_executor import create
# pylint: enable=import-outside-toplevel
mod = create(graph, rt_mod, dev)
mod.set_input(input_name, input_data)
graph_nodes = [n["name"] for n in json.loads(graph)["nodes"]]
graph_time = mod.run_individual(number=10, repeat=1, min_repeat_ms=5000)
print("|graph_nodes| = ", len(graph_nodes))
print("|graph_time| = ", len(graph_time))
graph_nodes_time = {k: float(v) for k, v in zip(graph_nodes, graph_time)}
for k, v in graph_nodes_time.items():
print(f"{k} : {v:.3f}")
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=rt_mod,
dev_type=ARGS.target.kind.name,
args=[input_data],
continuation=f_per_layer,
)
from tvm.meta_schedule.relay_integration import extract_task_from_relay
from tvm.meta_schedule import ApplyHistoryBest
from tvm.meta_schedule import schedule_rule, postproc
from tvm.meta_schedule.testing.tlcbench import load_quantized_bert_base
from tvm import meta_schedule as ms
from tvm.tir.tensor_intrin import (
VNNI_DOT_16x4_INTRIN as VNNI_INTRIN,
DP4A_INTRIN,
AMDGPU_SDOT4_INTRIN,
)
import tempfile
import tvm.topi.testing
config = ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=32,
max_trials_per_task=32,
max_trials_global=20000,
)
sch_rules_for_vnni = [
schedule_rule.AutoInline(
into_producer=False,
into_consumer=True,
inline_const_tensor=True,
disallow_if_then_else=True,
require_injective=True,
require_ordered=True,
disallow_op=["tir.exp"],
),
schedule_rule.AddRFactor(max_jobs_per_core=16, max_innermost_factor=64),
schedule_rule.MultiLevelTilingWithIntrin(
def manual_tir_common(do_tune=False):
M, N, K = 1024, 1024, 1024 # pylint: disable=invalid-name
data_shape = (M, K)
weight_shape = (N, K)
data_dtype = "uint8"
data = relay.var("data", shape=data_shape, dtype=data_dtype)
weight = relay.var("weight", shape=weight_shape, dtype="int8")
bias = relay.var("bias", shape=(weight_shape[0], ), dtype="int32")
# dense is tuned by the TIR schedule above, bmm is scheduled by TE (topi/x86/batch_matmul.py)
dense = relay.nn.dense(data, weight, out_dtype="int32")
bias_add = relay.nn.bias_add(dense, bias) + relay.const(1, dtype="int32")
out = relay.nn.batch_matmul(
relay.cast(relay.expand_dims(bias_add, 0), "uint8"),
relay.cast(relay.expand_dims(bias_add, 0), "int8"),
out_dtype="int32",
)
relay_mod = tvm.IRModule.from_expr(out)
target = "llvm -mcpu=cascadelake -num-cores 4"
dev = tvm.device(target, 0)
data = np.random.uniform(1, 10, size=(M, K)).astype("uint8")
weight_np = np.random.uniform(1, 10, size=weight_shape).astype("int8")
bias_np = np.random.uniform(1, 10,
size=(weight_shape[0], )).astype("int32")
ref = (relay.create_executor(
"vm", mod=relay_mod, device=dev,
target=target).evaluate()(*[data, weight_np, bias_np]).numpy())
params = {"weight": weight_np, "bias": bias_np}
if do_tune:
extracted_tasks = ms.extract_task_from_relay(relay_mod, target, params)
# Filter out tasks that we don't intend to schedule / tune with TIR.
tune_tasks = list(
filter(
lambda task: "dense" in task.task_name,
extracted_tasks,
))
config = ms.TuneConfig(
strategy="replay_trace",
num_trials_per_iter=64,
max_trials_per_task=20000,
max_trials_global=20000,
)
with tempfile.TemporaryDirectory() as work_dir:
# postprocs=lambda: [] is important to prevent default post processors from
# tampering with the manual schedule.
database = ms.tune_extracted_tasks(
tune_tasks,
config,
work_dir=work_dir,
postprocs=lambda: [],
)
else:
def schedule_fn(task, sch):
if "dense" not in task.task_name:
return False
block = sch.get_block("compute")
# Looks up schedule_rule annotation.
# See the comment in test_tune_relay_manual_tir_vnni().
schedule_rule = sch.get(block).annotations["schedule_rule"]
assert "dense_vnni" in schedule_rule
schedule_dense(block, M, False, sch)
return True
database = apply_fixed_schedules(relay_mod, target, params,
schedule_fn)
with ms.ApplyHistoryBest(database):
with tvm.transform.PassContext(
opt_level=3,
config={"relay.backend.use_meta_schedule": True},
):
# pylint: disable=W0105
"""
The log should say
Warning: Cannot find workload: tvmgen_default_fused_expand_dims
Warning: Cannot find workload: tvmgen_default_fused_cast
Warning: Cannot find workload: tvmgen_default_fused_cast_1
Warning: Cannot find workload: tvmgen_default_fused_nn_batch_matmul
This means batch matmul and others are scheduled by TE, and dense (the one not warned)
is found in the meta schedule tuning database during ApplyHistoryBest
"""
# pylint: enable=W0105
lib = relay.build(relay_mod, target=target, params=params)
runtime = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
runtime.set_input("data", data)
runtime.run()
out = runtime.get_output(0).numpy()
np.testing.assert_equal(out, ref)
def main():
describe()
print(f"Workload: {ARGS.model_name}")
onnx_model = onnx.load(ARGS.onnx_path)
shape_dict = {}
for item in ARGS.input_shape:
print(f" input_name: {item['name']}")
print(f" input_shape: {item['shape']}")
print(f" input_dtype: {item['dtype']}")
shape_dict[item["name"]] = item["shape"]
mod, params = from_onnx(onnx_model, shape_dict, freeze_params=True)
runner = ms.runner.RPCRunner(
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,
)
with ms.Profiler() as profiler:
lib = ms.tune_relay(
mod=mod,
target=ARGS.target,
config=ms.TuneConfig(
strategy="evolutionary",
num_trials_per_iter=64,
max_trials_per_task=ARGS.num_trials,
max_trials_global=ARGS.num_trials,
),
runner=runner, # type: ignore
work_dir=ARGS.work_dir,
params=params,
)
print("Tuning Time:")
print(profiler.table())
graph, rt_mod, params = lib.graph_json, lib.lib, lib.params
input_data = {}
for item in ARGS.input_shape:
input_name, input_shape, input_dtype = item["name"], item[
"shape"], item["dtype"]
if input_dtype.startswith("float"):
input_data[input_name] = np.random.uniform(
size=input_shape).astype(input_dtype)
else:
input_data[input_name] = np.random.randint(low=0,
high=10000,
size=input_shape,
dtype=input_dtype)
def f_timer(rt_mod, dev, input_data):
# pylint: disable=import-outside-toplevel
from tvm.contrib.graph_executor import GraphModule
# pylint: enable=import-outside-toplevel
mod = GraphModule(rt_mod["default"](dev))
for input_name, input_value in input_data.items():
mod.set_input(input_name, input_value)
ftimer = mod.module.time_evaluator(
"run",
dev,
min_repeat_ms=500,
repeat=3,
)
results = list(np.array(ftimer().results) * 1000.0) # type: ignore
print("Running time in time_evaluator: ", results)
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=lib,
dev_type=ARGS.target.kind.name,
args=input_data,
continuation=f_timer,
)
def f_per_layer(rt_mod, dev, input_data):
# pylint: disable=import-outside-toplevel
from tvm.contrib.debugger.debug_executor import create
# pylint: enable=import-outside-toplevel
mod = create(graph, rt_mod, dev)
for input_name, input_value in input_data.items():
mod.set_input(input_name, input_value)
graph_nodes = [n["name"] for n in json.loads(graph)["nodes"]]
graph_time = mod.run_individual(number=10,
repeat=1,
min_repeat_ms=5000)
print("|graph_nodes| = ", len(graph_nodes))
print("|graph_time| = ", len(graph_time))
graph_nodes_time = {
k: float(v)
for k, v in zip(graph_nodes, graph_time)
}
for k, v in graph_nodes_time.items():
print(f"{k} : {v:.3f}")
run_module_via_rpc(
rpc_config=ARGS.rpc_config,
lib=rt_mod,
dev_type=ARGS.target.kind.name,
args=input_data,
continuation=f_per_layer,
)
