chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,238 @@
|
||||
# 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 public namespace.
|
||||
|
||||
Dialect resolution mechanism
|
||||
----------------------------
|
||||
|
||||
``tvm.script`` is a virtual namespace: dialect names like ``tirx`` and
|
||||
``relax`` are not bound as static attributes here. Instead:
|
||||
|
||||
- ``register_dialect(name, module_path)`` writes an entry to
|
||||
``_DIALECT_REGISTRY: dict[str, str]``. Each in-tree dialect's
|
||||
``__init__.py`` calls this on import (e.g., ``tvm.tirx.__init__.py``
|
||||
calls ``tvm.script.register_dialect("tirx", "tvm.tirx.script")``).
|
||||
Out-of-tree dialects can register themselves the same way.
|
||||
|
||||
- ``__getattr__(name)`` (PEP 562) fires on missing attribute access.
|
||||
If ``name`` is in ``_DIALECT_REGISTRY``, the listed module is imported
|
||||
and cached as a normal module attribute. Subsequent accesses
|
||||
skip ``__getattr__`` (cached in ``globals()``).
|
||||
|
||||
- Subpackages ``tvm.script.parser``, ``tvm.script.ir_builder``, etc.
|
||||
each define their own ``__getattr__`` that consults the SAME
|
||||
``_DIALECT_REGISTRY`` and appends their suffix. So
|
||||
``tvm.script.parser.tirx`` resolves to ``tvm.tirx.script.parser`` via
|
||||
the dialect registry + ``.parser`` suffix.
|
||||
|
||||
- For deep statement-form imports like
|
||||
``from tvm.script.parser.tirx.entry import ObjectProxy``, PEP 562's
|
||||
``__getattr__`` is not enough — it only handles one-level
|
||||
``from X import Y``. A ``sys.meta_path`` finder (see
|
||||
``_DialectRedirectFinder``) intercepts the import machinery to
|
||||
register the real module under the legacy name in ``sys.modules``,
|
||||
so subsequent attribute walks resolve correctly.
|
||||
|
||||
Each dialect's ``tvm.<dialect>.script`` package MUST expose ``parser``,
|
||||
``ir_builder``, and (where applicable) ``printer`` as submodules. This
|
||||
convention is what makes the suffix-append redirect work uniformly.
|
||||
IR is foundational (script depends on ir) and is NOT a dialect; its
|
||||
script handlers live in the shared core, not via this registry.
|
||||
|
||||
Bootstrap order
|
||||
---------------
|
||||
|
||||
``python/tvm/__init__.py`` imports ``tvm.script`` BEFORE importing any
|
||||
dialect package (``tvm.tirx``, ``tvm.relax``, …). This guarantees that
|
||||
``tvm.script.register_dialect`` is reachable the moment a dialect's own
|
||||
``__init__.py`` runs and calls it. The ``tvm.script`` module itself
|
||||
stays dialect-agnostic at load time (no dialect submodules are eagerly
|
||||
imported here), so there is no circular dependency.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import importlib.util
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
_DIALECT_REGISTRY: dict[str, str] = {}
|
||||
|
||||
# Subpackages of `tvm.script` whose per-dialect children are redirected to a
|
||||
# matching subpackage under `tvm.<dialect>.script`. The values are the
|
||||
# subpackage name on the dialect side (e.g. `tvm.<dialect>.script.parser`).
|
||||
_REDIRECTED_SUBPACKAGES = {
|
||||
"tvm.script.parser": "parser",
|
||||
"tvm.script.ir_builder": "builder",
|
||||
}
|
||||
|
||||
|
||||
def register_dialect(name: str, module_path: str) -> None:
|
||||
"""Register a dialect's script package path.
|
||||
|
||||
Writes ``name -> module_path`` into ``_DIALECT_REGISTRY``. After
|
||||
registration, ``tvm.script.<name>`` resolves to ``module_path`` via
|
||||
``__getattr__``, and ``tvm.script.parser.<name>`` / ``tvm.script.ir_builder.<name>``
|
||||
resolve to ``module_path + ".parser"`` / ``module_path + ".builder"`` etc.
|
||||
via each subpackage's own ``__getattr__``. Deep statement-form imports
|
||||
(e.g., ``from tvm.script.parser.<name>.entry import X``) are handled
|
||||
by ``_DialectRedirectFinder`` on ``sys.meta_path``.
|
||||
|
||||
This function is idempotent — re-registering the same name with the same
|
||||
path is harmless.
|
||||
|
||||
Each in-tree dialect calls this from its own ``__init__.py``::
|
||||
|
||||
import tvm.script
|
||||
tvm.script.register_dialect("tirx", "tvm.tirx.script")
|
||||
|
||||
Out-of-tree dialects do the same in their own package init without
|
||||
editing any in-tree file.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The short name exposed under ``tvm.script.<name>`` (e.g. ``"tirx"``).
|
||||
module_path : str
|
||||
The full dotted module path of the dialect's script package, e.g.
|
||||
``"tvm.tirx.script"``. That package must expose ``parser`` and
|
||||
``ir_builder`` as submodules (and ``printer`` where applicable) so
|
||||
that the suffix-append redirect works uniformly.
|
||||
"""
|
||||
_DIALECT_REGISTRY[name] = module_path
|
||||
|
||||
|
||||
def _redirect_target(fullname: str) -> str | None:
|
||||
"""Return the target module path for a redirected ``tvm.script[...]`` name.
|
||||
|
||||
Returns ``None`` if ``fullname`` is not a redirected name.
|
||||
"""
|
||||
if fullname.startswith("tvm.script."):
|
||||
# tvm.script.<dialect>[.subpath]
|
||||
rest = fullname[len("tvm.script.") :]
|
||||
head, _, tail = rest.partition(".")
|
||||
if head in _DIALECT_REGISTRY and "." not in head:
|
||||
target = _DIALECT_REGISTRY[head]
|
||||
return f"{target}.{tail}" if tail else target
|
||||
# tvm.script.parser.<dialect>[.subpath] / tvm.script.ir_builder.<dialect>[.subpath]
|
||||
for prefix, sub in _REDIRECTED_SUBPACKAGES.items():
|
||||
if fullname == prefix or not fullname.startswith(prefix + "."):
|
||||
continue
|
||||
rest = fullname[len(prefix) + 1 :]
|
||||
head, _, tail = rest.partition(".")
|
||||
if head in _DIALECT_REGISTRY:
|
||||
target = f"{_DIALECT_REGISTRY[head]}.{sub}"
|
||||
return f"{target}.{tail}" if tail else target
|
||||
return None
|
||||
|
||||
|
||||
class _DialectRedirectFinder:
|
||||
"""``sys.meta_path`` finder that redirects ``tvm.script.<dialect>`` import paths.
|
||||
|
||||
PEP 562 ``__getattr__`` only handles one-level attribute lookups
|
||||
(``from tvm.script import tirx``). It cannot intercept deep
|
||||
statement-form imports such as::
|
||||
|
||||
from tvm.script.parser.tirx.entry import ObjectProxy
|
||||
import tvm.script.ir_builder.relax.ir
|
||||
|
||||
This finder is installed on ``sys.meta_path`` to cover those cases.
|
||||
When the import machinery asks for a module whose full name starts with
|
||||
``tvm.script.<dialect>`` (or ``tvm.script.parser.<dialect>``, etc.) and
|
||||
that dialect is in ``_DIALECT_REGISTRY``, :meth:`find_spec` imports the
|
||||
real target module (e.g. ``tvm.tirx.script.parser.entry``) and returns an
|
||||
alias spec whose loader hands back that module, so the import machinery
|
||||
registers it in ``sys.modules`` under the legacy name and all subsequent
|
||||
imports and attribute walks resolve without going through the redirect
|
||||
again.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def find_spec(cls, fullname, path, target=None):
|
||||
redirected = _redirect_target(fullname)
|
||||
if redirected is None:
|
||||
return None
|
||||
# Resolve the target module and return an alias spec for it. Do NOT
|
||||
# pre-register ``sys.modules[fullname]`` here: if ``fullname`` is in
|
||||
# ``sys.modules`` when find_spec returns, ``_bootstrap._find_spec``
|
||||
# discards the returned spec in favor of ``module.__spec__`` — the
|
||||
# target's original SourceFileLoader spec — and ``_load_unlocked``
|
||||
# then RE-EXECUTES the target module and replaces its canonical
|
||||
# ``sys.modules`` entry with the duplicate. The machinery registers
|
||||
# the alias under ``fullname`` itself when loading the spec below.
|
||||
module = importlib.import_module(redirected)
|
||||
return importlib.util.spec_from_loader(fullname, _AliasLoader(module))
|
||||
|
||||
|
||||
class _AliasLoader:
|
||||
"""Loader that returns an already-resolved module for an alias spec."""
|
||||
|
||||
def __init__(self, module):
|
||||
self._module = module
|
||||
# ``module_from_spec`` unconditionally stamps the alias spec onto the
|
||||
# module returned by ``create_module``; capture the canonical values
|
||||
# so ``exec_module`` can restore them.
|
||||
self._spec = getattr(module, "__spec__", None)
|
||||
self._loader = getattr(module, "__loader__", None)
|
||||
|
||||
def create_module(self, spec):
|
||||
return self._module
|
||||
|
||||
def exec_module(self, module):
|
||||
# Module is already populated by the redirect target; just restore
|
||||
# the canonical ``__spec__``/``__loader__`` that the import machinery
|
||||
# overwrote with the alias spec (a stale alias ``__spec__.parent``
|
||||
# breaks relative imports inside the module).
|
||||
if self._spec is not None:
|
||||
module.__spec__ = self._spec
|
||||
if self._loader is not None:
|
||||
module.__loader__ = self._loader
|
||||
|
||||
|
||||
# Install the redirect finder once. Re-importing tvm.script (e.g. during a
|
||||
# pytest reload) must not stack duplicates.
|
||||
if not any(isinstance(f, _DialectRedirectFinder) for f in sys.meta_path):
|
||||
sys.meta_path.append(_DialectRedirectFinder())
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
if name in _DIALECT_REGISTRY:
|
||||
module = importlib.import_module(_DIALECT_REGISTRY[name])
|
||||
globals()[name] = module
|
||||
return module
|
||||
if name == "ir":
|
||||
# IR is foundational — its parser is a real submodule under
|
||||
# tvm.script.parser.ir, exposed here as `tvm.script.ir` for the
|
||||
# legacy `from tvm.script import ir as I` pattern.
|
||||
ir_parser = importlib.import_module("tvm.script.parser.ir")
|
||||
globals()["ir"] = ir_parser
|
||||
return ir_parser
|
||||
if name in ("from_source", "parse"):
|
||||
from .parser._core import parse # pylint: disable=import-outside-toplevel
|
||||
|
||||
globals()["from_source"] = parse
|
||||
globals()["parse"] = parse
|
||||
return parse
|
||||
if name == "ir_module":
|
||||
# ir_module lives in the IR parser at tvm.script.parser.ir; the IR
|
||||
# layer is foundational, so we resolve it directly rather than via
|
||||
# the dialect registry.
|
||||
ir_parser = importlib.import_module("tvm.script.parser.ir")
|
||||
ir_module_value = ir_parser.ir_module
|
||||
globals()["ir_module"] = ir_module_value
|
||||
return ir_module_value
|
||||
raise AttributeError(f"module 'tvm.script' has no attribute {name!r}")
|
||||
@@ -0,0 +1,37 @@
|
||||
# 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.
|
||||
|
||||
# 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"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script", __name__)
|
||||
@@ -0,0 +1,311 @@
|
||||
# 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.
|
||||
# ruff: noqa: F821, RUF012
|
||||
"""Highlight printed TVM script."""
|
||||
|
||||
import functools
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import warnings
|
||||
from typing import Any, Union
|
||||
|
||||
|
||||
def cprint(
|
||||
printable: Any | str,
|
||||
style: str | None = None,
|
||||
black_format: bool = False,
|
||||
) -> None:
|
||||
"""Print TVMScript string with Pygments highlight and Black auto-formatting.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
printable : Union[IRModule, PrimFunc, str]
|
||||
|
||||
The TVMScript to be printed
|
||||
|
||||
style : str, optional
|
||||
|
||||
Pygmentize printing style, auto-detected if None.
|
||||
|
||||
black_format: bool
|
||||
|
||||
If true, use the formatter Black to format the TVMScript
|
||||
|
||||
Notes
|
||||
-----
|
||||
|
||||
The style parameter follows the Pygments style names or Style objects. Three
|
||||
built-in styles are extended: "light", "dark" and "ansi". By default, "light"
|
||||
will be used for notebook environment and terminal style will be "ansi" for
|
||||
better style consistency. As an fallback when the optional Pygment library is
|
||||
not installed, plain text will be printed with a one-time warning to suggest
|
||||
installing the Pygment library. Other Pygment styles can be found in
|
||||
https://pygments.org/styles/
|
||||
|
||||
The default pygmentize style can also be set with the environment
|
||||
variable "TVM_PYGMENTIZE_STYLE".
|
||||
"""
|
||||
if hasattr(printable, "script") and callable(getattr(printable, "script")):
|
||||
printable = printable.script()
|
||||
elif not isinstance(printable, str):
|
||||
raise TypeError(
|
||||
f"Only can print strings or objects with `script` method, but got: {type(printable)}"
|
||||
)
|
||||
|
||||
if black_format:
|
||||
printable = _format(printable)
|
||||
|
||||
is_in_notebook = "ipykernel" in sys.modules # in notebook env (support html display).
|
||||
|
||||
style = _get_pygments_style(style, is_in_notebook)
|
||||
|
||||
if style is None:
|
||||
print(printable)
|
||||
return
|
||||
|
||||
# pylint: disable=import-outside-toplevel
|
||||
from pygments import highlight
|
||||
from pygments.formatters import HtmlFormatter, Terminal256Formatter
|
||||
from pygments.lexers.python import Python3Lexer
|
||||
|
||||
if is_in_notebook:
|
||||
from IPython import display # pylint: disable=import-outside-toplevel
|
||||
|
||||
formatter = HtmlFormatter(style=style)
|
||||
formatter.noclasses = True # inline styles
|
||||
html = highlight(printable, Python3Lexer(), formatter)
|
||||
display.display(display.HTML(html))
|
||||
else:
|
||||
print(highlight(printable, Python3Lexer(), Terminal256Formatter(style=style)))
|
||||
|
||||
|
||||
@functools.lru_cache
|
||||
def _get_formatter(formatter: str | None = None):
|
||||
def get_ruff_formatter():
|
||||
if shutil.which("ruff") is None:
|
||||
return None
|
||||
|
||||
def formatter(code_str):
|
||||
proc = subprocess.Popen(
|
||||
["ruff", "format", "--stdin-filename=TVMScript"],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
encoding="utf-8",
|
||||
)
|
||||
stdout, _stderr = proc.communicate(code_str)
|
||||
return stdout
|
||||
|
||||
return formatter
|
||||
|
||||
def get_black_formatter():
|
||||
try:
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import black
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
def formatter(code_str):
|
||||
return black.format_str(code_str, mode=black.FileMode())
|
||||
|
||||
return formatter
|
||||
|
||||
def get_fallback_formatter():
|
||||
def formatter(code_str):
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("once", UserWarning)
|
||||
ruff_install_cmd = sys.executable + " -m pip install ruff"
|
||||
black_install_cmd = (
|
||||
sys.executable + ' -m pip install "black==22.3.0" --upgrade --user'
|
||||
)
|
||||
warnings.warn(
|
||||
f"Neither the 'ruff' formatter nor the 'black' formatter is available. "
|
||||
f"To print formatted TVM script, please a formatter. \n"
|
||||
f"To install ruff: {ruff_install_cmd}\n"
|
||||
f"To install black: {black_install_cmd}",
|
||||
category=UserWarning,
|
||||
)
|
||||
return code_str
|
||||
|
||||
return formatter
|
||||
|
||||
# formatter = "black"
|
||||
if formatter is None:
|
||||
options = [get_ruff_formatter, get_black_formatter]
|
||||
elif formatter == "ruff":
|
||||
options = [get_ruff_formatter]
|
||||
elif formatter == "black":
|
||||
options = [get_black_formatter]
|
||||
else:
|
||||
raise ValueError(f"Unknown formatter: {formatter}")
|
||||
|
||||
for option in options:
|
||||
func = option()
|
||||
if func is not None:
|
||||
return func
|
||||
return get_fallback_formatter()
|
||||
|
||||
|
||||
def _format(code_str: str, formatter: str | None = None) -> str:
|
||||
"""Format a code string using Black.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
code_str: str
|
||||
|
||||
The string containing Python/TVMScript code to format
|
||||
|
||||
formatter: Optional[str]
|
||||
|
||||
The formatter to use. Can specify `ruff`, `black`, or
|
||||
auto-select by passing `None`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
formatted: str
|
||||
|
||||
The formatted Python/TVMScript code
|
||||
|
||||
"""
|
||||
return _get_formatter(formatter)(code_str)
|
||||
|
||||
|
||||
def _get_pygments_style(
|
||||
style: str | None, is_in_notebook: bool
|
||||
) -> Union["pygments.style.Style", str] | None:
|
||||
"""Select a pygments style to use
|
||||
|
||||
Parameters
|
||||
----------
|
||||
style: str
|
||||
|
||||
The style specifier to use. If None, auto-select a style.
|
||||
|
||||
is_in_notebook: bool
|
||||
|
||||
Whether python is currently running in a jupyter notebook.
|
||||
Used for automatic selection.
|
||||
|
||||
Returns
|
||||
-------
|
||||
style: Optional[Union['pygments.style.Style',str]]
|
||||
|
||||
If pygments is installed, the style object or string, suitable
|
||||
for use as the "style" argument to pygments formatters. If
|
||||
pygments is not installed, returns None.
|
||||
|
||||
"""
|
||||
try:
|
||||
# pylint: disable=import-outside-toplevel
|
||||
import pygments
|
||||
from packaging import version
|
||||
from pygments.style import Style
|
||||
from pygments.token import Comment, Keyword, Name, Number, Operator, String
|
||||
|
||||
if version.parse(pygments.__version__) < version.parse("2.4.0"):
|
||||
raise ImportError("Required Pygments version >= 2.4.0 but got " + pygments.__version__)
|
||||
except ImportError as err:
|
||||
if err.name == "packaging":
|
||||
name = "packaging"
|
||||
elif err.name == "pygments":
|
||||
name = "Pygments>=2.4.0"
|
||||
else:
|
||||
raise ValueError(f'Package "{err.name}" should not be used')
|
||||
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("once", UserWarning)
|
||||
install_cmd = sys.executable + f' -m pip install "{name}" --upgrade --user'
|
||||
warnings.warn(
|
||||
str(err)
|
||||
+ "\n"
|
||||
+ f"To print highlighted TVM script, please install {name}:\n"
|
||||
+ install_cmd,
|
||||
category=UserWarning,
|
||||
)
|
||||
return None
|
||||
|
||||
class JupyterLight(Style):
|
||||
"""A Jupyter-Notebook-like Pygments style configuration (aka. "light")"""
|
||||
|
||||
background_color = ""
|
||||
styles = {
|
||||
Keyword: "bold #008000",
|
||||
Keyword.Type: "nobold #008000",
|
||||
Name.Function: "#0000FF",
|
||||
Name.Class: "bold #0000FF",
|
||||
Name.Decorator: "#AA22FF",
|
||||
String: "#BA2121",
|
||||
Number: "#008000",
|
||||
Operator: "bold #AA22FF",
|
||||
Operator.Word: "bold #008000",
|
||||
Comment: "italic #007979",
|
||||
}
|
||||
|
||||
class VSCDark(Style):
|
||||
"""A VSCode-Dark-like Pygments style configuration (aka. "dark")"""
|
||||
|
||||
background_color = ""
|
||||
styles = {
|
||||
Keyword: "bold #c586c0",
|
||||
Keyword.Type: "#82aaff",
|
||||
Keyword.Namespace: "#4ec9b0",
|
||||
Name.Class: "bold #569cd6",
|
||||
Name.Function: "bold #dcdcaa",
|
||||
Name.Decorator: "italic #fe4ef3",
|
||||
String: "#ce9178",
|
||||
Number: "#b5cea8",
|
||||
Operator: "#bbbbbb",
|
||||
Operator.Word: "#569cd6",
|
||||
Comment: "italic #6a9956",
|
||||
}
|
||||
|
||||
class AnsiTerminalDefault(Style):
|
||||
"""The default style for terminal display with ANSI colors (aka. "ansi")"""
|
||||
|
||||
background_color = ""
|
||||
styles = {
|
||||
Keyword: "bold ansigreen",
|
||||
Keyword.Type: "nobold ansigreen",
|
||||
Name.Class: "bold ansiblue",
|
||||
Name.Function: "bold ansiblue",
|
||||
Name.Decorator: "italic ansibrightmagenta",
|
||||
String: "ansiyellow",
|
||||
Number: "ansibrightgreen",
|
||||
Operator: "bold ansimagenta",
|
||||
Operator.Word: "bold ansigreen",
|
||||
Comment: "italic ansibrightblack",
|
||||
}
|
||||
|
||||
if style == "light":
|
||||
return JupyterLight
|
||||
elif style == "dark":
|
||||
return VSCDark
|
||||
elif style == "ansi":
|
||||
return AnsiTerminalDefault
|
||||
|
||||
if style is not None:
|
||||
return style
|
||||
|
||||
style_from_environment = os.environ.get("TVM_PYGMENTIZE_STYLE", "").strip()
|
||||
if style_from_environment:
|
||||
return style_from_environment
|
||||
|
||||
if is_in_notebook:
|
||||
return JupyterLight
|
||||
|
||||
return AnsiTerminalDefault
|
||||
@@ -0,0 +1,48 @@
|
||||
# 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.
|
||||
"""The ir_builder subpackage of TVMScript.
|
||||
|
||||
Per-dialect builder submodules (``tvm.script.ir_builder.tirx``, etc.) are
|
||||
resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`. When a dialect
|
||||
is accessed (e.g. ``tvm.script.ir_builder.tirx``), this subpackage's
|
||||
``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
|
||||
``<dialect_module_path>.builder`` (e.g. ``tvm.tirx.script.builder``),
|
||||
caching the result so subsequent accesses skip ``__getattr__``.
|
||||
|
||||
The IR layer is foundational and is NOT registered as a dialect — its
|
||||
builder lives as a real submodule ``tvm.script.ir_builder.ir``.
|
||||
|
||||
See :mod:`tvm.script` for a full description of the dialect resolution
|
||||
mechanism, including the ``_DialectRedirectFinder`` that handles
|
||||
deep statement-form imports.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
from .base import IRBuilder
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
# Lazy import to avoid loading tvm.script during dialect bootstrap.
|
||||
from tvm.script import _DIALECT_REGISTRY # pylint: disable=import-outside-toplevel
|
||||
|
||||
if name in _DIALECT_REGISTRY:
|
||||
module = importlib.import_module(f"{_DIALECT_REGISTRY[name]}.builder")
|
||||
globals()[name] = module
|
||||
return module
|
||||
raise AttributeError(f"module 'tvm.script.ir_builder' has no attribute {name!r}")
|
||||
@@ -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.ir_builder"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder", __name__) # pylint: disable=protected-access
|
||||
@@ -0,0 +1,202 @@
|
||||
# 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.
|
||||
"""A generic IRBuilder across the TVM stack"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from tvm_ffi import register_object as _register_object
|
||||
|
||||
from tvm.runtime import Object as _Object
|
||||
|
||||
from . import _ffi_api
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.IRBuilderFrame")
|
||||
class IRBuilderFrame(_Object):
|
||||
"""A stack frame of the IRBuilder used to keep track of the current scope.
|
||||
|
||||
Furthermore, the information stored in each stack frame can be useful for context-dependent
|
||||
IR construction.
|
||||
|
||||
Examples
|
||||
--------
|
||||
|
||||
The `T.match_buffer` below instead an element in the buffer map of `PrimFuncFrame`:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
|
||||
with IRBuilder() as builder:
|
||||
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
||||
# to `builder`'s stack of frames
|
||||
buffer = T.match_buffer(...)
|
||||
|
||||
|
||||
The `T.match_buffer` below instead generates `MatchBufferRegion` in a TIR block:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
|
||||
with IRBuilder() as builder:
|
||||
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
||||
# to `builder`'s stack of frames
|
||||
with T.sblock(...): # pushes a BlockFrame (subclass of IRBuilderFrame)
|
||||
# to `builder`'s stack of frames
|
||||
buffer = T.match_buffer(...)
|
||||
|
||||
"""
|
||||
|
||||
def __enter__(self) -> "IRBuilderFrame":
|
||||
_ffi_api.IRBuilderFrameEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, trace) -> None: # pylint: disable=unused-argument
|
||||
if exc_type is None and exc_value is None:
|
||||
# Do not execute `FrameExit` if the with scope exits because of exceptions
|
||||
_ffi_api.IRBuilderFrameExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
def add_callback(self, callback: Callable[[], None]) -> None:
|
||||
"""Add a callback method invoked when exiting the with-scope.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
callback : Callable[[], None]
|
||||
The callback method to be invoked.
|
||||
"""
|
||||
_ffi_api.IRBuilderFrameAddCallback( # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
self, callback
|
||||
)
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.IRBuilder")
|
||||
class IRBuilder(_Object):
|
||||
"""A dialect-agnostic IRBuilder that constructs any IR of TVM.
|
||||
|
||||
Examples
|
||||
--------
|
||||
An idiomatic use of this class is to put this inside the with-scope,
|
||||
call dialect-specific methods accordingly. Upon exiting the scope.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script.ir_builder import tirx as T
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
|
||||
with IRBuilder() as builder:
|
||||
with T.prim_func(...): # pushes a PrimFuncFrame (subclass of IRBuilderFrame)
|
||||
# to `builder`'s stack of frames
|
||||
buffer = T.match_buffer(...)
|
||||
|
||||
return builder.get() # returns the constructed IR, i.e. tirx.PrimFunc
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
"""Construct an IRBuilder."""
|
||||
self.__init_handle_by_constructor__(
|
||||
_ffi_api.IRBuilder # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
)
|
||||
|
||||
def __enter__(self) -> "IRBuilder":
|
||||
"""Enter the with-scope for IRBuilder, which allows the IRBuilder to be discoverable
|
||||
using `IRBuilder.current()`.
|
||||
|
||||
Examples
|
||||
--------
|
||||
.. code-block:: python
|
||||
|
||||
from tvm.script.ir_builder import IRBuilder
|
||||
|
||||
with IRBuilder() as builder:
|
||||
assert IRBuilder.current() == builder
|
||||
|
||||
"""
|
||||
_ffi_api.IRBuilderEnter(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
return self
|
||||
|
||||
def __exit__(self, ptype, value, trace) -> None: # pylint: disable=unused-argument
|
||||
_ffi_api.IRBuilderExit(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
@staticmethod
|
||||
def current() -> "IRBuilder":
|
||||
"""Get the current IRBuilder put in the with-scope.
|
||||
|
||||
Returns
|
||||
-------
|
||||
builder : IRBuilder
|
||||
The current IRBuilder.
|
||||
"""
|
||||
return _ffi_api.IRBuilderCurrent() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
@staticmethod
|
||||
def is_in_scope() -> bool:
|
||||
"""See if the current thread-local scope has an IRBuilder.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
Whether the current thread-local scope has an IRBuilder
|
||||
"""
|
||||
return _ffi_api.IRBuilderIsInScope() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
def get(self) -> _Object:
|
||||
"""Get the constructed IR."""
|
||||
return _ffi_api.IRBuilderGet(self) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
@staticmethod
|
||||
def name(s: str, v: Any) -> Any:
|
||||
"""Set the name of an object.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
s : str
|
||||
The name of the object.
|
||||
v : Any
|
||||
The object to name.
|
||||
|
||||
Returns
|
||||
-------
|
||||
v : Any
|
||||
The same object with the name set.
|
||||
"""
|
||||
return _ffi_api.IRBuilderName(s, v) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
@staticmethod
|
||||
def name_many( # pylint: disable=invalid-name
|
||||
s: list[str],
|
||||
vs: list[Any],
|
||||
) -> list[Any]:
|
||||
"""Set the name of a list of objects.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
s : List[str]
|
||||
The names of the objects.
|
||||
vs : List[Any]
|
||||
The objects to name.
|
||||
|
||||
Returns
|
||||
-------
|
||||
vs : List[Any]
|
||||
The same objects with the names set.
|
||||
"""
|
||||
assert len(s) == len(vs)
|
||||
return [IRBuilder.name(i, v) for i, v in zip(s, vs)]
|
||||
@@ -0,0 +1,33 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir"""
|
||||
|
||||
from .frame import IRModuleFrame
|
||||
from .ir import (
|
||||
decl_function,
|
||||
def_function,
|
||||
ir_module,
|
||||
module_attrs,
|
||||
module_get_attr,
|
||||
module_set_attr,
|
||||
module_global_infos,
|
||||
lookup_vdevice,
|
||||
lookup_name,
|
||||
vdevice,
|
||||
dummy_global_info,
|
||||
)
|
||||
@@ -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"""
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
tvm_ffi.init_ffi_api("script.ir_builder.ir", __name__) # pylint: disable=protected-access
|
||||
@@ -0,0 +1,25 @@
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir.frame"""
|
||||
|
||||
from tvm_ffi import register_object as _register_object
|
||||
|
||||
from ..base import IRBuilderFrame
|
||||
|
||||
|
||||
@_register_object("script.ir_builder.IRModuleFrame")
|
||||
class IRModuleFrame(IRBuilderFrame): ...
|
||||
@@ -0,0 +1,192 @@
|
||||
# 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.
|
||||
"""Package tvm.script.ir_builder.ir.ir"""
|
||||
|
||||
from tvm.ir import BaseFunc, DummyGlobalInfo, GlobalInfo, GlobalVar, VDevice
|
||||
from tvm.runtime import Object as tvm_Object
|
||||
|
||||
from . import _ffi_api
|
||||
from .frame import IRModuleFrame
|
||||
|
||||
|
||||
def ir_module() -> IRModuleFrame:
|
||||
"""Start a ir_module frame.
|
||||
Returns
|
||||
-------
|
||||
frame: IRModuleFrame
|
||||
The constructed frame.
|
||||
"""
|
||||
return _ffi_api.IRModule() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def decl_function(func_name: str, func_signature: BaseFunc) -> GlobalVar:
|
||||
"""Declare a Function without given the specific function implementation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_name : str
|
||||
The function unique name.
|
||||
|
||||
func_signature: BaseFunc
|
||||
A Function w/o body, which used to specify the function signature
|
||||
(i.e. func params and func return type/shape).
|
||||
|
||||
Note
|
||||
----
|
||||
It is usually used in cross-function call. And we can specify the function by `DefFunction`
|
||||
|
||||
Returns
|
||||
-------
|
||||
gv : GlobalVar
|
||||
The corresponding GlobalVar.
|
||||
"""
|
||||
if not isinstance(func_signature, BaseFunc):
|
||||
raise ValueError(
|
||||
"decl_function expects an instance of BaseFunc, "
|
||||
f"but {func_signature} is of type {type(func_signature)}"
|
||||
)
|
||||
return _ffi_api.DeclFunction( # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
func_name, func_signature
|
||||
)
|
||||
|
||||
|
||||
def def_function(func_name: str, func: BaseFunc) -> None:
|
||||
"""Define the function which is declared before.
|
||||
Parameters
|
||||
----------
|
||||
func_name : str
|
||||
The function unique name.
|
||||
func: BaseFunc
|
||||
The given function implementation
|
||||
"""
|
||||
return _ffi_api.DefFunction(func_name, func) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_attrs(attrs: dict[str, tvm_Object], allow_overwrite=False) -> None:
|
||||
"""Specify the attrs of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attrs: Dict[str, Object]
|
||||
The module attrs.
|
||||
allow_overwrite: bool
|
||||
Whether allow overwrite the existing attrs.
|
||||
"""
|
||||
return _ffi_api.ModuleAttrs(attrs, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_get_attr(attr_key: str) -> tvm_Object | None:
|
||||
"""Get the specified attr of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attr_key: str
|
||||
The key of the attr to be retrieved.
|
||||
Returns
|
||||
-------
|
||||
attr: Optional[Object]
|
||||
The specified module attr or None if not found.
|
||||
"""
|
||||
return _ffi_api.ModuleGetAttr(attr_key) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_set_attr(
|
||||
attr_key: str, attr_value: tvm_Object | None, allow_overwrite: bool = False
|
||||
) -> None:
|
||||
"""Set the specified attr of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
attr_key: str
|
||||
The key of the attr to be set.
|
||||
attr_value: Optional[Object]
|
||||
The value of the attr to be set.
|
||||
allow_overwrite: bool
|
||||
Whether allow overwrite the existing attr.
|
||||
"""
|
||||
return _ffi_api.ModuleSetAttr(attr_key, attr_value, allow_overwrite) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def module_global_infos(global_infos: dict[str, list[GlobalInfo]]) -> None:
|
||||
"""Specify the global infos of the ir_module frame.
|
||||
Parameters
|
||||
----------
|
||||
global_infos: Dict[str, List[GlobalInfo]]
|
||||
The module global infos.
|
||||
"""
|
||||
return _ffi_api.ModuleGlobalInfos(global_infos) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
############################### GlobalInfo ###############################
|
||||
|
||||
|
||||
def dummy_global_info() -> DummyGlobalInfo:
|
||||
"""Create a dummy global info expression.
|
||||
Returns
|
||||
-------
|
||||
res : DummyGlobalInfo
|
||||
The result dummy global info.
|
||||
"""
|
||||
return DummyGlobalInfo() # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def vdevice(target=None, vdevice_id: int = 0, memory_scope: str = "global") -> VDevice:
|
||||
"""Create a virtual device global info.
|
||||
Parameters
|
||||
----------
|
||||
target
|
||||
The target.
|
||||
vdevice_id: int
|
||||
The virtual device index.
|
||||
memory_scope: str
|
||||
The memory scope, default is "global"
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : VDevice
|
||||
The result virtual device.
|
||||
"""
|
||||
return VDevice(target, vdevice_id, memory_scope) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def lookup_vdevice(target_kind: str | None = None, device_index: int = -1) -> VDevice:
|
||||
"""Retrieve a virtual device from the globalinfo vdevice list.
|
||||
Parameters
|
||||
----------
|
||||
target_kind: str
|
||||
The target device kind, for example 'llvm' or 'cuda'.
|
||||
device_index: int
|
||||
The virtual device index.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : VDevice
|
||||
The result virtual device.
|
||||
"""
|
||||
return _ffi_api.LookupVDevice(target_kind, device_index) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
|
||||
|
||||
def lookup_name(name: str) -> bool:
|
||||
"""Check if a global variable with the given name exists.
|
||||
Parameters
|
||||
----------
|
||||
name: str
|
||||
The name of the global variable.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : bool
|
||||
True if the global variable exists, False otherwise.
|
||||
"""
|
||||
return _ffi_api.LookupName(name) # type: ignore[attr-defined] # pylint: disable=no-member
|
||||
@@ -0,0 +1,51 @@
|
||||
# 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.
|
||||
"""The parser subpackage of TVMScript.
|
||||
|
||||
Per-dialect parser submodules (``tvm.script.parser.tirx``, etc.) are
|
||||
resolved lazily via :data:`tvm.script._DIALECT_REGISTRY`. When a dialect
|
||||
is accessed (e.g. ``tvm.script.parser.tirx``), this subpackage's
|
||||
``__getattr__`` looks up the dialect in ``_DIALECT_REGISTRY`` and imports
|
||||
``<dialect_module_path>.parser`` (e.g. ``tvm.tirx.script.parser``),
|
||||
caching the result so subsequent accesses skip ``__getattr__``.
|
||||
|
||||
The IR layer is foundational and is NOT registered as a dialect — its
|
||||
parser lives as a real submodule ``tvm.script.parser.ir``, with
|
||||
``ir_module`` re-exported at this level for convenience.
|
||||
|
||||
See :mod:`tvm.script` for a full description of the dialect resolution
|
||||
mechanism, including the ``_DialectRedirectFinder`` that handles
|
||||
deep statement-form imports.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
from typing import Any
|
||||
|
||||
from . import _core, ir, tirx
|
||||
from ._core import parse
|
||||
from .ir import ir_module
|
||||
|
||||
|
||||
def __getattr__(name: str) -> Any:
|
||||
# Lazy import to avoid loading tvm.script during dialect bootstrap.
|
||||
from tvm.script import _DIALECT_REGISTRY # pylint: disable=import-outside-toplevel
|
||||
|
||||
if name in _DIALECT_REGISTRY:
|
||||
module = importlib.import_module(f"{_DIALECT_REGISTRY[name]}.parser")
|
||||
globals()[name] = module
|
||||
return module
|
||||
raise AttributeError(f"module 'tvm.script.parser' has no attribute {name!r}")
|
||||
@@ -0,0 +1,24 @@
|
||||
# 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 Licens.
|
||||
# ruff: noqa: F401
|
||||
"""The core parser infra"""
|
||||
|
||||
# pylint: disable=unused-import
|
||||
from .core import dispatch, doc, utils
|
||||
from .core.dispatch import OpMethod, register_op
|
||||
from .core.entry import parse, scan_macro
|
||||
from .core.parser import Parser
|
||||
@@ -0,0 +1,20 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""The core parser infra"""
|
||||
|
||||
from . import diagnostics, dispatch, doc, doc_core, entry, evaluator, parser, utils
|
||||
@@ -0,0 +1,305 @@
|
||||
# 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.
|
||||
# ruff: noqa: E741
|
||||
"""TVM Script Parser Source and diagnostics"""
|
||||
|
||||
import inspect
|
||||
import sys
|
||||
|
||||
from tvm.error import DiagnosticError
|
||||
|
||||
from . import doc
|
||||
|
||||
|
||||
class Source:
|
||||
"""Source code class for TVMScript.
|
||||
|
||||
It is constructed by source code str or doc AST tree.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_name : str
|
||||
The filename of the file where the source code locates.
|
||||
|
||||
start_line : int
|
||||
The first line number of the source code.
|
||||
|
||||
start_column : int
|
||||
The first column number of the first line of the source code.
|
||||
|
||||
source : str
|
||||
The source code str of source code.
|
||||
|
||||
full_source : str
|
||||
The complete source code of the file where the source code locates.
|
||||
"""
|
||||
|
||||
source_name: str
|
||||
start_line: int
|
||||
start_column: int
|
||||
source: str
|
||||
full_source: str
|
||||
|
||||
def __init__(self, program: str | doc.AST):
|
||||
if isinstance(program, str):
|
||||
self.source_name = "<str>"
|
||||
self.start_line = 1
|
||||
self.start_column = 0
|
||||
self.source = program
|
||||
self.full_source = program
|
||||
return
|
||||
|
||||
self.source_name = inspect.getsourcefile(program) # type: ignore
|
||||
lines, self.start_line = getsourcelines(program) # type: ignore
|
||||
if lines:
|
||||
self.start_column = len(lines[0]) - len(lines[0].lstrip())
|
||||
else:
|
||||
self.start_column = 0
|
||||
if self.start_column and lines:
|
||||
self.source = "\n".join([l[self.start_column :].rstrip() for l in lines])
|
||||
else:
|
||||
self.source = "".join(lines)
|
||||
try:
|
||||
# It will cause a problem when running in Jupyter Notebook.
|
||||
# `mod` will be <module '__main__'>, which is a built-in module
|
||||
# and `getsource` will throw a TypeError
|
||||
mod = inspect.getmodule(program)
|
||||
if mod:
|
||||
self.full_source = inspect.getsource(mod)
|
||||
else:
|
||||
self.full_source = self.source
|
||||
except TypeError:
|
||||
# It's a work around for Jupyter problem.
|
||||
# Since `findsource` is an internal API of inspect, we just use it
|
||||
# as a fallback method.
|
||||
src, _ = inspect.findsource(program) # type: ignore
|
||||
self.full_source = "".join(src)
|
||||
|
||||
def as_ast(self) -> doc.AST:
|
||||
"""Parse the source code into AST.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : doc.AST
|
||||
The AST of source code.
|
||||
"""
|
||||
return doc.parse(self.source)
|
||||
|
||||
|
||||
_getfile = inspect.getfile # pylint: disable=invalid-name
|
||||
_findsource = inspect.findsource # pylint: disable=invalid-name
|
||||
|
||||
|
||||
def _patched_inspect_getfile(obj):
|
||||
"""Work out which source or compiled file an object was defined in."""
|
||||
if not inspect.isclass(obj):
|
||||
return _getfile(obj)
|
||||
mod = getattr(obj, "__module__", None)
|
||||
if mod is not None:
|
||||
file = getattr(sys.modules[mod], "__file__", None)
|
||||
if file is not None:
|
||||
return file
|
||||
for _, member in inspect.getmembers(obj):
|
||||
if inspect.isfunction(member):
|
||||
if obj.__qualname__ + "." + member.__name__ == member.__qualname__:
|
||||
return inspect.getfile(member)
|
||||
raise TypeError(f"Source for {obj!r} not found")
|
||||
|
||||
|
||||
def findsource(obj):
|
||||
"""Return the entire source file and starting line number for an object."""
|
||||
import linecache # pylint: disable=import-outside-toplevel
|
||||
|
||||
if not inspect.isclass(obj):
|
||||
return _findsource(obj)
|
||||
|
||||
file = inspect.getsourcefile(obj)
|
||||
if file:
|
||||
linecache.checkcache(file)
|
||||
else:
|
||||
file = inspect.getfile(obj)
|
||||
if not (file.startswith("<") and file.endswith(">")):
|
||||
raise OSError("source code not available")
|
||||
|
||||
module = inspect.getmodule(obj, file)
|
||||
if module:
|
||||
lines = linecache.getlines(file, module.__dict__)
|
||||
else:
|
||||
lines = linecache.getlines(file)
|
||||
if not lines:
|
||||
raise OSError("could not get source code")
|
||||
qual_names = obj.__qualname__.replace(".<locals>", "<locals>").split(".")
|
||||
in_comment = 0
|
||||
scope_stack = []
|
||||
indent_info = {}
|
||||
for i, line in enumerate(lines):
|
||||
n_comment = line.count('"""')
|
||||
if n_comment:
|
||||
# update multi-line comments status
|
||||
in_comment = in_comment ^ (n_comment & 1)
|
||||
continue
|
||||
if in_comment:
|
||||
# skip lines within multi-line comments
|
||||
continue
|
||||
indent = len(line) - len(line.lstrip())
|
||||
tokens = line.split()
|
||||
if len(tokens) > 1:
|
||||
name = None
|
||||
if tokens[0] == "def":
|
||||
name = tokens[1].split(":")[0].split("(")[0] + "<locals>"
|
||||
elif tokens[0] == "class":
|
||||
name = tokens[1].split(":")[0].split("(")[0]
|
||||
# pop scope if we are less indented
|
||||
while scope_stack and indent_info[scope_stack[-1]] >= indent:
|
||||
scope_stack.pop()
|
||||
if name:
|
||||
scope_stack.append(name)
|
||||
indent_info[name] = indent
|
||||
if scope_stack == qual_names:
|
||||
return lines, i
|
||||
|
||||
raise OSError("could not find class definition")
|
||||
|
||||
|
||||
def getsourcelines(obj):
|
||||
"""Extract the block of code at the top of the given list of lines."""
|
||||
obj = inspect.unwrap(obj)
|
||||
lines, l_num = findsource(obj)
|
||||
return inspect.getblock(lines[l_num:]), l_num + 1
|
||||
|
||||
|
||||
inspect.getfile = _patched_inspect_getfile
|
||||
|
||||
|
||||
def _format_source_snippet(
|
||||
source_lines: list,
|
||||
lineno: int,
|
||||
col_offset: int,
|
||||
end_lineno: int,
|
||||
end_col_offset: int,
|
||||
) -> str:
|
||||
"""Format a source code snippet with a column/span marker.
|
||||
|
||||
Renders every source line spanned by the diagnostic (``lineno`` through
|
||||
``end_lineno``, inclusive) with a per-line gutter, followed by a caret
|
||||
underline covering the offending span. For a single-line span
|
||||
(``end_lineno == lineno``) only the columns ``col_offset`` (inclusive)
|
||||
through ``end_col_offset`` (exclusive) are underlined; for a multi-line
|
||||
span the underline covers the start column to end-of-line on the first
|
||||
line, the full text of interior lines, and the start of the final line up
|
||||
to ``end_col_offset``.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source_lines : list of str
|
||||
Lines of the source code.
|
||||
|
||||
lineno : int
|
||||
1-based starting line number in the source.
|
||||
|
||||
col_offset : int
|
||||
1-based starting column (inclusive) on ``lineno``.
|
||||
|
||||
end_lineno : int
|
||||
1-based ending line number in the source (>= ``lineno``).
|
||||
|
||||
end_col_offset : int
|
||||
1-based ending column (exclusive) on ``end_lineno``.
|
||||
|
||||
Returns
|
||||
-------
|
||||
snippet : str
|
||||
Formatted source snippet with caret-marker line(s).
|
||||
"""
|
||||
if end_lineno < lineno:
|
||||
end_lineno = lineno
|
||||
|
||||
# Determine the gutter width so that all line numbers line up.
|
||||
header_width = len(f" {end_lineno} ")
|
||||
no_line_header = " " * header_width
|
||||
|
||||
parts = [f"{no_line_header}| "]
|
||||
for cur_lineno in range(lineno, end_lineno + 1):
|
||||
idx = cur_lineno - 1
|
||||
if not 0 <= idx < len(source_lines):
|
||||
continue
|
||||
line_text = source_lines[idx].rstrip("\n")
|
||||
line_header = f" {cur_lineno} ".rjust(header_width)
|
||||
|
||||
# Compute the underline span [start_col, stop_col) for this line.
|
||||
start_col = col_offset if cur_lineno == lineno else 1
|
||||
stop_col = end_col_offset if cur_lineno == end_lineno else len(line_text) + 1
|
||||
|
||||
marker = ""
|
||||
for i in range(1, len(line_text) + 1):
|
||||
if start_col <= i < stop_col:
|
||||
marker += "^"
|
||||
else:
|
||||
marker += " "
|
||||
parts.append(f"{line_header}| {line_text}")
|
||||
parts.append(f"{no_line_header}| {marker}")
|
||||
|
||||
if len(parts) == 1:
|
||||
return ""
|
||||
return "\n".join(parts)
|
||||
|
||||
|
||||
class Diagnostics:
|
||||
"""Diagnostics class for error reporting in parser.
|
||||
|
||||
Formats parse errors with source location context and raises directly,
|
||||
without going through DiagnosticContext.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source : Source
|
||||
The source code.
|
||||
"""
|
||||
|
||||
source: Source
|
||||
|
||||
def __init__(self, source: Source):
|
||||
self.source = source
|
||||
|
||||
def error(self, node: doc.AST, message: str) -> None:
|
||||
"""Emit a diagnostic error by raising with source location context.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The node with diagnostic error.
|
||||
|
||||
message : str
|
||||
The diagnostic message.
|
||||
"""
|
||||
lineno = getattr(node, "lineno", 1)
|
||||
col_offset = getattr(node, "col_offset", self.source.start_column)
|
||||
end_lineno = getattr(node, "end_lineno", lineno)
|
||||
end_col_offset = getattr(node, "end_col_offset", col_offset)
|
||||
lineno += self.source.start_line - 1
|
||||
end_lineno += self.source.start_line - 1
|
||||
col_offset += self.source.start_column + 1
|
||||
end_col_offset += self.source.start_column + 1
|
||||
|
||||
source_lines = self.source.full_source.splitlines(keepends=True)
|
||||
snippet = _format_source_snippet(
|
||||
source_lines, lineno, col_offset, end_lineno, end_col_offset
|
||||
)
|
||||
|
||||
location = f"{self.source.source_name}:{lineno}:{col_offset}"
|
||||
formatted = f"error: {message}\n --> {location}\n{snippet}"
|
||||
raise DiagnosticError(formatted)
|
||||
@@ -0,0 +1,157 @@
|
||||
# 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.
|
||||
"""Parser dispatching infrastructure"""
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from .doc import AST
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .parser import Parser
|
||||
|
||||
|
||||
ParseMethod = Callable[["Parser", AST], None]
|
||||
ParseVTable: dict[tuple[str, str], ParseMethod] = {}
|
||||
|
||||
OpMethod = Callable[..., Any]
|
||||
OpVTable: dict[tuple[type, AST, int], OpMethod] = {}
|
||||
|
||||
|
||||
def register(token: str, type_name: str):
|
||||
"""Register a method for a dispatch token and type name.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
token : str
|
||||
The token for IR, e.g., T for TIR and R for Relax.
|
||||
|
||||
type_name : str
|
||||
The type name of AST node, e.g., FunctionDef, With, For.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : callable
|
||||
The function to register dispatched method of parsing
|
||||
corresponding token and AST node type.
|
||||
"""
|
||||
|
||||
def func(method: ParseMethod):
|
||||
"""Register a method in parser virtual table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
method : ParseMethod
|
||||
The dispatched method to be registered in parser virtual table.
|
||||
"""
|
||||
ParseVTable[(token, type_name)] = method
|
||||
|
||||
return func
|
||||
|
||||
|
||||
def get(
|
||||
token: str,
|
||||
type_name: str,
|
||||
default: ParseMethod | None = None,
|
||||
) -> ParseMethod | None:
|
||||
"""Get a registered method for a dispatch token and type name,
|
||||
or return a default method if no registered methods with this dispatch token and type name.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
token : str
|
||||
The token for IR, e.g., T for TIR and R for Relax.
|
||||
|
||||
type_name : str
|
||||
The type name of AST node, e.g., FunctionDef, With, For.
|
||||
|
||||
default : Optional[ParseMethod]
|
||||
The default method when no registered methods with this dispatch token and type name.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : Optional[ParseMethod]
|
||||
The dispatched method of parsing corresponding token and AST node type.
|
||||
"""
|
||||
return ParseVTable.get((token, type_name), default)
|
||||
|
||||
|
||||
def register_op(operand_type: type, op_node_type: AST, operand_index: int):
|
||||
"""Register a method for a operand type, AST operator node and operand index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
operand_type : Type
|
||||
The type of operands, e.g., tirx.Expr, tirx.IterVar.
|
||||
|
||||
op_node_type : AST
|
||||
The doc AST operator node type, e.g., doc.Add, doc.Eq.
|
||||
|
||||
operand_index : int
|
||||
The operand index, i.e., 0 for left operand and 1 for right operand.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : callable
|
||||
The function to register dispatched method of parsing
|
||||
corresponding a operand type, AST operator node and operand index.
|
||||
"""
|
||||
|
||||
def func(method: OpMethod):
|
||||
"""Register a method in parser operator virtual table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
method : ParseMethod
|
||||
The dispatched method to be registered in parser operator virtual table.
|
||||
"""
|
||||
OpVTable[(operand_type, op_node_type, operand_index)] = method
|
||||
|
||||
return func
|
||||
|
||||
|
||||
def get_op(
|
||||
operand_type: type,
|
||||
op_node_type: type,
|
||||
operand_index: int,
|
||||
default: OpMethod | None = None,
|
||||
) -> OpMethod | None:
|
||||
"""Register a method for a operand type, AST operator node and operand index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
operand_type : Type
|
||||
The type of operands, e.g., tirx.Expr, tirx.IterVar.
|
||||
|
||||
op_node_type : AST
|
||||
The doc AST operator node type, e.g., doc.Add, doc.Eq.
|
||||
|
||||
operand_index : int
|
||||
The operand index, i.e., 0 for left operand and 1 for right operand.
|
||||
|
||||
|
||||
default : Optional[OpMethod]
|
||||
The default method when no registered methods with this operand type,
|
||||
AST operator node and operand index.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : Optional[OpMethod]
|
||||
The function to register dispatched method of parsing
|
||||
corresponding a operand type, AST operator node and operand index.
|
||||
"""
|
||||
return OpVTable.get((operand_type, op_node_type, operand_index), default)
|
||||
@@ -0,0 +1,315 @@
|
||||
# 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.
|
||||
# ruff: noqa: E722, F403
|
||||
"""TVM Script Parser doc AST"""
|
||||
|
||||
import ast
|
||||
import inspect
|
||||
import typing
|
||||
from collections import defaultdict
|
||||
|
||||
from . import doc_core as doc
|
||||
from .doc_core import * # pylint: disable=unused-import,wildcard-import,redefined-builtin,W0614
|
||||
|
||||
FnToDoc = typing.Callable[[ast.AST], doc.AST]
|
||||
FnFromDoc = typing.Callable[[doc.AST], ast.AST]
|
||||
|
||||
|
||||
class Entry:
|
||||
"""Mapping entry between python AST node type str and doc AST.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
to_doc : typing.Optional[FnToDoc]
|
||||
The callable methods for converting python AST node to doc AST.
|
||||
|
||||
from_doc : typing.Optional[FnFromDoc]
|
||||
The callable methods for converting doc AST to python AST node.
|
||||
"""
|
||||
|
||||
to_doc: FnToDoc | None
|
||||
from_doc: FnFromDoc | None
|
||||
|
||||
def __init__(self):
|
||||
self.to_doc = None
|
||||
self.from_doc = None
|
||||
|
||||
|
||||
class Registry:
|
||||
"""Registration map from python AST node type str to methods of conversion
|
||||
between python AST node and doc AST node.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
_inst : typing.Optional["Registry"]
|
||||
The instance of Registry.
|
||||
|
||||
table : typing.Dict[str, Entry]
|
||||
The registration map from python AST node type str to methods of conversion
|
||||
between python AST node and doc AST node.
|
||||
"""
|
||||
|
||||
_inst: typing.Optional["Registry"] = None
|
||||
table: dict[str, Entry]
|
||||
|
||||
def __init__(self):
|
||||
self.table = defaultdict(Entry)
|
||||
|
||||
|
||||
def register_to_doc(name: str):
|
||||
"""Register the to_doc method for python AST node type.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The type of python AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
f : Callable[[FnToDoc], None]
|
||||
The function of registering the to_doc method for python AST node type.
|
||||
"""
|
||||
|
||||
def f(to_doc: FnToDoc): # pylint: disable=redefined-outer-name
|
||||
reg = Registry._inst # pylint: disable=protected-access
|
||||
reg.table[name].to_doc = to_doc
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def register_from_doc(name: str):
|
||||
"""Register the from_doc method for python AST node type.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The type of python AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
f : Callable[[FnFromDoc], None]
|
||||
The function of registering the from_doc method for python AST node type.
|
||||
"""
|
||||
|
||||
def f(to_doc: FnFromDoc): # pylint: disable=redefined-outer-name
|
||||
reg = Registry._inst # pylint: disable=protected-access
|
||||
reg.table[name].from_doc = to_doc
|
||||
|
||||
return f
|
||||
|
||||
|
||||
def _is_atomic_type(node):
|
||||
return (
|
||||
node is None
|
||||
or node in [..., True, False]
|
||||
or isinstance(
|
||||
node,
|
||||
int | float | str | bool | bytes | complex,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _get_registry_entry(cls_name, attr):
|
||||
cls_name = cls_name.split(".")[-1]
|
||||
reg = Registry._inst # pylint: disable=protected-access
|
||||
if cls_name in reg.table:
|
||||
entry = reg.table[cls_name]
|
||||
return getattr(entry, attr, None)
|
||||
return None
|
||||
|
||||
|
||||
def from_doc(node):
|
||||
"""Get original python AST node from doc AST node.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The doc AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : ast.AST
|
||||
The corresponding AST node.
|
||||
"""
|
||||
if _is_atomic_type(node):
|
||||
return node
|
||||
if isinstance(node, tuple):
|
||||
return tuple(from_doc(n) for n in node)
|
||||
if isinstance(node, list):
|
||||
return [from_doc(n) for n in node]
|
||||
func = _get_registry_entry(node.__class__.__name__, "from_doc")
|
||||
if not func:
|
||||
raise NotImplementedError(f"from_doc is not implemented for: {node.__class__.__name__}")
|
||||
return func(node)
|
||||
|
||||
|
||||
def to_doc(node):
|
||||
"""Get doc AST node from python AST node.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : ast.AST
|
||||
The AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : doc.AST
|
||||
The corresponding doc AST node.
|
||||
"""
|
||||
if _is_atomic_type(node):
|
||||
return node
|
||||
if isinstance(node, tuple):
|
||||
return tuple(to_doc(n) for n in node)
|
||||
if isinstance(node, list):
|
||||
return [to_doc(n) for n in node]
|
||||
func = _get_registry_entry(node.__class__.__name__, "to_doc")
|
||||
if not func:
|
||||
raise NotImplementedError(f"to_doc is not implemented for: {node.__class__.__name__}")
|
||||
return func(node)
|
||||
|
||||
|
||||
def parse(
|
||||
source: str,
|
||||
filename: str = "<unknown>",
|
||||
mode: str = "exec",
|
||||
) -> doc.AST:
|
||||
"""Parse TVMScript source code str to doc AST.
|
||||
|
||||
Its interface is consistent with python built-in ast.parse.
|
||||
And it will parse by python 3.8 first if possible,
|
||||
or it will parse with python version in current environment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source : str
|
||||
The TVMScript source code.
|
||||
|
||||
filename : str
|
||||
The optional filename of the file where source code locates.
|
||||
|
||||
mode : str
|
||||
The parsing mode for ast.parse.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : doc.AST
|
||||
The parsed doc AST.
|
||||
"""
|
||||
try:
|
||||
program = ast.parse( # pylint: disable=unexpected-keyword-arg
|
||||
source=source,
|
||||
filename=filename,
|
||||
mode=mode,
|
||||
feature_version=(3, 8),
|
||||
)
|
||||
except: # pylint: disable=bare-except
|
||||
program = ast.parse(
|
||||
source=source,
|
||||
filename=filename,
|
||||
mode=mode,
|
||||
)
|
||||
return to_doc(program)
|
||||
|
||||
|
||||
class NodeVisitor:
|
||||
"""Node visitor for doc AST"""
|
||||
|
||||
def visit(self, node: doc.AST) -> None:
|
||||
if isinstance(node, list | tuple):
|
||||
for item in node:
|
||||
self.visit(item)
|
||||
return
|
||||
if not isinstance(node, doc.AST):
|
||||
return
|
||||
getattr(
|
||||
self,
|
||||
"visit_" + node.__class__.__name__.split(".")[-1],
|
||||
self.generic_visit,
|
||||
)(node)
|
||||
|
||||
def generic_visit(self, node: doc.AST) -> None:
|
||||
for field in node.__class__._FIELDS: # pylint: disable=protected-access
|
||||
value = getattr(node, field, None)
|
||||
if value is None:
|
||||
pass
|
||||
elif isinstance(value, doc.AST | list | tuple):
|
||||
self.visit(value)
|
||||
|
||||
|
||||
class NodeTransformer:
|
||||
"""Node transformer for doc AST"""
|
||||
|
||||
def visit(self, node: doc.AST) -> doc.AST:
|
||||
if isinstance(node, list):
|
||||
return [self.visit(item) for item in node]
|
||||
if isinstance(node, tuple):
|
||||
return tuple(self.visit(item) for item in node)
|
||||
if not isinstance(node, doc.AST):
|
||||
return node
|
||||
return getattr(
|
||||
self,
|
||||
"visit_" + node.__class__.__name__.split(".")[-1],
|
||||
self.generic_visit,
|
||||
)(node)
|
||||
|
||||
def generic_visit(self, node: doc.AST) -> doc.AST:
|
||||
kv: dict[str, typing.Any] = {}
|
||||
for field in node.__class__._FIELDS: # pylint: disable=protected-access
|
||||
value = getattr(node, field, None)
|
||||
if value is None:
|
||||
pass
|
||||
elif isinstance(value, doc.AST | list | tuple):
|
||||
value = self.visit(value)
|
||||
kv[field] = value
|
||||
return node.__class__(**kv)
|
||||
|
||||
|
||||
def _register_default():
|
||||
class DefaultTranslator:
|
||||
def __init__(self, doc_cls, func, fields):
|
||||
self.doc_cls = doc_cls # getattr(doc, name)
|
||||
self.func = func
|
||||
self.fields = fields
|
||||
|
||||
def __call__(self, node):
|
||||
kv = {attr: self.func(getattr(node, attr, None)) for attr in self.fields}
|
||||
return self.doc_cls(**kv)
|
||||
|
||||
Registry._inst = Registry() # pylint: disable=protected-access
|
||||
for cls_name in dir(doc):
|
||||
doc_cls = getattr(doc, cls_name)
|
||||
if not hasattr(ast, cls_name):
|
||||
continue
|
||||
if inspect.isclass(doc_cls) and issubclass(doc_cls, doc.AST):
|
||||
assert "." not in cls_name
|
||||
register_to_doc(cls_name)(
|
||||
DefaultTranslator(
|
||||
getattr(doc, cls_name),
|
||||
to_doc,
|
||||
doc_cls._FIELDS, # pylint: disable=protected-access
|
||||
)
|
||||
)
|
||||
register_from_doc(cls_name)(
|
||||
DefaultTranslator(
|
||||
getattr(ast, cls_name),
|
||||
from_doc,
|
||||
doc_cls._FIELDS, # pylint: disable=protected-access
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
_register_default()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,210 @@
|
||||
# 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.
|
||||
"""The entry point of TVM parser."""
|
||||
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
import tvm
|
||||
|
||||
from ....ir.module import IRModule
|
||||
from ...ir_builder import IRBuilder
|
||||
from . import doc
|
||||
from .diagnostics import Source
|
||||
from .error import ParserError
|
||||
from .parser import Parser
|
||||
|
||||
WELL_FORMED_ERROR_MESSAGE = (
|
||||
"Program is not well-formed. If this is deliberate, consider "
|
||||
"setting check_well_formed in the top-level decorator to False "
|
||||
"(e.g., @I.ir_module(check_well_formed=False) or "
|
||||
"@R.function(check_well_formed=False))."
|
||||
)
|
||||
|
||||
|
||||
def _default_globals() -> dict[str, Any]:
|
||||
# lazy import here to avoid circular deps
|
||||
from tvm.script import tirx as _tirx_dsl # pylint: disable=import-outside-toplevel
|
||||
from tvm.script.parser import (
|
||||
ir, # pylint: disable=import-outside-toplevel
|
||||
relax, # pylint: disable=import-outside-toplevel
|
||||
)
|
||||
from tvm.script.parser import tirx as _tirx_parser # pylint: disable=import-outside-toplevel
|
||||
from tvm.script.tirx import tile as _tirx_tile # pylint: disable=import-outside-toplevel
|
||||
from tvm.tirx import layout as _tirx_layout # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Expose the layout `Axis` class so printed layout sugar like
|
||||
# `4 @ Axis.laneid` round-trips without per-script imports. Injecting just
|
||||
# `Axis` (one short symbol) avoids name collisions with common user shape
|
||||
# vars like `m`, `P`, `F` that registered axes happen to share names with.
|
||||
return {
|
||||
"tvm": tvm,
|
||||
"I": ir,
|
||||
"ir": ir,
|
||||
"T": _tirx_parser,
|
||||
"tir": _tirx_parser,
|
||||
"R": relax,
|
||||
"relax": relax,
|
||||
"Tx": _tirx_tile,
|
||||
"tirx": _tirx_dsl,
|
||||
"Axis": _tirx_layout.Axis,
|
||||
}
|
||||
|
||||
|
||||
def scan_macro(program: Any | str, extra_vars: dict[str, Any] | None = None) -> Any:
|
||||
"""Generate the AST, and the source code for __repr__."""
|
||||
# The AST will be converted into TIR at the time of expansion.
|
||||
source = Source(program)
|
||||
closure_vars = extra_vars or _default_globals()
|
||||
return source, closure_vars
|
||||
|
||||
|
||||
def parse(
|
||||
program: doc.AST | Any | str,
|
||||
extra_vars: dict[str, Any] | None = None,
|
||||
check_well_formed: bool = True,
|
||||
s_tir: bool = False,
|
||||
) -> Any:
|
||||
"""Register a method for a operand type, AST operator node and operand index.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
program : Union[doc.AST, Any, str]
|
||||
The TVMScript code to parse.
|
||||
|
||||
extra_vars : Dict[str, Any]
|
||||
The extra variable table for parsing.
|
||||
|
||||
check_well_formed : bool
|
||||
Whether to check well-formedness after parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : Any
|
||||
The parsed TVMScript program.
|
||||
"""
|
||||
if extra_vars is None:
|
||||
extra_vars = _default_globals()
|
||||
|
||||
ann = {}
|
||||
all_pyfuncs = {}
|
||||
if inspect.isfunction(program):
|
||||
ann = {program.__name__: program.__annotations__}
|
||||
elif inspect.isclass(program):
|
||||
for name, func in program.__dict__.items():
|
||||
if inspect.isfunction(func):
|
||||
ann[name] = func.__annotations__
|
||||
all_pyfuncs[name] = func
|
||||
|
||||
source = Source(program)
|
||||
parser = Parser(source, ann)
|
||||
with IRBuilder() as builder:
|
||||
try:
|
||||
parser.parse(extra_vars=extra_vars)
|
||||
except ParserError as err:
|
||||
parser.report_error(err.node, err.args[0])
|
||||
ret = builder.get()
|
||||
# Attach pyfuncs to the IRModule
|
||||
if inspect.isclass(program) and isinstance(ret, IRModule):
|
||||
_attach_pyfuncs_to_irmodule(ret, all_pyfuncs)
|
||||
|
||||
# check well-formedness in both Relax and TIR
|
||||
if check_well_formed:
|
||||
check_ret = ret
|
||||
if not isinstance(check_ret, IRModule):
|
||||
check_ret = IRModule.from_expr(ret)
|
||||
|
||||
source_ast = source.as_ast()
|
||||
|
||||
if isinstance(
|
||||
ret, IRModule | tvm.relax.Function
|
||||
) and not tvm.relax.analysis.check_well_formed(ret):
|
||||
parser.report_error(source_ast, err=WELL_FORMED_ERROR_MESSAGE)
|
||||
|
||||
try:
|
||||
if s_tir:
|
||||
tvm.tirx.analysis.verify_well_formed(check_ret)
|
||||
else:
|
||||
tvm.tirx.analysis.verify_tirx_well_formed(check_ret)
|
||||
except Exception as err: # pylint: disable=broad-exception-caught
|
||||
parser.report_error(
|
||||
source_ast,
|
||||
err=f"{WELL_FORMED_ERROR_MESSAGE}\n\nTraceback: {err!s}",
|
||||
)
|
||||
return ret
|
||||
|
||||
|
||||
def _create_python_packed_func(pyfunc):
|
||||
"""Create a PackedFunc wrapper for a Python function.
|
||||
|
||||
This function creates a PackedFunc that can be called from TVM runtime
|
||||
and will execute the original Python function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pyfunc : Callable
|
||||
The Python function to wrap.
|
||||
|
||||
Returns
|
||||
-------
|
||||
PackedFunc
|
||||
A PackedFunc that wraps the Python function.
|
||||
"""
|
||||
|
||||
def packed_func_wrapper(*args, **kwargs):
|
||||
"""Wrapper function that calls the original Python function."""
|
||||
try:
|
||||
result = pyfunc(*args, **kwargs)
|
||||
return result
|
||||
except Exception as error:
|
||||
print(f"Error calling Python function {pyfunc.__name__}: {error}")
|
||||
raise
|
||||
|
||||
return packed_func_wrapper
|
||||
|
||||
|
||||
def _attach_pyfuncs_to_irmodule(irmodule, all_pyfuncs):
|
||||
"""Attach Python functions to IRModule with reduced nesting."""
|
||||
if not all_pyfuncs:
|
||||
return
|
||||
|
||||
if not hasattr(irmodule, "pyfuncs"):
|
||||
irmodule.pyfuncs = {}
|
||||
|
||||
for global_var, func in irmodule.functions_items():
|
||||
if not isinstance(func, tvm.relax.ExternFunc):
|
||||
continue
|
||||
if not func.attrs.get("is_pyfunc", False):
|
||||
continue
|
||||
|
||||
pyfunc_name = global_var.name_hint
|
||||
if pyfunc_name not in all_pyfuncs:
|
||||
continue
|
||||
|
||||
pyfunc = all_pyfuncs[pyfunc_name]
|
||||
irmodule.pyfuncs[pyfunc_name] = pyfunc
|
||||
|
||||
try:
|
||||
source_code = inspect.getsource(pyfunc)
|
||||
func = func.with_attr("python_source", source_code)
|
||||
except (OSError, TypeError):
|
||||
func = func.with_attr("python_source", f"# Source unavailable for {pyfunc_name}")
|
||||
|
||||
packed_func = _create_python_packed_func(pyfunc)
|
||||
func = func.with_attr("python_packed_func", packed_func)
|
||||
|
||||
irmodule[global_var] = func
|
||||
@@ -0,0 +1,27 @@
|
||||
# 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.
|
||||
"""Error classes for diagnostics."""
|
||||
|
||||
from . import doc
|
||||
|
||||
|
||||
class ParserError(Exception):
|
||||
"""Error class for diagnostics."""
|
||||
|
||||
def __init__(self, node: doc.AST, msg: str):
|
||||
super().__init__(msg)
|
||||
self.node = node
|
||||
@@ -0,0 +1,604 @@
|
||||
# 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.
|
||||
"""AST Evaluation"""
|
||||
|
||||
import ast
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import tvm
|
||||
|
||||
from . import dispatch, doc
|
||||
from .error import ParserError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .parser import Parser
|
||||
|
||||
DEFAULT_OP: dict[type, Callable[..., Any]] = {
|
||||
doc.Add: lambda a, b: a + b,
|
||||
doc.Sub: lambda a, b: a - b,
|
||||
doc.Mult: lambda a, b: a * b,
|
||||
doc.Div: lambda a, b: a / b,
|
||||
doc.FloorDiv: lambda a, b: a // b,
|
||||
doc.Mod: lambda a, b: a % b,
|
||||
doc.LShift: lambda a, b: a << b,
|
||||
doc.RShift: lambda a, b: a >> b,
|
||||
doc.BitOr: lambda a, b: a | b,
|
||||
doc.BitXor: lambda a, b: a ^ b,
|
||||
doc.BitAnd: lambda a, b: a & b,
|
||||
doc.MatMult: lambda a, b: a @ b,
|
||||
doc.Pow: lambda a, b: a**b,
|
||||
doc.Eq: lambda a, b: a == b,
|
||||
doc.NotEq: lambda a, b: a != b,
|
||||
doc.Lt: lambda a, b: a < b,
|
||||
doc.LtE: lambda a, b: a <= b,
|
||||
doc.Gt: lambda a, b: a > b,
|
||||
doc.GtE: lambda a, b: a >= b,
|
||||
doc.Is: lambda a, b: a is b,
|
||||
doc.IsNot: lambda a, b: a is not b,
|
||||
doc.In: lambda a, b: a in b,
|
||||
doc.NotIn: lambda a, b: a not in b,
|
||||
doc.And: lambda a, b: a and b,
|
||||
doc.Or: lambda a, b: a or b,
|
||||
doc.Invert: lambda a: ~a,
|
||||
doc.Not: lambda a: not a,
|
||||
doc.UAdd: lambda a: +a,
|
||||
doc.USub: lambda a: -a,
|
||||
}
|
||||
|
||||
|
||||
def _get_builtin_or_none(name: str):
|
||||
builtins = globals().get("__builtins__")
|
||||
if not builtins:
|
||||
return None
|
||||
if not isinstance(builtins, dict) and hasattr(builtins, "__dict__"):
|
||||
builtins = builtins.__dict__
|
||||
if isinstance(builtins, dict):
|
||||
return builtins.get(name)
|
||||
return None
|
||||
|
||||
|
||||
class ExprEvaluator:
|
||||
"""Expression evaluator for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser bound with the evaluator.
|
||||
|
||||
value_table : Dict[str, Any]
|
||||
The value table for expression evaluation.
|
||||
|
||||
new_value_count : int
|
||||
The count for intermediate result added during evaluation.
|
||||
"""
|
||||
|
||||
parser: "Parser"
|
||||
value_table: dict[str, Any]
|
||||
new_value_count: int
|
||||
|
||||
def __init__(self, parser: "Parser", value_table: dict[str, Any]) -> None:
|
||||
super().__init__()
|
||||
self.parser = parser
|
||||
self.value_table = value_table
|
||||
self.new_value_count = 0
|
||||
|
||||
@staticmethod
|
||||
def eval(parser: "Parser", value_table: dict[str, Any], node: doc.AST) -> Any:
|
||||
"""Expression evaluation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser bound with the evaluator.
|
||||
|
||||
value_table : Dict[str, Any]
|
||||
The value table for expression evaluation.
|
||||
|
||||
node : doc.AST
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
self = ExprEvaluator(parser, value_table)
|
||||
result = self._visit(node) # pylint: disable=protected-access
|
||||
if isinstance(result, doc.Name):
|
||||
if result.id in self.value_table:
|
||||
return self.value_table[result.id]
|
||||
else:
|
||||
builtin = _get_builtin_or_none(result.id)
|
||||
if builtin:
|
||||
return builtin
|
||||
raise ParserError(result, f"Undefined variable: {result.id}")
|
||||
if isinstance(result, doc.Constant):
|
||||
return result.value
|
||||
raise TypeError(f"Unexpected result type: {type(result)}")
|
||||
|
||||
def _add_intermediate_result(self, value: Any) -> doc.Name:
|
||||
"""Add intermediate result during evaluation into value table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : Any
|
||||
The intermediate result.
|
||||
|
||||
Returns
|
||||
-------
|
||||
name : doc.Name
|
||||
The doc AST name node with intermediate name for intermediate result.
|
||||
"""
|
||||
name = f"__tvm_tmp_value_{self.new_value_count}"
|
||||
self.new_value_count += 1
|
||||
self.value_table[name] = value
|
||||
lineno = 0
|
||||
col_offset = 0
|
||||
return doc.Name(
|
||||
id=name,
|
||||
ctx=doc.Load(),
|
||||
lineno=lineno,
|
||||
col_offset=col_offset,
|
||||
end_lineno=None,
|
||||
end_col_offset=None,
|
||||
)
|
||||
|
||||
def _visit(self, node: doc.AST) -> Any:
|
||||
"""General doc AST node visiting method for expression evaluation.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
args = []
|
||||
if (
|
||||
isinstance(node, doc.Call)
|
||||
and hasattr(node.func, "attr")
|
||||
and node.func.attr not in ["reads", "writes", "match_buffer"]
|
||||
) or isinstance(node, doc.BinOp | doc.UnaryOp | doc.Compare | doc.BoolOp | doc.IfExp):
|
||||
if isinstance(node, doc.BinOp):
|
||||
args = [node.left, node.right]
|
||||
elif isinstance(node, doc.UnaryOp):
|
||||
args = [node.operand]
|
||||
elif isinstance(node, doc.Compare):
|
||||
args = [node.left, *node.comparators]
|
||||
elif isinstance(node, doc.IfExp):
|
||||
args = [node.test, node.body, node.orelse]
|
||||
elif isinstance(node, doc.Call):
|
||||
args = node.args
|
||||
elif isinstance(node, doc.BoolOp):
|
||||
args = node.values
|
||||
for arg in args:
|
||||
if isinstance(arg, doc.Subscript) and isinstance(arg.slice, doc.Slice | doc.Tuple):
|
||||
if isinstance(arg.slice, doc.Slice):
|
||||
check_slices = [arg.slice]
|
||||
else:
|
||||
check_slices = []
|
||||
for p in arg.slice.elts:
|
||||
if isinstance(p, doc.Slice):
|
||||
check_slices.append(p)
|
||||
for s in check_slices:
|
||||
if not s.step and s.upper and s.lower:
|
||||
s.step = doc.Constant(
|
||||
1,
|
||||
None,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 1,
|
||||
s.upper.lineno,
|
||||
s.upper.end_col_offset + 2,
|
||||
)
|
||||
if isinstance(node, list):
|
||||
return [self._visit(n) for n in node]
|
||||
if isinstance(node, tuple):
|
||||
return tuple(self._visit(n) for n in node)
|
||||
assert isinstance(node, doc.AST)
|
||||
if isinstance(node, doc.Name):
|
||||
if node.id not in self.value_table and not _get_builtin_or_none(node.id):
|
||||
raise ParserError(node, f"Undefined variable: {node.id}")
|
||||
return node
|
||||
if isinstance(
|
||||
node,
|
||||
doc.Constant | doc.expr_context | doc.operator | doc.boolop | doc.unaryop | doc.cmpop,
|
||||
):
|
||||
return node
|
||||
if isinstance(node, doc.keyword):
|
||||
return doc.keyword(arg=node.arg, value=self._visit(node.value))
|
||||
if not isinstance(node, doc.expr | doc.Slice):
|
||||
return node
|
||||
if isinstance(node, doc.Lambda):
|
||||
return self._eval_lambda(node)
|
||||
if isinstance(node, doc.Starred):
|
||||
value = self._visit(node.value)
|
||||
return doc.Starred(
|
||||
value=value,
|
||||
ctx=node.ctx,
|
||||
lineno=node.lineno,
|
||||
col_offset=node.col_offset,
|
||||
end_lineno=node.end_lineno,
|
||||
end_col_offset=node.end_col_offset,
|
||||
)
|
||||
|
||||
if isinstance(node, doc.ListComp | doc.SetComp | doc.DictComp):
|
||||
value = self._eval_expr(node)
|
||||
return self._add_intermediate_result(value)
|
||||
|
||||
fields = {}
|
||||
for field in node.__class__._FIELDS: # pylint: disable=protected-access
|
||||
attr = getattr(node, field)
|
||||
if isinstance(attr, doc.AST | tuple | list):
|
||||
fields[field] = self._visit(attr)
|
||||
else:
|
||||
fields[field] = attr
|
||||
try:
|
||||
if isinstance(node, doc.BoolOp):
|
||||
value = self._eval_bool_op(fields)
|
||||
elif isinstance(node, doc.Compare):
|
||||
value = self._eval_compare(fields)
|
||||
elif isinstance(node, doc.UnaryOp):
|
||||
value = self._eval_unary_op(fields)
|
||||
elif isinstance(node, doc.BinOp):
|
||||
value = self._eval_bin_op(fields)
|
||||
elif isinstance(node, doc.IfExp):
|
||||
value = self._eval_if_exp(fields)
|
||||
elif isinstance(node, doc.Slice):
|
||||
value = self._eval_slice(fields)
|
||||
else:
|
||||
value = self._eval_expr(node.__class__(**fields))
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.parser.report_error(node, err)
|
||||
return self._add_intermediate_result(value)
|
||||
|
||||
def _eval_lambda(self, node: doc.Lambda) -> Any:
|
||||
"""The doc AST lambda node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Lambda
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
try:
|
||||
value = self._eval_expr(node)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.parser.report_error(node, err)
|
||||
return self._add_intermediate_result(value)
|
||||
|
||||
def _eval_bool_op(self, fields: dict[str, Any]) -> Any:
|
||||
"""The doc AST boolean operator node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of boolean operation information,
|
||||
e.g., operator types, operand values.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
op = fields["op"]
|
||||
if not isinstance(op, doc.And | doc.Or):
|
||||
raise TypeError(f"Unexpected operator: {op}")
|
||||
value = self._eval_expr(fields["values"][0])
|
||||
for rhs in fields["values"][1:]:
|
||||
value = _eval_op(op, values=[value, self._eval_expr(rhs)])
|
||||
return value
|
||||
|
||||
def _eval_compare(self, fields: dict[str, Any]) -> Any:
|
||||
"""The doc AST comparison operation node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of comparison operation information,
|
||||
e.g., operator types, operand values.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
values = [self._eval_expr(fields["left"])]
|
||||
values.extend([self._eval_expr(rhs) for rhs in fields["comparators"]])
|
||||
result = None
|
||||
assert len(fields["ops"]) == len(values) - 1
|
||||
|
||||
for index, op in enumerate(fields["ops"]):
|
||||
sub_result = _eval_op(op, values=[values[index], values[index + 1]])
|
||||
if result is None:
|
||||
result = sub_result
|
||||
else:
|
||||
result = _eval_op(doc.And(), values=[result, sub_result])
|
||||
return result
|
||||
|
||||
def _eval_unary_op(self, fields: dict[str, Any]) -> Any:
|
||||
"""The doc AST unary operation node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of unary operation information,
|
||||
e.g., operator types, operand values.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
value = self._eval_expr(fields["operand"])
|
||||
value = _eval_op(fields["op"], values=[value])
|
||||
return value
|
||||
|
||||
def _eval_bin_op(self, fields: dict[str, Any]) -> Any:
|
||||
"""The doc AST binary operation node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of binary operation information,
|
||||
e.g., operator types, operand values.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
return _eval_op(
|
||||
fields["op"],
|
||||
values=[
|
||||
self._eval_expr(fields["left"]),
|
||||
self._eval_expr(fields["right"]),
|
||||
],
|
||||
)
|
||||
|
||||
def _eval_if_exp(self, fields: dict[str, Any]) -> Any:
|
||||
"""The doc AST if-else expression node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of if-else expression information,
|
||||
e.g., test, body, orelse.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
test = self._eval_expr(fields["test"])
|
||||
body = self._eval_expr(fields["body"])
|
||||
orelse = self._eval_expr(fields["orelse"])
|
||||
if isinstance(test, bool):
|
||||
return body if test else orelse
|
||||
elif tvm.ir.is_prim_expr(test) and test.ty.matches_code(tvm.DataTypeCode.BOOL):
|
||||
return tvm.tirx.op.if_then_else(test, body, orelse)
|
||||
else:
|
||||
raise TypeError(f"Expected Python bool or TIR bool, but got {type(test)}")
|
||||
|
||||
def _eval_slice(self, fields: dict[str, Any]) -> slice:
|
||||
"""The doc AST slice node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fields : Dict[str, Any]
|
||||
The dictionary of slice information,
|
||||
e.g., lower bound, upper bound, step.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : slice
|
||||
The evaluation result.
|
||||
"""
|
||||
lower, upper, step = fields["lower"], fields["upper"], fields["step"]
|
||||
|
||||
lower = self._eval_expr(lower) if lower is not None else None
|
||||
upper = self._eval_expr(upper) if upper is not None else None
|
||||
step = self._eval_expr(step) if step is not None else None
|
||||
|
||||
return slice(lower, upper, step)
|
||||
|
||||
def _eval_expr(self, v: Any) -> Any:
|
||||
"""The doc AST expression node evaluating method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
v : Any
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
return _eval_expr(v, self.value_table)
|
||||
|
||||
|
||||
def eval_expr(
|
||||
parser: "Parser",
|
||||
node: doc.expr | doc.Expression,
|
||||
dict_globals: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
"""Expression evaluation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser bound with the evaluator.
|
||||
|
||||
node : Union[doc.expr, doc.Expression]
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
dict_globals : Optional[Dict[str, Any]]
|
||||
The optional global value table for expression evaluation.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
value_table = {}
|
||||
if dict_globals is not None:
|
||||
value_table.update(dict_globals)
|
||||
return ExprEvaluator.eval(parser, value_table, node)
|
||||
|
||||
|
||||
def eval_assign(
|
||||
parser: "Parser",
|
||||
target: doc.expr,
|
||||
source: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Expression assignment evaluation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser bound with the evaluator.
|
||||
|
||||
target : doc.expr
|
||||
The root node of AST tree node of assigned expression to evaluate.
|
||||
|
||||
source : Any
|
||||
The source to be assigned with evaluated expression.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
try:
|
||||
return _eval_assign(target, source)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
parser.report_error(target, err)
|
||||
raise
|
||||
|
||||
|
||||
def _eval_expr(
|
||||
node: doc.expr | doc.Expression,
|
||||
dict_globals: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
"""Expression evaluation implementation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : Union[doc.expr, doc.Expression]
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
dict_globals : Optional[Dict[str, Any]]
|
||||
The optional global value table for expression evaluation.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
node = doc.from_doc(node)
|
||||
if isinstance(node, ast.expr):
|
||||
node = ast.Expression(body=node)
|
||||
assert isinstance(node, ast.Expression), "Expects an ast.Expression, but gets: " + str(node)
|
||||
if dict_globals is None:
|
||||
dict_globals = {}
|
||||
node = ast.fix_missing_locations(node)
|
||||
exe = compile(node, filename="<ast>", mode="eval")
|
||||
return eval(exe, dict_globals) # pylint: disable=eval-used
|
||||
|
||||
|
||||
def _eval_op(
|
||||
op: doc.AST,
|
||||
values: list[Any],
|
||||
):
|
||||
"""Operation expression evaluation implementation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
op : doc.AST
|
||||
The root node of AST tree node of operation expression to evaluate.
|
||||
|
||||
values : List[Any]
|
||||
The list of values of operands.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
op_type = type(op) # pylint: disable=protected-access
|
||||
for i, v in enumerate(values):
|
||||
v_type = getattr(type(v), "_dispatch_type", None)
|
||||
if v_type is None:
|
||||
continue
|
||||
f = dispatch.get_op(
|
||||
operand_type=v_type, op_node_type=op_type, operand_index=i, default=None
|
||||
)
|
||||
if f is not None:
|
||||
return f(*values)
|
||||
return DEFAULT_OP[op_type](*values)
|
||||
|
||||
|
||||
def _eval_assign(
|
||||
target: doc.expr,
|
||||
source: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""Expression assignment evaluation implementation for TVMScript parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : doc.expr
|
||||
The root node of AST tree node of assigned expression to evaluate.
|
||||
|
||||
source : Any
|
||||
The source to be assigned with evaluated expression.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
target = doc.from_doc(target)
|
||||
assert isinstance(target, ast.expr)
|
||||
RHS_VAR_NAME = "__tvm_rhs_var__" # pylint: disable=invalid-name
|
||||
rhs_var_name = RHS_VAR_NAME
|
||||
dict_locals = {rhs_var_name: source}
|
||||
mod = ast.fix_missing_locations(
|
||||
ast.Module(
|
||||
body=[
|
||||
ast.Assign(
|
||||
targets=[target],
|
||||
value=ast.Name(
|
||||
id=rhs_var_name,
|
||||
ctx=ast.Load(),
|
||||
),
|
||||
)
|
||||
],
|
||||
type_ignores=[],
|
||||
)
|
||||
)
|
||||
exe = compile(mod, filename="<ast>", mode="exec")
|
||||
exec(exe, {}, dict_locals) # pylint: disable=exec-used
|
||||
del dict_locals[rhs_var_name]
|
||||
return dict_locals
|
||||
@@ -0,0 +1,952 @@
|
||||
# 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.
|
||||
"""The core parser"""
|
||||
|
||||
import abc
|
||||
import inspect
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
from tvm.error import DiagnosticError
|
||||
from tvm.ir import GlobalVar
|
||||
|
||||
from . import dispatch, doc
|
||||
from .diagnostics import Diagnostics, Source
|
||||
from .evaluator import eval_assign, eval_expr
|
||||
|
||||
DEFAULT_VISIT = {
|
||||
"Interactive",
|
||||
"Module",
|
||||
"Expression",
|
||||
"Pass",
|
||||
}
|
||||
|
||||
|
||||
def _deferred(exit_f: Callable[[], None]):
|
||||
"""Created context with certain exit function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
exit_f : Callable[[], None]
|
||||
The function to call when exiting the context.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The created context.
|
||||
"""
|
||||
|
||||
@contextmanager
|
||||
def context():
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
exit_f()
|
||||
|
||||
return context()
|
||||
|
||||
|
||||
def _do_nothing(*args, **kwargs): # pylint: disable=unused-argument
|
||||
pass
|
||||
|
||||
|
||||
class ScriptMacro(abc.ABC):
|
||||
"""Representation of a script macro.
|
||||
|
||||
This is a callable object, intended to be called from the expression evaluator.
|
||||
The evaluator is expected to insert the current parser into the environment
|
||||
undef the name given by "parser_object_name".
|
||||
|
||||
Once called, the ScriptMacro object will locate the current parser, and use it
|
||||
to parse the macro's body and produce the result.
|
||||
|
||||
There were two major considerations for this design:
|
||||
1. Implementing hygienic and non-hygienic macros.
|
||||
2. Implementing macros that return values.
|
||||
|
||||
Macro uses in TIR are only allowed at a statement-level, and they don't produce
|
||||
any values. Parsing of such macros could easily be done by intercepting doc.Call
|
||||
nodes in the TIR parser. If a macro is a value-producing expression, then there
|
||||
may not be a direct way to intercept calls to it if it's embedded in a complex
|
||||
expression. Because macros use function-call syntax, the evaluator will try to
|
||||
call the macro object, which this design relies on to parse and evaluate the macro.
|
||||
"""
|
||||
|
||||
parser_object_name = "__current_script_parser__"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Source,
|
||||
closure_vars: dict[str, Any],
|
||||
func: Callable,
|
||||
hygienic: bool,
|
||||
) -> None:
|
||||
self.source = source
|
||||
self.closure_vars = closure_vars
|
||||
self.func = func
|
||||
self.hygienic = hygienic
|
||||
|
||||
def __repr__(self):
|
||||
return self.source.source
|
||||
|
||||
@abc.abstractmethod
|
||||
def parse_macro(self, parser: "Parser") -> Any:
|
||||
"""The main macro parsing function. Different scripts may have different
|
||||
ways to parse a macro, and to return a value to the evaluator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
parser : Parser
|
||||
The parser with the appropriate frame already created and populated depending
|
||||
macro's hygiene settings,
|
||||
|
||||
Returns
|
||||
-------
|
||||
The return value depends on the specifics of the particular script. It can be
|
||||
"None" or any other value or any type.
|
||||
"""
|
||||
|
||||
def _find_parser_def(self):
|
||||
outer_frame_infos = inspect.getouterframes(inspect.currentframe())
|
||||
for finfo in outer_frame_infos:
|
||||
parser = finfo.frame.f_globals.get(ScriptMacro.parser_object_name)
|
||||
if parser is not None:
|
||||
return parser
|
||||
raise RuntimeError(f"{ScriptMacro.parser_object_name} not available")
|
||||
|
||||
def get_macro_def(self):
|
||||
ast_module = self.source.as_ast()
|
||||
for decl in ast_module.body:
|
||||
if isinstance(decl, doc.FunctionDef) and decl.name == self.func.__name__:
|
||||
return decl
|
||||
raise RuntimeError(f"cannot find macro definition for {self.func.__name__}")
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
param_binding = inspect.signature(self.func).bind(*args, **kwargs)
|
||||
param_binding.apply_defaults()
|
||||
local_vars = param_binding.arguments
|
||||
parser = self._find_parser_def()
|
||||
|
||||
with parser.with_diag_source(self.source):
|
||||
if self.hygienic:
|
||||
saved_var_table = parser.var_table
|
||||
parser.var_table = VarTable()
|
||||
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in self.closure_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
for k, v in local_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
|
||||
parse_result = self.parse_macro(parser)
|
||||
|
||||
parser.var_table = saved_var_table
|
||||
|
||||
else:
|
||||
with parser.var_table.with_frame():
|
||||
for k, v in local_vars.items():
|
||||
parser.var_table.add(k, v)
|
||||
|
||||
parse_result = self.parse_macro(parser)
|
||||
|
||||
return parse_result
|
||||
|
||||
|
||||
class VarTableFrame:
|
||||
"""The variable table frame.
|
||||
A frame of variable table stores the variables created in one block or scope.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vars : Set[str]
|
||||
The set of variable names in the variable table frame.
|
||||
"""
|
||||
|
||||
vars: set[str]
|
||||
|
||||
def __init__(self):
|
||||
self.vars = set()
|
||||
|
||||
def add(self, var: str):
|
||||
"""Add a new variable into variable table frame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
var : str
|
||||
The name of new variable.
|
||||
"""
|
||||
if var in self.vars:
|
||||
raise ValueError(f"Variable {var} already defined in current scope")
|
||||
self.vars.add(var)
|
||||
|
||||
def pop_all(self, fn_pop: Callable[[str], None]):
|
||||
"""Pop out all variable in variable table frame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
fn_pop : Callable[[str], None]
|
||||
The methods to call when popping each variable.
|
||||
"""
|
||||
for var in self.vars:
|
||||
fn_pop(var)
|
||||
self.vars.clear()
|
||||
|
||||
|
||||
class VarTable:
|
||||
"""The variable table.
|
||||
A variable table stores the all variables when parsing TVMScript.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[VarTableFrame]
|
||||
The list or stack of variable table frame.
|
||||
|
||||
name2value : Dict[str, List[Any]]
|
||||
The dictionary for variable table name-based query.
|
||||
"""
|
||||
|
||||
frames: list[VarTableFrame]
|
||||
name2value: dict[str, list[Any]]
|
||||
|
||||
def __init__(self):
|
||||
self.frames = []
|
||||
self.name2value = defaultdict(list)
|
||||
|
||||
def with_frame(self):
|
||||
"""Create a new variable table frame as with statement.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new variable table frame.
|
||||
"""
|
||||
|
||||
def pop_frame():
|
||||
frame = self.frames.pop()
|
||||
frame.pop_all(lambda name: self.name2value[name].pop())
|
||||
|
||||
self.frames.append(VarTableFrame())
|
||||
return _deferred(pop_frame)
|
||||
|
||||
def add(self, var: str, value: Any, allow_shadowing: bool = False):
|
||||
"""Add a new variable to variable table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
var : str
|
||||
The name of variable.
|
||||
|
||||
value : Any
|
||||
The value of variable.
|
||||
|
||||
allow_shadowing : bool
|
||||
The options of whether variable shadowing allowed for this variable.
|
||||
"""
|
||||
# Skip if the key and value are equal to those in the var_table
|
||||
if self.name2value[var] and isinstance(self.name2value[var][-1], type(value)):
|
||||
if isinstance(value, np.ndarray) and (self.name2value[var][-1] == value).all():
|
||||
return
|
||||
elif self.name2value[var][-1] == value:
|
||||
return
|
||||
if allow_shadowing and var in self.frames[-1].vars:
|
||||
# Shadowing
|
||||
self.name2value[var][-1] = value
|
||||
else:
|
||||
self.frames[-1].add(var)
|
||||
self.name2value[var].append(value)
|
||||
|
||||
def get(self) -> dict[str, Any]:
|
||||
"""Get a variable dictionary of latest variables.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The variable dictionary copy of latest variables.
|
||||
"""
|
||||
return {key: values[-1] for key, values in self.name2value.items() if values}
|
||||
|
||||
def get_at_depth(self, depth: int) -> dict[str, Any]:
|
||||
"""Get variables visible at the given frame depth, using current values.
|
||||
|
||||
For each variable name that appears in frames 0..depth-1, count how many
|
||||
times it was pushed (to handle shadowing), then index into name2value at
|
||||
count-1 to retrieve the latest value visible at that depth.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
depth : int
|
||||
The frame depth (number of frames visible).
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : dict[str, Any]
|
||||
Variable dictionary of values visible at the given depth.
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
name_count: dict[str, int] = defaultdict(int)
|
||||
for frame_idx in range(min(depth, len(self.frames))):
|
||||
for name in self.frames[frame_idx].vars:
|
||||
name_count[name] += 1
|
||||
for name, count in name_count.items():
|
||||
if self.name2value[name]:
|
||||
result[name] = self.name2value[name][count - 1]
|
||||
return result
|
||||
|
||||
def exist(self, value: Any) -> bool:
|
||||
"""Check if any value exists in variable table.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : Any
|
||||
The value of variable.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : bool
|
||||
The existence of the value.
|
||||
"""
|
||||
return any(
|
||||
value.same_as(known_value)
|
||||
for known_value_stack in self.name2value.values()
|
||||
for known_value in known_value_stack
|
||||
)
|
||||
|
||||
|
||||
def _dispatch_wrapper(func: dispatch.ParseMethod) -> dispatch.ParseMethod:
|
||||
def _wrapper(self: "Parser", node: doc.AST) -> None:
|
||||
try:
|
||||
return func(self, node)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
raise
|
||||
|
||||
return _wrapper
|
||||
|
||||
|
||||
def _dispatch(self: "Parser", type_name: str) -> dispatch.ParseMethod:
|
||||
for token in [self.dispatch_tokens[-1], "default"]:
|
||||
func = dispatch.get(token=token, type_name=type_name, default=None)
|
||||
if func is not None:
|
||||
return _dispatch_wrapper(func)
|
||||
return _dispatch_wrapper(lambda self, node: self.generic_visit(node))
|
||||
|
||||
|
||||
class Parser(doc.NodeVisitor):
|
||||
"""The TVMScript parser
|
||||
|
||||
Parameters
|
||||
----------
|
||||
diag : Diagnostics
|
||||
The diagnostics for error reporting.
|
||||
|
||||
dispatch_tokens : List[str]
|
||||
The list of dispatching tokens to dispatching parsing method
|
||||
of different IRs and different doc AST structure.
|
||||
|
||||
var_table : VarTable
|
||||
The variable table for parsing.
|
||||
"""
|
||||
|
||||
diag: Diagnostics
|
||||
dispatch_tokens: list[str]
|
||||
function_annotations: dict[str, dict[str, Any]] | None
|
||||
var_table: VarTable
|
||||
inside_function: bool # whether we are within a function
|
||||
current_class: str | None = None # current class being parsed
|
||||
base_py_module_context: bool = False # whether current class inherits from BasePyModule
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
source: Source,
|
||||
function_annotations: dict[str, dict[str, Any]],
|
||||
) -> None:
|
||||
self.diag = Diagnostics(source)
|
||||
self.dispatch_tokens = ["default"]
|
||||
self.function_annotations = function_annotations
|
||||
self.var_table = VarTable()
|
||||
self.inside_function = False
|
||||
|
||||
def parse(self, extra_vars: dict[str, Any] | None = None) -> Any:
|
||||
"""The main parse method for parser.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
extra_vars : Optional[Dict[str, Any]]
|
||||
The optional global value table for parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The doc AST node visiting result.
|
||||
"""
|
||||
if extra_vars is None:
|
||||
extra_vars = {}
|
||||
with self.var_table.with_frame():
|
||||
for k, v in extra_vars.items():
|
||||
self.var_table.add(k, v)
|
||||
node = self.diag.source.as_ast()
|
||||
self.visit(node)
|
||||
|
||||
def get_dispatch_token(self, node: doc.FunctionDef) -> str:
|
||||
if not isinstance(node, doc.FunctionDef):
|
||||
self.report_error(node, "Only can get dispatch token for function.")
|
||||
if not node.decorator_list:
|
||||
self.report_error(node, "Function must be decorated")
|
||||
# TODO: only the last decorator is parsed
|
||||
decorator = self.eval_expr(node.decorator_list[-1])
|
||||
if not hasattr(decorator, "dispatch_token"):
|
||||
self.report_error(node, "The parser does not understand the decorator")
|
||||
return decorator.dispatch_token
|
||||
|
||||
def with_dispatch_token(self, token: str):
|
||||
"""Add a new dispatching token as with statement.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
token : str
|
||||
The dispatching token.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new dispatching token.
|
||||
"""
|
||||
|
||||
self.dispatch_tokens.append(token)
|
||||
enter_func = dispatch.get(token=token, type_name="enter_token", default=lambda *args: None)
|
||||
context = enter_func(self)
|
||||
|
||||
def pop_token():
|
||||
exit_func = dispatch.get(
|
||||
token=token, type_name="exit_token", default=lambda *args: None
|
||||
)
|
||||
exit_func(self, context)
|
||||
self.dispatch_tokens.pop()
|
||||
|
||||
return _deferred(pop_token)
|
||||
|
||||
def set_class_context(self, class_name: str, is_base_py_module: bool = False):
|
||||
"""Set the current class context for parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
class_name : str
|
||||
The name of the current class being parsed.
|
||||
is_base_py_module : bool
|
||||
Whether the current class inherits from BasePyModule.
|
||||
"""
|
||||
self.current_class = class_name
|
||||
self.base_py_module_context = is_base_py_module
|
||||
|
||||
def _get_current_class_context(self) -> str | None:
|
||||
"""Get the current class context.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Optional[str]
|
||||
The name of the current class, or None if not in a class context.
|
||||
"""
|
||||
return self.current_class
|
||||
|
||||
def _is_base_py_module_context(self) -> bool:
|
||||
"""Check if the current class context allows Python functions.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if Python functions are allowed in the current context.
|
||||
"""
|
||||
return self.base_py_module_context
|
||||
|
||||
def with_diag_source(self, source: Source):
|
||||
"""Add a new source as with statement.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
source : Source
|
||||
The source for diagnostics.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The context with new source.
|
||||
"""
|
||||
|
||||
last_diag = self.diag
|
||||
self.diag = Diagnostics(source)
|
||||
|
||||
def pop_source():
|
||||
self.diag = last_diag
|
||||
|
||||
return _deferred(pop_source)
|
||||
|
||||
def eval_expr(
|
||||
self,
|
||||
node: doc.Expression | doc.expr,
|
||||
extra_vars: dict[str, Any] | None = None,
|
||||
) -> Any:
|
||||
"""Expression evaluation when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : Union[doc.expr, doc.Expression]
|
||||
The root node of AST tree node of expression to evaluate.
|
||||
|
||||
extra_vars : Optional[Dict[str, Any]]
|
||||
The optional global value table for expression evaluation.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The evaluation result.
|
||||
"""
|
||||
var_values = self.var_table.get()
|
||||
if extra_vars is not None:
|
||||
for k, v in extra_vars.items():
|
||||
var_values[k] = v
|
||||
var_values[ScriptMacro.parser_object_name] = self
|
||||
return eval_expr(self, node, var_values)
|
||||
|
||||
def _duplicate_lhs_check(self, target: doc.expr) -> bool | set[str]:
|
||||
"""Check whether duplicate lhs exists in assignment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : doc.expr
|
||||
The doc AST expr node for lhs.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Union[bool, Set[str]]
|
||||
The result of true if duplicate lhs exists,
|
||||
or the set of lhs names if no duplicate lhs exists.
|
||||
"""
|
||||
if isinstance(target, doc.Tuple | doc.List):
|
||||
vars: set[str] = set() # pylint: disable=redefined-builtin
|
||||
for i in target.elts:
|
||||
res = self._duplicate_lhs_check(i)
|
||||
if isinstance(res, bool) and res:
|
||||
return True
|
||||
assert isinstance(res, set)
|
||||
if vars & res:
|
||||
return True
|
||||
vars = vars.union(res)
|
||||
return vars
|
||||
elif isinstance(target, doc.Name):
|
||||
return {target.id}
|
||||
elif isinstance(target, doc.Starred):
|
||||
return self._duplicate_lhs_check(target.value)
|
||||
elif isinstance(target, doc.Attribute):
|
||||
# Attribute assignment like packedB.data = ..., treated as rebinding.
|
||||
return {target.attr}
|
||||
else:
|
||||
self.report_error(target, "Invalid type in assign statement")
|
||||
raise NotImplementedError
|
||||
|
||||
def eval_assign(
|
||||
self,
|
||||
target: doc.expr,
|
||||
source: Any,
|
||||
bind_value: Callable[["Parser", doc.expr, str, Any], Any],
|
||||
allow_shadowing: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
"""Expression assignment evaluation when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
target : doc.expr
|
||||
The root node of AST tree node of assigned expression to evaluate.
|
||||
|
||||
source : Any
|
||||
The source to be assigned with evaluated expression.
|
||||
|
||||
bind_value : Callable[["Parser", doc.expr, str, Any], Any]
|
||||
The value binding method when assigning the values to variables.
|
||||
|
||||
allow_shadowing : bool
|
||||
The options of whether variable shadowing allowed for assignment.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Dict[str, Any]
|
||||
The dictionary of assignment result.
|
||||
"""
|
||||
if self._duplicate_lhs_check(target) is True:
|
||||
self.report_error(target, "Duplicate vars assigned.")
|
||||
var_values = eval_assign(self, target, source)
|
||||
for k, v in var_values.items():
|
||||
var = bind_value(self, target, k, v)
|
||||
self.var_table.add(k, var, allow_shadowing)
|
||||
return var_values
|
||||
|
||||
def report_error(self, node: doc.AST, err: Exception | str) -> None: # pylint: disable=no-self-use
|
||||
"""The error reporting when parsing.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The doc AST node with errors.
|
||||
|
||||
err: Union[Exception, str]
|
||||
The error to report.
|
||||
"""
|
||||
|
||||
# If the error is already being raised as a DiagnosticError,
|
||||
# re-raise it without wrapping it in a DiagnosticContext.
|
||||
if isinstance(err, DiagnosticError):
|
||||
raise err
|
||||
|
||||
# Only take the last line of the error message
|
||||
if isinstance(err, RuntimeError):
|
||||
lines = list(filter(None, str(err).split("\n")))
|
||||
msg = lines[-1] if lines else (str(err) or type(err).__name__)
|
||||
elif isinstance(err, KeyError):
|
||||
msg = "KeyError: " + str(err)
|
||||
else:
|
||||
msg = str(err)
|
||||
|
||||
try:
|
||||
self.diag.error(node, msg)
|
||||
except Exception as diag_err:
|
||||
# Calling self.diag.error is guaranteed to throw an
|
||||
# exception. When shown to a user, this error should
|
||||
# reference the point of error within the provided
|
||||
# TVMScript. However, when caught in pdb, the full
|
||||
# traceback should be available for debugging.
|
||||
if isinstance(err, Exception):
|
||||
diag_err = diag_err.with_traceback(err.__traceback__)
|
||||
raise diag_err
|
||||
|
||||
def visit(self, node: doc.AST) -> None:
|
||||
"""The general visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AST
|
||||
The doc AST node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
if isinstance(node, list | tuple):
|
||||
for item in node:
|
||||
self.visit(item)
|
||||
return
|
||||
if not isinstance(node, doc.AST):
|
||||
return
|
||||
name = node.__class__.__name__.split(".")[-1]
|
||||
if name in DEFAULT_VISIT:
|
||||
func = self.generic_visit
|
||||
else:
|
||||
func = getattr(self, "visit_" + name, None)
|
||||
if func is None:
|
||||
raise NotImplementedError(f"Visitor of AST node is not implemented: {name}")
|
||||
try:
|
||||
func(node)
|
||||
except Exception as err: # pylint: disable=broad-except
|
||||
self.report_error(node, err)
|
||||
|
||||
def visit_body(self, node: list[doc.stmt]) -> Any:
|
||||
"""The general body visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : List[doc.stmt]
|
||||
The list of statements in body.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
for stmt in node:
|
||||
self.visit(stmt)
|
||||
|
||||
def visit_tvm_annotation(self, node: doc.expr) -> Any:
|
||||
"""The general TVM annotation visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.expr
|
||||
The doc AST expr node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "tvm_annotation")(self, node)
|
||||
|
||||
def visit_FunctionDef(self, node: doc.FunctionDef) -> None: # pylint: disable=invalid-name
|
||||
"""The general function definition visit method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.FunctionDef
|
||||
The doc FunctionDef node.
|
||||
"""
|
||||
token = self.get_dispatch_token(node)
|
||||
func = dispatch.get(token=token, type_name="FunctionDef", default=None)
|
||||
if func is None:
|
||||
self.report_error(
|
||||
node,
|
||||
"""The parser does not understand the decorator,
|
||||
or visit_FunctionDef is not implemented for the decorator with token: """
|
||||
+ token,
|
||||
)
|
||||
_dispatch(self, "pre_visit_local_function")(self, node)
|
||||
_dispatch_wrapper(func)(self, node)
|
||||
_dispatch(self, "post_visit_local_function")(self, node)
|
||||
|
||||
def visit_tvm_declare_function(self, node: doc.FunctionDef) -> GlobalVar:
|
||||
token = self.get_dispatch_token(node)
|
||||
with self.with_dispatch_token(token):
|
||||
return _dispatch(self, "tvm_declare_function")(self, node)
|
||||
|
||||
def visit_ClassDef(self, node: doc.ClassDef) -> Any: # pylint: disable=invalid-name
|
||||
"""The general class definition visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.ClassDef
|
||||
The doc AST class definition node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
func = dispatch.get(token="ir", type_name="ClassDef", default=None)
|
||||
if func is None:
|
||||
self.report_error(node, "The parser does not understand the decorator")
|
||||
_dispatch_wrapper(func)(self, node)
|
||||
|
||||
def visit_arguments(self, node: doc.arguments) -> Any:
|
||||
"""The general arguments visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.arguments
|
||||
The doc AST arguments node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "arguments")(self, node)
|
||||
|
||||
def visit_For(self, node: doc.For) -> Any: # pylint: disable=invalid-name
|
||||
"""The general for visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.For
|
||||
The doc AST for node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "For")(self, node)
|
||||
|
||||
def visit_While(self, node: doc.While) -> Any: # pylint: disable=invalid-name
|
||||
"""The general while visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.While
|
||||
The doc AST while node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "While")(self, node)
|
||||
|
||||
def visit_With(self, node: doc.With) -> Any: # pylint: disable=invalid-name
|
||||
"""The general with visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.With
|
||||
The doc AST with node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "With")(self, node)
|
||||
|
||||
def visit_Assign(self, node: doc.Assign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general assign visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assign
|
||||
The doc AST assign node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Assign")(self, node)
|
||||
|
||||
def visit_AnnAssign(self, node: doc.AnnAssign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general annotated assign visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assign
|
||||
The doc AST annotated assign node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "AnnAssign")(self, node)
|
||||
|
||||
def visit_Expr(self, node: doc.Expr) -> Any: # pylint: disable=invalid-name
|
||||
"""The general expression visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Expr
|
||||
The doc AST expression node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Expr")(self, node)
|
||||
|
||||
def visit_If(self, node: doc.If) -> Any: # pylint: disable=invalid-name
|
||||
"""The general if visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.If
|
||||
The doc AST if node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "If")(self, node)
|
||||
|
||||
def visit_AugAssign(self, node: doc.AugAssign) -> Any: # pylint: disable=invalid-name
|
||||
"""The general augmented assignment visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.AugAssign
|
||||
The doc AST augmented assignment node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "AugAssign")(self, node)
|
||||
|
||||
def visit_Assert(self, node: doc.Assert) -> Any: # pylint: disable=invalid-name
|
||||
"""The general assert visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Assert
|
||||
The doc AST assert node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Assert")(self, node)
|
||||
|
||||
def visit_Return(self, node: doc.Return) -> Any: # pylint: disable=invalid-name
|
||||
"""The general return visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Return
|
||||
The doc AST return node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Return")(self, node)
|
||||
|
||||
def visit_Continue(self, node: doc.Continue) -> Any: # pylint: disable=invalid-name
|
||||
"""The general continue visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Continue
|
||||
The doc AST continue node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Continue")(self, node)
|
||||
|
||||
def visit_Break(self, node: doc.Break) -> Any: # pylint: disable=invalid-name
|
||||
"""The general break visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Break
|
||||
The doc AST break node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Break")(self, node)
|
||||
|
||||
def visit_Nonlocal(self, node: doc.Nonlocal) -> Any: # pylint: disable=invalid-name
|
||||
"""The general nonlocal visiting method.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.Nonlocal
|
||||
The doc AST nonlocal node.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Any
|
||||
The visiting result.
|
||||
"""
|
||||
return _dispatch(self, "Nonlocal")(self, node)
|
||||
@@ -0,0 +1,212 @@
|
||||
# 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.
|
||||
"""TVM Script Parser utils"""
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from types import FrameType
|
||||
from typing import Any
|
||||
|
||||
from .diagnostics import findsource
|
||||
|
||||
|
||||
def get_func_nonlocals(func):
|
||||
"""A modified version of `inspect.getclosurevars`"""
|
||||
|
||||
if inspect.ismethod(func):
|
||||
func = func.__func__
|
||||
|
||||
if not inspect.isfunction(func):
|
||||
raise TypeError(f"{func!r} is not a Python function")
|
||||
|
||||
code = func.__code__
|
||||
# Nonlocal references are named in co_freevars and resolved
|
||||
# by looking them up in __closure__ by positional index
|
||||
nonlocal_vars = {}
|
||||
if func.__closure__ is not None:
|
||||
for var, cell in zip(code.co_freevars, func.__closure__):
|
||||
try:
|
||||
nonlocal_vars[var] = cell.cell_contents
|
||||
except ValueError as err:
|
||||
# cell_contents may raise ValueError if the cell is empty.
|
||||
if "empty" not in str(err):
|
||||
raise
|
||||
return nonlocal_vars
|
||||
|
||||
|
||||
def inspect_function_capture(func: Callable) -> dict[str, Any]:
|
||||
"""Capture function non-locals and global variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func : Callable
|
||||
The function to inspect.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Dict[str, Any]
|
||||
The function variables map with non-local or global variables.
|
||||
"""
|
||||
captured = {
|
||||
**func.__globals__, # type: ignore
|
||||
**get_func_nonlocals(func),
|
||||
}
|
||||
return captured
|
||||
|
||||
|
||||
def inspect_class_capture(cls: type) -> dict[str, Any]:
|
||||
"""Capture class non-locals and global variables.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
cls : type
|
||||
The class to inspect.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : Dict[str, Any]
|
||||
The class variables map with non-local or global variables.
|
||||
"""
|
||||
result: dict[str, Any] = {}
|
||||
for _, v in cls.__dict__.items():
|
||||
if inspect.isfunction(v):
|
||||
func_vars = inspect_function_capture(v)
|
||||
result.update(**func_vars)
|
||||
return result
|
||||
|
||||
|
||||
def _collect_annotation_names(source_obj: type | Callable) -> set[str]:
|
||||
"""Parse source AST to find names used in function annotations.
|
||||
|
||||
Returns the set of ``ast.Name`` identifiers found inside argument
|
||||
annotations and return annotations of any function definitions in
|
||||
*source_obj*.
|
||||
"""
|
||||
import ast
|
||||
import textwrap
|
||||
|
||||
try:
|
||||
source = textwrap.dedent(inspect.getsource(source_obj))
|
||||
tree = ast.parse(source)
|
||||
except (OSError, TypeError):
|
||||
return set()
|
||||
|
||||
names: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
|
||||
for arg in node.args.args + node.args.posonlyargs + node.args.kwonlyargs:
|
||||
if arg.annotation:
|
||||
for n in ast.walk(arg.annotation):
|
||||
if isinstance(n, ast.Name):
|
||||
names.add(n.id)
|
||||
if node.returns:
|
||||
for n in ast.walk(node.returns):
|
||||
if isinstance(n, ast.Name):
|
||||
names.add(n.id)
|
||||
return names
|
||||
|
||||
|
||||
def _has_string_annotations(source_obj: type | Callable) -> bool:
|
||||
"""Check if *source_obj* has stringified annotations (PEP 563)."""
|
||||
if inspect.isclass(source_obj):
|
||||
return any(
|
||||
isinstance(a, str)
|
||||
for v in source_obj.__dict__.values()
|
||||
if inspect.isfunction(v)
|
||||
for a in v.__annotations__.values()
|
||||
)
|
||||
return any(isinstance(a, str) for a in getattr(source_obj, "__annotations__", {}).values())
|
||||
|
||||
|
||||
def _get_enclosing_scope_names(qualname: str) -> set[str]:
|
||||
"""Extract lexically enclosing scope names from ``__qualname__``.
|
||||
|
||||
For ``outer.<locals>.inner.<locals>.func`` this returns ``{"outer", "inner"}``.
|
||||
"""
|
||||
parts = qualname.split(".")
|
||||
return {p for p in parts[:-1] if p != "<locals>"}
|
||||
|
||||
|
||||
def resolve_closure_vars(
|
||||
source_obj: type | Callable, extra_vars: dict[str, Any], outer_stack: list
|
||||
) -> None:
|
||||
"""Resolve closure variables hidden by PEP 563.
|
||||
|
||||
With ``from __future__ import annotations``, variables used only in
|
||||
annotations are not captured in ``__closure__``. This function parses
|
||||
the source AST to find names used in function annotations, then looks
|
||||
them up in lexically enclosing scope frames identified via
|
||||
``__qualname__``.
|
||||
|
||||
Only triggered when annotations are actually strings (PEP 563 active).
|
||||
Only annotation-referenced names are added, and only from enclosing
|
||||
scopes — not from arbitrary caller frames.
|
||||
|
||||
Works for both classes (``@I.ir_module``) and functions (``@T.prim_func``).
|
||||
"""
|
||||
if not _has_string_annotations(source_obj):
|
||||
return
|
||||
ann_names = _collect_annotation_names(source_obj)
|
||||
enclosing = _get_enclosing_scope_names(source_obj.__qualname__)
|
||||
for name in ann_names:
|
||||
if name not in extra_vars:
|
||||
for frame_info in outer_stack[1:]:
|
||||
if frame_info.frame.f_code.co_name in enclosing:
|
||||
if name in frame_info.frame.f_locals:
|
||||
extra_vars[name] = frame_info.frame.f_locals[name]
|
||||
break
|
||||
|
||||
|
||||
def is_defined_in_class(frames: list[FrameType], obj: Any) -> bool:
|
||||
"""Check whether a object is defined in a class scope.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames : List[FrameType]
|
||||
The frame stack of the object, obtained by `inspect.stack()`.
|
||||
|
||||
Returns
|
||||
-------
|
||||
res : bool
|
||||
The result if the object is defined in a class scope.
|
||||
"""
|
||||
|
||||
def _is_tvmscript_class_annotator(line: str) -> bool:
|
||||
"""Checks if the line contains a TVMScript annotator for a class
|
||||
|
||||
These match either `@I.ir_module` or `@R.rewriter`, or their
|
||||
imported names `@ir_module` or `@rewriter`.
|
||||
"""
|
||||
|
||||
return line.startswith("@") and ("ir_module" in line or "rewriter" in line)
|
||||
|
||||
if len(frames) > 2:
|
||||
frame_info = frames[2]
|
||||
code_context = frame_info.code_context
|
||||
if code_context is None:
|
||||
return False
|
||||
line = code_context[0].strip()
|
||||
if _is_tvmscript_class_annotator(line):
|
||||
return True
|
||||
if line.startswith("class"):
|
||||
lineno = frame_info.lineno
|
||||
if lineno >= 2:
|
||||
source, _ = findsource(obj)
|
||||
line = source[lineno - 2].strip()
|
||||
if _is_tvmscript_class_annotator(line):
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,35 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""The ir module parser"""
|
||||
|
||||
from tvm.ir import Range
|
||||
from ...ir_builder.ir import * # pylint: disable=redefined-builtin
|
||||
from . import parser as _parser
|
||||
from .entry import ir_module, pyfunc
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Range",
|
||||
"dummy_global_info",
|
||||
"ir_module",
|
||||
"lookup_vdevice",
|
||||
"module_attrs",
|
||||
"module_global_infos",
|
||||
"pyfunc",
|
||||
"vdevice",
|
||||
]
|
||||
@@ -0,0 +1,163 @@
|
||||
# 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.
|
||||
# pylint: disable=import-outside-toplevel
|
||||
"""The entry point of TVM parser for ir module."""
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
|
||||
from tvm import cpu, ir
|
||||
from tvm.ir import GlobalVar, IRModule
|
||||
|
||||
from .._core import parse, utils
|
||||
|
||||
|
||||
# this formulation allows us to support having @I.ir_module
|
||||
# appear as a decorator by itself or to have optional arguments
|
||||
# like @I.ir_module(check_well_formed=False)
|
||||
def ir_module(
|
||||
mod: type | None = None, check_well_formed: bool = True, s_tir: bool = False
|
||||
) -> IRModule:
|
||||
"""The parsing method for ir module, by using `@ir_module` as decorator.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
mod : Type
|
||||
The class to be parsed as ir module.
|
||||
|
||||
check_well_formed : bool
|
||||
Whether to check well-formedness during parsing.
|
||||
|
||||
Returns
|
||||
-------
|
||||
ir_module : IRModule
|
||||
The parsed ir module.
|
||||
"""
|
||||
|
||||
# Capture stack outside wrapper (wrapper adds to the stack)
|
||||
outer_stack = inspect.stack()
|
||||
|
||||
def decorator_wrapper(mod):
|
||||
if not inspect.isclass(mod):
|
||||
raise TypeError(f"Expect a class, but got: {mod}")
|
||||
|
||||
# Check BasePyModule inheritance
|
||||
base_py_module_inherited = any(base.__name__ == "BasePyModule" for base in mod.__bases__)
|
||||
|
||||
extra_vars = utils.inspect_class_capture(mod)
|
||||
# Resolve closure variables hidden by PEP 563 (annotation-only names)
|
||||
utils.resolve_closure_vars(mod, extra_vars, outer_stack)
|
||||
m = parse(mod, extra_vars, check_well_formed=check_well_formed, s_tir=s_tir)
|
||||
|
||||
if base_py_module_inherited:
|
||||
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
|
||||
# because tvm.script is loaded before tvm.relax during tvm initialization.
|
||||
from tvm.relax.base_py_module import BasePyModule
|
||||
from tvm.relax.expr import ExternFunc # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Collect pyfunc methods
|
||||
pyfunc_methods = [
|
||||
name
|
||||
for name, attr in mod.__dict__.items()
|
||||
if hasattr(attr, "dispatch_token") and attr.dispatch_token == "pyfunc"
|
||||
]
|
||||
|
||||
mod._pyfunc_methods = pyfunc_methods
|
||||
|
||||
# Create ExternFunc nodes
|
||||
|
||||
for method_name in pyfunc_methods:
|
||||
try:
|
||||
existing_gvars = [
|
||||
global_var
|
||||
for global_var in m.get_global_vars()
|
||||
if global_var.name_hint == method_name
|
||||
]
|
||||
|
||||
extern_func = ExternFunc(method_name)
|
||||
extern_func = extern_func.with_attr("is_pyfunc", True)
|
||||
extern_func = extern_func.with_attr("function_type", "python")
|
||||
extern_func = extern_func.with_attr("python_function_name", method_name)
|
||||
extern_func = extern_func.with_attr(
|
||||
"python_source", f"# Source for {method_name}"
|
||||
)
|
||||
extern_func = extern_func.with_attr("python_packed_func", None)
|
||||
|
||||
if existing_gvars:
|
||||
m[existing_gvars[0]] = extern_func
|
||||
else:
|
||||
m[GlobalVar(method_name)] = extern_func
|
||||
|
||||
except Exception: # pylint: disable=broad-exception-caught
|
||||
continue
|
||||
|
||||
class ModuleFactory:
|
||||
"""Factory class for creating BasePyModule instances with Python functions."""
|
||||
|
||||
def __init__(self, module, pyfunc_methods, original_class):
|
||||
self.ir_module = module
|
||||
self.pyfunc_methods = pyfunc_methods
|
||||
self.original_class = original_class
|
||||
|
||||
def __call__(self, device=None, target=None):
|
||||
if device is None:
|
||||
device = cpu(0)
|
||||
|
||||
instance_ir_mod = ir.IRModule()
|
||||
for global_var, func in self.ir_module.functions_items():
|
||||
instance_ir_mod[global_var] = func
|
||||
|
||||
instance = BasePyModule(instance_ir_mod, device, target)
|
||||
|
||||
for method_name in self.pyfunc_methods:
|
||||
if hasattr(self.original_class, method_name):
|
||||
method = getattr(self.original_class, method_name)
|
||||
instance.add_python_function(method_name, method)
|
||||
|
||||
return instance
|
||||
|
||||
def __getattr__(self, name):
|
||||
if hasattr(self.ir_module, name):
|
||||
return getattr(self.ir_module, name)
|
||||
raise AttributeError(
|
||||
f"'{self.__class__.__name__}' object has no attribute '{name}'"
|
||||
)
|
||||
|
||||
factory = ModuleFactory(m, pyfunc_methods, mod)
|
||||
setattr(factory, "__name__", mod.__name__)
|
||||
return factory
|
||||
|
||||
setattr(m, "__name__", mod.__name__)
|
||||
return m
|
||||
|
||||
if mod is not None:
|
||||
# if there are no optional args given, this will directly invoke the wrapper
|
||||
return decorator_wrapper(mod)
|
||||
else:
|
||||
# if there is a optional arg given, it returns the wrapper function
|
||||
# as a new decorator and applies it
|
||||
setattr(decorator_wrapper, "dispatch_token", "ir")
|
||||
return decorator_wrapper
|
||||
|
||||
|
||||
def pyfunc(func: Callable):
|
||||
# Set the dispatch_token on the decorated function
|
||||
setattr(func, "dispatch_token", "pyfunc")
|
||||
return func
|
||||
|
||||
|
||||
setattr(pyfunc, "dispatch_token", "pyfunc")
|
||||
@@ -0,0 +1,212 @@
|
||||
# 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.
|
||||
# pylint: disable=unused-argument
|
||||
"""The base parser for ir module"""
|
||||
|
||||
from tvm.ir import GlobalVar
|
||||
|
||||
from ...ir_builder import ir as I
|
||||
from .._core import Parser, dispatch, doc
|
||||
|
||||
|
||||
class ModuleWithGlobalVars:
|
||||
"""A Module that can add global vars during parsing, to support `Module.function` syntax."""
|
||||
|
||||
def __getattr__(self, attr):
|
||||
# Customize the error message.
|
||||
# NOTE: `__getattr__` is only called when the attribute access fails with an AttributeError
|
||||
raise AttributeError(f"Cannot find the function `{attr}` in the current IRModule")
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="ClassDef")
|
||||
def _visit_class_def(self: Parser, node: doc.ClassDef) -> None:
|
||||
"""The class definition visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST class definition node.
|
||||
"""
|
||||
|
||||
with self.var_table.with_frame():
|
||||
with I.ir_module():
|
||||
# Step 0. Add the class name to the var table
|
||||
fake_module = ModuleWithGlobalVars()
|
||||
self.var_table.add(node.name, fake_module)
|
||||
|
||||
# Step 1: Check if this class inherits from BasePyModule
|
||||
is_base_py_module = _check_base_py_module_inheritance(node)
|
||||
if is_base_py_module:
|
||||
# Store this information in the IRModule for later use
|
||||
I.module_attrs({"base_py_module": True})
|
||||
# Set the parser context to allow Python functions
|
||||
self.set_class_context(node.name, True)
|
||||
else:
|
||||
# Set the parser context to disallow Python functions
|
||||
self.set_class_context(node.name, False)
|
||||
|
||||
# Step 2. Visit non-function stmts, including but not limited to
|
||||
# 1. `I.module_attrs`
|
||||
# 2. `I.module_global_infos`
|
||||
with self.with_dispatch_token("ir"):
|
||||
for stmt in node.body:
|
||||
if not isinstance(stmt, doc.FunctionDef):
|
||||
self.visit(stmt)
|
||||
|
||||
# Step 3. Visit function stmts to declare the global vars
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.FunctionDef):
|
||||
global_var = self.visit_tvm_declare_function(stmt)
|
||||
fake_module.__setattr__(stmt.name, global_var)
|
||||
|
||||
# Step 4. Visit and parse the functions
|
||||
with self.with_dispatch_token("ir"):
|
||||
for stmt in node.body:
|
||||
if isinstance(stmt, doc.FunctionDef):
|
||||
self.visit(stmt)
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="Assign")
|
||||
def _visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
"""The assign visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST assign node.
|
||||
"""
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0].id
|
||||
rhs = self.eval_expr(node.value)
|
||||
|
||||
I.decl_function(lhs, rhs)
|
||||
I.def_function(lhs, rhs)
|
||||
|
||||
|
||||
@dispatch.register(token="ir", type_name="Expr")
|
||||
def _visit_expr(self: Parser, node: doc.Expr) -> None:
|
||||
"""The expression visiting method for ir module.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
self : Parser
|
||||
The visiting parser.
|
||||
|
||||
node : doc.ClassDef
|
||||
The doc AST expression node.
|
||||
"""
|
||||
self.eval_expr(node.value)
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="Assign")
|
||||
def visit_assign(self: Parser, node: doc.Assign) -> None:
|
||||
if len(node.targets) != 1:
|
||||
self.report_error(node, "Consequential assignments like 'a = b = c' are not supported.")
|
||||
lhs = node.targets[0]
|
||||
rhs = self.eval_expr(node.value)
|
||||
self.eval_assign(
|
||||
target=lhs, source=rhs, bind_value=lambda _a, _b, _c, value: value, allow_shadowing=True
|
||||
)
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="pre_visit_local_function")
|
||||
def pre_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@dispatch.register(token="default", type_name="post_visit_local_function")
|
||||
def post_visit_local_function(self: Parser, node: doc.Expr) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@dispatch.register(token="pyfunc", type_name="tvm_declare_function")
|
||||
def visit_tvm_declare_function(self: Parser, node: doc.FunctionDef) -> GlobalVar:
|
||||
"""Declare a Python function as an ExternFunc in the IRModule."""
|
||||
# Check if Python functions are allowed in this context
|
||||
# We need to check if we're in a class that inherits from BasePyModule
|
||||
current_class = self._get_current_class_context()
|
||||
if current_class and not self._is_base_py_module_context():
|
||||
self.report_error(
|
||||
node,
|
||||
"@I.pyfunc are only allowed in classes that inherit from BasePyModule. "
|
||||
f"Class '{current_class}' does not inherit from BasePyModule.",
|
||||
)
|
||||
|
||||
# Lazy import: tvm.relax cannot be imported at module level in tvm.script.parser
|
||||
# because tvm.script is loaded before tvm.relax during tvm initialization.
|
||||
from tvm.relax import ExternFunc # pylint: disable=import-outside-toplevel
|
||||
|
||||
# Create ExternFunc with proper attributes for Python functions
|
||||
func = ExternFunc(node.name)
|
||||
func = func.with_attr("is_pyfunc", True)
|
||||
func = func.with_attr("function_type", "python")
|
||||
func = func.with_attr("python_function_name", node.name)
|
||||
|
||||
# Add placeholder attributes that will be filled in later
|
||||
func = func.with_attr("python_source", f"# Source will be filled for {node.name}")
|
||||
func = func.with_attr("python_packed_func", None) # Will be filled in entry.py
|
||||
|
||||
# Store the function name for later retrieval
|
||||
return I.decl_function(node.name, func)
|
||||
|
||||
|
||||
@dispatch.register(token="pyfunc", type_name="FunctionDef")
|
||||
def visit_function_def(self: Parser, node: doc.FunctionDef) -> None:
|
||||
"""Visit Python function definition - no need to parse the body."""
|
||||
# Python function body is not parsed in TVMScript
|
||||
|
||||
|
||||
def _check_base_py_module_inheritance(node: doc.ClassDef) -> bool:
|
||||
"""Check if a class inherits from BasePyModule.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
node : doc.ClassDef
|
||||
The class definition node to check.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if the class inherits from BasePyModule, False otherwise.
|
||||
"""
|
||||
if not node.bases:
|
||||
return False
|
||||
|
||||
# Check each base class
|
||||
for base in node.bases:
|
||||
if hasattr(base, "id"):
|
||||
if base.id == "BasePyModule":
|
||||
return True
|
||||
elif hasattr(base, "attr"):
|
||||
if base.attr == "BasePyModule":
|
||||
return True
|
||||
elif hasattr(base, "value") and hasattr(base.value, "id"):
|
||||
if (
|
||||
base.value.id in ["BasePyModule", "tvm", "relax"]
|
||||
and hasattr(base, "attr")
|
||||
and base.attr == "BasePyModule"
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
@@ -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.
|
||||
"""
|
||||
@@ -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
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user