chore: import upstream snapshot with attribution

python/tvm/script/printer/__init__.py

@@ -0,0 +1,21 @@
+# 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.
+"""
+TVMScript Unified Printer
+This package provides a set of APIs to print supported TVM IR into TVMScript
+in a roundtrippable way.
+"""

python/tvm/script/printer/_ffi_api.py

@@ -0,0 +1,21 @@
+# 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.
+"""FFI APIs for tvm.script.printer"""
+
+import tvm_ffi
+
+tvm_ffi.init_ffi_api("script.printer", __name__)  # pylint: disable=protected-access

python/tvm/script/printer/doc.py

@@ -0,0 +1,520 @@
+# 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.
+"""Doc types for TVMScript Unified Printer"""
+
+from collections.abc import Sequence
+from enum import IntEnum, unique
+from typing import Union
+
+from tvm_ffi import register_object
+from tvm_ffi.access_path import AccessPath
+
+from tvm.runtime import Object
+from tvm.tirx import FloatImm, IntImm
+
+from . import _ffi_api
+
+
+@register_object("script.printer.Doc")
+class Doc(Object):
+    """Base class of all Docs"""
+
+
+@register_object("script.printer.ExprDoc")
+class ExprDoc(Doc):
+    """Base class of all expression Docs"""
+
+    def attr(self, name: str) -> "AttrAccessDoc":
+        """
+        Create a doc that represents attribute access on self.
+
+        Parameters
+        ----------
+        name : str
+            The attribute name to access
+
+        Returns
+        -------
+        doc : AttrAccessDoc
+        """
+        return _ffi_api.ExprDocAttr(self, name)  # type: ignore # pylint: disable=no-member
+
+    def call(self, *args: tuple["ExprDoc"], **kwargs: dict[str, "ExprDoc"]) -> "CallDoc":
+        """
+        Create a doc that represents function call, with self as callee.
+
+        Parameters
+        ----------
+        *args : ExprDoc
+            The positional arguments of the function call.
+        **kwargs
+            The keyword arguments of the function call.
+
+        Returns
+        -------
+        doc : CallDoc
+        """
+        kwargs_keys = list(kwargs.keys())
+        kwargs_values = list(kwargs.values())
+        return _ffi_api.ExprDocCall(self, args, kwargs_keys, kwargs_values)  # type: ignore # pylint: disable=no-member
+
+    _IndexType = Union["ExprDoc", "SliceDoc"]
+
+    def __getitem__(self, indices: tuple[_IndexType] | _IndexType) -> "IndexDoc":
+        """
+        Create a doc that represents index access on self.
+
+        Parameters
+        ----------
+        indices : Union[Tuple[Union["ExprDoc", "SliceDoc"]], Union["ExprDoc", "SliceDoc"]]
+            The indices to access
+
+        Returns
+        -------
+        doc : IndexDoc
+        """
+        if not isinstance(indices, tuple):
+            indices = (indices,)
+        return _ffi_api.ExprDocIndex(self, indices)  # type: ignore # pylint: disable=no-member
+
+    def __iter__(self):
+        """
+        This is implemented to prevent confusing error message when trying to use ExprDoc
+        as iterable. According to PEP-234, An object can be iterated over if it
+        implements __iter__() or __getitem__(). If an object has only __getitem__
+        but not __iter__, interpreter will iterate the object by calling
+        __getitem__ with 0, 1, 2, ..., until an IndexError is raised.
+
+        https://peps.python.org/pep-0234/#python-api-specification
+        """
+        raise RuntimeError(f"{self.__class__} cannot be used as iterable.")
+
+
+@register_object("script.printer.StmtDoc")
+class StmtDoc(Doc):
+    """Base class of statement doc"""
+
+
+@register_object("script.printer.StmtBlockDoc")
+class StmtBlockDoc(Doc):
+    """The container doc that holds a list of StmtDoc.
+
+    Note: `StmtBlockDoc` is never used in the IR, but a temporary container that allows holding a
+    list of StmtDoc.
+    """
+
+    stmts: Sequence[StmtDoc]
+
+    def __init__(self, stmts: list[StmtDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.StmtBlockDoc, stmts)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.LiteralDoc")
+class LiteralDoc(ExprDoc):
+    """Doc that represents literal value"""
+
+    value: str | IntImm | FloatImm | None
+
+    def __init__(
+        self,
+        value: str | float | bool | int | None,
+        path: AccessPath | None = None,
+    ):
+        if value is None:
+            self.__init_handle_by_constructor__(_ffi_api.LiteralDocNone, path)  # type: ignore # pylint: disable=no-member
+        elif isinstance(value, str):
+            self.__init_handle_by_constructor__(
+                _ffi_api.LiteralDocStr,  # type: ignore # pylint: disable=no-member
+                value,
+                path,
+            )
+        elif isinstance(value, float):
+            self.__init_handle_by_constructor__(
+                _ffi_api.LiteralDocFloat,  # type: ignore # pylint: disable=no-member
+                value,
+                path,
+            )
+        elif isinstance(value, bool):
+            self.__init_handle_by_constructor__(
+                _ffi_api.LiteralDocBoolean,  # type: ignore # pylint: disable=no-member
+                value,
+                path,
+            )
+        elif isinstance(value, int):
+            self.__init_handle_by_constructor__(
+                _ffi_api.LiteralDocInt,  # type: ignore # pylint: disable=no-member
+                value,
+                path,
+            )
+        else:
+            raise TypeError(f"Unsupported type {type(value)} for LiteralDoc")
+
+
+@register_object("script.printer.IdDoc")
+class IdDoc(ExprDoc):
+    """Doc that represents identifier"""
+
+    name: str
+
+    def __init__(self, name: str):
+        self.__init_handle_by_constructor__(_ffi_api.IdDoc, name)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.AttrAccessDoc")
+class AttrAccessDoc(ExprDoc):
+    """Doc that represents attribute access on an expression"""
+
+    value: ExprDoc
+    name: str
+
+    def __init__(self, value: ExprDoc, name: str):
+        self.__init_handle_by_constructor__(_ffi_api.AttrAccessDoc, value, name)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.IndexDoc")
+class IndexDoc(ExprDoc):
+    """Doc that represents index access on an expression"""
+
+    value: ExprDoc
+    indices: Sequence[Union[ExprDoc, "SliceDoc"]]
+
+    def __init__(self, value: ExprDoc, indices: list[Union[ExprDoc, "SliceDoc"]]):
+        self.__init_handle_by_constructor__(_ffi_api.IndexDoc, value, indices)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.CallDoc")
+class CallDoc(ExprDoc):
+    """Doc that represents function call"""
+
+    callee: ExprDoc
+    args: Sequence[ExprDoc]
+    kwargs_keys: Sequence[str]
+    kwargs_values: Sequence[ExprDoc]
+
+    def __init__(self, callee: ExprDoc, *args: tuple[ExprDoc], **kwargs: dict[str, ExprDoc]):
+        kwargs_keys = list(kwargs.keys())
+        kwargs_values = list(kwargs.values())
+        self.__init_handle_by_constructor__(
+            _ffi_api.CallDoc,  # type: ignore # pylint: disable=no-member
+            callee,
+            args,
+            kwargs_keys,
+            kwargs_values,
+        )
+
+
+@unique
+class OperationKind(IntEnum):
+    """
+    This enum represents the kind of operation (operator) in OperationDoc
+
+    It's mirrored from OperationDocNode::Kind at include/tvm/script/printer/doc.h
+    """
+
+    # The name convention follows https://docs.python.org/3/library/ast.html
+    # pylint: disable=invalid-name
+
+    _UnaryStart = 0
+    USub = 1
+    Invert = 2
+    Not = 3
+    _UnaryEnd = 4
+
+    _BinaryStart = 5
+    Add = 6
+    Sub = 7
+    Mult = 8
+    Div = 9
+    FloorDiv = 10
+    Mod = 11
+    Pow = 12
+    LShift = 13
+    RShift = 14
+    BitAnd = 15
+    BitOr = 16
+    BitXor = 17
+    Lt = 18
+    LtE = 19
+    Eq = 20
+    NotEq = 21
+    Gt = 22
+    GtE = 23
+    And = 24
+    Or = 25
+    MatMul = 26
+    _BinaryEnd = 27
+
+    _SpecialStart = 28
+    IfThenElse = 29
+    _SpecialEnd = 30
+
+    # pylint: enable=invalid-name
+
+
+@register_object("script.printer.OperationDoc")
+class OperationDoc(ExprDoc):
+    """
+    Doc that represents operation
+
+    It can be unary, binary and other special operators (for example, the
+    if-then-else expression).
+    """
+
+    kind: OperationKind
+    operands: Sequence[ExprDoc]
+
+    def __init__(self, kind: OperationKind, operands: list[ExprDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.OperationDoc, kind, operands)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.LambdaDoc")
+class LambdaDoc(ExprDoc):
+    """Doc that represents lambda function"""
+
+    args: Sequence[IdDoc]
+    body: ExprDoc
+
+    def __init__(self, args: list[IdDoc], body: ExprDoc):
+        self.__init_handle_by_constructor__(_ffi_api.LambdaDoc, args, body)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.TupleDoc")
+class TupleDoc(ExprDoc):
+    """Doc that represents tuple literal"""
+
+    elements: Sequence[ExprDoc]
+
+    def __init__(self, elements: list[ExprDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.TupleDoc, elements)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.ListDoc")
+class ListDoc(ExprDoc):
+    """Doc that represents list literal"""
+
+    elements: Sequence[ExprDoc]
+
+    def __init__(self, elements: list[ExprDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.ListDoc, elements)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.DictDoc")
+class DictDoc(ExprDoc):
+    """Doc that represents dict literal"""
+
+    keys: Sequence[ExprDoc]
+    values: Sequence[ExprDoc]
+
+    def __init__(self, content: dict[ExprDoc, ExprDoc]):
+        keys = list(content.keys())
+        values = list(content.values())
+        self.__init_handle_by_constructor__(_ffi_api.DictDoc, keys, values)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.SliceDoc")
+class SliceDoc(ExprDoc):
+    """
+    Doc that represents slice in Index expression
+
+    This doc can only appear in `IndexDoc.indices`.
+    """
+
+    start: ExprDoc | None
+    stop: ExprDoc | None
+    step: ExprDoc | None
+
+    def __init__(
+        self,
+        start: ExprDoc | None = None,
+        stop: ExprDoc | None = None,
+        step: ExprDoc | None = None,
+    ):
+        self.__init_handle_by_constructor__(_ffi_api.SliceDoc, start, stop, step)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.AssignDoc")
+class AssignDoc(StmtDoc):
+    """Doc that represents assign statement."""
+
+    lhs: ExprDoc
+    rhs: ExprDoc | None
+    annotation: ExprDoc | None
+
+    def __init__(self, lhs: ExprDoc, rhs: ExprDoc | None, annotation: ExprDoc | None = None):
+        self.__init_handle_by_constructor__(
+            _ffi_api.AssignDoc,  # type: ignore # pylint: disable=no-member
+            lhs,
+            rhs,
+            annotation,
+        )
+
+
+@register_object("script.printer.IfDoc")
+class IfDoc(StmtDoc):
+    """Doc that represent if-then-else statement."""
+
+    predicate: ExprDoc
+    then_branch: Sequence[StmtDoc]
+    else_branch: Sequence[StmtDoc]
+
+    def __init__(self, predicate: ExprDoc, then_branch: list[StmtDoc], else_branch: list[StmtDoc]):
+        self.__init_handle_by_constructor__(
+            _ffi_api.IfDoc,  # type: ignore # pylint: disable=no-member
+            predicate,
+            then_branch,
+            else_branch,
+        )
+
+
+@register_object("script.printer.WhileDoc")
+class WhileDoc(StmtDoc):
+    """Doc that represents while statement."""
+
+    predicate: ExprDoc
+    body: Sequence[StmtDoc]
+
+    def __init__(self, predicate: ExprDoc, body: list[StmtDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.WhileDoc, predicate, body)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.ForDoc")
+class ForDoc(StmtDoc):
+    """Doc that represents for statement."""
+
+    lhs: ExprDoc
+    rhs: ExprDoc
+    body: Sequence[StmtDoc]
+
+    def __init__(self, lhs: ExprDoc, rhs: ExprDoc, body: list[StmtDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.ForDoc, lhs, rhs, body)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.ScopeDoc")
+class ScopeDoc(StmtDoc):
+    """
+    Doc that represents special scopes.
+
+    Specifically, this means the with statement in Python:
+
+    with <rhs> as <lhs>:
+        <body...>
+    """
+
+    lhs: ExprDoc | None
+    rhs: ExprDoc
+    body: Sequence[StmtDoc]
+
+    def __init__(self, lhs: ExprDoc | None, rhs: ExprDoc, body: list[StmtDoc]):
+        self.__init_handle_by_constructor__(_ffi_api.ScopeDoc, lhs, rhs, body)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.ExprStmtDoc")
+class ExprStmtDoc(StmtDoc):
+    """Doc that represents an expression as statement."""
+
+    expr: ExprDoc
+
+    def __init__(self, expr: ExprDoc):
+        self.__init_handle_by_constructor__(_ffi_api.ExprStmtDoc, expr)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.AssertDoc")
+class AssertDoc(StmtDoc):
+    """Doc that represents assert statement."""
+
+    test: ExprDoc
+    msg: ExprDoc | None
+
+    def __init__(self, test: ExprDoc, msg: ExprDoc | None = None):
+        self.__init_handle_by_constructor__(_ffi_api.AssertDoc, test, msg)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.ReturnDoc")
+class ReturnDoc(StmtDoc):
+    """Doc that represents return statement."""
+
+    value: ExprDoc
+
+    def __init__(self, value: ExprDoc):
+        self.__init_handle_by_constructor__(_ffi_api.ReturnDoc, value)  # type: ignore # pylint: disable=no-member
+
+
+@register_object("script.printer.FunctionDoc")
+class FunctionDoc(StmtDoc):
+    """Doc that represents function definition."""
+
+    name: IdDoc
+    args: Sequence[AssignDoc]
+    decorators: Sequence[ExprDoc]
+    return_type: ExprDoc | None
+    body: Sequence[StmtDoc]
+
+    def __init__(
+        self,
+        name: IdDoc,
+        args: list[AssignDoc],
+        decorators: list[ExprDoc],
+        return_type: ExprDoc | None,
+        body: list[StmtDoc],
+    ):
+        self.__init_handle_by_constructor__(
+            _ffi_api.FunctionDoc,  # type: ignore # pylint: disable=no-member
+            name,
+            args,
+            decorators,
+            return_type,
+            body,
+        )
+
+
+@register_object("script.printer.ClassDoc")
+class ClassDoc(StmtDoc):
+    """Doc that represents class definition."""
+
+    name: IdDoc
+    decorators: Sequence[ExprDoc]
+    body: Sequence[StmtDoc]
+
+    def __init__(self, name: IdDoc, decorators: list[ExprDoc], body: list[StmtDoc]):
+        self.__init_handle_by_constructor__(
+            _ffi_api.ClassDoc,  # type: ignore # pylint: disable=no-member
+            name,
+            decorators,
+            body,
+        )
+
+
+@register_object("script.printer.CommentDoc")
+class CommentDoc(StmtDoc):
+    """Doc that represents comment."""
+
+    def __init__(self, comment: str):
+        self.__init_handle_by_constructor__(
+            _ffi_api.CommentDoc,
+            comment,  # type: ignore # pylint: disable=no-member
+        )
+
+
+@register_object("script.printer.DocStringDoc")
+class DocStringDoc(StmtDoc):
+    """Doc that represents docstring."""
+
+    def __init__(self, docs: str):
+        self.__init_handle_by_constructor__(
+            _ffi_api.DocStringDoc,
+            docs,  # type: ignore # pylint: disable=no-member
+        )

python/tvm/script/printer/doc_printer.py

@@ -0,0 +1,60 @@
+# 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.
+"""Functions to print doc into text format"""
+
+from tvm_ffi.access_path import AccessPath
+
+from tvm.runtime.script_printer import PrinterConfig
+
+from . import _ffi_api
+from .doc import Doc
+
+
+def to_python_script(
+    doc: Doc,
+    indent_spaces: int = 4,
+    print_line_numbers: bool = False,
+    num_context_lines: int | None = None,
+    path_to_underline: list[AccessPath] | None = None,
+) -> str:
+    """Convert Doc into Python script.
+
+    Parameters
+    ----------
+    doc : Doc
+        The doc to convert into Python script
+    indent_spaces : int
+        The number of indent spaces to use in the output
+    print_line_numbers: bool
+        Whether to print line numbers
+    num_context_lines : Optional[int]
+        Number of context lines to print around the underlined text
+    path_to_underline : Optional[AccessPath]
+        Object path to be underlined
+
+    Returns
+    -------
+    script : str
+        The text representation of Doc in Python syntax
+    """
+    cfg = PrinterConfig(
+        indent_spaces=indent_spaces,
+        print_line_numbers=print_line_numbers,
+        num_context_lines=num_context_lines,
+        path_to_underline=path_to_underline,
+    )
+    return _ffi_api.DocToPythonScript(doc, cfg)  # type: ignore # pylint: disable=no-member