chore: import upstream snapshot with attribution

python/tvm/ir/utils.py

@@ -0,0 +1,161 @@
+# 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.
+"""Utilities shared across TVM IR packages."""
+
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def derived_object(cls: type[T]) -> type[T]:
+    """A decorator to register derived subclasses for TVM objects.
+
+    Parameters
+    ----------
+    cls : type
+        The derived class to be registered.
+
+    Returns
+    -------
+    cls : type
+        The decorated TVM object.
+
+    Example
+    -------
+    .. code-block:: python
+
+        @register_object("s_tir.meta_schedule.PyRunner")
+        class _PyRunner(meta_schedule.Runner):
+            def __init__(self, f_run: Callable = None):
+                self.__init_handle_by_constructor__(_ffi_api.RunnerPyRunner, f_run)
+
+        class PyRunner:
+            _tvm_metadata = {
+                "cls": _PyRunner,
+                "methods": ["run"]
+            }
+            def run(self, runner_inputs):
+                raise NotImplementedError
+
+        @derived_object
+        class LocalRunner(PyRunner):
+            def run(self, runner_inputs):
+                ...
+    """
+
+    import functools  # pylint: disable=import-outside-toplevel
+    import weakref  # pylint: disable=import-outside-toplevel
+
+    def _extract(inst: type, name: str):
+        """Extract function from intrinsic class."""
+
+        def method(*args, **kwargs):
+            return getattr(inst, name)(*args, **kwargs)
+
+        for inherit_cls, base_cls in zip(cls.__mro__, cls.__mro__[1:]):
+            # extract functions that differ from the base class
+            if not hasattr(base_cls, name):
+                continue
+            if getattr(base_cls, name) is getattr(inherit_cls, name):
+                continue
+            return method
+
+        # for task scheduler return None means calling default function
+        # otherwise it will trigger a RuntimeError of method not implemented
+        # on the c++ side when you call the method
+        return None
+
+    assert isinstance(cls.__base__, type)
+    if hasattr(cls, "_type") and cls._type == "TVMDerivedObject":  # type: ignore
+        raise TypeError(
+            f"Inheritance from a decorated object `{cls.__name__}` is not allowed. "
+            f"Please inherit from `{cls.__name__}._cls`."
+        )
+    assert hasattr(cls, "_tvm_metadata"), (
+        "Please use the user-facing method overriding class, i.e., PyRunner."
+    )
+
+    base = cls.__base__
+    metadata = getattr(base, "_tvm_metadata")
+    fields = metadata.get("fields", [])
+    methods = metadata.get("methods", [])
+
+    base_cls = metadata["cls"]
+    slots = []
+    if getattr(base_cls, "__dictoffset__", 0) == 0:
+        slots.append("__dict__")
+    if getattr(base_cls, "__weakrefoffset__", 0) == 0:
+        slots.append("__weakref__")
+
+    class TVMDerivedObject(base_cls):  # type: ignore
+        """The derived object to avoid cyclic dependency."""
+
+        __slots__ = tuple(slots)
+
+        _cls = cls
+        _type = "TVMDerivedObject"
+
+        def __init__(self, *args, **kwargs):
+            """Constructor."""
+            self._inst = cls(*args, **kwargs)
+
+            super().__init__(
+                # the constructor's parameters, builder, runner, etc.
+                *[getattr(self._inst, name) for name in fields],
+                # the function methods, init_with_tune_context, build, run, etc.
+                *[_extract(self._inst, name) for name in methods],
+            )
+
+            # for task scheduler hybrid funcs in c++ & python side
+            # using weakref to avoid cyclic dependency
+            self._inst._outer = weakref.ref(self)
+
+        def __getattr__(self, name):
+            import inspect  # pylint: disable=import-outside-toplevel
+
+            try:
+                # fall back to instance attribute if there is not any
+                # return self._inst.__getattribute__(name)
+                result = self._inst.__getattribute__(name)
+            except AttributeError:
+                result = super().__getattr__(name)
+
+            if inspect.ismethod(result):
+
+                def method(*args, **kwargs):
+                    return result(*args, **kwargs)
+
+                # set __own__ to aviod implicit deconstruction
+                setattr(method, "__own__", self)
+                return method
+
+            return result
+
+        def __setattr__(self, name, value):
+            if name not in ["_inst", "key", "handle"]:
+                self._inst.__setattr__(name, value)
+            else:
+                super().__setattr__(name, value)
+
+    functools.update_wrapper(TVMDerivedObject.__init__, cls.__init__)  # type: ignore
+    TVMDerivedObject.__name__ = cls.__name__
+    TVMDerivedObject.__doc__ = cls.__doc__
+    TVMDerivedObject.__module__ = cls.__module__
+    for key, value in cls.__dict__.items():
+        if isinstance(value, classmethod | staticmethod):
+            setattr(TVMDerivedObject, key, value)
+    return TVMDerivedObject