onnxslim 0.1.80__py3-none-any.whl

onnxslim/third_party/onnx_graphsurgeon/logger/logger.py

@@ -0,0 +1,261 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+#
+
+import enum
+import inspect
+import os
+import sys
+import time
+
+from onnxslim.third_party.onnx_graphsurgeon.util.exception import (
+    OnnxGraphSurgeonException,
+)
+
+
+# Context manager to apply indentation to messages
+class LoggerIndent:
+    def __init__(self, logger, indent):
+        """Initialize the LoggerIndent context manager with the specified logger and indentation level."""
+        self.logger = logger
+        self.old_indent = self.logger.logging_indent
+        self.indent = indent
+
+    def __enter__(self):
+        """Set logger indentation level on entering the context."""
+        self.logger.logging_indent = self.indent
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Reset logger indentation level on exiting the context."""
+        self.logger.logging_indent = self.old_indent
+
+
+# Context manager to suppress messages
+class LoggerSuppress:
+    def __init__(self, logger, severity):
+        """Initialize a LoggerSuppress object with a logger and severity level."""
+        self.logger = logger
+        self.old_severity = self.logger.severity
+        self.severity = severity
+
+    def __enter__(self):
+        """Set logger severity to a specified level when entering the context."""
+        self.logger.severity = self.severity
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Reset logger severity to its original level when exiting the context."""
+        self.logger.severity = self.old_severity
+
+
+class LogMode(enum.IntEnum):
+    EACH = 0  # Log the message each time
+    ONCE = 1  # Log the message only once. The same message will not be logged again.
+
+
+class Logger:
+    ULTRA_VERBOSE = -10
+    VERBOSE = 0
+    DEBUG = 10
+    INFO = 20
+    WARNING = 30
+    ERROR = 40
+    CRITICAL = 50
+
+    SEVERITY_LETTER_MAPPING = {
+        ULTRA_VERBOSE: "[UV]",
+        VERBOSE: "[V]",
+        DEBUG: "[D]",
+        INFO: "[I]",
+        WARNING: "[W]",
+        ERROR: "[E]",
+        CRITICAL: "[C]",
+    }
+
+    SEVERITY_COLOR_MAPPING = {
+        ULTRA_VERBOSE: "cyan",
+        VERBOSE: "dark_gray",
+        DEBUG: "light_gray",
+        INFO: "light_green",
+        WARNING: "light_yellow",
+        ERROR: "red_1",
+        CRITICAL: "red_1",
+    }
+
+    def __init__(self, severity=INFO, colors=True, letter=True, timestamp=False, line_info=False):
+        """
+        Logger.
+
+        Args:
+            severity (Logger.Severity): Messages below this severity are ignored.
+            colors (bool): Whether to use colored output.
+            letter (bool): Whether to prepend each logging message with a letter indicating it's severity. Defaults to True.
+            timestamp (bool): Whether to include a timestamp in the logging output. Defaults to False.
+            line_info (bool): Whether to include file and line number information in the logging output. Defaults to False.
+        """
+        self._severity = severity
+        self.logging_indent = 0
+        self.root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, os.pardir))
+        self.once_logged = set()
+        self.colors = colors
+        self.letter = letter
+        self.timestamp = timestamp
+        self.line_info = line_info
+        self.logger_callbacks = []
+
+    @property
+    def severity(self):
+        """Returns the logging severity level."""
+        return self._severity
+
+    @severity.setter
+    def severity(self, value):
+        """Returns or sets the logging severity level with callback updates."""
+        self._severity = value
+        for callback in self.logger_callbacks:
+            callback(self._severity)
+
+    def register_callback(self, callback):
+        """
+        Registers a callback with the logger, which will be invoked when the logging severity is modified. The callback
+        is guaranteed to be called at least once in the register_callback function.
+
+        Args:
+            callback (Callable(Logger.Severity)): A callback that accepts the current logger severity.
+        """
+        callback(self._severity)
+        self.logger_callbacks.append(callback)
+
+    def indent(self, level=1):
+        """Returns a context manager that indents all strings logged by the specified amount."""
+        return LoggerIndent(self, level + self.logging_indent)
+
+    def suppress(self, severity=CRITICAL):
+        """
+        Returns a context manager that temporarily changes the severity of the logger for its duration.
+
+        Args:
+            severity (Logger.Severity): The severity to set the logger to. Defaults to Logger.CRITICAL, which will suppress all messages.
+        """
+        return LoggerSuppress(self, severity)
+
+    # If once is True, the logger will only log this message a single time. Useful in loops.
+    # message may be a callable which returns a message. This way, only if the message needs to be logged is it ever generated.
+    def log(self, message, severity, mode=LogMode.EACH, stack_depth=2):
+        """Logs a message with a specified severity and mode, supporting both single and repeated logging based on
+        conditions.
+        """
+
+        def process_message(message, stack_depth):
+            """Generates a log message prefix with file name and line number based on the specified stack depth."""
+
+            def get_prefix():
+                def get_line_info():
+                    module = inspect.getmodule(sys._getframe(stack_depth + 3)) or inspect.getmodule(
+                        sys._getframe(stack_depth + 2)
+                    )
+                    filename = module.__file__
+                    filename = os.path.relpath(filename, self.root_dir)
+                    # If the file is not located in trt_smeagol, use its basename instead.
+                    if os.pardir in filename:
+                        filename = os.path.basename(filename)
+                    return f"[{filename}:{sys._getframe(stack_depth).f_lineno}] "
+
+                prefix = ""
+                if self.letter:
+                    prefix += f"{Logger.SEVERITY_LETTER_MAPPING[severity]} "
+                if self.timestamp:
+                    prefix += "({:}) ".format(time.strftime("%X"))
+                if self.line_info:
+                    prefix += get_line_info()
+                return prefix
+
+            def apply_indentation(message):
+                """Indent each line in the message by the specified logging_indent level."""
+                message_lines = str(message).splitlines()
+                return "\n".join(["\t" * self.logging_indent + line for line in message_lines])
+
+            def apply_color(message):
+                """Apply color formatting to the message if color support is enabled."""
+                if self.colors:
+                    try:
+                        import colored
+
+                        color = Logger.SEVERITY_COLOR_MAPPING[severity]
+                        return colored.stylize(message, [colored.fg(color)])
+                    except ImportError:
+                        self.colors = False
+                        self.warning(
+                            "colored module is not installed, will not use colors when logging. To enable colors, please install the colored module: python3 -m pip install colored"
+                        )
+                        self.colors = True
+                return message
+
+            prefix = get_prefix()
+            message = apply_indentation(message)
+            return apply_color(f"{prefix}{message}")
+
+        def should_log(message):
+            """Determines if a message should be logged based on the severity level and logging mode."""
+            should = severity >= self._severity
+            if mode == LogMode.ONCE:
+                message_hash = hash(message)
+                should &= message_hash not in self.once_logged
+                self.once_logged.add(message_hash)
+            return should
+
+        if not should_log(message):
+            return
+
+        if callable(message):
+            message = message()
+        message = str(message)
+        print(process_message(message, stack_depth=stack_depth))
+
+    def ultra_verbose(self, message, mode=LogMode.EACH):
+        """Logs an ultra-verbose message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.ULTRA_VERBOSE, mode=mode, stack_depth=3)
+
+    def verbose(self, message, mode=LogMode.EACH):
+        """Logs a verbose message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.VERBOSE, mode=mode, stack_depth=3)
+
+    def debug(self, message, mode=LogMode.EACH):
+        """Logs a debug message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.DEBUG, mode=mode, stack_depth=3)
+
+    def info(self, message, mode=LogMode.EACH):
+        """Logs an informational message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.INFO, mode=mode, stack_depth=3)
+
+    def warning(self, message, mode=LogMode.EACH):
+        """Logs a warning message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.WARNING, mode=mode, stack_depth=3)
+
+    def error(self, message, mode=LogMode.EACH):
+        """Logs an error message with a specified mode and stack depth of 3."""
+        self.log(message, Logger.ERROR, mode=mode, stack_depth=3)
+
+    # Like error, but immediately exits.
+    def critical(self, message):
+        """Logs a critical message with a stack depth of 3 and raises an OnnxGraphSurgeonException."""
+        self.log(message, Logger.CRITICAL, stack_depth=3)
+        raise OnnxGraphSurgeonException(message) from None  # Erase exception chain
+
+
+global G_LOGGER
+G_LOGGER = Logger()

onnxslim/third_party/onnx_graphsurgeon/util/exception.py

@@ -0,0 +1,20 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+#
+
+
+class OnnxGraphSurgeonException(Exception):
+    """An exception raised by ONNX-GraphSurgeon."""

onnxslim/third_party/onnx_graphsurgeon/util/misc.py

@@ -0,0 +1,252 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed 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.
+#
+from __future__ import annotations
+
+from collections import OrderedDict
+from collections.abc import Sequence
+
+import numpy as np
+from onnx import AttributeProto
+
+from onnxslim.third_party.onnx_graphsurgeon.logger import G_LOGGER
+
+
+# default_value exists to solve issues that might result from Python's normal default argument behavior.
+# Specifically, consider the following class:
+#
+# class MyClass(object):
+#     def __init__(self, value=[]):
+#         self.value = value
+#
+# This leads to unwanted behavior when the default value is used:
+#
+# >>> x = MyClass()
+# >>> x.value.append("SHOULD NOT BE IN Y")
+# >>> y = MyClass()
+# >>> y.value
+# ['SHOULD NOT BE IN Y']
+#
+# If we rewrite the class using default value:
+#
+# class MyClass(object):
+#     def __init__(self, value=None):
+#         self.value = default_value(value, [])
+#
+# Then we get the desired behavior:
+#
+# >>> x = MyClass()
+# >>> x.value.append("SHOULD NOT BE IN Y")
+# >>> y = MyClass()
+# >>> y.value
+# []
+def default_value(value, default):
+    """Return the value if not None, otherwise return the default value."""
+    return value if value is not None else default
+
+
+def combine_dicts(dict0, dict1):
+    """
+    Combine two dictionaries.
+
+    Values in the second will overwrite values in the first.
+    """
+    if dict1 is None:
+        return dict0
+    combined = OrderedDict()
+    combined.update(dict0)
+    combined.update(dict1)
+    return combined
+
+
+def unique_dicts(dict0, dict1):
+    """
+    Subtract two dictionaries.
+
+    Values in the second will be subtracted from the first.
+    """
+    return {k: v for k, v in dict0.items() if k not in dict1} if dict1 else dict0
+
+
+def is_dynamic_dimension(dim):
+    """Check if a dimension is dynamic (non-integer or negative)."""
+    return not isinstance(dim, int) or dim < 0
+
+
+def is_dynamic_shape(shape):
+    """Determine if any dimension in the given shape is dynamic (non-integer or negative)."""
+    return any(is_dynamic_dimension(dim) for dim in shape)
+
+
+def volume(obj):
+    """Calculate the volume by multiplying the elements of an iterable object."""
+    vol = 1
+    for elem in obj:
+        vol *= elem
+    return vol
+
+
+_ONNX_ATTR_TYPE_TO_GS_TYPE = {}
+_GS_TYPE_TO_ONNX_ATTR_TYPE = {}
+
+
+# This method prevents circular import of Tensor and Graph
+def _init_dicts():
+    """Initialize mapping dictionaries to prevent circular imports of Tensor and Graph."""
+    global _ONNX_ATTR_TYPE_TO_GS_TYPE
+    global _GS_TYPE_TO_ONNX_ATTR_TYPE
+    if _ONNX_ATTR_TYPE_TO_GS_TYPE and _GS_TYPE_TO_ONNX_ATTR_TYPE:
+        return
+
+    from onnxslim.third_party.onnx_graphsurgeon.ir.graph import Graph
+    from onnxslim.third_party.onnx_graphsurgeon.ir.tensor import Tensor
+
+    _ONNX_ATTR_TYPE_TO_GS_TYPE = {
+        AttributeProto.UNDEFINED: None,
+        AttributeProto.FLOAT: float,
+        AttributeProto.INT: int,
+        AttributeProto.STRING: str,
+        AttributeProto.TENSOR: Tensor,
+        AttributeProto.GRAPH: Graph,
+        AttributeProto.SPARSE_TENSOR: AttributeProto.SPARSE_TENSOR,
+        AttributeProto.TYPE_PROTO: AttributeProto.TYPE_PROTO,
+        AttributeProto.FLOATS: list[float],
+        AttributeProto.INTS: list[int],
+        AttributeProto.STRINGS: list[str],
+        AttributeProto.TENSORS: list[Tensor],
+        AttributeProto.GRAPHS: list[Graph],
+        AttributeProto.SPARSE_TENSORS: AttributeProto.SPARSE_TENSORS,
+        AttributeProto.TYPE_PROTOS: AttributeProto.TYPE_PROTOS,
+    }
+    _GS_TYPE_TO_ONNX_ATTR_TYPE = {v: k for k, v in _ONNX_ATTR_TYPE_TO_GS_TYPE.items()}
+
+
+def convert_from_onnx_attr_type(onnx_attr_type):
+    """Converts an ONNX attribute type to its corresponding GS attribute type."""
+    _init_dicts()
+    return _ONNX_ATTR_TYPE_TO_GS_TYPE[onnx_attr_type]
+
+
+def convert_to_onnx_attr_type(any_type):
+    """Converts a given type to its corresponding ONNX attribute type."""
+    _init_dicts()
+    if any_type in _GS_TYPE_TO_ONNX_ATTR_TYPE:
+        return _GS_TYPE_TO_ONNX_ATTR_TYPE[any_type]
+    if np.issubdtype(any_type, np.floating):
+        return AttributeProto.FLOAT
+    if np.issubdtype(any_type, np.integer):
+        return AttributeProto.INT
+    G_LOGGER.warning(f"Unable to convert {any_type} into an ONNX AttributeType")
+
+
+# Special type of list that synchronizes contents with another list.
+# Concrete example: Assume some node, n, contains an input tensor, t. If we remove t from n.inputs,
+# we also need to remove n from t.outputs. To avoid having to do this manually, we use SynchronizedList,
+# which takes an attribute name as a parameter, and then synchronizes to that attribute of each of its elements.
+# So, in the example above, we can make n.inputs a synchronized list whose field_name is set to "outputs".
+# See test_ir.TestNodeIO for functional tests
+class SynchronizedList(list):
+    def __init__(self, parent_obj, field_name, initial):
+        """Initialize a SynchronizedList with a parent object, a field name, and an initial set of elements."""
+        self.parent_obj = parent_obj
+        self.field_name = field_name
+        self.extend(initial)
+
+    def _add_to_elem(self, elem):
+        """Append the parent_obj to the list attribute defined by field_name in the provided elem object."""
+        list.append(getattr(elem, self.field_name), self.parent_obj)
+
+    def _remove_from_elem(self, elem):
+        """Remove the parent_obj from the list attribute defined by field_name in the provided elem object."""
+        list.remove(getattr(elem, self.field_name), self.parent_obj)
+
+    def __delitem__(self, index):
+        """Remove the element at the specified index and update the corresponding list attribute in the parent
+        object.
+        """
+        self._remove_from_elem(self[index])
+        super().__delitem__(index)
+
+    def __setitem__(self, index, elem):
+        """Update the element at the specified index and modify the corresponding list attribute in the parent
+        object.
+        """
+        self._remove_from_elem(self[index])
+        super().__setitem__(index, elem)
+        self._add_to_elem(elem)
+
+    def append(self, x):
+        """Append an element to the list and update the parent object's corresponding list attribute."""
+        super().append(x)
+        self._add_to_elem(x)
+
+    def extend(self, iterable: Sequence[object]):
+        """Extend the list with elements from an iterable and update the parent object's corresponding list
+        attribute.
+        """
+        super().extend(iterable)
+        for elem in iterable:
+            self._add_to_elem(elem)
+
+    def insert(self, i, x):
+        """Insert an element at a given position and update the parent object's corresponding list attribute."""
+        super().insert(i, x)
+        self._add_to_elem(x)
+
+    def remove(self, x):
+        """Remove an element from the list and update the parent object's corresponding list attribute."""
+        super().remove(x)
+        self._remove_from_elem(x)
+
+    def pop(self, i=-1):
+        """Remove and return the element at index i (default last) from the list and update the parent object's
+        corresponding list attribute.
+        """
+        elem = super().pop(i)
+        self._remove_from_elem(elem)
+        return elem
+
+    def clear(self):
+        """Clear all elements from the list and update the parent object's corresponding list attribute."""
+        for elem in self:
+            self._remove_from_elem(elem)
+        super().clear()
+
+    def __add__(self, other_list: list[object]):
+        """Concatenate the current list with another list and return the resulting list."""
+        return list(self) + list(other_list)
+
+    def __iadd__(self, other_list: list[object]):
+        """Append elements from another list to the current list and return the modified list."""
+        self.extend(other_list)
+        return self
+
+    def __copy__(self):
+        """Return a shallow copy of the current list."""
+        return list(self)
+
+    def __deepcopy__(self, memo):
+        """Return a deep copy of the current list."""
+        return list(self)
+
+
+def sequences_equal(seq1, seq2):
+    """Check if two sequences are equal by comparing their lengths and elements."""
+    length_match = len(seq1) == len(seq2)
+    if not length_match:
+        return False
+
+    return all(elem1 == elem2 for elem1, elem2 in zip(seq1, seq2))