def search_common(workload=matmul_auto_scheduler_test,
target="llvm",
search_policy='empty',
seed=random.randint(1, 1 << 30),
runner='local',
cost_model=auto_scheduler.RandomModel(),
num_measure_trials=2,
init_search_callbacks=None):
print("Test %s schedule search with the default search policy" % (target))
random.seed(seed)
N = 128
workload_key = auto_scheduler.make_workload_key(workload, (N, N, N))
dag = auto_scheduler.ComputeDAG(workload_key)
target = tvm.target.create(target)
task = auto_scheduler.SearchTask(dag, workload_key, target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
init_search_callbacks = init_search_callbacks or []
init_search_callbacks.append(
auto_scheduler.PreloadMeasuredStates(log_file))
if search_policy == 'empty':
search_policy = auto_scheduler.EmptyPolicy(task)
elif search_policy == 'sketch':
search_policy = auto_scheduler.SketchPolicy(
task, init_search_callbacks=init_search_callbacks)
tuning_options = auto_scheduler.TuningOptions(
num_measure_trials=num_measure_trials,
runner=runner,
verbose=1,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)])
sch, args = auto_scheduler.auto_schedule(task, search_policy,
tuning_options)
inp, res = auto_scheduler.load_best(log_file, workload_key, target)
print("==== Python Code ====")
print(dag.print_python_code_from_state(inp.state))
try:
print("==== Lowered Stmt ====")
print(tvm.lower(sch, args, simple_mode=True))
mod = tvm.build(sch, args, target)
ctx = tvm.context(str(target), 0)
dtype = dag.tensors[0].dtype
a = tvm.nd.array(np.random.uniform(size=(N, N)).astype(dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=(N, N)).astype(dtype), ctx)
c = tvm.nd.array(np.zeros((N, N), dtype=dtype), ctx)
mod(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(),
np.dot(a.asnumpy(), b.asnumpy()),
rtol=1e-5)
print("==== Verification passed ====")
except Exception:
raise Exception("Error encountered with seed: %d" % (seed))
print()
def resume_search(task, log_file):
print("Resume search:")
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(log_file)
search_policy = auto_scheduler.SketchPolicy(
task, cost_model, init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file)]
)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file)]
)
task.tune(tune_option, search_policy=search_policy)
def resume_search(task, logfile_name):
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(logfile_name)
search_policy = auto_scheduler.SketchPolicy(
task,
cost_model,
init_search_callbacks=[
auto_scheduler.PreloadMeasuredStates(logfile_name)
])
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5,
measure_callbacks=[auto_scheduler.RecordToFile(logfile_name)])
sch, args = auto_scheduler.auto_schedule(task,
search_policy,
tuning_options=tune_option)
def resume_search(task, log_file):
print("Resume search:")
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(log_file)
search_policy = auto_scheduler.SketchPolicy(
task, cost_model, init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file)]
)
measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5,
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
task.tune(tune_option, search_policy=search_policy)
# Kill the measurement process
del measure_ctx
# learning the behavior of the auto-scheduler.
print("Equivalent python schedule:")
print(task.compute_dag.print_python_code_from_state(inp.state))
# Rebuild the binary. This shows how you can apply the best schedule from a
# log file without reruning the search again.
sch, args = task.compute_dag.apply_steps_from_state(inp.state)
func = tvm.build(sch, args, target)
######################################################################
# A more complicated example is to resume the search.
# In this case, we need to create the search policy and cost model by ourselves
# and resume the status of search policy and cost model with the log file.
# In the example below we resume the status and do more 5 trials.
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(log_file)
search_policy = auto_scheduler.SketchPolicy(
task, cost_model, init_search_callbacks=[auto_scheduler.PreloadMeasuredStates(log_file)]
)
measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5,
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
sch, args = auto_scheduler.auto_schedule(task, search_policy, tuning_options=tune_option)
# Kill the measurement process
del measure_ctx
print(task.compute_dag.print_python_code_from_state(inp.state))
# Rebuild the binary. This shows how you can apply the best schedule from a
# log file without reruning the search again.
sch, args = task.compute_dag.apply_steps_from_state(inp.state)
func = tvm.build(sch, args, target)
######################################################################
# A more complicated example is to resume the search.
# In this case, we need to create the search policy and cost model by ourselves
# and resume the status of search policy and cost model with the log file.
# In the example below we resume the status and do more 5 trials.
cost_model = auto_scheduler.XGBModel()
cost_model.update_from_file(log_file)
search_policy = auto_scheduler.SketchPolicy(
task, cost_model, init_search_callbacks=[
auto_scheduler.PreloadMeasuredStates(log_file)]
)
measure_ctx = auto_scheduler.LocalRPCMeasureContext(min_repeat_ms=300)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=5,
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
sch, args = auto_scheduler.auto_schedule(
task, search_policy, tuning_options=tune_option)
# Kill the measurement process
del measure_ctx
def search_common(
workload=matmul_auto_scheduler_test,
target="llvm",
search_policy="sketch",
seed=0,
runner="local",
num_measure_trials=100,
cost_model=auto_scheduler.RandomModel(),
init_search_callbacks=None,
):
print("Test search policy '%s' for '%s'" % (search_policy, target))
random.seed(seed)
N = 128
target = tvm.target.Target(target)
task = auto_scheduler.create_task(workload, (N, N, N), target)
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
init_search_callbacks = init_search_callbacks or []
init_search_callbacks.append(
auto_scheduler.PreloadMeasuredStates(log_file))
if search_policy == "empty":
search_policy = auto_scheduler.EmptyPolicy(task)
elif search_policy == "sketch":
search_policy = auto_scheduler.SketchPolicy(
task,
program_cost_model=cost_model,
init_search_callbacks=init_search_callbacks)
else:
raise ValueError("Invalid policy: " + search_policy)
tuning_options = auto_scheduler.TuningOptions(
num_measure_trials=num_measure_trials,
num_measures_per_round=2,
early_stopping=1,
runner=runner,
verbose=2,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
sch, args = auto_scheduler.auto_schedule(task, search_policy,
tuning_options)
inp, res = auto_scheduler.load_best(log_file, task.workload_key,
target)
print("==== Python Code ====")
print(task.compute_dag.print_python_code_from_state(inp.state))
try:
print("==== Lowered Stmt ====")
print(tvm.lower(sch, args, simple_mode=True))
mod = tvm.build(sch, args, target)
ctx = tvm.context(str(target), 0)
dtype = task.compute_dag.tensors[0].dtype
a = tvm.nd.array(np.random.uniform(size=(N, N)).astype(dtype), ctx)
b = tvm.nd.array(np.random.uniform(size=(N, N)).astype(dtype), ctx)
c = tvm.nd.array(np.zeros((N, N), dtype=dtype), ctx)
mod(a, b, c)
tvm.testing.assert_allclose(c.asnumpy(),
np.dot(a.asnumpy(), b.asnumpy()),
rtol=1e-5)
print("==== Verification passed ====")
except Exception:
raise Exception("Error encountered with seed: %d" % (seed))
print()
def search_common(
task=None,
target="llvm",
search_policy="sketch",
runner="local",
num_measure_trials=100,
cost_model=auto_scheduler.RandomModel(),
init_search_callbacks=None,
):
if task is None:
task = auto_scheduler.SearchTask(func=matmul_auto_scheduler_test,
args=(64, 64, 64),
target=target)
target = task.target
print("Test search policy '%s' for '%s'" % (search_policy, target))
with tempfile.NamedTemporaryFile() as fp:
log_file = fp.name
init_search_callbacks = init_search_callbacks or []
init_search_callbacks.append(
auto_scheduler.PreloadMeasuredStates(log_file))
if search_policy == "empty":
search_policy = auto_scheduler.EmptyPolicy(task)
elif search_policy == "sketch":
search_policy = auto_scheduler.SketchPolicy(
task,
program_cost_model=cost_model,
init_search_callbacks=init_search_callbacks)
else:
raise ValueError("Invalid policy: " + search_policy)
# Tune
tuning_options = auto_scheduler.TuningOptions(
num_measure_trials=num_measure_trials,
num_measures_per_round=2,
early_stopping=1,
runner=runner,
measure_callbacks=[
auto_scheduler.RecordToFile(log_file),
CustomMeasureCallback()
],
)
task.tune(tuning_options=tuning_options, search_policy=search_policy)
# Compile with the best schedule
sch, args = task.apply_best(log_file)
mod = tvm.build(sch, args, target)
# Compile with naive schedule for correctness check
sch, args = task.compute_dag.apply_steps_from_state(
task.compute_dag.init_state)
mod_ref = tvm.build(sch, args, "llvm")
ctx = tvm.device(str(target), 0)
np_arrays = [
np.random.uniform(size=get_const_tuple(x.shape)).astype(x.dtype)
for x in args
]
tvm_arrays = [tvm.nd.array(x, ctx) for x in np_arrays]
mod(*tvm_arrays)
actual = [x.numpy() for x in tvm_arrays]
tvm_arrays = [tvm.nd.array(x) for x in np_arrays]
mod_ref(*tvm_arrays)
expected = [x.numpy() for x in tvm_arrays]
for x, y in zip(actual, expected):
tvm.testing.assert_allclose(x, y, rtol=1e-5)
