From 77998824c614328c681410b130009080379e97fb Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Sat, 5 Nov 2016 14:31:33 -0500 Subject: [PATCH] Use arrows to visualize deps in kernel printing --- loopy/kernel/__init__.py | 40 ++++++++---- loopy/kernel/tools.py | 137 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 12 deletions(-) diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index 83119d770..d618d4b0d 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -1140,18 +1140,36 @@ class LoopKernel(RecordWithoutPickling): lines.append("INSTRUCTIONS:") loop_list_width = 35 - printed_insn_ids = set() + # {{{ topological sort - Fore = self.options._fore - Style = self.options._style + printed_insn_ids = set() + printed_insn_order = [] - def print_insn(insn): + def insert_insn_into_order(insn): if insn.id in printed_insn_ids: return printed_insn_ids.add(insn.id) for dep_id in sorted(insn.depends_on): - print_insn(kernel.id_to_insn[dep_id]) + insert_insn_into_order(kernel.id_to_insn[dep_id]) + + printed_insn_order.append(insn) + + for insn in kernel.instructions: + insert_insn_into_order(insn) + + # }}} + + import loopy as lp + + Fore = self.options._fore + Style = self.options._style + + from loopy.kernel.tools import draw_dependencies_as_unicode_arrows + for insn, (arrows, extender) in zip( + printed_insn_order, + draw_dependencies_as_unicode_arrows( + printed_insn_order, fore=Fore, style=Style)): if isinstance(insn, lp.MultiAssignmentBase): lhs = ", ".join(str(a) for a in insn.assignees) @@ -1206,13 +1224,15 @@ class LoopKernel(RecordWithoutPickling): core = Fore.MAGENTA+rhs+Style.RESET_ALL if len(loop_list) > loop_list_width: - lines.append("[%s]" % loop_list) - lines.append("%s%s # %s" % ( + lines.append("%s[%s]" % (arrows, loop_list)) + lines.append("%s%s%s # %s" % ( + extender, (loop_list_width+2)*" ", core, ", ".join(options))) else: - lines.append("[%s]%s%s # %s" % ( + lines.append("%s[%s]%s%s # %s" % ( + arrows, loop_list, " "*(loop_list_width-len(loop_list)), core, ",".join(options))) @@ -1223,10 +1243,6 @@ class LoopKernel(RecordWithoutPickling): lines.append(10*" " + "if (%s)" % " && ".join( [str(x) for x in insn.predicates])) - import loopy as lp - for insn in kernel.instructions: - print_insn(insn) - dep_lines = [] for insn in kernel.instructions: if insn.depends_on: diff --git a/loopy/kernel/tools.py b/loopy/kernel/tools.py index 54236efca..0a69ac4f7 100644 --- a/loopy/kernel/tools.py +++ b/loopy/kernel/tools.py @@ -1,4 +1,5 @@ """Operations on the kernel object.""" +# coding=utf-8 from __future__ import division, absolute_import, print_function @@ -1083,4 +1084,140 @@ def guess_var_shape(kernel, var_name): # }}} + +# {{{ find_recursive_dependencies + +def find_recursive_dependencies(kernel, insn_ids): + queue = list(insn_ids) + + result = set(insn_ids) + + while queue: + new_queue = [] + + for insn_id in queue: + insn = kernel.id_to_insn[insn_id] + additionals = insn.depends_on - result + result.update(additionals) + new_queue.extend(additionals) + + queue = new_queue + + return result + +# }}} + + +# {{{ find_reverse_dependencies + +def find_reverse_dependencies(kernel, insn_ids): + """Finds a set of IDs of instructions that depend on one of the insn_ids. + + :arg insn_ids: a set of instruction IDs + """ + return frozenset( + insn.id + for insn in kernel.instructions + if insn.depends_on & insn_ids) + +# }}} + + +# {{{ draw_dependencies_as_unicode_arrows + +def draw_dependencies_as_unicode_arrows( + instructions, fore, style, flag_downward=True): + """ + :arg instructions: an ordered iterable of :class:`loopy.InstructionBase` + instances + :arg fore: if given, will be used like a :mod:`colorama` ``Fore`` object + to color-code dependencies. (E.g. red for downward edges) + :returns: A list of tuples (arrows, extender) with Unicode-drawn dependency + arrows, one per entry of *instructions*. *extender* can be used to + extend arrows below the line of an instruction. + """ + reverse_deps = {} + + for insn in instructions: + for dep in insn.depends_on: + reverse_deps.setdefault(dep, []).append(insn.id) + + # mapping of (from_id, to_id) tuples to column_index + dep_to_column = {} + + # {{{ find column assignments + + # mapping from column indices to (end_insn_id, updown) + columns_in_use = {} + + n_columns = [0] + + def find_free_column(): + i = 0 + while i in columns_in_use: + i += 1 + if i+1 > n_columns[0]: + n_columns[0] = i+1 + row.append(" ") + return i + + def do_flag_downward(s, updown): + if flag_downward and updown == "down": + return fore.RED+s+style.RESET_ALL + else: + return s + + def make_extender(): + result = n_columns[0] * [" "] + for col, (_, updown) in six.iteritems(columns_in_use): + result[col] = do_flag_downward("│", updown) + + return result + + rows = [] + for insn in instructions: + row = make_extender() + + for rdep in reverse_deps.get(insn.id, []): + assert rdep != insn.id + + dep_key = (rdep, insn.id) + if dep_key not in dep_to_column: + col = dep_to_column[dep_key] = find_free_column() + columns_in_use[col] = (rdep, "up") + row[col] = "↱" + + for dep in insn.depends_on: + assert dep != insn.id + dep_key = (insn.id, dep) + if dep_key not in dep_to_column: + col = dep_to_column[dep_key] = find_free_column() + columns_in_use[col] = (dep, "down") + row[col] = do_flag_downward("┌", "down") + + for col, (end, updown) in list(six.iteritems(columns_in_use)): + if insn.id == end: + del columns_in_use[col] + if updown == "up": + row[col] = "└" + else: + row[col] = do_flag_downward("↳", updown) + + extender = make_extender() + + rows.append(("".join(row), "".join(extender))) + + # }}} + + def extend_to_uniform_length(s): + return s + " "*(n_columns[0]-len(s)) + + return [ + (extend_to_uniform_length(row), + extend_to_uniform_length(extender)) + for row, extender in rows] + +# }}} + + # vim: foldmethod=marker -- GitLab