chore: import upstream snapshot with attribution

python/tvm/ir/instrument.py

@@ -0,0 +1,341 @@
+# 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=invalid-name,unused-argument
+"""Common pass instrumentation across IR variants."""
+
+import functools
+import inspect
+import re
+import shutil
+from pathlib import Path
+
+import tvm_ffi
+
+import tvm.runtime
+
+from . import _ffi_instrument_api
+
+
+@tvm_ffi.register_object("instrument.PassInstrument")
+class PassInstrument(tvm.runtime.Object):
+    """A pass instrument implementation.
+
+    To use, a user class can either subclass from PassInstrument
+    directly, or can apply the :py:func:`pass_instrument` wrapper.  In
+    either case, the `enter_pass_ctx`, `exit_pass_ctx`, `should_run`,
+    `run_before_pass`, and `run_after_pass` methods can be defined to
+    adjust the instrument's behavior.  See the no-op implementations
+    in this class definition for more information on each.
+
+    """
+
+    def __init__(self):
+        # initialize handle in case pi_cls creation failed.
+        cls = type(self)
+
+        # If the child class declared the method, then use it.
+        # Otherwise, pass None to avoid a C++ -> Python round trip for
+        # a no-op.
+        def get_child_method(name):
+            if getattr(cls, name) is getattr(PassInstrument, name):
+                return None
+
+            return getattr(self, name)
+
+        # Create runtime pass instrument object.
+        # register instance's enter_pass_ctx,exit_pass_ctx, should_run, run_before_pass and
+        # run_after_pass methods to it if present.
+        self.__init_handle_by_constructor__(
+            _ffi_instrument_api.PassInstrument,
+            cls.__name__,
+            get_child_method("enter_pass_ctx"),
+            get_child_method("exit_pass_ctx"),
+            get_child_method("should_run"),
+            get_child_method("run_before_pass"),
+            get_child_method("run_after_pass"),
+        )
+
+    def enter_pass_ctx(self):
+        """Called when entering the instrumented context.
+
+        Returns
+        -------
+        None
+        """
+
+    def exit_pass_ctx(self):
+        """Called when exiting the instrumented context.
+
+        Returns
+        -------
+        None
+        """
+
+    def should_run(self, mod, info):
+        """Determine whether to run the pass or not.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        should_run : bool
+
+            True to run the pass, or False to skip the pass.
+        """
+
+    def run_before_pass(self, mod, info):
+        """Instrument before the pass runs.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        None
+        """
+
+    def run_after_pass(self, mod, info):
+        """Instrument after the pass runs.
+
+        Called once for each pass that is run while the instrumented
+        context is active.
+
+        Parameters
+        ----------
+        mod : tvm.ir.module.IRModule
+
+            The module on which an optimization pass is being run.
+
+        info : tvm.transform.PassInfo
+
+            The pass information.
+
+        Returns
+        -------
+        None
+        """
+
+
+def _wrap_class_pass_instrument(pi_cls):
+    """Wrap a python class as pass instrument"""
+
+    # No additional wrapping needed if the user class already
+    # inherits.
+    if issubclass(pi_cls, PassInstrument):
+        return pi_cls
+
+    class PyPassInstrument(pi_cls, PassInstrument):
+        """Internal wrapper class to create a class instance."""
+
+        def __init__(self, *args, **kwargs):
+            # initialize handle in case pi_cls creation failed.
+            pi_cls.__init__(self, *args, **kwargs)
+            PassInstrument.__init__(self)
+
+    functools.update_wrapper(PyPassInstrument.__init__, pi_cls.__init__)
+    PyPassInstrument.__name__ = pi_cls.__name__
+    PyPassInstrument.__doc__ = pi_cls.__doc__
+    PyPassInstrument.__module__ = pi_cls.__module__
+    return PyPassInstrument
+
+
+def pass_instrument(pi_cls=None):
+    """Decorate a pass instrument.
+
+    Parameters
+    ----------
+    pi_class : class
+        Instrument class. See example below.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        @tvm.instrument.pass_instrument
+        class SkipPass:
+            def __init__(self, skip_pass_name):
+                self.skip_pass_name = skip_pass_name
+
+            # Uncomment to customize
+            # def enter_pass_ctx(self):
+            #    pass
+
+            # Uncomment to customize
+            # def exit_pass_ctx(self):
+            #    pass
+
+            # If pass name contains keyword, skip it by return False. (return True: not skip)
+            def should_run(self, mod, pass_info)
+                if self.skip_pass_name in pass_info.name:
+                    return False
+                return True
+
+            # Uncomment to customize
+            # def run_before_pass(self, mod, pass_info):
+            #    pass
+
+            # Uncomment to customize
+            # def run_after_pass(self, mod, pass_info):
+            #    pass
+
+        skip_annotate = SkipPass("AnnotateSpans")
+        with tvm.transform.PassContext(instruments=[skip_annotate]):
+            tvm.compile(mod, "llvm")
+    """
+
+    def create_pass_instrument(pi_cls):
+        if not inspect.isclass(pi_cls):
+            raise TypeError("pi_cls must be a class")
+
+        return _wrap_class_pass_instrument(pi_cls)
+
+    if pi_cls:
+        return create_pass_instrument(pi_cls)
+    return create_pass_instrument
+
+
+@tvm_ffi.register_object("instrument.PassInstrument")
+class PassTimingInstrument(tvm.runtime.Object):
+    """A wrapper to create a passes time instrument that implemented in C++"""
+
+    def __init__(self):
+        self.__init_handle_by_constructor__(_ffi_instrument_api.MakePassTimingInstrument)
+
+    @staticmethod
+    def render():
+        """Retrieve rendered time profile result
+        Returns
+        -------
+        string : string
+            The rendered string result of time profiles
+
+        Examples
+        --------
+
+        .. code-block:: python
+
+            timing_inst = PassTimingInstrument()
+            with tvm.transform.PassContext(instruments=[timing_inst]):
+                relax_mod = relax.transform.FuseOps()(relax_mod)
+                # before exiting the context, get profile results.
+                profiles = timing_inst.render()
+        """
+        return _ffi_instrument_api.RenderTimePassProfiles()
+
+
+@pass_instrument
+class PassPrintingInstrument:
+    """A pass instrument to print if before or
+    print ir after each element of a named pass."""
+
+    def __init__(self, print_before_pass_names, print_after_pass_names):
+        self.print_before_pass_names = print_before_pass_names
+        self.print_after_pass_names = print_after_pass_names
+
+    def run_before_pass(self, mod, pass_info):
+        if pass_info.name in self.print_before_pass_names:
+            print(f"Print IR before: {pass_info.name}\n{mod}\n\n")
+
+    def run_after_pass(self, mod, pass_info):
+        if pass_info.name in self.print_after_pass_names:
+            print(f"Print IR after: {pass_info.name}\n{mod}\n\n")
+
+
+@pass_instrument
+class PrintAfterAll:
+    """Print the name of the pass, the IR, only after passes execute."""
+
+    def run_after_pass(self, mod, info):
+        print(f"After Running Pass: {info}")
+        print(mod)
+
+
+@pass_instrument
+class PrintBeforeAll:
+    """Print the name of the pass, the IR, only before passes execute."""
+
+    def run_before_pass(self, mod, info):
+        print(f"Before Running Pass: {info}")
+        print(mod)
+
+
+@pass_instrument
+class DumpIR:
+    """Dump the IR after the pass runs."""
+
+    def __init__(self, dump_dir: Path | str, refresh: bool = False):
+        if isinstance(dump_dir, Path):
+            self.dump_dir = dump_dir
+        else:
+            self.dump_dir = Path(dump_dir)
+        self.counter = 0
+        if refresh and self.dump_dir.is_dir():
+            self._safe_remove_dump_dir()
+
+    def _safe_remove_dump_dir(self):
+        """Remove dump directory only if it contains only dumped IR files."""
+        # Pattern for dumped files: {counter:03d}_{pass_name}.py
+        dump_pattern = re.compile(r"^\d{3}_.*\.py$")
+
+        # Check all files in the directory
+        for item in self.dump_dir.iterdir():
+            # If there's a subdirectory or a file that doesn't match the pattern, abort
+            if item.is_dir() or not dump_pattern.match(item.name):
+                print(
+                    f"WARNING: Skipping removal of {self.dump_dir} as it contains "
+                    f"non-dumped files or directories. Please clean it manually."
+                )
+                return
+
+        # Safe to remove - only contains dumped files
+        try:
+            shutil.rmtree(self.dump_dir)
+        except OSError as e:
+            print(f"WARNING: Failed to remove directory {self.dump_dir}: {e}")
+
+    def run_after_pass(self, mod, info):
+        self.dump_dir.mkdir(parents=True, exist_ok=True)
+        try:
+            sanitized_pass_name = re.sub(r'[<>:"/\\|?*]', "_", info.name)
+            with open(self.dump_dir / f"{self.counter:03d}_{sanitized_pass_name}.py", "w") as f:
+                f.write(mod.script())
+        except Exception:  # pylint: disable=broad-exception-caught
+            print(f"WARNING: Failed to dump IR for pass {info.name}")
+        finally:
+            self.counter += 1