def test_meta_schedule_database_reload():
mod: IRModule = Matmul
with tempfile.TemporaryDirectory() as tmpdir:
database = _create_tmp_database(tmpdir)
token = database.commit_workload(mod)
trace = _create_schedule(mod, _schedule_matmul).trace
records = [
TuningRecord(
trace,
[7.0, 8.0, 9.0],
token,
tvm.target.Target("llvm"),
ArgInfo.from_prim_func(func=mod["main"]), # pylint: disable=unsubscriptable-object
),
TuningRecord(
trace,
[1.0, 2.0, 3.0],
token,
tvm.target.Target("llvm"),
ArgInfo.from_prim_func(func=mod["main"]), # pylint: disable=unsubscriptable-object
),
TuningRecord(
trace,
[4.0, 5.0, 6.0],
token,
tvm.target.Target("llvm"),
ArgInfo.from_prim_func(func=mod["main"]), # pylint: disable=unsubscriptable-object
),
]
for record in records:
database.commit_tuning_record(record)
new_database = JSONDatabase( # pylint: disable=unused-variable
path_workload=database.path_workload,
path_tuning_record=database.path_tuning_record,
)
token = new_database.commit_workload(mod)
ret = new_database.get_top_k(token, 2)
assert len(ret) == 2
try:
_equal_record(ret[0], records[2])
_equal_record(ret[1], records[1])
except AssertionError:
_equal_record(ret[0], records[1])
_equal_record(ret[1], records[2])
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:
rt_mod1: tvm.runtime.Module = tune_relay(
mod=mod,
params=params,
target=target,
config=TuneConfig(
strategy="evolutionary",
num_trials_per_iter=32,
max_trials_per_task=20000,
max_trials_global=20000,
search_strategy_config={
"genetic_num_iters": 10,
},
),
work_dir=work_dir,
database=JSONDatabase(
osp.join(work_dir, "workload.json"),
osp.join(work_dir, "records.json"),
),
)
# Compile without meta-scheduler 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 test_meta_schedule_relay_lowering():
data_shape = (1, 3, 16, 16)
weight_shape = (8, 3, 5, 5)
data = relay.var("data", relay.TensorType(data_shape, "float32"))
weight = relay.var("weight", relay.TensorType(weight_shape, "float32"))
y = relay.nn.conv2d(
data,
weight,
padding=(2, 2),
kernel_size=(5, 5),
kernel_layout="OIHW",
out_dtype="float32",
)
f = relay.Function([data, weight], y)
mod = tvm.IRModule.from_expr(f)
mod = relay.transform.InferType()(mod)
data_sample = np.random.rand(*data_shape).astype("float32")
weight_sample = np.random.rand(*weight_shape).astype("float32")
params = {mod["main"].params[1].name_hint: weight_sample}
input_name = "data"
dev = tvm.cpu()
target = Target("llvm --num-cores=16")
data = tvm.nd.array(data_sample, dev)
with tempfile.TemporaryDirectory() as work_dir:
database = JSONDatabase(osp.join(work_dir, "workload.json"),
osp.join(work_dir, "records.json"))
database.commit_tuning_record(
TuningRecord(
Trace([], {}),
[0.0],
database.commit_workload(
tvmgen_default_fused_nn_contrib_conv2d_NCHWc),
target=target,
args_info=[],
))
with ApplyHistoryBest(database):
with tvm.transform.PassContext(
opt_level=3,
config={"relay.backend.use_meta_schedule": True},
):
rt_mod1 = relay.build(mod, target=target, params=params)
# Compile without meta-scheduler 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 _create_tmp_database(tmpdir: str) -> JSONDatabase:
path_workload = osp.join(tmpdir, "workloads.json")
path_tuning_record = osp.join(tmpdir, "tuning_records.json")
return JSONDatabase(path_workload, path_tuning_record)