Hzfengsy commented on a change in pull request #4651: Tensor Expression Debug Display (TEDD) URL: https://github.com/apache/incubator-tvm/pull/4651#discussion_r365149671
########## File path: python/tvm/contrib/tedd.py ########## @@ -0,0 +1,506 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tensor Expression Debug Display (TEDD), visualizing Tensor Expression""" +import html +from graphviz import Digraph +from graphviz import Source +from IPython.display import display +from IPython.display import SVG +import tvm + +TVMDD_TABLE_BODY_WIDTH = 30 +# Must match enum IterVarType defined in include/tvm/expr.h +ITERVAR_TYPE_STRING_MAP = { + 0: ('kDataPar', '#FFFFFF'), + 1: ('kThreadIndex', '#2980B9'), + 2: ('kCommReduce', '#FAD7A0'), + 3: ('kOrdered', '#D35400'), + 4: ('kOpaque', '#ABB2B9'), + 5: ('kUnrolled', '#D2B4DE'), + 6: ('kVectorized', '#AED6F1'), + 7: ('kParallelized', '#F5B7B1'), + 8: ('kTensorized', '#A9DFBF'), +} + + +def get_or_create_dot_id(obj, prefix="", assert_on_missing=False): + """If obj's ID has been registered, return it. + If not, either assert or create a unique and legal ID, register and + return it, according to assert_on_missing. + ID must be a unique and legal Dotty ID. + + Parameters + ---------- + obj : objet + Serve as the key to the ID. + + prefix : string + Prefix to attach to the ID. Usually use obj's non-unique + name as prefix. + + assert_on_missing : bool + Asserts or not if object doesn't have a registered ID. + """ + prefix = prefix.replace('.', '_') + if not hasattr(get_or_create_dot_id, "obj_id_dict"): + get_or_create_dot_id.obj_id_dict = {} + if obj not in get_or_create_dot_id.obj_id_dict: + if assert_on_missing: + assert False, 'dot_id ' + str(obj) + ' has not been registered.' + else: + get_or_create_dot_id.obj_id_dict[obj] = prefix + hex(id(obj)) + return get_or_create_dot_id.obj_id_dict[obj] + + +def get_port_id(is_input, index): + return 'I_' + str(index) if is_input else 'O_' + str(index) + + +def get_itervar_type_info(iter_type): + assert iter_type < len( + ITERVAR_TYPE_STRING_MAP), 'Unknown IterVar type: ' + str(iter_type) + return ITERVAR_TYPE_STRING_MAP[iter_type] + + +def get_itervar_label_color(itervar, iv_type): + type_info = get_itervar_type_info(iv_type) + return str(itervar.var) + '(' + type_info[0] + ')', type_info[1] + + +def linebrk(s, n): + """ Break input string s with <br/> for every n charactors.""" + result = '' + j = 0 + for i, c in enumerate(s): + if j == n and i != len(s) - 1: + result = result + '\n' + j = 0 + j = j + 1 + result = result + c + result = html.escape(str(result), quote=True) + result = result.replace('\n', '<br/>') + return result + + +def create_graph(name="", rankdir='BT'): + graph = Digraph(name=name) + graph.graph_attr['rankdir'] = rankdir + return graph + + +def itervar_label(itervar, index, index_color, label): + return '<TR><TD PORT="' + get_or_create_dot_id(itervar, str( + itervar.var)) + '" BGCOLOR="' + index_color + '">' + str( + index + ) + '</TD><TD BGCOLOR="white" PORT="itervar">' + label + '</TD></TR>' + + +def stage_label(stage): + return stage.op.name + '<br/>Scope: ' + stage.scope + + +def legend_label(): + label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">' + for iter_type in ITERVAR_TYPE_STRING_MAP: + name, color = ITERVAR_TYPE_STRING_MAP[iter_type] + label += '<TR><TD BGCOLOR="' + color + '"></TD>' \ + + '<TD BGCOLOR="white">' + name + '</TD></TR>' + label += '</TABLE>>' + return label + + +def legend_dot(g): + with g.subgraph(name='cluster_legend') as subgraph: + subgraph.attr(label='Legend') + label = legend_label() + subgraph.node('legend', label, shape='none', margin='0') + + +def dump_graph(dot_string, + showsvg=True, + dotfilepath='', + outputdotstring=False): + """Output dot_string in various formats.""" + if dotfilepath: + try: + dot_file = open(dotfilepath, "w+") + dot_file.write(dot_string) + dot_file.close() + except IOError: + print('Cannot open file: ' + dotfilepath) + if showsvg: + src = Source(dot_string) + display(SVG(src.pipe(format='svg'))) + if outputdotstring: + return dot_string + return None + + +def viz_schedule_tree(sch, + showsvg=False, Review comment: Sounds great! As for the CI test, I would like to skip this test on the CI until the CI updating. Also, is there a way to allow contributors to modify the docker image on CI? @tqchen Many new features depend on different environments. It is difficult to change the CI docker every time by only one person. ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: us...@infra.apache.org With regards, Apache Git Services
