wehub-resource-sync
521 lines
15 KiB
Python
521 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.
|
|
"""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
|
|
)
|