From ea02d453ce4df68c6bebfe7e6d84f6b1f42bd025 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 27 Jul 2020 21:42:10 -0500 Subject: [PATCH 1/7] Add graph visualization support in pytools.visualization --- .gitignore | 3 + doc/reference.rst | 1 + examples/visualization.py | 40 ++++++ pytato/__init__.py | 3 + pytato/codegen.py | 37 ++++-- pytato/visualization.py | 261 ++++++++++++++++++++++++++++++++++++++ requirements.txt | 2 +- 7 files changed, 332 insertions(+), 15 deletions(-) create mode 100755 examples/visualization.py create mode 100644 pytato/visualization.py diff --git a/.gitignore b/.gitignore index 1b2f54d..8e4ee63 100644 --- a/.gitignore +++ b/.gitignore @@ -18,3 +18,6 @@ distribute*tar.gz doc/_build .mypy_cache + +*.dot +*.svg diff --git a/doc/reference.rst b/doc/reference.rst index 43f6555..581f7b6 100644 --- a/doc/reference.rst +++ b/doc/reference.rst @@ -7,3 +7,4 @@ Reference .. automodule:: pytato.program .. automodule:: pytato.target .. automodule:: pytato.codegen +.. automodule:: pytato.visualization diff --git a/examples/visualization.py b/examples/visualization.py new file mode 100755 index 0000000..d0150dd --- /dev/null +++ b/examples/visualization.py @@ -0,0 +1,40 @@ +#!/usr/bin/env python +"""Demonstrates graph visualization with Graphviz.""" + +import logging +import numpy as np +import subprocess + +import pytato as pt + + +logger = logging.getLogger(__name__) + + +GRAPH_DOT = "graph.dot" +GRAPH_SVG = "graph.svg" + + +def main(): + ns = pt.Namespace() + + pt.make_size_param(ns, "n") + array = pt.make_placeholder(ns, "array", shape="n", dtype=np.float) + stack = pt.stack([array, 2*array, array + 6]) + ns.assign("stack", stack) + result = stack @ stack.T + + from pytato.visualization import get_dot_graph + dot_code = get_dot_graph(result) + + with open(GRAPH_DOT, "w") as outf: + outf.write(dot_code) + logger.info("wrote '%s'", GRAPH_DOT) + + subprocess.run(["dot", "-Tsvg", GRAPH_DOT, "-o", GRAPH_SVG], check=True) + logger.info("wrote '%s'", GRAPH_SVG) + + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + main() diff --git a/pytato/__init__.py b/pytato/__init__.py index 5bc498c..06a3d83 100644 --- a/pytato/__init__.py +++ b/pytato/__init__.py @@ -36,6 +36,7 @@ from pytato.array import ( from pytato.codegen import generate_loopy from pytato.target import Target, PyOpenCLTarget +from pytato.visualization import get_dot_graph, show_dot_graph __all__ = ( "DottedName", "Namespace", "Array", "DictOfNamedArrays", @@ -49,4 +50,6 @@ __all__ = ( "generate_loopy", "Target", "PyOpenCLTarget", + + "get_dot_graph", "show_dot_graph", ) diff --git a/pytato/codegen.py b/pytato/codegen.py index f74051c..3a742b8 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -72,6 +72,7 @@ Code Generation Internals .. autofunction:: get_loopy_temporary .. autofunction:: add_store .. autofunction:: rename_reductions +.. autofunction:: normalize_outputs """ @@ -689,6 +690,26 @@ def rename_reductions( loopy_expr_context.reduction_bounds = new_reduction_bounds return result + +def normalize_outputs(result: Union[Array, DictOfNamedArrays]) -> DictOfNamedArrays: + """Convert outputs of a computation to the canonical form. + + Performs a conversion to :class:`~pytato.DictOfNamedArrays` if necessary. + + :param result: Outputs of the computation. + """ + if not isinstance(result, (Array, DictOfNamedArrays)): + raise TypeError("outputs of the computation should be " + "either an Array or a DictOfNamedArrays") + + if isinstance(result, Array): + outputs = DictOfNamedArrays({"_pt_out": result}) + else: + assert isinstance(outputs, DictOfNamedArrays) + outputs = result + + return outputs + # }}} @@ -702,21 +723,9 @@ def generate_loopy(result: Union[Array, DictOfNamedArrays], :param options: Code generation options for the kernel. :returns: A wrapped generated :mod:`loopy` kernel """ - # {{{ get namespace and outputs - - outputs: DictOfNamedArrays - - if isinstance(result, Array): - outputs = DictOfNamedArrays({"_pt_out": result}) - namespace = outputs.namespace - else: - assert isinstance(result, DictOfNamedArrays) - outputs = result - - namespace = outputs.namespace + outputs: DictOfNamedArrays = normalize_outputs(result) del result - - # }}} + namespace = outputs.namespace if target is None: target = PyOpenCLTarget() diff --git a/pytato/visualization.py b/pytato/visualization.py new file mode 100644 index 0000000..be264df --- /dev/null +++ b/pytato/visualization.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +__copyright__ = """ +Copyright (C) 2020 Matt Wala +""" + +__license__ = """ +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +""" + + +import contextlib +import dataclasses +import html +import itertools +from typing import Callable, Dict, Union, Iterator, List, Mapping, Union + +from pytools import UniqueNameGenerator +from pytools.codegen import CodeGenerator as CodeGeneratorBase + +from pytato.array import ( + Array, DictOfNamedArrays, IndexLambda, InputArgumentBase, + Stack, ShapeType, TagsType) +from pytato.codegen import normalize_outputs +import pytato.transform + + +__doc__ = """ +.. currentmodule:: pytato + +Graph Visualization +------------------- + +.. autofunction:: get_dot_graph +.. autofunction:: show_dot_graph + +""" + + +# {{{ array -> dot node converter + +@dataclasses.dataclass +class DotNodeInfo: + title: str + fields: Dict[str, str] + edges: Dict[str, Array] + + +def stringify_tags(tags: TagsType) -> str: + components = sorted(str(elem) for elem in tags) + return "{" + ", ".join(components) + "}" + + +def stringify_shape(shape: ShapeType) -> str: + components = [str(elem) for elem in shape] + if not components: + components = [","] + elif len(components) == 1: + components[0] += "," + return "(" + ", ".join(components) + ")" + + +class ArrayToDotNodeInfoMapper(pytato.transform.Mapper): + def get_common_dot_info(self, expr: Array) -> DotNodeInfo: + title = type(expr).__name__ + fields = dict(shape=stringify_shape(expr.shape), + dtype=str(expr.dtype), + tags=stringify_tags(expr.tags)) + edges: Dict[str, Array] = {} + return DotNodeInfo(title, fields, edges) + + def handle_unsupported_array(self, expr: Array, # type: ignore + nodes: Dict[Array, DotNodeInfo]) -> None: + # Default handler, does its best to guess how to handle fields. + if expr in nodes: + return + info = self.get_common_dot_info(expr) + + for field in expr.fields: + if field in info.fields: + continue + attr = getattr(expr, field) + + if isinstance(attr, Array): + self.rec(attr, nodes) + info.edges[field] = attr + elif isinstance(attr, tuple): + info.fields[field] = stringify_shape(attr) + else: + info.fields[field] = str(attr) + + nodes[expr] = info + + def map_index_lambda(self, expr: IndexLambda, + nodes: Dict[Array, DotNodeInfo]) -> None: + if expr in nodes: + return + + info = self.get_common_dot_info(expr) + info.fields["expr"] = str(expr.expr) + + for name, val in expr.bindings.items(): + self.rec(val, nodes) + info.edges[name] = val + + nodes[expr] = info + + def map_stack(self, expr: Stack, nodes: Dict[Array, DotNodeInfo]) -> None: + if expr in nodes: + return + + info = self.get_common_dot_info(expr) + info.fields["axis"] = str(expr.axis) + + for i, array in enumerate(expr.arrays): + self.rec(array, nodes) + info.edges[str(i)] = array + + nodes[expr] = info + + +def dot_escape(s: str) -> str: + # "\" and HTML are significant in graphviz. + return html.escape(s.replace("\\", "\\\\")) + + +class DotEmitter(CodeGeneratorBase): + @contextlib.contextmanager + def block(self, name: str) -> Iterator[None]: + self(name + " {") + self.indent() + yield + self.dedent() + self("}") + + +def _emit_array(emit: DotEmitter, info: DotNodeInfo, id: str) -> None: + td_attrib = "border=\"0\"" + table_attrib = "border=\"0\" cellborder=\"1\" cellspacing=\"0\"" + + rows = ["%s" + % (td_attrib, dot_escape(info.title))] + + for name, field in info.fields.items(): + rows.append( + "%s:%s" + % (td_attrib, dot_escape(name), td_attrib, dot_escape(field))) + + table = "\n%s
" % (table_attrib, "".join(rows)) + emit("%s [label=<%s>]" % (id, table)) + + +def _emit_name_cluster(emit: DotEmitter, names: Mapping[str, Array], + array_to_id: Mapping[Array, str], id_gen: Callable[[str], str], + label: str) -> None: + edges = [] + + with emit.block("subgraph cluster_%s" % label): + emit("node [shape=ellipse]") + emit("label=\"%s\"" % label) + + for name, array in names.items(): + name_id = id_gen(label) + emit("%s [label=\"%s\"]" % (name_id, dot_escape(name))) + array_id = array_to_id[array] + # Edges must be outside the cluster. + edges.append((name_id, array_id)) + + for name_id, array_id in edges: + emit("%s -> %s" % (name_id, array_id)) + +# }}} + + +def get_dot_graph(result: Union[Array, DictOfNamedArrays]) -> str: + r"""Return a string in the `dot `_ language depicting the + graph of the computation of *result*. + + :arg result: Outputs of the computation (cf. :func:`pytato.generate_loopy`). + """ + outputs: DictOfNamedArrays = normalize_outputs(result) + del result + namespace = outputs.namespace + + nodes: Dict[Array, DotNodeInfo] = {} + mapper = ArrayToDotNodeInfoMapper() + for elem in itertools.chain(namespace.values(), outputs.values()): + mapper(elem, nodes) + + input_arrays: List[Array] = [] + internal_arrays: List[Array] = [] + array_to_id: Dict[Array, str] = {} + + id_gen = UniqueNameGenerator() + for array in nodes: + array_to_id[array] = id_gen("array") + if isinstance(array, InputArgumentBase): + input_arrays.append(array) + else: + internal_arrays.append(array) + + emit = DotEmitter() + + with emit.block("digraph computation"): + emit("node [shape=rectangle]") + + # Emit inputs. + with emit.block("subgraph cluster_Inputs"): + emit("label=\"Inputs\"") + for array in input_arrays: + _emit_array(emit, nodes[array], array_to_id[array]) + + # Emit non-inputs. + for array in internal_arrays: + _emit_array(emit, nodes[array], array_to_id[array]) + + # Emit edges. + for array, node in nodes.items(): + for label, tail_array in node.edges.items(): + tail = array_to_id[tail_array] + head = array_to_id[array] + emit("%s -> %s [label=\"%s\"]" % (tail, head, dot_escape(label))) + + # Emit output/namespace name mappings. + _emit_name_cluster(emit, outputs, array_to_id, id_gen, label="Outputs") + _emit_name_cluster(emit, namespace, array_to_id, id_gen, label="Namespace") + + output: str = emit.get() + return output + + +def show_dot_graph(result: Union[Array, DictOfNamedArrays]) -> None: + """Show a graph representing the computation of *result* in a browser. + + :arg result: Outputs of the computation (cf. :func:`pytato.generate_loopy`). + """ + dot_code: str + + if isinstance(result, str): + dot_code = result + else: + dot_code = get_dot_graph(result) + + from pymbolic.imperative.utils import show_dot + show_dot(dot_code) diff --git a/requirements.txt b/requirements.txt index 8900766..e489999 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -git+https://github.com/inducer/pytools.git +git+https://github.com/mattwala/pytools.git@code-generator git+https://github.com/inducer/loopy.git -- GitLab From 8d157c33c93009c7c058a79ed56016a4eb694c76 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 27 Jul 2020 21:47:25 -0500 Subject: [PATCH 2/7] CI fixes --- pytato/codegen.py | 2 +- pytato/visualization.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index 3a742b8..55587e2 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -705,7 +705,7 @@ def normalize_outputs(result: Union[Array, DictOfNamedArrays]) -> DictOfNamedArr if isinstance(result, Array): outputs = DictOfNamedArrays({"_pt_out": result}) else: - assert isinstance(outputs, DictOfNamedArrays) + assert isinstance(result, DictOfNamedArrays) outputs = result return outputs diff --git a/pytato/visualization.py b/pytato/visualization.py index be264df..b8ac339 100644 --- a/pytato/visualization.py +++ b/pytato/visualization.py @@ -29,7 +29,7 @@ import contextlib import dataclasses import html import itertools -from typing import Callable, Dict, Union, Iterator, List, Mapping, Union +from typing import Callable, Dict, Union, Iterator, List, Mapping from pytools import UniqueNameGenerator from pytools.codegen import CodeGenerator as CodeGeneratorBase -- GitLab From ce7ed9dc3ea39f9477ae0bc95706762574ead5fa Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 27 Jul 2020 21:52:08 -0500 Subject: [PATCH 3/7] Skip conversion to svg if dot not present --- examples/visualization.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/examples/visualization.py b/examples/visualization.py index d0150dd..05b06db 100755 --- a/examples/visualization.py +++ b/examples/visualization.py @@ -3,6 +3,7 @@ import logging import numpy as np +import shutil import subprocess import pytato as pt @@ -31,8 +32,12 @@ def main(): outf.write(dot_code) logger.info("wrote '%s'", GRAPH_DOT) - subprocess.run(["dot", "-Tsvg", GRAPH_DOT, "-o", GRAPH_SVG], check=True) - logger.info("wrote '%s'", GRAPH_SVG) + dot_path = shutil.which("dot") + if dot_path is not None: + subprocess.run([dot_path, "-Tsvg", GRAPH_DOT, "-o", GRAPH_SVG], check=True) + logger.info("wrote '%s'", GRAPH_SVG) + else: + logger.info("'dot' executable not found; cannot convert to SVG") if __name__ == "__main__": -- GitLab From 53094dfad0f9cb77cb78b19bbe311f82abbd110b Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 27 Jul 2020 22:04:36 -0500 Subject: [PATCH 4/7] No need for return statement variable annotation --- pytato/visualization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pytato/visualization.py b/pytato/visualization.py index b8ac339..421ad2a 100644 --- a/pytato/visualization.py +++ b/pytato/visualization.py @@ -241,8 +241,7 @@ def get_dot_graph(result: Union[Array, DictOfNamedArrays]) -> str: _emit_name_cluster(emit, outputs, array_to_id, id_gen, label="Outputs") _emit_name_cluster(emit, namespace, array_to_id, id_gen, label="Namespace") - output: str = emit.get() - return output + return emit.get() def show_dot_graph(result: Union[Array, DictOfNamedArrays]) -> None: -- GitLab From 9ae21074c53346bab98962150c040714f35857b1 Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Tue, 28 Jul 2020 05:06:29 +0200 Subject: [PATCH 5/7] Use pt.get_dot_graph() --- examples/visualization.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/visualization.py b/examples/visualization.py index 05b06db..f41f33b 100755 --- a/examples/visualization.py +++ b/examples/visualization.py @@ -25,8 +25,7 @@ def main(): ns.assign("stack", stack) result = stack @ stack.T - from pytato.visualization import get_dot_graph - dot_code = get_dot_graph(result) + dot_code = pt.get_dot_graph(result) with open(GRAPH_DOT, "w") as outf: outf.write(dot_code) -- GitLab From a25ea9d077808d9fde604da7d37c0c4fa79c41dc Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Mon, 27 Jul 2020 22:12:02 -0500 Subject: [PATCH 6/7] Re-dangle requirements.txt --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index e489999..8900766 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -git+https://github.com/mattwala/pytools.git@code-generator +git+https://github.com/inducer/pytools.git git+https://github.com/inducer/loopy.git -- GitLab From 3eeb32fbac8fd62586c881d2eb994585c0dae8ce Mon Sep 17 00:00:00 2001 From: Matt Wala Date: Sun, 2 Aug 2020 19:16:20 -0500 Subject: [PATCH 7/7] Fix variable name --- pytato/codegen.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pytato/codegen.py b/pytato/codegen.py index c6884e5..de93b1e 100644 --- a/pytato/codegen.py +++ b/pytato/codegen.py @@ -795,9 +795,8 @@ def generate_loopy(result: Union[Array, DictOfNamedArrays], :param options: Code generation options for the kernel. :returns: A wrapped generated :mod:`loopy` kernel """ - outputs: DictOfNamedArrays = normalize_outputs(result) + orig_outputs: DictOfNamedArrays = normalize_outputs(result) del result - namespace = outputs.namespace if target is None: target = PyOpenCLTarget() -- GitLab