chore: import upstream snapshot with attribution

python/tvm/backend/loader.py

@@ -0,0 +1,187 @@
+# 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.
+"""Backend loading and public alias support."""
+
+from __future__ import annotations
+
+import importlib
+import importlib.util
+import sys
+import types
+from typing import Any
+
+_LOADED_BACKENDS: dict[str, Any] = {}
+
+
+class _AliasModule(types.ModuleType):
+    """Module object that exposes a backend module under a public alias."""
+
+    def __init__(self, fullname: str, module):
+        super().__init__(fullname, getattr(module, "__doc__", None))
+        self.__dict__["__tvm_backend_module__"] = module
+        self.__dict__["__package__"] = fullname.rpartition(".")[0]
+        if hasattr(module, "__all__"):
+            self.__dict__["__all__"] = module.__all__
+        if hasattr(module, "__path__"):
+            self.__dict__["__path__"] = []
+
+    def __getattr__(self, name: str):
+        return getattr(self.__dict__["__tvm_backend_module__"], name)
+
+    def __setattr__(self, name: str, value):
+        setattr(self.__dict__["__tvm_backend_module__"], name, value)
+
+    def __delattr__(self, name: str):
+        delattr(self.__dict__["__tvm_backend_module__"], name)
+
+    def __dir__(self):
+        return sorted(set(super().__dir__()) | set(dir(self.__dict__["__tvm_backend_module__"])))
+
+
+class _AliasLoader:
+    """Loader that returns an already-resolved module for an alias spec."""
+
+    def __init__(self, fullname: str, module):
+        self._fullname = fullname
+        self._module = module
+
+    def create_module(self, spec):
+        return _get_alias_module(self._fullname, self._module)
+
+    def exec_module(self, module):
+        _set_module_alias(self._fullname, self._module)
+        return None
+
+    def is_package(self, fullname):
+        return hasattr(self._module, "__path__")
+
+
+def _redirect_tirx_backend_alias(fullname: str) -> str | None:
+    prefix = "tvm.tirx."
+    if not fullname.startswith(prefix):
+        return None
+    rest = fullname[len(prefix) :]
+    backend_name, sep, tail = rest.partition(".")
+    if not sep or backend_name not in _LOADED_BACKENDS:
+        return None
+    return f"tvm.backend.{backend_name}.{tail}"
+
+
+class _BackendAliasFinder:
+    """Redirect ``tvm.tirx.<backend>.*`` imports to ``tvm.backend.<backend>.*``."""
+
+    @classmethod
+    def find_spec(cls, fullname, path, target=None):
+        redirected = _redirect_tirx_backend_alias(fullname)
+        if redirected is None:
+            return None
+        module = importlib.import_module(redirected)
+        _set_module_alias(fullname, module)
+        loader = _AliasLoader(fullname, module)
+        spec = importlib.util.spec_from_loader(
+            fullname, loader, is_package=hasattr(module, "__path__")
+        )
+        if spec is not None and hasattr(module, "__path__"):
+            spec.submodule_search_locations = []
+        return spec
+
+
+if not any(isinstance(finder, _BackendAliasFinder) for finder in sys.meta_path):
+    sys.meta_path.insert(0, _BackendAliasFinder())
+
+
+def _get_alias_module(alias: str, module):
+    existing = sys.modules.get(alias)
+    if (
+        isinstance(existing, _AliasModule)
+        and existing.__dict__.get("__tvm_backend_module__") is module
+    ):
+        return existing
+    return _AliasModule(alias, module)
+
+
+def _set_module_alias(alias: str, module, *, direct: bool = False) -> None:
+    alias_module = module if direct else _get_alias_module(alias, module)
+    sys.modules[alias] = alias_module
+    parent_name, _, child_name = alias.rpartition(".")
+    parent = sys.modules.get(parent_name)
+    if parent is not None:
+        setattr(parent, child_name, alias_module)
+
+
+def _alias_loaded_backend_modules(name: str) -> None:
+    backend_prefix = f"tvm.backend.{name}"
+    public_prefix = f"tvm.tirx.{name}"
+    for module_name, module in sorted(list(sys.modules.items())):
+        if module_name == backend_prefix or module_name.startswith(f"{backend_prefix}."):
+            public_name = f"{public_prefix}{module_name[len(backend_prefix) :]}"
+            _set_module_alias(public_name, module, direct=module_name == backend_prefix)
+
+
+def _import_backend(name: str):
+    module_name = f"tvm.backend.{name}"
+    try:
+        return importlib.import_module(module_name)
+    except ModuleNotFoundError as err:
+        if err.name == module_name:
+            raise ImportError(
+                f"Cannot load TVM backend {name!r}: expected Python package {module_name!r}. "
+                "Install the backend package or check the backend name."
+            ) from err
+        raise
+
+
+def load(name: str) -> None:
+    """Load a backend's Python registration hooks.
+
+    Loading is idempotent.  A backend package must live at ``tvm.backend.<name>``
+    and expose ``register_backend()``.
+    """
+
+    if name in _LOADED_BACKENDS:
+        return None
+
+    module = _import_backend(name)
+    register_backend = getattr(module, "register_backend", None)
+    if register_backend is None:
+        raise AttributeError(f"Backend package 'tvm.backend.{name}' has no register_backend()")
+
+    import tvm.tirx as tirx  # pylint: disable=import-outside-toplevel
+
+    setattr(tirx, name, module)
+    sys.modules[f"tvm.tirx.{name}"] = module
+    _LOADED_BACKENDS[name] = module
+    try:
+        register_backend()
+        _alias_loaded_backend_modules(name)
+    except Exception:
+        _LOADED_BACKENDS.pop(name, None)
+        if getattr(tirx, name, None) is module:
+            delattr(tirx, name)
+        if sys.modules.get(f"tvm.tirx.{name}") is module:
+            sys.modules.pop(f"tvm.tirx.{name}", None)
+        raise
+    return None
+
+
+def is_loaded(name: str) -> bool:
+    """Return whether a backend has been loaded."""
+
+    return name in _LOADED_BACKENDS
+
+
+__all__ = ["is_loaded", "load"]