Tianqi Chen
7504e3ed1a
## Summary Restructure TVMScript to be dialect-agnostic at the script-core layer while letting each extension dialect (TIRX, Relax) own its own per-dialect script subtree. IR is below script in the dependency stack and is NOT a peer dialect — its script handlers stay in the shared core. This PR folds together two coupled refactors that were initially opened as separate PRs (#19478 and the original #19479); they share rename / relocation surface so they ship as one cohesive change. ## What this PR does ### Per-dialect script subtree (originally #19479) - Moves per-dialect printer + builder from `src/script/{printer,ir_builder}/{tirx,relax}/` to `src/{tirx,relax}/script/{printer,builder}/`. - Tightens `src/script/*.cc` CMake glob to the dialect-free core. - Refactors `IRBuilder::DeclFunction` to dispatch via FFI registry (`script.ir_builder.decl_function.<type-key>`); removes cross-dialect includes from the shared core. - Adds `tvm.script.register_dialect` API + `__getattr__` + a `sys.meta_path` finder for Python-side dialect discovery. In-tree dialects (tirx, relax) registered centrally in `python/tvm/__init__.py`. - Drops the obsolete static re-export shims at `python/tvm/script/{parser,ir_builder}/{tirx,relax}/`. ### Dialect-agnostic printer config (originally #19478) - Relocates `include/tvm/ir/script_printer.h` → `include/tvm/script/printer/config.h` next to the rest of the printer's public surface. The header is not IR-specific. - Renames `TVM_SCRIPT_REPR` → `TVM_REGISTER_SCRIPT_AS_REPR` for clarity (the macro registers Script as the kRepr callback + per-type vtable dispatch). Aligns with the `TVM_REGISTER_*` family. - Drops dialect-hardcoded `PrinterConfig` fields (`tir_prefix`, `relax_prefix`, `show_all_struct_info`, `buffer_dtype`) in favor of a generic `ffi::Map<String, Any> extra_config` keyed by `"<dialect>.<knob>"`. Each call site reads via the templated accessor `config->GetExtraConfig<T>("...", default)`. - Promotes `std::string` config fields to `ffi::String`. After this lands, the script-printer core knows nothing specific about any dialect — new dialects plug in via the registry pattern with zero core edits. Public Python API surface unchanged.
405 lines
15 KiB
Python
405 lines
15 KiB
Python
# 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
|
|
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",
|
|
tir_prefix: str = "T",
|
|
relax_prefix: str = "R",
|
|
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_struct_info: bool = True,
|
|
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 = {
|
|
"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,
|
|
"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
|
|
"tirx.prefix": tir_prefix,
|
|
"tirx.buffer_dtype": buffer_dtype,
|
|
"relax.prefix": relax_prefix,
|
|
"relax.show_all_struct_info": show_all_struct_info,
|
|
}
|
|
|
|
if name is not None:
|
|
cfg["name"] = name
|
|
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",
|
|
tir_prefix: str = "T",
|
|
relax_prefix: str = "R",
|
|
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 = -1,
|
|
syntax_sugar: bool = True,
|
|
show_object_address: bool = False,
|
|
show_all_struct_info: bool = True,
|
|
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
|
|
tir_prefix : str = "T"
|
|
The prefix of AST nodes from tvm.tirx
|
|
relax_prefix : str = "R"
|
|
The prefix of AST nodes from tvm.relax
|
|
module_alias : str = "cls"
|
|
The alias of the current module at cross-function call,
|
|
Directly use module name if it's empty.
|
|
buffer_dtype : str = "float32"
|
|
The default data type of buffer
|
|
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_struct_info: 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.
|
|
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
|
|
|
|
"""
|
|
return _script(
|
|
self,
|
|
PrinterConfig(
|
|
name=name,
|
|
show_meta=show_meta,
|
|
ir_prefix=ir_prefix,
|
|
tir_prefix=tir_prefix,
|
|
relax_prefix=relax_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,
|
|
show_all_struct_info=show_all_struct_info,
|
|
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",
|
|
tir_prefix: str = "T",
|
|
relax_prefix: str = "R",
|
|
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 = -1,
|
|
syntax_sugar: bool = True,
|
|
show_object_address: bool = False,
|
|
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,
|
|
tir_prefix=tir_prefix,
|
|
relax_prefix=relax_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,
|
|
),
|
|
)
|
|
|
|
def show(
|
|
self,
|
|
style: str | None = None,
|
|
black_format: bool | None = None,
|
|
*,
|
|
name: str | None = None,
|
|
show_meta: bool = False,
|
|
ir_prefix: str = "I",
|
|
tir_prefix: str = "T",
|
|
relax_prefix: str = "R",
|
|
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 = -1,
|
|
syntax_sugar: bool = True,
|
|
show_object_address: bool = False,
|
|
show_all_struct_info: bool = True,
|
|
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
|
|
tir_prefix : str = "T"
|
|
The prefix of AST nodes from tvm.tirx
|
|
relax_prefix : str = "R"
|
|
The prefix of AST nodes from tvm.relax
|
|
module_alias : str = "cls"
|
|
The alias of the current module at cross-function call,
|
|
Directly use module name if it's empty.
|
|
buffer_dtype : str = "float32"
|
|
The default data type of buffer
|
|
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_struct_info: 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.
|
|
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 ( # pylint: disable=import-outside-toplevel
|
|
cprint,
|
|
)
|
|
|
|
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,
|
|
tir_prefix=tir_prefix,
|
|
relax_prefix=relax_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,
|
|
show_all_struct_info=show_all_struct_info,
|
|
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,
|
|
)
|