Now I only have a te.Schedule
, is it possible to optimize it with auto_scheduler
? I tried converting it to ComputeDAG
and using auto_scheduler to optimize, but got a KeyError:
def get_optimized_code_from_schedule(sche: te.Schedule, target):
compute_dag = auto_scheduler.ComputeDAG(sche)
workload_key = compute_dag.workload_key()
task = auto_scheduler.SearchTask(compute_dag=compute_dag, workload_key=workload_key, target=target)
num_trials = 400
log_file = "ansor_log.json"
measure_ctx = auto_scheduler.LocalRPCMeasureContext(repeat=1, min_repeat_ms=300, timeout=10)
tune_option = auto_scheduler.TuningOptions(
num_measure_trials=num_trials,
runner=measure_ctx.runner,
measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
)
task.tune(tune_option)
sch, args = task.apply_best(log_file)
return tvm.build(sch, args, target=target).import_modules[0].get_source()
# File "/share/tvm/python/tvm/auto_scheduler/measure.py", line 141, in serialize
# serialize_workload_registry_entry(self.task.workload_key),
# File "/share/tvm/python/tvm/auto_scheduler/workload_registry.py", line 229, in serialize_workload_registry_entry
# svalue = WORKLOAD_FUNC_REGISTRY[sname]
# KeyError: '86cc3050675f461b1e799b2935158a97''