We would like to share a visualization tool (Tensor Expression Debug Display, or TEDD) for displaying scheduling results in Tensor Expression. We use three graphs (Schedule tree, IterVar relationship graph, and dataflow graph) to represent Tensor Expression schedules, as described in the operation model document linked from An operational model of schedule primitives in Tensor Expression.
The tool is written in Python, extracts information from a given schedule with Tensor Expression python APIs, and renders them in Graphviz dot and SV. It depends on these modules: tvm, graphviz, IPython, and html.
We’d like to hear feedback from the community and to make a PR, if it is useful.
The following are some example figures.
The final schedule from tutorial opt_conv_cuda.py:
The final schedule from verify_softmax in topi/tests/python/test_topi_softmax.py:
Some example usages:
import tedd
tedd.viz_dataflow_graph(s, False, '/tmp/dfg.dot')
tedd.viz_schedule_tree(s, False, '/tmp/scheduletree.dot')
tedd.viz_itervar_relationship_graph(s, False, '/tmp/itervar.dot')
