[autoscheduler] Print tuned Python schedule

Questions
#1

The documentation lists that as a method of tvm.auto_scheduler.ComputeDAG we can get a Python code representation of the schedule with print_python_code_from_state().

Described as:

Print transform steps in the history of a State as TVM’s python schedule code.

This is used to print transformation steps for debugging. Use apply_steps_from_state if you want to get a schedule for code generation.

However, I am trying to figure out to use it. I have a simple model with a single operation, and have loaded it into the two variables: mod and params.

Additionally, I have performed auto-scheduling, and have a tuned log-file in the json format.

From this, how would I get this printed schedule?

If I try building a tvm.auto_scheduler.ComputeDAG object by passing mod to it I get a RecursionError: maximum recursion depth exceeded while calling a Python object error.

Same if I try and get the Ansor workload from it:

    tasks, task_weights = auto_scheduler.extract_tasks(
        mod["main"], params, device_info['target_string'], device_info['target_host']
    )

    for i, t in enumerate(tasks):
         tgt_dag = str(t.compute_dag)
         tvm.auto_scheduler.ComputeDAG(tgt_dag)

Even if they did work however, I’m not sure how I’d link it to my auto-schedule JSON file.

The file python/tvm/auto_scheduler/relay_integration.py has a function auto_schedule_topi() instantiates a ComputeDAG, though the docstring says:

Note: This is used internally for relay integration. Do not use this as a general user-facing API.

Any pointers on getting this Python schedule for a tuned model?

#2

Did you try the API mentioned in this tutorial? Specifically, it uses print(task.print_best(log_file)). The API you used is for getting the schedule for compilation, not for printing.

https://tvm.apache.org/docs/tutorials/auto_scheduler/tune_matmul_x86.html#using-the-record-file

1 Like
#3

Exactly what I’m looking for, thank you.

1 Like