From 3586d05ba4b38e374ce3c2261cd77587c6af2ab5 Mon Sep 17 00:00:00 2001 From: ganler Date: Tue, 26 Jul 2022 21:06:12 -0700 Subject: [PATCH 1/8] feat(ux): highlight tvm script --- python/gen_requirements.py | 1 + python/tvm/ir/module.py | 13 +++ python/tvm/script/highlight.py | 86 +++++++++++++++++++ python/tvm/tir/function.py | 13 +++ ...st_tvmscript_printer_python_doc_printer.py | 26 ++++++ 5 files changed, 139 insertions(+) create mode 100644 python/tvm/script/highlight.py diff --git a/python/gen_requirements.py b/python/gen_requirements.py index 7e2c3e218618..e7f550b25c93 100755 --- a/python/gen_requirements.py +++ b/python/gen_requirements.py @@ -64,6 +64,7 @@ ( "Base requirements needed to install tvm", [ + "Pygments", "attrs", "cloudpickle", "decorator", diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py index b94d50dbf20d..49944729fafb 100644 --- a/python/tvm/ir/module.py +++ b/python/tvm/ir/module.py @@ -276,6 +276,19 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str: self, tir_prefix, show_meta ) # type: ignore + def show(self, style: str = "light") -> None: + """ + A sugar for print highlighted TVM script. + Parameters + ---------- + style : str, optional + Pygments styles extended by "light" (default) and "dark", by default "light" + """ + from tvm.script.highlight import cprint # pylint: disable=import-outside-toplevel + + # Use deferred import to avoid circular import while keeping cprint under tvm/script + cprint(self, style=style) + def get_attr(self, attr_key): """Get the IRModule attribute. diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py new file mode 100644 index 000000000000..f37eb4be35f7 --- /dev/null +++ b/python/tvm/script/highlight.py @@ -0,0 +1,86 @@ +# 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. +"""Highlight printed TVM script. +""" + +from typing import Union + +from pygments import highlight +from pygments.lexers import Python3Lexer +from pygments.formatters import Terminal256Formatter +from pygments.style import Style +from pygments.token import Keyword, Name, Comment, String, Number, Operator + +from tvm.ir import IRModule +from tvm.tir import PrimFunc + + +class VSCDark(Style): + """A VSCode-Dark-like Pygments style configuration""" + + styles = { + Keyword: "bold #c586c0", + Keyword.Namespace: "#4ec9b0", + Keyword.Type: "#82aaff", + Name.Function: "bold #dcdcaa", + Name.Class: "bold #569cd6", + Name.Decorator: "italic #fe4ef3", + String: "#ce9178", + Number: "#b5cea8", + Operator: "#bbbbbb", + Operator.Word: "#569cd6", + Comment: "italic #6a9956", + } + + +class JupyterLight(Style): + """A Jupyter-Notebook-like Pygments style configuration""" + + styles = { + Keyword: "bold #008000", + Keyword.Type: "nobold #008000", + Name.Function: "#0000FF", + Name.Class: "bold #0000FF", + Name.Decorator: "#AA22FF", + String: "#BA2121", + Number: "#008000", + Operator: "bold #AA22FF", + Operator.Word: "bold #008000", + Comment: "italic #007979", + } + + +def cprint(printable: Union[IRModule, PrimFunc], style="light") -> None: + """ + Print highlighted TVM script string with Pygments + Parameters + ---------- + printable : Union[IRModule, PrimFunc] + The TVM script to be printed + style : str, optional + Style of the printed script + Notes + ----- + The style parameter follows the Pygments style names or Style objects. Two + built-in styles are extended: "light" (default) and "dark". Other styles + can be found in https://pygments.org/styles/ + """ + if style == "light": + style = JupyterLight + elif style == "dark": + style = VSCDark + print(highlight(printable.script(), Python3Lexer(), Terminal256Formatter(style=style))) diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py index a921c5b9fc40..dff2aab76372 100644 --- a/python/tvm/tir/function.py +++ b/python/tvm/tir/function.py @@ -195,6 +195,19 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str: self, tir_prefix, show_meta ) # type: ignore + def show(self, style: str = "light") -> None: + """ + A sugar for print highlighted TVM script. + Parameters + ---------- + style : str, optional + Pygments styles extended by "light" (default) and "dark", by default "light" + """ + from tvm.script.highlight import cprint # pylint: disable=import-outside-toplevel + + # Use deferred import to avoid circular import while keeping cprint under tvm/script + cprint(self, style=style) + @tvm._ffi.register_object("tir.TensorIntrin") class TensorIntrin(Object): diff --git a/tests/python/unittest/test_tvmscript_printer_python_doc_printer.py b/tests/python/unittest/test_tvmscript_printer_python_doc_printer.py index 55b5e88c88c8..05b6103e51db 100644 --- a/tests/python/unittest/test_tvmscript_printer_python_doc_printer.py +++ b/tests/python/unittest/test_tvmscript_printer_python_doc_printer.py @@ -16,6 +16,8 @@ # under the License. import pytest +import tvm +from tvm.script import tir as T from tvm.script.printer.doc_printer import to_python_script from tvm.script.printer.doc import LiteralDoc @@ -51,3 +53,27 @@ def format_script(s: str) -> str: ) def test_print_literal_doc(doc, expected): assert to_python_script(doc).rstrip("\n") == format_script(expected) + + +def test_highlight_script(): + @tvm.script.ir_module + class Module: + @T.prim_func + def main( # type: ignore + a: T.handle, + b: T.handle, + c: T.handle, + ) -> None: # pylint: disable=no-self-argument + T.func_attr({"global_symbol": "main", "tir.noalias": True}) + A = T.match_buffer(a, [16, 128, 128]) + B = T.match_buffer(b, [16, 128, 128]) + C = T.match_buffer(c, [16, 128, 128]) + for n, i, j, k in T.grid(16, 128, 128, 128): + with T.block("matmul"): + vn, vi, vj, vk = T.axis.remap("SSSR", [n, i, j, k]) + with T.init(): + C[vn, vi, vj] = 0.0 # type: ignore + C[vn, vi, vj] = C[vn, vi, vj] + A[vn, vi, vk] * B[vn, vj, vk] + + Module.show() + Module["main"].show() From 939647554991554746a669daccf4ce2c98c5f770 Mon Sep 17 00:00:00 2001 From: ganler Date: Wed, 27 Jul 2022 14:04:22 -0700 Subject: [PATCH 2/8] resolve dependency --- docker/install/ubuntu_install_python_package.sh | 1 + python/gen_requirements.py | 1 - python/tvm/script/highlight.py | 2 +- 3 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/install/ubuntu_install_python_package.sh b/docker/install/ubuntu_install_python_package.sh index d97bd647cda5..729598a8c256 100755 --- a/docker/install/ubuntu_install_python_package.sh +++ b/docker/install/ubuntu_install_python_package.sh @@ -22,6 +22,7 @@ set -o pipefail # install libraries for python package on ubuntu pip3 install --upgrade \ + Pygments \ attrs \ cloudpickle \ cython \ diff --git a/python/gen_requirements.py b/python/gen_requirements.py index e7f550b25c93..7e2c3e218618 100755 --- a/python/gen_requirements.py +++ b/python/gen_requirements.py @@ -64,7 +64,6 @@ ( "Base requirements needed to install tvm", [ - "Pygments", "attrs", "cloudpickle", "decorator", diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index f37eb4be35f7..d99c975c9186 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -20,7 +20,7 @@ from typing import Union from pygments import highlight -from pygments.lexers import Python3Lexer +from pygments.lexers.python import Python3Lexer from pygments.formatters import Terminal256Formatter from pygments.style import Style from pygments.token import Keyword, Name, Comment, String, Number, Operator From 9761f727b86dcd2a141da172fe167fc5050afc54 Mon Sep 17 00:00:00 2001 From: ganler Date: Wed, 27 Jul 2022 16:10:00 -0700 Subject: [PATCH 3/8] refact(tvmscript): highlight fallback as plain text with warning; put ansi colors as default terminal style --- python/tvm/ir/module.py | 4 +- python/tvm/script/highlight.py | 146 +++++++++++------- python/tvm/tir/function.py | 2 +- .../test_tvmscript_printer_highlight.py | 3 + 4 files changed, 101 insertions(+), 54 deletions(-) diff --git a/python/tvm/ir/module.py b/python/tvm/ir/module.py index 49944729fafb..06537e2cdc4d 100644 --- a/python/tvm/ir/module.py +++ b/python/tvm/ir/module.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. """IRModule that holds the functions and type definitions.""" +from typing import Optional + from tvm._ffi.base import string_types import tvm._ffi @@ -276,7 +278,7 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str: self, tir_prefix, show_meta ) # type: ignore - def show(self, style: str = "light") -> None: + def show(self, style: Optional[str] = None) -> None: """ A sugar for print highlighted TVM script. Parameters diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index d99c975c9186..57c75e8ce8b1 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -17,54 +17,15 @@ """Highlight printed TVM script. """ -from typing import Union - -from pygments import highlight -from pygments.lexers.python import Python3Lexer -from pygments.formatters import Terminal256Formatter -from pygments.style import Style -from pygments.token import Keyword, Name, Comment, String, Number, Operator +from typing import Union, Optional +import warnings +import sys from tvm.ir import IRModule from tvm.tir import PrimFunc -class VSCDark(Style): - """A VSCode-Dark-like Pygments style configuration""" - - styles = { - Keyword: "bold #c586c0", - Keyword.Namespace: "#4ec9b0", - Keyword.Type: "#82aaff", - Name.Function: "bold #dcdcaa", - Name.Class: "bold #569cd6", - Name.Decorator: "italic #fe4ef3", - String: "#ce9178", - Number: "#b5cea8", - Operator: "#bbbbbb", - Operator.Word: "#569cd6", - Comment: "italic #6a9956", - } - - -class JupyterLight(Style): - """A Jupyter-Notebook-like Pygments style configuration""" - - styles = { - Keyword: "bold #008000", - Keyword.Type: "nobold #008000", - Name.Function: "#0000FF", - Name.Class: "bold #0000FF", - Name.Decorator: "#AA22FF", - String: "#BA2121", - Number: "#008000", - Operator: "bold #AA22FF", - Operator.Word: "bold #008000", - Comment: "italic #007979", - } - - -def cprint(printable: Union[IRModule, PrimFunc], style="light") -> None: +def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) -> None: """ Print highlighted TVM script string with Pygments Parameters @@ -72,15 +33,96 @@ def cprint(printable: Union[IRModule, PrimFunc], style="light") -> None: printable : Union[IRModule, PrimFunc] The TVM script to be printed style : str, optional - Style of the printed script + Printing style, auto-detected if None. Notes ----- - The style parameter follows the Pygments style names or Style objects. Two - built-in styles are extended: "light" (default) and "dark". Other styles - can be found in https://pygments.org/styles/ + The style parameter follows the Pygments style names or Style objects. Three + built-in styles are extended: "light", "dark" and "ansi". By default, "light" + will be used for notebook environment and terminal style will be "ansi" for + better style consistency. As an fallback when the optional Pygment library is + not installed, plain text will be printed with a one-time warning to suggest + installing the Pygment library. Other Pygment styles can be found in + https://pygments.org/styles/ """ - if style == "light": - style = JupyterLight - elif style == "dark": - style = VSCDark - print(highlight(printable.script(), Python3Lexer(), Terminal256Formatter(style=style))) + + try: + from pygments import highlight + from pygments.lexers.python import Python3Lexer + from pygments.formatters import Terminal256Formatter + from pygments.style import Style + from pygments.token import Keyword, Name, Comment, String, Number, Operator + except ImportError: + with warnings.catch_warnings(): + warnings.simplefilter("once", UserWarning) + install_cmd = sys.executable + " -m pip install Pygments --user" + warnings.warn( + "To print highlighted TVM script, please install Pygments:\n" + install_cmd, + category=UserWarning, + ) + print(printable.script()) + else: + + class JupyterLight(Style): + """A Jupyter-Notebook-like Pygments style configuration (aka. "dark")""" + + styles = { + Keyword: "bold #008000", + Keyword.Type: "nobold #008000", + Name.Function: "#0000FF", + Name.Class: "bold #0000FF", + Name.Decorator: "#AA22FF", + String: "#BA2121", + Number: "#008000", + Operator: "bold #AA22FF", + Operator.Word: "bold #008000", + Comment: "italic #007979", + } + + class VSCDark(Style): + """A VSCode-Dark-like Pygments style configuration (aka. "dark")""" + + styles = { + Keyword: "bold #c586c0", + Keyword.Type: "#82aaff", + Keyword.Namespace: "#4ec9b0", + Name.Class: "bold #569cd6", + Name.Function: "bold #dcdcaa", + Name.Decorator: "italic #fe4ef3", + String: "#ce9178", + Number: "#b5cea8", + Operator: "#bbbbbb", + Operator.Word: "#569cd6", + Comment: "italic #6a9956", + } + + class AnsiTerminalDefault(Style): + """The default style for terminal display with ANSI colors (aka. "ansi")""" + + styles = { + Keyword: "bold ansimagenta", + Keyword.Type: "ansibrightblue", + Keyword.Namespace: "ansicyan", + Name.Class: "bold ansiblue", + Name.Function: "bold ansiyellow", + Name.Decorator: "italic ansibrightmagenta", + String: "ansibrightyellow", + Number: "ansibrightgreen", + Operator: "ansired", + Operator.Word: "ansibrightcyan", + Comment: "italic ansibrightblack", + } + + if style is None: + # choose style automatically according to the environment: + if "ipykernel" in sys.modules: # in notebook env. + style = JupyterLight + else: # in a terminal or something. + style = AnsiTerminalDefault + elif style == "light": + style = JupyterLight + elif style == "dark": + style = VSCDark + elif style == "ansi": + style = AnsiTerminalDefault + + print(highlight(printable.script(), Python3Lexer(), Terminal256Formatter(style=style))) diff --git a/python/tvm/tir/function.py b/python/tvm/tir/function.py index dff2aab76372..f06376147b9a 100644 --- a/python/tvm/tir/function.py +++ b/python/tvm/tir/function.py @@ -195,7 +195,7 @@ def script(self, tir_prefix: str = "T", show_meta: bool = False) -> str: self, tir_prefix, show_meta ) # type: ignore - def show(self, style: str = "light") -> None: + def show(self, style: Optional[str] = None) -> None: """ A sugar for print highlighted TVM script. Parameters diff --git a/tests/python/unittest/test_tvmscript_printer_highlight.py b/tests/python/unittest/test_tvmscript_printer_highlight.py index 94ce8dc6cac9..cc3469a2ceea 100644 --- a/tests/python/unittest/test_tvmscript_printer_highlight.py +++ b/tests/python/unittest/test_tvmscript_printer_highlight.py @@ -42,3 +42,6 @@ def main( # type: ignore Module.show() Module["main"].show() + Module["main"].show(style="light") + Module["main"].show(style="dark") + Module["main"].show(style="ansi") From fd7315582d27ca1dbc0dcf48a9c93194ccc5e120 Mon Sep 17 00:00:00 2001 From: ganler Date: Wed, 27 Jul 2022 16:20:48 -0700 Subject: [PATCH 4/8] refact(tvmscript): make terminal style close to the default notebook style --- python/tvm/script/highlight.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index 57c75e8ce8b1..b2da62d11dea 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -99,16 +99,15 @@ class AnsiTerminalDefault(Style): """The default style for terminal display with ANSI colors (aka. "ansi")""" styles = { - Keyword: "bold ansimagenta", - Keyword.Type: "ansibrightblue", - Keyword.Namespace: "ansicyan", + Keyword: "bold ansigreen", + Keyword.Type: "nobold ansigreen", Name.Class: "bold ansiblue", - Name.Function: "bold ansiyellow", + Name.Function: "bold ansiblue", Name.Decorator: "italic ansibrightmagenta", - String: "ansibrightyellow", + String: "ansiyellow", Number: "ansibrightgreen", - Operator: "ansired", - Operator.Word: "ansibrightcyan", + Operator: "bold ansimagenta", + Operator.Word: "bold ansigreen", Comment: "italic ansibrightblack", } From 10b4a72c85c36446ea6ce86c8c8983212d1a0168 Mon Sep 17 00:00:00 2001 From: ganler Date: Wed, 27 Jul 2022 17:32:21 -0700 Subject: [PATCH 5/8] fix(pylint): disable=import-outside-toplevel --- python/tvm/script/highlight.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index b2da62d11dea..d0e610292ea7 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -46,6 +46,7 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) -> """ try: + # pylint: disable=import-outside-toplevel from pygments import highlight from pygments.lexers.python import Python3Lexer from pygments.formatters import Terminal256Formatter From dd85ba8a5cd54c0a61182331eb01e41952cf2e62 Mon Sep 17 00:00:00 2001 From: ganler Date: Thu, 28 Jul 2022 00:38:30 -0500 Subject: [PATCH 6/8] fix(ci-dep): Pygments>=2.4.0 to support ansicolors w/o # --- docker/install/ubuntu_install_python_package.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/install/ubuntu_install_python_package.sh b/docker/install/ubuntu_install_python_package.sh index 729598a8c256..3018dd9a5b5d 100755 --- a/docker/install/ubuntu_install_python_package.sh +++ b/docker/install/ubuntu_install_python_package.sh @@ -22,7 +22,7 @@ set -o pipefail # install libraries for python package on ubuntu pip3 install --upgrade \ - Pygments \ + Pygments>=2.4.0 \ attrs \ cloudpickle \ cython \ From 242a376f1f271fccee895372d82d3ef4f5ed1bef Mon Sep 17 00:00:00 2001 From: ganler Date: Thu, 28 Jul 2022 11:02:30 -0700 Subject: [PATCH 7/8] refact: making Pygments versioning most robust and user-friendly --- docker/install/ubuntu_install_python_package.sh | 2 +- python/tvm/script/highlight.py | 14 +++++++++++--- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/docker/install/ubuntu_install_python_package.sh b/docker/install/ubuntu_install_python_package.sh index 3018dd9a5b5d..3fc310c47e34 100755 --- a/docker/install/ubuntu_install_python_package.sh +++ b/docker/install/ubuntu_install_python_package.sh @@ -22,7 +22,7 @@ set -o pipefail # install libraries for python package on ubuntu pip3 install --upgrade \ - Pygments>=2.4.0 \ + "Pygments>=2.4.0" \ attrs \ cloudpickle \ cython \ diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index d0e610292ea7..db3e1acb3912 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -47,17 +47,25 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) -> try: # pylint: disable=import-outside-toplevel + import pygments from pygments import highlight from pygments.lexers.python import Python3Lexer from pygments.formatters import Terminal256Formatter from pygments.style import Style from pygments.token import Keyword, Name, Comment, String, Number, Operator - except ImportError: + from packaging import version + + if version.parse(pygments.__version__) < version.parse("2.4.0"): + raise ImportError("Required Pygments version >= 2.4.0 but got " + pygments.__version__) + except ImportError as e: with warnings.catch_warnings(): warnings.simplefilter("once", UserWarning) - install_cmd = sys.executable + " -m pip install Pygments --user" + install_cmd = sys.executable + ' -m pip install "Pygments>=2.4.0" --upgrade --user' warnings.warn( - "To print highlighted TVM script, please install Pygments:\n" + install_cmd, + str(e) + + "\n" + + "To print highlighted TVM script, please install Pygments:\n" + + install_cmd, category=UserWarning, ) print(printable.script()) From ed91f58d0f6288c7bd3e8cb69614298db60ee3a2 Mon Sep 17 00:00:00 2001 From: ganler Date: Thu, 28 Jul 2022 15:12:14 -0700 Subject: [PATCH 8/8] fix: pylint var naming --- python/tvm/script/highlight.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/script/highlight.py b/python/tvm/script/highlight.py index db3e1acb3912..03476ba60cd2 100644 --- a/python/tvm/script/highlight.py +++ b/python/tvm/script/highlight.py @@ -57,12 +57,12 @@ def cprint(printable: Union[IRModule, PrimFunc], style: Optional[str] = None) -> if version.parse(pygments.__version__) < version.parse("2.4.0"): raise ImportError("Required Pygments version >= 2.4.0 but got " + pygments.__version__) - except ImportError as e: + except ImportError as err: with warnings.catch_warnings(): warnings.simplefilter("once", UserWarning) install_cmd = sys.executable + ' -m pip install "Pygments>=2.4.0" --upgrade --user' warnings.warn( - str(e) + str(err) + "\n" + "To print highlighted TVM script, please install Pygments:\n" + install_cmd,