This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-tvm.git
The following commit(s) were added to refs/heads/master by this push:
new b91dbca Add colors to compute_at edges and thread/block indices.
(#5111)
b91dbca is described below
commit b91dbca6fd63468d615fe3e63b8b4e019fa956b0
Author: yongfeng-nv <[email protected]>
AuthorDate: Fri Mar 20 12:21:45 2020 -0400
Add colors to compute_at edges and thread/block indices. (#5111)
---
python/tvm/contrib/tedd.py | 23 ++++++++++++++++++++---
tests/python/contrib/test_tedd.py | 2 +-
2 files changed, 21 insertions(+), 4 deletions(-)
diff --git a/python/tvm/contrib/tedd.py b/python/tvm/contrib/tedd.py
index 68e15f2..ae57211 100644
--- a/python/tvm/contrib/tedd.py
+++ b/python/tvm/contrib/tedd.py
@@ -37,6 +37,19 @@ ITERVAR_TYPE_STRING_MAP = {
8: ('kTensorized', '#A9DFBF'),
}
+PALETTE = {
+ 0: '#000000',
+ 1: '#922B21',
+ 2: '#76448A',
+ 3: '#1F618D',
+ 4: '#148F77',
+ 5: '#B7950B',
+ 6: '#AF601A',
+ 7: '#F5B7B1',
+ 8: '#A9DFBF',
+}
+
+PALETTE_SIZE = 9
def dom_path_to_string(dom_path, prefix=""):
path_string = prefix
@@ -458,8 +471,8 @@ def viz_schedule_tree(sch,
var_attr_label = ''
if "thread" in leafiv["properties"] and \
leafiv["properties"]["thread"] is not None:
- var_attr_label = var_attr_label + "<br/>(" + str(
- leafiv["properties"]["thread"]) + ")"
+ var_attr_label = var_attr_label + "<br/><font
color=\"#2980B9\">(" + str(
+ leafiv["properties"]["thread"]) + ")</font>"
if "intrin" in leafiv["properties"] and \
leafiv["properties"]["intrin"] is not None:
var_attr_label = var_attr_label + "<br/>" + \
@@ -483,7 +496,11 @@ def viz_schedule_tree(sch,
[stage["attaching_to"][0]], "Stage") + ":" + dom_path_to_string(
stage["attaching_to"],
"IterVar") if stage["attaching_to"] is not None else "ROOT"
- g.edge(src, dst)
+ color = PALETTE[
+ stage["attaching_to"][1] +
+ 1] if stage["attaching_to"] is not None and stage["attaching_to"][
+ 1] < PALETTE_SIZE - 1 else PALETTE[0]
+ g.edge(src, dst, color=color)
graph = create_schedule_tree_graph("Schedule Tree")
s = extract_dom_for_viz(sch)
diff --git a/tests/python/contrib/test_tedd.py
b/tests/python/contrib/test_tedd.py
index 6e5f3a4..58ff064 100644
--- a/tests/python/contrib/test_tedd.py
+++ b/tests/python/contrib/test_tedd.py
@@ -125,7 +125,7 @@ def test_schedule_tree():
findany(r"r.outer\(kCommReduce\)", str)
findany(r"label=ROOT", str)
# Check the compute_at edge
- findany(r"Stage_1", str)
+ findany(r"Stage_1.*\[color\=\"\#000000\"\]", str)
if checkdepdency():
verify()
