# 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. """Configuration of TVMScript printer""" import os from collections.abc import Sequence from tvm_ffi import get_global_func, register_object from tvm_ffi.access_path import AccessPath from tvm.runtime import Object from . import _ffi_node_api @register_object("script.PrinterConfig") class PrinterConfig(Object): """Configuration of TVMScript printer""" binding_names: Sequence[str] show_meta: bool ir_prefix: str module_alias: str buffer_dtype: str int_dtype: str float_dtype: str verbose_expr: bool indent_spaces: int print_line_numbers: bool num_context_lines: int syntax_sugar: bool show_object_address: bool 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 def __init__( self, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = "I", module_alias: str = "cls", buffer_dtype: str = "float32", int_dtype: str = "int32", float_dtype: str = "void", verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int | None = None, syntax_sugar: bool = True, show_object_address: bool = False, 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[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None, ) -> None: if num_context_lines is None: num_context_lines = -1 cfg: dict = { "show_meta": show_meta, "ir_prefix": ir_prefix, "module_alias": module_alias, "buffer_dtype": buffer_dtype, "int_dtype": int_dtype, "float_dtype": float_dtype, "verbose_expr": verbose_expr, "indent_spaces": indent_spaces, "print_line_numbers": print_line_numbers, "num_context_lines": num_context_lines, "syntax_sugar": syntax_sugar, "show_object_address": show_object_address, "path_to_underline": path_to_underline, "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: cfg["name"] = name if extra_config is not None: cfg["extra_config"] = extra_config self.__init_handle_by_constructor__( _ffi_node_api.PrinterConfig, cfg, # type: ignore # pylint: disable=no-member ) def _script(obj: Object, config: PrinterConfig) -> str: return _ffi_node_api.TVMScriptPrinterScript(obj, config) # type: ignore # pylint: disable=no-member def _relax_script(obj: Object, config: PrinterConfig) -> str: func = get_global_func("script.printer.ReprPrintRelax") return func(obj, config) class Scriptable: """A base class that enables the script() and show() method.""" def script( self, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = "I", module_alias: str = "cls", int_dtype: str = "int32", float_dtype: str = "void", verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, show_object_address: bool = False, 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[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None, ) -> str: """Print TVM IR into TVMScript text format Parameters ---------- name : Optional[str] = None The name of the object show_meta : bool = False Whether to print the meta data of the object ir_prefix : str = "I" The prefix of AST nodes from tvm.ir module_alias : str = "cls" The alias of the current module at cross-function call, Directly use module name if it's empty. int_dtype : str = "int32" The default data type of integer float_dtype : str = "void" The default data type of float verbose_expr : bool = False Whether to print the detailed definition of each variable in the expression indent_spaces : int = 4 The number of spaces for indentation print_line_numbers : bool = False 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. syntax_sugar: bool = True Whether to output with syntax sugar, set false for complete printing. show_object_address: bool = False Whether to include the object's address as part of the TVMScript name 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. Keys are conventionally namespaced as ".", 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 Object path to be annotated obj_to_underline : Optional[List[Object]] = None Object to be underlined obj_to_annotate : Optional[Dict[Object, str]] = None Object to be annotated Returns ------- script : str The TVM Script of the given TVM IR """ # Auto-switch to tirx (`T`/`tirx`) flavor only when explicitly # printing a PrimFunc / IRModule that has no s_tir-tagged content. # Free objects (Buffer, BufferRegion, ...) keep the default `T`/`tir` # flavor -- they have no enclosing function to indicate tirx vs s_tir. merged_extra: dict = {} if extra_config is not None: merged_extra.update(extra_config) # Only auto-switch if the caller has not already set a tirx.prefix override. if "tirx.prefix" not in merged_extra: from tvm.ir import IRModule # pylint: disable=import-outside-toplevel from tvm.tirx import PrimFunc # pylint: disable=import-outside-toplevel switch_to_tirx = False if isinstance(self, PrimFunc): attrs = getattr(self, "attrs", None) if attrs is None or not attrs.get("s_tir", False): switch_to_tirx = True elif isinstance(self, IRModule): any_prim = False any_s_tir = False for _, base_func in self.functions.items(): if isinstance(base_func, PrimFunc): any_prim = True if getattr(base_func, "attrs", None) and base_func.attrs.get( "s_tir", False ): any_s_tir = True break if any_prim and not any_s_tir: switch_to_tirx = True if switch_to_tirx: merged_extra["tirx.prefix"] = "T" return _script( self, PrinterConfig( name=name, show_meta=show_meta, ir_prefix=ir_prefix, module_alias=module_alias, int_dtype=int_dtype, float_dtype=float_dtype, verbose_expr=verbose_expr, indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, show_object_address=show_object_address, 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, ), ) def _relax_script( self, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = "I", module_alias: str = "cls", int_dtype: str = "int32", float_dtype: str = "void", verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, show_object_address: bool = False, extra_config: dict | None = None, path_to_underline: list[AccessPath] | None = None, path_to_annotate: dict[AccessPath, str] | None = None, obj_to_underline: list[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None, ) -> str: return _relax_script( self, PrinterConfig( name=name, show_meta=show_meta, ir_prefix=ir_prefix, module_alias=module_alias, int_dtype=int_dtype, float_dtype=float_dtype, verbose_expr=verbose_expr, indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, show_object_address=show_object_address, 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, ), ) def show( self, style: str | None = None, black_format: bool | None = None, *, name: str | None = None, show_meta: bool = False, ir_prefix: str = "I", module_alias: str = "cls", int_dtype: str = "int32", float_dtype: str = "void", verbose_expr: bool = False, indent_spaces: int = 4, print_line_numbers: bool = False, num_context_lines: int = -1, syntax_sugar: bool = True, show_object_address: bool = False, 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[Object] | None = None, obj_to_annotate: dict[Object, str] | None = None, ) -> None: """A sugar for print highlighted TVM script. Parameters ---------- style : str, optional Pygmentize printing style, auto-detected if None. See `tvm.script.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 Whether to print the meta data of the object ir_prefix : str = "I" The prefix of AST nodes from tvm.ir module_alias : str = "cls" The alias of the current module at cross-function call, Directly use module name if it's empty. int_dtype : str = "int32" The default data type of integer float_dtype : str = "void" The default data type of float verbose_expr : bool = False Whether to print the detailed definition of each variable in the expression indent_spaces : int = 4 The number of spaces for indentation print_line_numbers : bool = False 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. syntax_sugar: bool = True Whether to output with syntax sugar, set false for complete printing. show_object_address: bool = False Whether to include the object's address as part of the TVMScript name 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. path_to_underline : Optional[List[AccessPath]] = None Object path to be underlined path_to_annotate : Optional[Dict[AccessPath, str]] = None Object path to be annotated obj_to_underline : Optional[List[Object]] = None Object to be underlined obj_to_annotate : Optional[Dict[Object, str]] = None Object to be annotated """ from tvm.script.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, show_meta=show_meta, ir_prefix=ir_prefix, module_alias=module_alias, int_dtype=int_dtype, float_dtype=float_dtype, verbose_expr=verbose_expr, indent_spaces=indent_spaces, print_line_numbers=print_line_numbers, num_context_lines=num_context_lines, syntax_sugar=syntax_sugar, show_object_address=show_object_address, 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, ), style=style, black_format=black_format, )