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}")
|
||||
Reference in New Issue
Block a user