This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit c6d6448a9b7a63e5c5683cda594e31e8c899fe1f Author: Tianqi Chen <[email protected]> AuthorDate: Sat Oct 17 17:06:14 2026 -0400 [REFACTOR][Script] Simplify printer configuration and display options (#20778) Display TVMScript directly with syntax highlighting by removing the obsolete Black formatting option and environment switch from the shared printer and its wrappers. Keep dialect-specific type-annotation settings in `extra_config`, place that configuration last in printer signatures, and keep TIRx printer documentation independent of parent dialects. Preserve the Relax-owned Python-module convenience through explicit configuration adaptation. --- tirx/include/tvm/script/printer/config.h | 4 +- tirx/python/tvm_tirx/script/printer/highlight.py | 103 +--------------------- tirx/python/tvm_tirx/script/printer/scriptable.py | 75 ++++------------ 3 files changed, 19 insertions(+), 163 deletions(-) diff --git a/tirx/include/tvm/script/printer/config.h b/tirx/include/tvm/script/printer/config.h index dfb9dbbe40..440ca25b6b 100644 --- a/tirx/include/tvm/script/printer/config.h +++ b/tirx/include/tvm/script/printer/config.h @@ -73,8 +73,6 @@ class PrinterConfigObj : public ffi::Object { * Keys are conventionally namespaced as "<dialect>.<knob>", e.g.: * "ir.prefix" — the IR prefix (default "I") * "tirx.prefix" — the TIR prefix (default "T") - * "relax.prefix" — the Relax prefix (default "R") - * "relax.show_all_ty" — whether to show all Relax type annotations (default true) * * Use GetExtraConfig<T>(key, fallback) to read values with a typed fallback. */ @@ -84,7 +82,7 @@ class PrinterConfigObj : public ffi::Object { * \brief Look up a value in extra_config with type cast and fallback. * * Keys are conventionally namespaced as "<dialect>.<knob>" - * (e.g. "tirx.prefix", "relax.show_all_ty"). + * (e.g. "ir.prefix", "tirx.prefix"). */ template <typename T> T GetExtraConfig(const ffi::String& key, T fallback) const { diff --git a/tirx/python/tvm_tirx/script/printer/highlight.py b/tirx/python/tvm_tirx/script/printer/highlight.py index 0abb3d82e4..603b949793 100644 --- a/tirx/python/tvm_tirx/script/printer/highlight.py +++ b/tirx/python/tvm_tirx/script/printer/highlight.py @@ -17,10 +17,7 @@ # ruff: noqa: F821, RUF012 """Highlight printed TVM script.""" -import functools import os -import shutil -import subprocess import sys import warnings from typing import Any, Union @@ -29,9 +26,8 @@ from typing import Any, Union def cprint( printable: Any | str, style: str | None = None, - black_format: bool = False, ) -> None: - """Print TVMScript string with Pygments highlight and Black auto-formatting. + """Print TVMScript string with Pygments highlighting. Parameters ---------- @@ -43,10 +39,6 @@ def cprint( Pygmentize printing style, auto-detected if None. - black_format: bool - - If true, use the formatter Black to format the TVMScript - Notes ----- @@ -68,9 +60,6 @@ def cprint( f"Only can print strings or objects with `script` method, but got: {type(printable)}" ) - if black_format: - printable = _format(printable) - is_in_notebook = "ipykernel" in sys.modules # in notebook env (support html display). style = _get_pygments_style(style, is_in_notebook) @@ -95,96 +84,6 @@ def cprint( print(highlight(printable, Python3Lexer(), Terminal256Formatter(style=style))) [email protected]_cache -def _get_formatter(formatter: str | None = None): - def get_ruff_formatter(): - if shutil.which("ruff") is None: - return None - - def formatter(code_str): - proc = subprocess.Popen( - ["ruff", "format", "--stdin-filename=TVMScript"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - encoding="utf-8", - ) - stdout, _stderr = proc.communicate(code_str) - return stdout - - return formatter - - def get_black_formatter(): - try: - # pylint: disable=import-outside-toplevel - import black - except ImportError: - return None - - def formatter(code_str): - return black.format_str(code_str, mode=black.FileMode()) - - return formatter - - def get_fallback_formatter(): - def formatter(code_str): - with warnings.catch_warnings(): - warnings.simplefilter("once", UserWarning) - ruff_install_cmd = sys.executable + " -m pip install ruff" - black_install_cmd = ( - sys.executable + ' -m pip install "black==22.3.0" --upgrade --user' - ) - warnings.warn( - f"Neither the 'ruff' formatter nor the 'black' formatter is available. " - f"To print formatted TVM script, please a formatter. \n" - f"To install ruff: {ruff_install_cmd}\n" - f"To install black: {black_install_cmd}", - category=UserWarning, - ) - return code_str - - return formatter - - # formatter = "black" - if formatter is None: - options = [get_ruff_formatter, get_black_formatter] - elif formatter == "ruff": - options = [get_ruff_formatter] - elif formatter == "black": - options = [get_black_formatter] - else: - raise ValueError(f"Unknown formatter: {formatter}") - - for option in options: - func = option() - if func is not None: - return func - return get_fallback_formatter() - - -def _format(code_str: str, formatter: str | None = None) -> str: - """Format a code string using Black. - - Parameters - ---------- - code_str: str - - The string containing Python/TVMScript code to format - - formatter: Optional[str] - - The formatter to use. Can specify `ruff`, `black`, or - auto-select by passing `None`. - - Returns - ------- - formatted: str - - The formatted Python/TVMScript code - - """ - return _get_formatter(formatter)(code_str) - - def _get_pygments_style( style: str | None, is_in_notebook: bool ) -> Union["pygments.style.Style", str] | None: diff --git a/tirx/python/tvm_tirx/script/printer/scriptable.py b/tirx/python/tvm_tirx/script/printer/scriptable.py index c6823ccbad..e958253b35 100644 --- a/tirx/python/tvm_tirx/script/printer/scriptable.py +++ b/tirx/python/tvm_tirx/script/printer/scriptable.py @@ -16,7 +16,6 @@ # under the License. """Configuration of TVMScript printer""" -import os from collections.abc import Sequence import tvm_ffi @@ -43,11 +42,11 @@ class PrinterConfig(tvm_ffi.Object): indent_spaces: int print_line_numbers: bool num_context_lines: int - extra_config: dict path_to_underline: list[AccessPath] | None path_to_annotate: dict[AccessPath, str] | None - obj_to_underline: list[AccessPath] | None - obj_to_annotate: dict[AccessPath, str] | None + obj_to_underline: list[tvm_ffi.Object] | None + obj_to_annotate: dict[tvm_ffi.Object, str] | None + extra_config: dict def __init__( self, @@ -60,12 +59,11 @@ class PrinterConfig(tvm_ffi.Object): indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int | None = None, - show_all_ty: bool = True, - extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[tvm_ffi.Object] | None = None, obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, + extra_config: dict | None = None, ) -> None: if num_context_lines is None: num_context_lines = -1 @@ -81,8 +79,6 @@ class PrinterConfig(tvm_ffi.Object): "path_to_annotate": path_to_annotate, "obj_to_underline": obj_to_underline, "obj_to_annotate": obj_to_annotate, - # Dialect-specific config via dotted keys in extra_config - "relax.show_all_ty": show_all_ty, } if name is not None: @@ -113,12 +109,11 @@ class _ScriptableMethods: indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, - show_all_ty: bool = True, - extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[tvm_ffi.Object] | None = None, obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, + extra_config: dict | None = None, ) -> str: """Print TVM IR into TVMScript text format @@ -140,15 +135,6 @@ class _ScriptableMethods: Whether to print line numbers num_context_lines : int = -1 The number of lines of context to print before and after the line to underline. - show_all_ty: bool = True - If True (default), annotate all variable bindings with the struct - info of that variable. If False, only add annotations where - required for unambiguous round-trip of Relax -> TVMScript -> Relax. - extra_config : Optional[dict] = None - Dialect-specific configuration passed through to PrinterConfig.extra_config, - including ``ir.prefix`` for the IR namespace alias (default ``"I"``). - Keys are conventionally namespaced as "<dialect>.<knob>", e.g. - ``{"tirx.prefix": "T"}``. path_to_underline : Optional[List[AccessPath]] = None Object path to be underlined path_to_annotate : Optional[Dict[AccessPath, str]] = None @@ -157,6 +143,11 @@ class _ScriptableMethods: Object to be underlined obj_to_annotate : Optional[Dict[tvm_ffi.Object, str]] = None Object to be annotated + extra_config : Optional[dict] = None + Dialect-specific configuration passed through to PrinterConfig.extra_config, + including ``ir.prefix`` for the IR namespace alias (default ``"I"``). + Keys are conventionally namespaced as "<dialect>.<knob>", e.g. + ``{"tirx.prefix": "T"}``. Returns ------- @@ -182,19 +173,17 @@ class _ScriptableMethods: indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, - show_all_ty=show_all_ty, - extra_config=merged_extra if merged_extra else None, path_to_underline=path_to_underline, path_to_annotate=path_to_annotate, obj_to_underline=obj_to_underline, obj_to_annotate=obj_to_annotate, + extra_config=merged_extra, ), ) def show( self, style: str | None = None, - black_format: bool | None = None, *, name: str | None = None, show_meta: bool = False, @@ -204,14 +193,13 @@ class _ScriptableMethods: indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, - show_all_ty: bool = True, - extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[tvm_ffi.Object] | None = None, obj_to_annotate: dict[tvm_ffi.Object, str] | None = None, + extra_config: dict | None = None, ) -> None: - """A sugar for print highlighted TVM script. + """Print highlighted TVMScript. Parameters ---------- @@ -219,25 +207,6 @@ class _ScriptableMethods: Pygmentize printing style, auto-detected if None. See `tvm_tirx.script.printer.highlight.cprint` for more details. - black_format: Optional[bool] - - If true, use the formatter Black to format the TVMScript. - If false, do not apply the auto-formatter. - - If None (default), determine the behavior based on the - environment variable "TVM_BLACK_FORMAT". If this - environment variable is unset, set to the empty string, or - set to the integer zero, black auto-formatting will be - disabled. If the environment variable is set to a - non-zero integer, black auto-formatting will be enabled. - - Note that the "TVM_BLACK_FORMAT" environment variable only - applies to the `.show()` method, and not the underlying - `.script()` method. The `.show()` method is intended for - human-readable output based on individual user - preferences, while the `.script()` method is intended to - provided a consistent output regardless of environment. - name : Optional[str] = None The name of the object show_meta : bool = False @@ -254,13 +223,6 @@ class _ScriptableMethods: Whether to print line numbers num_context_lines : int = -1 The number of lines of context to print before and after the line to underline. - show_all_ty: bool = True - If True (default), annotate all variable bindings with the struct - info of that variable. If False, only add annotations where - required for unambiguous round-trip of Relax -> TVMScript -> Relax. - extra_config : Optional[dict] = None - Dialect-specific configuration passed through to PrinterConfig.extra_config, - including ``ir.prefix`` for the IR namespace alias (default ``"I"``). path_to_underline : Optional[List[AccessPath]] = None Object path to be underlined path_to_annotate : Optional[Dict[AccessPath, str]] = None @@ -269,16 +231,15 @@ class _ScriptableMethods: Object to be underlined obj_to_annotate : Optional[Dict[tvm_ffi.Object, str]] = None Object to be annotated + extra_config : Optional[dict] = None + Dialect-specific configuration passed through to PrinterConfig.extra_config, + including ``ir.prefix`` for the IR namespace alias (default ``"I"``). """ from tvm_tirx.script.printer.highlight import ( cprint, # pylint: disable=import-outside-toplevel ) - if black_format is None: - env = os.environ.get("TVM_BLACK_FORMAT") - black_format = env and int(env) - cprint( self.script( name=name, @@ -289,15 +250,13 @@ class _ScriptableMethods: indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, - show_all_ty=show_all_ty, - extra_config=extra_config, path_to_underline=path_to_underline, path_to_annotate=path_to_annotate, obj_to_underline=obj_to_underline, obj_to_annotate=obj_to_annotate, + extra_config=extra_config, ), style=style, - black_format=black_format, )
