Bohan Hou
859498dc01
## Summary This PR adds the initial TIRx support needed for low-level programming of Blackwell-class GPU architectures. As part of the ongoing TIRx refactor, it introduces TVMScript support for directly scripting advanced hardware features without relying on scheduling as the primary programming interface. The change keeps existing `s_tir` script support intact while making direct scripting a first-class path for TIRx programs. ## Main Changes - Add TIRx operator dispatch and layout infrastructure. - Add TVMScript support for new low-level TIRx operations. - Add analysis, transform, and lowering support for TIRx IR nodes. - Add CUDA/Blackwell-oriented codegen and intrinsic coverage. - Add Python and C++ integration points for TIRx scripting and runtime support. ## Validation - `pre-commit run --all-files` - `ninja -C build -j32` - `CUDA_VISIBLE_DEVICES=2 pytest tests/python/tirx/ -n 16` - `1723 passed, 47 skipped, 32 warnings` - `CUDA_VISIBLE_DEVICES=2 python -m pytest -v tests/python/all-platform-minimal-test` - `37 passed, 105 skipped` - `TVM_TEST_TARGETS=llvm python -m pytest -v tests/python/tirx-analysis tests/python/tirx-base tests/python/tirx-transform -n 16` - `664 passed, 25 skipped, 9 xfailed, 1 xpassed` ## Local CI Notes Some full CI-equivalent jobs were not locally reproducible because this machine is missing parts of the Apache TVM CI environment, including `llvm-config-15/17`, Vulkan, ROCm, Maven, Sphinx, Doxygen, Emscripten, and ARM/QEMU cross-toolchain components. Metal-specific tests were skipped locally because no Metal runtime is available.
450 lines
18 KiB
Python
450 lines
18 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
|
|
tir_prefix: str
|
|
tir_import_module: str
|
|
relax_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",
|
|
tir_import_module: str = "tir",
|
|
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,
|
|
"tir_prefix": tir_prefix,
|
|
"tir_import_module": tir_import_module,
|
|
"relax_prefix": relax_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",
|
|
tir_import_module: str = "tir",
|
|
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.tir
|
|
tir_import_module : str = "tir"
|
|
The module name in the printed import (e.g. \"tir\" or \"tirx\").
|
|
Use tir_import_module=\"tirx\" with tir_prefix=\"Tx\" for all-Tx output.
|
|
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
|
|
|
|
"""
|
|
# Auto-switch to tirx (`Tx`/`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.
|
|
tir_prefix_val = tir_prefix
|
|
tir_import_module_val = tir_import_module
|
|
if tir_prefix == "T" and tir_import_module == "tir":
|
|
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:
|
|
tir_prefix_val = "Tx"
|
|
tir_import_module_val = "tirx"
|
|
return _script(
|
|
self,
|
|
PrinterConfig(
|
|
name=name,
|
|
show_meta=show_meta,
|
|
ir_prefix=ir_prefix,
|
|
tir_prefix=tir_prefix_val,
|
|
tir_import_module=tir_import_module_val,
|
|
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",
|
|
tir_import_module: str = "tir",
|
|
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,
|
|
tir_import_module=tir_import_module,
|
|
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",
|
|
tir_import_module: str = "tir",
|
|
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 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,
|
|
tir_prefix=tir_prefix,
|
|
tir_import_module=tir_import_module,
|
|
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,
|
|
)
|