chore: import upstream snapshot with attribution

python/paddle/jit/sot/utils/__init__.py

@@ -0,0 +1,128 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 .call_ast_utils import get_static_function, try_ast_func  # noqa: F401
+from .envs import (  # noqa: F401
+    ENV_MIN_GRAPH_SIZE,
+    ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_CE_DEBUG_MODE,
+    ENV_SOT_COLLECT_INFO,
+    ENV_SOT_ENABLE_0_SIZE_FALLBACK,
+    ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
+    ENV_SOT_ENABLE_FASTER_GUARD,
+    ENV_SOT_ENABLE_GUARD_TREE,
+    ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
+    ENV_SOT_EXPORT,
+    ENV_SOT_FORCE_FALLBACK_SIR_IDS,
+    ENV_SOT_LOG_LEVEL,
+    ENV_SOT_SERIALIZE_INFO,
+    ENV_SOT_TRACE_NUMPY,
+    ENV_SOT_UNSAFE_CACHE_FASTPATH,
+    ENV_SOT_WITH_CONTROL_FLOW,
+    ENV_STRICT_MODE,
+    PEP508LikeEnvironmentVariable,
+    allow_dynamic_shape_guard,
+    enable_0_size_fallback_guard,
+    export_guard,
+    faster_guard_guard,
+    guard_tree_guard,
+    min_graph_size_guard,
+    sot_step_profiler_guard,
+    specialized_dim_numbers_guard,
+    strict_mode_guard,
+    with_control_flow_guard,
+)
+from .exceptions import (  # noqa: F401
+    BreakGraphError,
+    BreakGraphReasonBase,
+    BuiltinFunctionBreak,
+    ConditionalFallbackError,
+    DataDependencyControlFlowBreak,
+    DataDependencyDynamicShapeBreak,
+    DataDependencyOperationBreak,
+    ExportError,
+    FallbackError,
+    InnerError,
+    PsdbBreakReason,
+    SotCapturedException,
+    SotCapturedExceptionFactory,
+    SotErrorBase,
+    UnsupportedIteratorBreak,
+    UnsupportedOperationBreak,
+    inner_error_default_handler,
+)
+from .info_collector import (  # noqa: F401
+    BreakGraphReasonInfo,
+    CompileCountInfo,
+    InfoCollector,
+    NewSymbolHitRateInfo,
+    SubGraphInfo,
+    SubGraphRelationInfo,
+)
+from .magic_methods import magic_method_builtin_dispatch  # noqa: F401
+from .numpy_utils import (  # noqa: F401
+    NUMPY_API_SUPPORTED_DICT,
+)
+from .paddle_api_config import (  # noqa: F401
+    get_tensor_methods,
+    is_break_graph_tensor_methods,
+    is_directly_run_api,
+    is_inplace_api,
+    is_not_supported_paddle_layer,
+)
+from .utils import (  # noqa: F401
+    Cache,
+    ConstTypes,
+    NameGenerator,
+    ResumeFnNameFactory,
+    Singleton,
+    SIRToCodeMap,
+    SotUndefinedVar,
+    StepInfoManager,
+    already_unified_in_dynamic_and_static_graph,
+    count_if,
+    current_symbol_registry,
+    do_until_stop_iteration,
+    execute_time,
+    flatten,
+    flatten_extend,
+    get_api_fullname,
+    get_min_non_specialized_number,
+    get_numpy_ufuncs,
+    get_obj_stable_repr,
+    get_unbound_method,
+    hashable,
+    in_paddle_module,
+    is_break_graph_api,
+    is_builtin_fn,
+    is_comprehensive_name,
+    is_namedtuple_class,
+    is_paddle_api,
+    is_strict_mode,
+    list_contain_by_id,
+    list_find_index_by_id,
+    log,
+    log_do,
+    log_enabled,
+    log_format,
+    log_once,
+    map_if,
+    map_if_extend,
+    meta_str,
+    need_capture_control_flow,
+    no_eval_frame,
+    printable,
+    switch_symbol_registry,
+    update_list_inplace,
+)

python/paddle/jit/sot/utils/call_ast_utils.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 inspect
+import types
+
+import paddle
+
+from .envs import ENV_SOT_WITH_CONTROL_FLOW
+from .exceptions import InnerError
+from .utils import Singleton
+
+try_ast_codes = set()
+
+
+def try_ast_func(func):
+    def _is_wrapped(f):
+        return hasattr(f, '__wrapped__')
+
+    unwrapped_f = func
+    if hasattr(unwrapped_f, "__code__"):
+        try_ast_codes.add(func.__code__)
+
+    while _is_wrapped(unwrapped_f):
+        unwrapped_f = unwrapped_f.__wrapped__
+        if hasattr(unwrapped_f, "__code__"):
+            try_ast_codes.add(func.__code__)
+
+    return func
+
+
+class StaticFunctionManager(metaclass=Singleton):
+    def __init__(self):
+        self.code_map = {}
+
+    def ast_transform_with_frame(self, frame):
+        code = frame.f_code
+        if code not in try_ast_codes:
+            return None
+        if code not in self.code_map:
+            if code.co_name.startswith("#") or code.co_name.startswith("$"):
+                self.code_map[code] = None
+            elif len(code.co_cellvars) + len(code.co_freevars) != 0:
+                self.code_map[code] = None
+            else:
+                function = types.FunctionType(
+                    code,
+                    frame.f_globals,
+                    code.co_name,
+                    (),
+                    (),
+                )
+                function = paddle.jit.to_static(function, full_graph=True)
+                self.code_map[code] = function
+
+        return self.code_map[code]
+
+    def ast_transform_with_callable(self, fn):
+        if not inspect.isfunction(fn) or not hasattr(fn, "__code__"):
+            return None
+
+        code = fn.__code__
+        if code not in try_ast_codes:
+            return None
+        if code not in self.code_map:
+            if code.co_name.startswith("#") or code.co_name.startswith("$"):
+                self.code_map[code] = None
+            elif len(code.co_cellvars) + len(code.co_freevars) != 0:
+                self.code_map[code] = None
+            else:
+                self.code_map[code] = paddle.jit.to_static(fn, full_graph=True)
+
+        return self.code_map[code]
+
+
+def get_static_function(obj, type_):
+    if ENV_SOT_WITH_CONTROL_FLOW.get():
+        if type_ == "eval_frame":
+            return StaticFunctionManager().ast_transform_with_frame(obj)
+        elif type_ == "inline_call":
+            return StaticFunctionManager().ast_transform_with_callable(obj)
+        else:
+            raise InnerError(f"Can not get static function with type {type_}.")
+    return None

python/paddle/jit/sot/utils/envs.py

@@ -0,0 +1,238 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+
+import os
+from contextlib import contextmanager
+
+from paddle.utils.environments import (
+    BooleanEnvironmentVariable,
+    EnvironmentVariable,
+    EnvironmentVariableGuard,
+    IntegerEnvironmentVariable,
+    StringEnvironmentVariable,
+)
+
+
+class PEP508LikeEnvironmentVariable(EnvironmentVariable[dict[str, list[str]]]):
+    """
+    Environment variable parser following PEP 508 extras specification syntax.
+    https://peps.python.org/pep-0508/
+
+    Processes strings using PEP 508-style bracket notation for optional components:
+    "feat1[opt1,opt2], feat2[opt3,opt4]" -> {'feat1': ['opt1', 'opt2'], 'feat2': ['opt3', 'opt4']}
+    """
+
+    def __init__(self, name: str, default: dict[str, list[str]]):
+        super().__init__(name, default)
+        assert isinstance(default, dict), "default must be a dict"
+
+    def parse_from_string(self) -> dict[str, list[str]]:
+        env_var = os.getenv(self.name)
+        if env_var is None or env_var == "":
+            return self.default
+        items = self.split_by_unbracketed_commas(env_var)
+        ret = {}
+        for item in items:
+            ret.update(self.parse_parameterized_key(item))
+        return ret
+
+    def convert_to_string(self, value: dict[str, list[str]]) -> str:
+        assert isinstance(value, dict), "The input must be a dict"
+        assert all(isinstance(x, str) for x in value.keys()), (
+            "Keys must be a string"
+        )
+        assert all(isinstance(x, list) for x in value.values()), (
+            "Values must be a list"
+        )
+
+        env_list = []
+        for k, v in value.items():
+            env_list.append(f"{k}" + (f"[{','.join(v)}]" if len(v) else ""))
+
+        return ",".join(env_list)
+
+    @staticmethod
+    def split_by_unbracketed_commas(input_str: str) -> list[str]:
+        """Split string by commas that are not enclosed in square brackets"""
+        # "feat1[opt1,opt2], feat2[opt3], feat3" -> ["feat1[opt1,opt2]", "feat2[opt3]", "feat3"]
+        bracket_depth = 0
+        split_parts = []
+        _start = 0
+
+        for _current, char in enumerate(input_str):
+            if char == "[":
+                bracket_depth += 1
+            elif char == "]":
+                bracket_depth = max(
+                    0, bracket_depth - 1
+                )  # Prevent negative depth
+
+            if char == "," and bracket_depth == 0:
+                split_parts.append(input_str[_start:_current].strip())
+                _start = _current + 1  # Skip comma
+
+        # Add remaining content after last comma
+        if remaining := input_str[_start:].strip():
+            split_parts.append(remaining)
+
+        return split_parts
+
+    @staticmethod
+    def parse_parameterized_key(input_str: str) -> dict[str, list[str]]:
+        """Parse key with parameters in brackets into a dictionary."""
+
+        start_bracket = input_str.find("[")
+        end_bracket = input_str.rfind("]")
+
+        if start_bracket == -1 or end_bracket == -1:
+            return {input_str: []}
+
+        parameter_key = input_str[:start_bracket].strip()
+
+        # Extract and clean parameters
+        parameters_str = input_str[start_bracket + 1 : end_bracket]
+        parameter_values = [
+            v.strip() for v in parameters_str.split(",") if v.strip()
+        ]
+
+        return {parameter_key: parameter_values}
+
+
+ENV_MIN_GRAPH_SIZE = IntegerEnvironmentVariable("MIN_GRAPH_SIZE", 10)
+ENV_SOT_LOG_LEVEL = IntegerEnvironmentVariable("SOT_LOG_LEVEL", 0)
+ENV_STRICT_MODE = BooleanEnvironmentVariable("STRICT_MODE", False)
+ENV_SOT_WITH_CONTROL_FLOW = BooleanEnvironmentVariable(
+    "SOT_WITH_CONTROL_FLOW", True
+)
+ENV_SOT_EXPORT = StringEnvironmentVariable("SOT_EXPORT", "")
+ENV_SOT_ALLOW_DYNAMIC_SHAPE = BooleanEnvironmentVariable(
+    "SOT_ALLOW_DYNAMIC_SHAPE",
+    # Enable SOT dynamic shape as default in PIR mode
+    True,
+)
+ENV_SOT_ENABLE_FASTER_GUARD = BooleanEnvironmentVariable(
+    "SOT_ENABLE_FASTER_GUARD",
+    False,
+)
+ENV_SOT_ENABLE_STRICT_GUARD_CHECK = BooleanEnvironmentVariable(
+    "SOT_ENABLE_STRICT_GUARD_CHECK",
+    False,
+)
+ENV_SOT_ENABLE_GUARD_TREE = BooleanEnvironmentVariable(
+    "SOT_ENABLE_GUARD_TREE",
+    False,
+)
+ENV_ENABLE_SOT_STEP_PROFILER = BooleanEnvironmentVariable(
+    "ENABLE_SOT_STEP_PROFILER", False
+)
+ENV_SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE = BooleanEnvironmentVariable(
+    "SOT_BREAK_GRAPH_ON_GET_SYMBOLIC_VALUE", False
+)
+ENV_SOT_COLLECT_INFO = PEP508LikeEnvironmentVariable("SOT_COLLECT_INFO", {})
+ENV_SOT_SERIALIZE_INFO = BooleanEnvironmentVariable("SOT_SERIALIZE_INFO", False)
+ENV_SOT_CE_DEBUG_MODE = BooleanEnvironmentVariable("SOT_CE_DEBUG_MODE", False)
+ENV_SOT_FORCE_FALLBACK_SIR_IDS = StringEnvironmentVariable(
+    "SOT_FORCE_FALLBACK_SIR_IDS", ""
+)
+ENV_SOT_TRACE_NUMPY = BooleanEnvironmentVariable("ENV_SOT_TRACE_NUMPY", True)
+ENV_SOT_UNSAFE_CACHE_FASTPATH = BooleanEnvironmentVariable(
+    "SOT_UNSAFE_CACHE_FASTPATH", False
+)
+ENV_SOT_ENABLE_0_SIZE_FALLBACK = BooleanEnvironmentVariable(
+    "SOT_ENABLE_0_SIZE_FALLBACK", True
+)
+ENV_SOT_SPECIALIZED_DIM_NUMBERS = StringEnvironmentVariable(
+    "SOT_SPECIALIZED_DIM_NUMBERS", "0"
+)
+ENV_SOT_ENABLE_COMPILE_TIME_LIMIT = BooleanEnvironmentVariable(
+    "SOT_ENABLE_COMPILE_TIME_LIMIT", True
+)
+
+
+def update_ce_flags():
+    if not ENV_SOT_CE_DEBUG_MODE.get():
+        return
+    # Enable information collection flags to facilitate debugging and analysis
+
+    collected_info_item: dict[str, list[str]] = ENV_SOT_COLLECT_INFO.get()
+    collected_info_item.setdefault("breakgraph_reason", [])
+    collected_info_item.setdefault("subgraph_info", [])
+
+    ENV_SOT_COLLECT_INFO.set(collected_info_item)
+    ENV_SOT_SERIALIZE_INFO.set(True)
+
+
+update_ce_flags()
+
+
+@contextmanager
+def strict_mode_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_STRICT_MODE, value):
+        yield
+
+
+@contextmanager
+def min_graph_size_guard(value: int):
+    with EnvironmentVariableGuard(ENV_MIN_GRAPH_SIZE, value):
+        yield
+
+
+@contextmanager
+def with_control_flow_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_WITH_CONTROL_FLOW, value):
+        yield
+
+
+@contextmanager
+def export_guard(value: str):
+    with EnvironmentVariableGuard(ENV_SOT_EXPORT, value):
+        yield
+
+
+@contextmanager
+def allow_dynamic_shape_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_ALLOW_DYNAMIC_SHAPE, value):
+        yield
+
+
+@contextmanager
+def faster_guard_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_ENABLE_FASTER_GUARD, value):
+        yield
+
+
+@contextmanager
+def guard_tree_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_ENABLE_GUARD_TREE, value):
+        yield
+
+
+@contextmanager
+def sot_step_profiler_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_ENABLE_SOT_STEP_PROFILER, value):
+        yield
+
+
+@contextmanager
+def specialized_dim_numbers_guard(value: str):
+    with EnvironmentVariableGuard(ENV_SOT_SPECIALIZED_DIM_NUMBERS, value):
+        yield
+
+
+@contextmanager
+def enable_0_size_fallback_guard(value: bool):
+    with EnvironmentVariableGuard(ENV_SOT_ENABLE_0_SIZE_FALLBACK, value):
+        yield

python/paddle/jit/sot/utils/exceptions.py

@@ -0,0 +1,488 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+
+import traceback
+from typing import TYPE_CHECKING
+
+from .info_collector import BreakGraphReasonInfo
+
+if TYPE_CHECKING:
+    from ..opcode_translator.executor.variables.base import VariableBase
+
+
+class BreakGraphReasonBase:
+    """Base class for representing reasons why graph execution was interrupted.
+
+    Attributes:
+        reason_str (str): Description of the break reason
+        file_path (str): Path to the file where break occurred
+        line_number (int): Line number where break occurred
+    """
+
+    def __init__(
+        self,
+        reason_str,
+        file_path="",
+        line_number=-1,
+    ):
+        self.reason_str = reason_str
+        self.file_path = file_path
+        self.line_number = line_number
+
+    def __repr__(self) -> str:
+        return f"{self.reason_str}"
+
+
+class DataDependencyBreak(BreakGraphReasonBase):
+    pass
+
+
+class DataDependencyControlFlowBreak(DataDependencyBreak):
+    """Break reason for control flow execution."""
+
+    def __init__(self, reason_str=None, file_path="", line_number=-1):
+        if reason_str is None:
+            reason_str = "OpcodeInlineExecutor want break graph when simulate control flow."
+
+        super().__init__(
+            reason_str,
+            file_path,
+            line_number,
+        )
+
+
+class DataDependencyDynamicShapeBreak(DataDependencyBreak):
+    pass
+
+
+class DataDependencyOperationBreak(DataDependencyBreak):
+    pass
+
+
+class UnsupportedOperationBreak(BreakGraphReasonBase):
+    def __init__(
+        self,
+        *,
+        left_type=None,
+        right_type=None,
+        operator=None,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = f"Unsupported operator '{operator}' between {left_type} and {right_type}"
+        super().__init__(reason_str, file_path, line_number)
+
+
+class UnsupportedPaddleAPIBreak(UnsupportedOperationBreak):
+    def __init__(
+        self,
+        *,
+        fn_name=None,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = f"Not support Paddlepaddle API: {fn_name}"
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
+class UnsupportedNumPyAPIBreak(UnsupportedOperationBreak):
+    def __init__(
+        self,
+        *,
+        fn_name=None,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = f"Not support NumPy API: {fn_name}"
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
+class UnsupportedRandomAPIBreak(UnsupportedOperationBreak):
+    def __init__(
+        self,
+        *,
+        fn_name=None,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = f"Random function {fn_name} is not supported."
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
+class ForceBreak(UnsupportedOperationBreak):
+    def __init__(
+        self,
+        *,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = "Force break graph execution"
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
+class BuiltinFunctionBreak(UnsupportedOperationBreak):
+    """Break reason for unsupported built-in function calls.
+
+    Args:
+        fn_name (str): Name of the builtin function
+        arg_types (list): Types of the arguments passed to the function
+        file_path (str): Path to the file where break occurred
+        line_number (int): Line number where break occurred
+    """
+
+    def __init__(
+        self,
+        *,
+        fn_name=None,
+        arg_types=None,
+        reason_str=None,
+        file_path="",
+        line_number=-1,
+    ):
+        if reason_str is None:
+            reason_str = f"Not support builtin function: {fn_name} with args: Args({arg_types})"
+
+        super().__init__(
+            reason_str=reason_str,
+            file_path=file_path,
+            line_number=line_number,
+        )
+
+
+class SideEffectBreak(BreakGraphReasonBase):
+    pass
+
+
+class UnsupportedIteratorBreak(SideEffectBreak):
+    pass
+
+
+class InlineCallBreak(BreakGraphReasonBase):
+    pass
+
+
+class FallbackInlineCallBreak(InlineCallBreak):
+    pass
+
+
+class BreakGraphInlineCallBreak(InlineCallBreak):
+    pass
+
+
+class OtherInlineCallBreak(InlineCallBreak):
+    pass
+
+
+class DygraphInconsistentWithStaticBreak(BreakGraphReasonBase):
+    pass
+
+
+class PsdbBreakReason(BreakGraphReasonBase):
+    pass
+
+
+class InferMetaBreak(BreakGraphReasonBase):
+    """Break reason during meta information inference phase."""
+
+    pass
+
+
+class NullMetaBreak(BreakGraphReasonBase):
+    def __init__(
+        self,
+        *,
+        file_path="",
+        line_number=-1,
+    ):
+        super().__init__(
+            "Access attribute from null meta", file_path, line_number
+        )
+
+
+class SotErrorBase(Exception):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        from ..opcode_translator.breakpoint import BreakpointManager
+
+        BreakpointManager().on_event(f"{self.__class__.__name__}")
+
+    def print(self):
+        lines = traceback.format_tb(self.__traceback__)
+        print("".join(lines))
+
+
+class InnerError(SotErrorBase):
+    pass
+
+
+class HasNoAttributeError(InnerError):
+    pass
+
+
+class FallbackError(SotErrorBase):
+    def __init__(self, msg, disable_eval_frame=False):
+        super().__init__(msg)
+        self.disable_eval_frame = disable_eval_frame
+
+
+class ConditionalFallbackError(FallbackError): ...
+
+
+# raise in inline function call strategy.
+class BreakGraphError(SotErrorBase):
+    def __init__(self, reason: BreakGraphReasonBase = None):
+        super().__init__(str(reason))
+
+        if not isinstance(reason, BreakGraphReasonBase):
+            raise ValueError(
+                "reason must be a subclass of BreakGraphReasonBase"
+            )
+
+        self.reason = reason
+        BreakGraphReasonInfo.collect_break_graph_reason(reason)
+
+
+def inner_error_default_handler(func, message_fn):
+    """Wrap function and an error handling function and throw an InnerError."""
+
+    def impl(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except SotErrorBase as e:
+            raise e
+        except Exception as e:
+            message = message_fn(*args, **kwargs)
+            origin_exception_message = "\n".join(
+                traceback.format_exception(type(e), e, e.__traceback__)
+            )
+            raise InnerError(
+                f"{message}\nOrigin Exception is: \n {origin_exception_message}"
+            ) from e
+
+    return impl
+
+
+class ExportError(SotErrorBase):
+    pass
+
+
+class SotExtraInfo:
+    SOT_EXTRA_INFO_ATTR_NAME = "__SOT_EXTRA_INFO__"
+
+    def __init__(self, *, need_breakgraph: bool = False):
+        self.need_breakgraph = need_breakgraph
+
+    def set_need_breakgraph(self, need_breakgraph: bool):
+        self.need_breakgraph = need_breakgraph
+
+    def attach(self, err: BaseException):
+        setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, self)
+
+    @staticmethod
+    def default() -> SotExtraInfo:
+        return SotExtraInfo()
+
+    @staticmethod
+    def from_exception(err: BaseException) -> SotExtraInfo:
+        info = getattr(
+            err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, SotExtraInfo.default()
+        )
+        setattr(err, SotExtraInfo.SOT_EXTRA_INFO_ATTR_NAME, info)
+        return info
+
+
+class SotCapturedException(SotErrorBase):
+    # Represents an exception encountered during bytecode execution simulation.
+    # This exception is used by SOT to handle Python exceptions by mapping them to
+    # SotCapturedException for consistent exception handling in the simulation process.
+    ...
+
+
+class SotCapturedLookupError(SotCapturedException): ...
+
+
+class SotCapturedIndexError(SotCapturedLookupError): ...
+
+
+class SotCapturedKeyError(SotCapturedLookupError): ...
+
+
+class SotCapturedArithmeticError(SotCapturedException): ...
+
+
+class SotCapturedFloatingPointError(SotCapturedArithmeticError): ...
+
+
+class SotCapturedOverflowError(SotCapturedArithmeticError): ...
+
+
+class SotCapturedZeroDivisionError(SotCapturedArithmeticError): ...
+
+
+class SotCapturedImportError(SotCapturedException): ...
+
+
+class SotCapturedModuleNotFoundError(SotCapturedImportError): ...
+
+
+class SotCapturedRuntimeError(SotCapturedException): ...
+
+
+class SotCapturedNotImplementedError(SotCapturedRuntimeError): ...
+
+
+class SotCapturedRecursionError(SotCapturedRuntimeError): ...
+
+
+class SotCapturedNameError(SotCapturedException): ...
+
+
+class SotCapturedUnboundLocalError(SotCapturedNameError): ...
+
+
+class SotCapturedSyntaxError(SotCapturedException): ...
+
+
+class SotCapturedIndentationError(SotCapturedSyntaxError): ...
+
+
+class SotCapturedTabError(SotCapturedIndentationError): ...
+
+
+class SotCapturedOSError(SotCapturedException): ...
+
+
+class SotCapturedFileExistsError(SotCapturedOSError): ...
+
+
+class SotCapturedFileNotFoundError(SotCapturedOSError): ...
+
+
+class SotCapturedIsADirectoryError(SotCapturedOSError): ...
+
+
+class SotCapturedNotADirectoryError(SotCapturedOSError): ...
+
+
+class SotCapturedPermissionError(SotCapturedOSError): ...
+
+
+class SotCapturedTimeoutError(SotCapturedOSError): ...
+
+
+class SotCapturedStopIteration(SotCapturedOSError): ...
+
+
+class SotCapturedExceptionFactory:
+    # This dictionary maps common built-in Python Exception types to their corresponding SotCapturedException
+    # types, preserving the original exception hierarchy for proper inheritance behavior.
+    # Reference: https://docs.python.org/3/library/exceptions.html#exception-hierarchy
+    MAPPING = {
+        Exception: SotCapturedException,
+        LookupError: SotCapturedLookupError,
+        IndexError: SotCapturedIndexError,
+        KeyError: SotCapturedKeyError,
+        ArithmeticError: SotCapturedArithmeticError,
+        FloatingPointError: SotCapturedFloatingPointError,
+        OverflowError: SotCapturedOverflowError,
+        ZeroDivisionError: SotCapturedZeroDivisionError,
+        ImportError: SotCapturedImportError,
+        ModuleNotFoundError: SotCapturedModuleNotFoundError,
+        RuntimeError: SotCapturedRuntimeError,
+        NotImplementedError: SotCapturedNotImplementedError,
+        NameError: SotCapturedNameError,
+        UnboundLocalError: SotCapturedUnboundLocalError,
+        SyntaxError: SotCapturedSyntaxError,
+        IndentationError: SotCapturedIndentationError,
+        TabError: SotCapturedTabError,
+        OSError: SotCapturedOSError,
+        FileExistsError: SotCapturedFileExistsError,
+        FileNotFoundError: SotCapturedFileNotFoundError,
+        IsADirectoryError: SotCapturedIsADirectoryError,
+        NotADirectoryError: SotCapturedNotADirectoryError,
+        PermissionError: SotCapturedPermissionError,
+        TimeoutError: SotCapturedTimeoutError,
+        StopIteration: SotCapturedStopIteration,
+    }
+
+    @classmethod
+    def get(
+        cls,
+        exc_type: type[Exception],
+    ) -> type[SotCapturedException]:
+        if isinstance(exc_type, type) and issubclass(
+            exc_type, SotCapturedException
+        ):
+            return exc_type
+
+        if exc_type not in cls.MAPPING:
+            name = getattr(exc_type, "__name__", str(exc_type))
+            cls.MAPPING[exc_type] = type(
+                f"SotCaptured{name}", (SotCapturedException,), {}
+            )
+        return cls.MAPPING[exc_type]
+
+    @classmethod
+    def create(
+        cls,
+        origin_exc: Exception,
+        tracked_args: list[VariableBase] | None = None,
+    ) -> SotCapturedException:
+        # transform an Exception to SotCapturedException
+        exc_type = origin_exc.__class__
+
+        new_exc_type = cls.get(exc_type)
+        new_exc = new_exc_type(*origin_exc.args)
+        new_exc.__cause__ = origin_exc.__cause__
+        new_exc.__context__ = origin_exc.__context__
+        new_exc.__suppress_context__ = origin_exc.__suppress_context__
+        new_exc.__traceback__ = origin_exc.__traceback__
+
+        # Propagating Exception Parameters through SotCapturedException
+        if tracked_args is None:
+            tracked_args = []
+        new_exc.tracked_args = tracked_args
+
+        return new_exc

python/paddle/jit/sot/utils/info_collector.py

@@ -0,0 +1,468 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+
+import atexit
+import base64
+import json
+import sys
+from abc import ABC, abstractmethod
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple
+
+from typing_extensions import Self
+
+from .envs import ENV_SOT_COLLECT_INFO, ENV_SOT_SERIALIZE_INFO
+from .utils import Singleton
+
+if TYPE_CHECKING:
+    import types
+
+    from .exceptions import BreakGraphReasonBase
+
+PREFIX = "<sot>"
+SUFFIX = "</sot>"
+ENCODING = "utf-8"
+
+
+def try_import_graphviz():
+    try:
+        import graphviz
+
+        return graphviz
+    except ImportError:
+        return None
+
+
+class InfoType(Enum):
+    STEP_INFO = 0
+    E2E_INFO = 1
+
+
+class InfoCollector(metaclass=Singleton):
+    def __init__(self):
+        self._step_info: dict[str, list[InfoBase]] = {}
+        self._e2e_info: dict[str, list[InfoBase]] = {}
+
+    def get_info_dict(self, info_type: InfoType) -> dict[str, list[InfoBase]]:
+        if info_type == InfoType.STEP_INFO:
+            return self._step_info
+        else:
+            return self._e2e_info
+
+    def attach(self, cls: type[InfoBase], *args, **kwargs) -> None:
+        if self.need_collect(cls):
+            info = cls(*args, **kwargs)
+            self.register(info)
+
+    def register(self, info: InfoBase) -> None:
+        info_class_name = info.__class__.__name__
+        info_type = info.TYPE
+        info_dict = self.get_info_dict(info_type)
+        info_dict.setdefault(info_class_name, [])
+        info_dict[info_class_name].append(info)
+
+    def need_collect(self, cls: type[InfoBase]) -> bool:
+        return cls.SHORT_NAME in ENV_SOT_COLLECT_INFO.get()
+
+    def clear_step_info(self):
+        self._step_info.clear()
+
+    def clear_e2e_info(self):
+        self._e2e_info.clear()
+
+    def clear(self):
+        self.clear_step_info()
+        self.clear_e2e_info()
+
+    def print_step_report(self):
+        self.print_report(InfoType.STEP_INFO)
+
+    def print_e2e_info_atexit(self) -> None:
+        def atexit_hook():
+            self.print_report(InfoType.E2E_INFO)
+            sys.stdout.flush()
+            self.clear()
+
+        atexit.register(atexit_hook)
+
+    def print_report(self, info_type: InfoType) -> None:
+        if info_dict := self.get_info_dict(info_type):
+            print(self.generate_report(info_dict))
+
+    def generate_report(self, info_dict: dict[str, list[InfoBase]]) -> str:
+        report = ""
+        for info_class_name, info_list in info_dict.items():
+            cls = info_list[0].__class__
+            report += f"{info_class_name} ({cls.SHORT_NAME}):\n"
+            if ENV_SOT_SERIALIZE_INFO.get():
+                report += cls.json_report(info_list)
+            else:
+                report += cls.summary(info_list)
+            report += "\n"
+        return report
+
+
+InfoCollector().print_e2e_info_atexit()
+
+
+class InfoBase(ABC):
+    SHORT_NAME: ClassVar[str]
+    TYPE: ClassVar[InfoType]
+
+    def __init__(self): ...
+
+    @classmethod
+    @abstractmethod
+    def summary(cls, history: list[Self]) -> str: ...
+
+    @classmethod
+    def serialize(cls, obj: dict[str:Any]) -> str:
+        json_data = json.dumps(obj)
+        b64_bytes = base64.b64encode(json_data.encode(ENCODING))
+
+        return b64_bytes.decode(ENCODING)
+
+    @classmethod
+    def deserialize(cls, data: bytes | str) -> dict:
+        if isinstance(data, str):
+            data = data.encode(ENCODING)
+        json_str = base64.b64decode(data).decode(ENCODING)
+
+        return json.loads(json_str)
+
+
+class NewSymbolHitRateInfo(InfoBase):
+    SHORT_NAME = "new_symbol_hit_rate"
+    TYPE = InfoType.STEP_INFO
+
+    def __init__(
+        self, input_tensor_ids: list[int], output_tensor_ids: list[int]
+    ):
+        super().__init__()
+        self.input_tensor_ids = input_tensor_ids
+        self.output_tensor_ids = output_tensor_ids
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        if len(history) == 0:
+            return f"No {cls.SHORT_NAME} info"
+        if len(history) == 1:
+            return "Only one subgraph is generated"
+        known_tensor_ids = set()
+        hit_count = 0
+        all_count = sum([len(info.input_tensor_ids) for info in history[1:]])
+        for i, info in enumerate(history):
+            for tensor_id in info.input_tensor_ids:
+                # Skip the first graph
+                if i == 0:
+                    continue
+                if tensor_id in known_tensor_ids:
+                    hit_count += 1
+            for tensor_id in info.output_tensor_ids:
+                known_tensor_ids.add(tensor_id)
+        summary = f"All tensor count: {all_count}, hit count: {hit_count}\n"
+        summary += f"Hit rate: {hit_count / all_count:.2f}"
+        return summary
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
+
+class SubGraphRelationInfo(InfoBase):
+    SHORT_NAME = "subgraph_relation"
+    TYPE = InfoType.STEP_INFO
+    STEP_UNIQUE_ID = 0
+
+    class ConcreteShapeInfo(NamedTuple):
+        id: int
+        ir_shape: list[int]
+        real_shape: list[int]
+
+    def __init__(
+        self,
+        subgraph_name: str,
+        input_shape_infos: list[SubGraphRelationInfo.ConcreteShapeInfo],
+        output_shape_infos: list[SubGraphRelationInfo.ConcreteShapeInfo],
+        is_first_call: bool,
+        graph_size: int,
+    ):
+        super().__init__()
+        self.subgraph_name = subgraph_name
+        self.input_shape_infos = input_shape_infos
+        self.output_shape_infos = output_shape_infos
+        self.is_first_call = is_first_call
+        self.graph_size = graph_size
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        # TODO: attach input shape (with dynamic shape info)
+        cls.STEP_UNIQUE_ID += 1
+        if len(history) == 0:
+            return f"No {cls.SHORT_NAME} info"
+        if all(not subgraph_info.is_first_call for subgraph_info in history):
+            return "All subgraph are not the first call"
+        graphviz = try_import_graphviz()
+        if graphviz is None:
+            return "Please install graphviz to show the subgraph relation"
+        dot = graphviz.Digraph()
+        shape_infos = [
+            shape_info
+            for info in history
+            for shape_info in info.input_shape_infos + info.output_shape_infos
+        ]
+
+        def to_tensor_node_name(
+            shape_info: SubGraphRelationInfo.ConcreteShapeInfo,
+        ):
+            return f"tensor_{shape_info.id}"
+
+        visited_shape = set()
+        for shape_info in shape_infos:
+            if shape_info.id in visited_shape:
+                continue
+            visited_shape.add(shape_info.id)
+            dot.node(
+                to_tensor_node_name(shape_info),
+                f"Tensor {shape_info.id} shape={shape_info.real_shape}",
+                shape="rect",
+            )
+        for i, info in enumerate(history):
+            subgraph_id = f"subgraph_{i}"
+            dot.node(
+                subgraph_id,
+                f"Subgraph {i} ({info.subgraph_name}, size={info.graph_size})",
+                shape="oval",
+                fillcolor="cyan" if info.is_first_call else None,
+                style="filled" if info.is_first_call else None,
+            )
+            for shape_info in info.input_shape_infos:
+                dot.edge(
+                    to_tensor_node_name(shape_info),
+                    subgraph_id,
+                    label=str(shape_info.ir_shape),
+                )
+            for shape_info in info.output_shape_infos:
+                dot.edge(
+                    subgraph_id,
+                    to_tensor_node_name(shape_info),
+                    label=str(shape_info.ir_shape),
+                )
+
+        directory = Path(".") / "subgraph_relation"
+        directory.mkdir(exist_ok=True, parents=True)
+        filename = f"subgraph_relation_{cls.STEP_UNIQUE_ID}"
+        dot.render(directory / filename, format="svg", cleanup=True)
+        return f"Please check {directory / filename}.svg for subgraph relation"
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
+
+class CompileCountInfo(InfoBase):
+    SHORT_NAME = "compile_count"
+    TYPE = InfoType.E2E_INFO
+
+    def __init__(self, code: types.CodeType):
+        super().__init__()
+        self.code = code
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        if len(history) == 0:
+            return f"No {cls.SHORT_NAME} info"
+        code_count = {}
+        for info in history:
+            code_count[info.code] = code_count.get(info.code, 0) + 1
+        summary_lines = []
+        for code, count in sorted(
+            code_count.items(), key=lambda x: x[1], reverse=True
+        ):
+            filename, lineno = code.co_filename, code.co_firstlineno
+            summary_lines.append(
+                f"    {code.co_name} ({filename}:{lineno}): {count}"
+            )
+        summary = "\n".join(summary_lines)
+        return summary
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        # TODO: need to support serialize the output
+        return cls.summary(history)
+
+
+class BreakGraphReasonInfo(InfoBase):
+    SHORT_NAME = "breakgraph_reason"
+    TYPE = InfoType.E2E_INFO
+
+    def __init__(self, reason: BreakGraphReasonBase):
+        super().__init__()
+        self.reason = reason
+
+    @classmethod
+    def classify(cls, history: list[Self]) -> str:
+        reasons_dict = {}
+
+        for info in history:
+            name = info.reason.__class__.__name__
+            if name not in reasons_dict:
+                reasons_dict[name] = []
+            reasons_dict[name].append(str(info.reason))
+
+        sorted_reasons = list(reasons_dict.items())
+        sorted_reasons.sort(key=lambda x: len(x[1]), reverse=True)
+
+        return reasons_dict, sorted_reasons
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        reason_dict, reason_list = cls.classify(history)
+
+        return "\n".join(
+            [
+                f"{name} ({len(reasons)}):\n\t" + "\n\t".join(reasons)
+                for name, reasons in reason_list
+            ]
+        )
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        reason_dict, sorted_reasons = cls.classify(history)
+        reason_dict["count"] = {k: len(v) for k, v in sorted_reasons}
+        serialized = cls.serialize({cls.SHORT_NAME: reason_dict})
+
+        return f"{PREFIX}{serialized}{SUFFIX}"
+
+    @classmethod
+    def restore_from_string(cls, serialized: str) -> list[Self]:
+        # This method is the inverse of json_report
+
+        from paddle.jit.sot.utils import exceptions
+
+        history = []
+        obj = cls.deserialize(serialized)[cls.SHORT_NAME]
+        obj.pop("count")
+
+        for classname in obj:
+            ReasonClass = getattr(exceptions, classname, None)
+            for reason in obj[classname]:
+                history.append(cls(ReasonClass(reason_str=reason)))
+
+        return history
+
+    @staticmethod
+    def collect_break_graph_reason(reason: BreakGraphReasonBase):
+        if not InfoCollector().need_collect(BreakGraphReasonInfo):
+            return
+
+        InfoCollector().attach(BreakGraphReasonInfo, reason)
+
+
+class SubGraphInfo(InfoBase):
+    SHORT_NAME = "subgraph_info"
+    TYPE = InfoType.STEP_INFO
+
+    def __init__(self, graph: str, op_num: int, sir_name: str):
+        # NOTE: All data should be serializable
+        super().__init__()
+        self.graph = graph
+        self.op_num = op_num
+        self.sir_name = sir_name
+
+    def __str__(self):
+        return (
+            f"[SIR Name] {self.sir_name}   [OpNum] {self.op_num}\n{self.graph}"
+        )
+
+    @classmethod
+    def summary(cls, history: list[Self]) -> str:
+        num_of_subgraph = len(history)
+        sum_of_op_num = sum(item.op_num for item in history)
+
+        need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            cls.SHORT_NAME, []
+        )
+
+        details = ""
+        if need_details:
+            details = "\n".join(
+                [
+                    f"[SubGraphIdx] {idx}   {info}"
+                    for idx, info in enumerate(map(str, history))
+                ]
+            )
+
+        summary = f"[Number of subgraph] {num_of_subgraph} [Sum of opnum] {sum_of_op_num}"
+
+        return f"{summary}\n{details}"
+
+    @classmethod
+    def json_report(cls, history: list[Self]) -> str:
+        need_details = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            cls.SHORT_NAME, []
+        )
+
+        aggregated_info_list = []
+        for idx, record in enumerate(history):
+            entry_data = {}
+
+            entry_data["SIR_name"] = record.sir_name
+            entry_data["OpNum"] = record.op_num
+            entry_data["Graph"] = ""
+            if need_details:
+                entry_data["Graph"] = str(record.graph)
+            aggregated_info_list.append(entry_data)
+
+        serialized = cls.serialize({cls.SHORT_NAME: aggregated_info_list})
+
+        return f"{PREFIX}{serialized}{SUFFIX}"
+
+    @classmethod
+    def restore_from_string(cls, serialized: str) -> list[Self]:
+        # This method is the inverse of json_report
+
+        history = []
+        obj = cls.deserialize(serialized)[cls.SHORT_NAME]
+
+        for entry in obj:
+            history.append(
+                SubGraphInfo(
+                    graph=entry["Graph"],
+                    op_num=entry["OpNum"],
+                    sir_name=entry["SIR_name"],
+                )
+            )
+
+        return history
+
+    def __eq__(self, other):
+        need_graph_equal = "details" in ENV_SOT_COLLECT_INFO.get().get(
+            self.SHORT_NAME, []
+        )
+
+        graph_equal_or_not = True
+        if need_graph_equal:
+            graph_equal_or_not = self.graph == other.graph
+
+        return (
+            graph_equal_or_not
+            and self.op_num == other.op_num
+            and self.sir_name == other.sir_name
+        )

python/paddle/jit/sot/utils/magic_methods.py

@@ -0,0 +1,155 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+
+import operator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+from .utils import hashable
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    BinaryOp = Callable[[Any, Any], Any]
+    UnaryOp = Callable[[Any], Any]
+
+
+INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[BinaryOp, tuple[str, BinaryOp]] = {
+    # inplace op fn: (magic name, non-inplace op fn)
+    operator.iadd: ("__iadd__", operator.add),
+    operator.iand: ("__iand__", operator.and_),
+    operator.iconcat: ("__iconcat__", operator.concat),
+    operator.ifloordiv: ("__ifloordiv__", operator.floordiv),
+    operator.ilshift: ("__ilshift__", operator.lshift),
+    operator.imatmul: ("__imatmul__", operator.matmul),
+    operator.imod: ("__imod__", operator.mod),
+    operator.imul: ("__imul__", operator.mul),
+    operator.ior: ("__ior__", operator.or_),
+    operator.ipow: ("__ipow__", operator.pow),
+    operator.irshift: ("__irshift__", operator.rshift),
+    operator.isub: ("__isub__", operator.sub),
+    operator.itruediv: ("__itruediv__", operator.truediv),
+    operator.ixor: ("__ixor__", operator.xor),
+}
+
+NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES: dict[
+    BinaryOp, tuple[str, str | None]
+] = {
+    # op fn: (magic name, reverse magic name)
+    operator.add: ("__add__", "__radd__"),
+    operator.and_: ("__and__", "__rand__"),
+    operator.contains: ("__contains__", None),
+    operator.delitem: ("__delitem__", None),
+    operator.eq: ("__eq__", "__eq__"),
+    operator.floordiv: ("__floordiv__", "__rfloordiv__"),
+    operator.ge: ("__ge__", "__le__"),
+    operator.getitem: ("__getitem__", None),
+    operator.gt: ("__gt__", "__lt__"),
+    operator.le: ("__le__", "__ge__"),
+    operator.lshift: ("__lshift__", "__rlshift__"),
+    operator.lt: ("__lt__", "__gt__"),
+    operator.matmul: ("__matmul__", "__rmatmul__"),
+    operator.mod: ("__mod__", "__rmod__"),
+    operator.mul: ("__mul__", "__rmul__"),
+    operator.ne: ("__ne__", "__ne__"),
+    operator.or_: ("__or__", "__ror__"),
+    operator.pow: ("__pow__", "__rpow__"),
+    operator.rshift: ("__rshift__", "__rrshift__"),
+    operator.sub: ("__sub__", "__rsub__"),
+    operator.truediv: ("__truediv__", "__rtruediv__"),
+    operator.xor: ("__xor__", "__rxor__"),
+}
+
+UNARY_OPS_TO_MAGIC_NAMES: dict[UnaryOp, str] = {
+    operator.neg: "__neg__",
+    operator.invert: "__invert__",
+    operator.pos: "__pos__",
+    operator.abs: "__abs__",
+    operator.index: "__index__",
+    operator.inv: "__inv__",
+    operator.truth: "__bool__",
+    bool: "__bool__",
+    abs: "__abs__",
+    float: "__float__",
+    len: "__len__",
+    int: "__int__",
+    complex: "__complex__",
+}
+# TODO(SigureMo): support any, all, sum
+
+
+INPLACE_BINARY_OPS = set(INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
+NON_INPLACE_BINARY_OPS = set(NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES.keys())
+BINARY_OPS = INPLACE_BINARY_OPS | NON_INPLACE_BINARY_OPS
+UNARY_OPS = set(UNARY_OPS_TO_MAGIC_NAMES.keys())
+
+
+# NOTE: Both operator.pow and operator.ipow should be considered for inclusion in this list,
+# as they raise ZeroDivisionError when evaluating 0^n where n < 0 (division by zero).
+NEED_GUARD_ZERO_DIVISION_ERROR_OPS: list[BinaryOp] = [
+    operator.floordiv,
+    operator.truediv,
+    operator.mod,
+    operator.ifloordiv,
+    operator.itruediv,
+    operator.imod,
+]
+
+
+@dataclass
+class MagicMethod:
+    name: str
+    is_inplace: bool = False
+    is_reverse: bool = False
+
+
+def magic_method_builtin_dispatch(fn: BinaryOp | UnaryOp) -> list[MagicMethod]:
+    if not hashable(fn):
+        return []
+    if fn in INPLACE_BINARY_OPS:
+        inplace_magic_name, non_inplace_op = INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
+            fn
+        ]
+        return [
+            MagicMethod(inplace_magic_name, is_inplace=True),
+            *magic_method_builtin_dispatch(non_inplace_op),
+        ]
+    elif fn in NON_INPLACE_BINARY_OPS:
+        magic_name, reverse_magic_name = NON_INPLACE_BINARY_OPS_TO_MAGIC_NAMES[
+            fn
+        ]
+        magic_methods = [MagicMethod(magic_name)]
+        if reverse_magic_name is not None:
+            magic_methods.append(
+                MagicMethod(reverse_magic_name, is_reverse=True)
+            )
+        return magic_methods
+    elif fn in UNARY_OPS:
+        magic_name = UNARY_OPS_TO_MAGIC_NAMES[fn]
+        return [MagicMethod(magic_name)]
+    return []
+
+
+def non_inplace_op_to_inplace_op(
+    fn: BinaryOp,
+) -> BinaryOp | None:
+    for inplace_op, (
+        _,
+        non_inplace_op,
+    ) in INPLACE_BINARY_OPS_TO_MAGIC_NAMES.items():
+        if fn is non_inplace_op:
+            return inplace_op
+    return None

python/paddle/jit/sot/utils/numpy_utils.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 numpy as np
+
+import paddle
+
+NUMPY_API_SUPPORTED_DICT = {
+    np.add: paddle.add,
+    np.subtract: paddle.subtract,
+    np.multiply: paddle.multiply,
+    np.divide: paddle.divide,
+    np.equal: paddle.equal,
+}

python/paddle/jit/sot/utils/paddle_api_config.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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 inspect
+
+import paddle
+
+
+def is_inplace_api(func):
+    inplace_apis = {paddle.static.setitem}
+    return func in inplace_apis
+
+
+def get_tensor_methods():
+    return [
+        member_name
+        for member_name, member in inspect.getmembers(paddle.pir.Value)
+        if inspect.isfunction(member) or inspect.ismethoddescriptor(member)
+    ]
+
+
+def get_paddle_api():
+    modules = [
+        paddle,
+        paddle.nn.functional,
+        paddle.nn.quant,
+        paddle.incubate.nn.functional,
+        paddle.linalg,
+        paddle.signal,
+        paddle.fft,
+        paddle.vision.ops,
+        paddle.metric,
+        paddle.geometric,
+    ]
+    distributed_apis = [
+        paddle.distributed.all_reduce,
+        paddle.distributed.shard_tensor,
+        paddle.distributed.reshard,
+        paddle.distributed.all_gather,
+        paddle.distributed.alltoall,
+        paddle.distributed.barrier,
+        paddle.distributed.recv,
+        paddle.distributed.send,
+        paddle.distributed.broadcast,
+        paddle.distributed.unshard_dtensor,
+        paddle.distributed.auto_parallel.api.dtensor_to_local,
+        paddle.distributed.auto_parallel.api.dtensor_from_local,
+        paddle.distributed.auto_parallel.api.moe_global_mesh_tensor,
+        paddle.distributed.auto_parallel.api.moe_sub_mesh_tensors,
+    ]
+    special_paddle_apis = [
+        paddle.tensor.fill_constant,
+        paddle.tensor.top_p_sampling,
+    ]
+    non_operator_related_apis = [
+        paddle.in_dynamic_mode,
+        paddle.save,
+        paddle.load,
+        paddle.get_cuda_rng_state,
+        paddle.set_rng_state,
+        paddle.set_cuda_rng_state,
+        paddle.get_rng_state,
+        paddle.set_default_dtype,
+        paddle.check_shape,
+        paddle.summary,
+        paddle.finfo,
+        paddle.iinfo,
+        paddle.enable_static,
+        paddle.disable_static,
+        paddle.is_grad_enabled,
+    ]
+    # TODO: users should not call static_apis, but we need to use, so add static_apis here temporary
+    static_apis = [paddle.static.setitem, paddle.static.accuracy]
+    paddle_api_list = []
+    for module in modules:
+        for fn_name in getattr(module, "__all__", []):
+            fn = getattr(module, fn_name)
+            if inspect.isfunction(fn):
+                paddle_api_list.append(fn)
+    return list(
+        set(special_paddle_apis)
+        | set(distributed_apis)
+        | set(static_apis)
+        | set(paddle_api_list) - set(non_operator_related_apis)
+    )
+
+
+paddle_api_list = get_paddle_api()
+
+# TODO(Aurelius84): It seems that we use it to judge 'in_paddle_module()'.
+# Bug what does 'is_paddle_module' really means? Is all paddle.xx sub module
+# considered as paddle module？
+paddle_api_module_prefix = {
+    "paddle.nn.functional",
+}
+
+break_graph_functions = set()
+break_graph_layer_classes = set()
+break_graph_tensor_method = {
+    'register_hook',
+    'numpy',
+    'clear_gradient',
+    'tolist',
+    'item',
+    # TODO: Browse all possible functions and make prior judgments.
+}
+
+not_supported_paddle_layer = {paddle.nn.RNN}
+
+
+def is_not_supported_paddle_layer(layer_class):
+    return layer_class in not_supported_paddle_layer
+
+
+def is_break_graph_tensor_methods(method_name):
+    return method_name in break_graph_tensor_method
+
+
+def add_break_graph_function(fn):
+    break_graph_functions.add(fn)
+
+
+def add_break_graph_layer_class(layer_class: type[paddle.nn.Layer]):
+    break_graph_layer_classes.add(layer_class)
+
+
+def is_directly_run_api(api):
+    from .utils import hashable
+
+    if not hashable(api):
+        return False
+    NATIVE_CODE_PURE_FUNCTIONS = {
+        paddle.base.libpaddle.is_compiled_with_avx,
+        paddle.base.libpaddle.is_compiled_with_cuda,
+        paddle.base.libpaddle.is_compiled_with_cudnn_frontend,
+        paddle.base.libpaddle.is_compiled_with_rocm,
+        paddle.base.libpaddle.is_compiled_with_custom_device,
+        paddle.base.libpaddle.is_compiled_with_ipu,
+        paddle.base.libpaddle.is_compiled_with_xpu,
+        paddle.base.libpaddle.is_compiled_with_mkldnn,
+        paddle.base.libpaddle.is_compiled_with_onednn,
+        paddle.base.libpaddle.is_compiled_with_nccl,
+        paddle.base.libpaddle.is_compiled_with_mpi,
+        paddle.base.libpaddle.is_compiled_with_mpi_aware,
+        paddle.base.libpaddle.is_compiled_with_cinn,
+        paddle.base.libpaddle.is_compiled_with_distribute,
+        paddle.base.libpaddle.is_compiled_with_brpc,
+        paddle.base.libpaddle.is_compiled_with_dist,
+        paddle.base.libpaddle.is_compiled_with_flagcx,
+    }
+
+    if hasattr(paddle.base.libpaddle, "get_device_properties"):
+        NATIVE_CODE_PURE_FUNCTIONS.add(
+            paddle.base.libpaddle.get_device_properties
+        )
+
+    return api in NATIVE_CODE_PURE_FUNCTIONS

python/paddle/jit/sot/utils/utils.py

@@ -0,0 +1,541 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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
+
+import builtins
+import copy
+import inspect
+import sys
+import time
+import types
+import weakref
+from collections import OrderedDict
+from collections.abc import Callable
+from contextlib import contextmanager
+from dataclasses import is_dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING, Any, TypeVar
+from weakref import WeakValueDictionary
+
+import numpy as np
+
+import paddle
+from paddle.jit.dy2static.utils import (
+    TransformOptions,
+    dataclass_as_dict,
+    dataclass_from_dict,
+)
+from paddle.utils import flatten, map_structure
+
+from .envs import (
+    ENV_SOT_LOG_LEVEL,
+    ENV_SOT_SPECIALIZED_DIM_NUMBERS,
+    ENV_STRICT_MODE,
+)
+from .paddle_api_config import (
+    break_graph_functions,
+    paddle_api_list,
+    paddle_api_module_prefix,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle._typing import NestedStructure
+
+T = TypeVar("T")
+T1 = TypeVar("T1")
+T2 = TypeVar("T2")
+T3 = TypeVar("T3")
+ConstTypes = (int, float, str, bool, type(None), bytes)
+
+
+class Singleton(type):
+    _instances: dict[Any, Any] = {}
+
+    def __call__(cls, *args: Any, **kwargs: Any):
+        if cls not in cls._instances:
+            cls._instances[cls] = super().__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+
+class NameGenerator:
+    def __init__(self, prefix):
+        self.counter = 0
+        self.prefix = prefix
+
+    def next(self):
+        name = self.prefix + str(self.counter)
+        self.counter += 1
+        return name
+
+    def match_name(self, name: str) -> bool:
+        return name.startswith(self.prefix)
+
+
+class SymbolRegistry:
+    def __init__(self):
+        self.symbol_generator = NameGenerator(prefix="___t_")
+        self.tmp_names_record = OrderedDict()
+        self.declared_symbols: set[str] = set()
+        self.symbol_table = {}
+
+    def next_symbol(self) -> str:
+        return self.symbol_generator.next()
+
+    def request_symbol(self, expr: str) -> str:
+        if expr in self.symbol_table:
+            return self.symbol_table[expr]
+        symbol = self.next_symbol()
+        self.symbol_table[expr] = symbol
+        return symbol
+
+    def gen_expr(self, expr: str, gen_expr_fn):
+        symbol = self.symbol_table[expr]
+        if symbol in self.declared_symbols:
+            return symbol
+        self.declared_symbols.add(symbol)
+        return f"({symbol} := ({gen_expr_fn()}))"
+
+
+_symbol_registry = SymbolRegistry()
+
+
+@contextmanager
+def switch_symbol_registry():
+    global _symbol_registry
+    original_registry = _symbol_registry
+    _symbol_registry = SymbolRegistry()
+    yield
+    _symbol_registry = original_registry
+
+
+def current_symbol_registry():
+    global _symbol_registry
+    return _symbol_registry
+
+
+class ResumeFnNameFactory(metaclass=Singleton):
+    def __init__(self) -> None:
+        self.gen = NameGenerator('resume_')
+
+    def next(self):
+        name = self.gen.next()
+        return name
+
+
+class SIRToCodeMap(metaclass=Singleton):
+    def __init__(self):
+        self._map = {}
+
+    def register(self, sir, code):
+        self._map[sir.name] = code
+
+    def get(self, sir):
+        return self._map.get(sir.name)
+
+
+def log(level, *args):
+    cur_level = ENV_SOT_LOG_LEVEL.get()
+    if level <= cur_level:
+        print(*args, end="", flush=True)
+
+
+def log_do(level, fn):
+    cur_level = ENV_SOT_LOG_LEVEL.get()
+    if level <= cur_level:
+        fn()
+
+
+def log_format(level, str, *args):
+    cur_level = ENV_SOT_LOG_LEVEL.get()
+    if level <= cur_level:
+        print(str.format(*args), end="", flush=True)
+
+
+def log_enabled(level):
+    return level <= ENV_SOT_LOG_LEVEL.get()
+
+
+@lru_cache
+def log_once(msg):
+    print(msg, flush=True)
+
+
+def no_eval_frame(func):
+    def no_eval_frame_func(*args, **kwargs):
+        old_cb = paddle.framework.core.set_eval_frame(None)
+        try:
+            retval = func(*args, **kwargs)
+        except:
+            raise
+        finally:
+            paddle.framework.core.set_eval_frame(old_cb)
+        return retval
+
+    return no_eval_frame_func
+
+
+def is_comprehensive_name(name):
+    return name in ["<listcomp>", "<dictcomp>", "<setcomp>", "<genexpr>"]
+
+
+def is_paddle_api(func):
+    if isinstance(func, paddle.nn.Layer):  # ignore all the classes
+        return False
+    if hasattr(func, "__self__"):  # ignore all the methods
+        return False
+    if inspect.isclass(
+        func
+    ):  # paddle.Tensor should not be wrapped, but how about other situations?
+        return False
+    return in_paddle_module(func) or func in paddle_api_list
+
+
+def already_unified_in_dynamic_and_static_graph(fn):
+    if is_paddle_api(fn):
+        return True
+    return not TransformOptions.check_fn_need_transform(
+        fn, TransformOptions.ToStaticMode.SOT
+    )
+
+
+def need_capture_control_flow(fn):
+    return TransformOptions.check_fn_need_capture_control_flow(fn)
+
+
+def is_builtin_fn(fn):
+    special_builtin_fns = [weakref.ref]
+    if fn in special_builtin_fns:
+        return True
+    if isinstance(fn, types.BuiltinFunctionType):
+        return True
+    for member_name, member in inspect.getmembers(builtins):
+        if member is fn and isinstance(member, type):
+            return True
+    return False
+
+
+def in_paddle_module(func):
+    if hasattr(func, "__module__"):
+        module_str = func.__module__
+        if module_str is None:
+            return False
+        log(5, "find paddle function with __module__: ", module_str, "\n")
+        if hasattr(func, "__name__"):
+            log(
+                5, "                     with __name__  : ", func.__name__, "\n"
+            )
+        log(5, "                     with results   : ")
+        for prefix in paddle_api_module_prefix:
+            if module_str.startswith(prefix):
+                log(5, " True\n")
+                return True
+    log(5, " False\n")
+    return False
+
+
+def is_break_graph_api(func):
+    return func in break_graph_functions
+
+
+def is_namedtuple_class(cls):
+    if not inspect.isclass(cls):
+        return False
+    if not issubclass(cls, tuple):
+        return False
+    # The signature created by nametuple function
+    namedtuple_attrs = {"_make", "_asdict", "_fields", "_replace"}
+    cls_attrs = set(dir(cls))
+    return namedtuple_attrs.issubset(cls_attrs)
+
+
+def map_if(
+    *structures: NestedStructure[T1],
+    pred: Callable[[T1], bool],
+    true_fn: Callable[[T1], T2],
+    false_fn: Callable[[T1], T3],
+) -> NestedStructure[T2 | T3]:
+    def replace(*args):
+        if pred(*args):
+            return true_fn(*args)
+        return false_fn(*args)
+
+    return map_structure(replace, *structures)
+
+
+def flatten_extend(structure):
+    for item in flatten(structure):
+        if isinstance(item, slice):
+            yield item.start
+            yield item.stop
+            yield item.step
+        else:
+            yield item
+
+
+def map_if_extend(structure, pred, true_fn, false_fn):
+    """support extended structures like slice and SliceVariable"""
+
+    def wrapped_pred(x):
+        if isinstance(x, slice):
+            return True
+        if is_dataclass(x) and not isinstance(x, type):
+            return True
+        return pred(x)
+
+    def wrapped_true_fn(x):
+        if isinstance(x, (slice)):
+            l = [x.start, x.stop, x.step]
+            l = map_if_extend(l, pred, true_fn, false_fn)
+            return slice(*l)
+
+        if is_dataclass(x) and not isinstance(x, type):
+            dt_dict = dataclass_as_dict(x)
+            dt_dict = map_if_extend(dt_dict, pred, true_fn, false_fn)
+            return dataclass_from_dict(type(x), dt_dict)
+
+        return true_fn(x)
+
+    return map_if(
+        structure, pred=wrapped_pred, true_fn=wrapped_true_fn, false_fn=false_fn
+    )
+
+
+def count_if(*structures, pred):
+    def is_true(*args):
+        if pred(*args):
+            return 1
+        return 0
+
+    return sum(flatten(map_structure(is_true, *structures)))
+
+
+class Cache:
+    def __init__(self, weak=False, copy=False):
+        if not weak:
+            self.cache = {}
+        else:
+            self.cache = WeakValueDictionary()
+        self.hit_num = 0
+        self.copy = copy
+
+    def __call__(self, *args, **kwargs):
+        cache_key = self.key_fn(*args, **kwargs)
+        if not hashable(cache_key):
+            return self.value_fn(*args, **kwargs)
+        if cache_key in self.cache:
+            log(5, "cache hit: ", cache_key, "\n")
+            self.hit_num += 1
+            cache_item = self.cache[cache_key]
+            if self.copy:
+                cache_item = copy.deepcopy(cache_item)
+            return cache_item
+        value = self.value_fn(*args, **kwargs)
+        self.cache[cache_key] = value
+        return value
+
+    def clear(self):
+        self.cache.clear()
+        self.hit_num = 0
+
+    def key_fn(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def value_fn(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+def execute_time(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        print("Execute time:", execution_time)
+        return result
+
+    return wrapper
+
+
+def meta_str(shape, dtype, stop_gradient):
+    return f"(shape: {shape}, dtype: {dtype}, stop_gradient: {stop_gradient})"
+
+
+def is_strict_mode():
+    return ENV_STRICT_MODE.get()
+
+
+def list_find_index_by_id(li: list[Any], item: Any) -> int:
+    return [id(it) for it in li].index(id(item))
+
+
+def list_contain_by_id(li: list[Any], item: Any) -> int:
+    return id(item) in [id(it) for it in li]
+
+
+def get_unbound_method(obj, name):
+    # TODO(dev): Consider the case of patching methods to instances
+    return getattr(obj.__class__, name)
+
+
+class SotUndefinedVar(metaclass=Singleton):
+    pass
+
+
+def hashable(obj):
+    try:
+        hash(obj)
+        return True
+    except TypeError as e:
+        return False
+
+
+def printable(obj):
+    try:
+        str(obj)
+        return True
+    except Exception as e:
+        return False
+
+
+class StepInfo:
+    BACK_TRACE_STEPS = 20
+
+    def __init__(self):
+        self.step_count = -1
+
+    def need_back_trace(self):
+        return self.step_count < self.BACK_TRACE_STEPS
+
+
+class StepInfoManager(metaclass=Singleton):
+    def __init__(self):
+        self.step_record = {}
+        self.current_code = None
+        self.current_step_info = None
+
+    @contextmanager
+    def step_guard(self, code):
+        try:
+            old_code = self.current_code
+            old_info = self.current_step_info
+
+            self.current_code = code
+            if code not in self.step_record:
+                self.step_record[code] = StepInfo()
+            self.current_step_info = self.step_record[code]
+
+            self.current_step_info.step_count += 1
+            yield
+        finally:
+            self.current_code = old_code
+            self.current_step_info = old_info
+
+    @property
+    def need_back_trace(self):
+        return (
+            self.current_step_info is not None
+            and self.current_step_info.need_back_trace()
+        )
+
+    @property
+    def current_step(self):
+        return self.current_step_info.step_count
+
+    def clear(self):
+        self.step_record.clear()
+        self.current_code = None
+        self.current_step = -1
+
+
+def get_api_fullname(api):
+    api_name = api.__name__
+    module_str = api.__module__
+    while len(module_str) > 0:
+        if module_str not in sys.modules:
+            return api_name
+        module = sys.modules[module_str]
+        if hasattr(module, api_name):
+            return module_str + "." + api_name
+        module_str = module_str.rpartition(".")[0]
+    return None
+
+
+def get_numpy_ufuncs():
+    ufuncs = [
+        ufunc
+        for _, ufunc in inspect.getmembers(
+            np, lambda member: isinstance(member, np.ufunc)
+        )
+    ]
+    unary_ufuncs = filter(lambda ufunc: ufunc.nin == 1, ufuncs)
+    binary_ufuncs = filter(lambda ufunc: ufunc.nin == 2, ufuncs)
+    return list(unary_ufuncs), list(binary_ufuncs)
+
+
+def do_until_stop_iteration(fn: Callable[[], T]) -> list[T]:
+    from paddle.jit.sot.utils.exceptions import SotCapturedStopIteration
+
+    res = []
+    while True:
+        try:
+            res.append(fn())
+        except SotCapturedStopIteration:
+            break
+    return res
+
+
+def update_list_inplace(
+    original_list: list[T], new_contents: list[T]
+) -> list[T]:
+    original_list.clear()
+    original_list.extend(new_contents)
+    return original_list
+
+
+def get_obj_stable_repr(obj) -> str:
+    if hasattr(obj, '__qualname__'):
+        return obj.__qualname__
+    if hasattr(obj, '__name__'):
+        return obj.__name__
+
+    class_name = obj.__class__.__name__
+
+    # If module is available and not __main__, include it
+    if hasattr(obj, "__class__") and hasattr(obj.__class__, "__module__"):
+        module = obj.__class__.__module__
+        if module not in ("__main__", "builtins"):
+            return f"{module}.{class_name}()"
+
+    return f"{class_name}()"
+
+
+def get_min_non_specialized_number() -> int:
+    specialized_dim_numbers_raw_str = (
+        ENV_SOT_SPECIALIZED_DIM_NUMBERS.get().lower()
+    )
+    assert specialized_dim_numbers_raw_str in [
+        "no",
+        "0",
+        "01",
+    ], f"Unsupported specialized_dim_numbers: {specialized_dim_numbers_raw_str}"
+    to_min_non_specialized_number = {
+        # specialized numbers, minimum non-specialized number
+        "no": 0,
+        "0": 1,
+        "01": 2,
+    }
+    return to_min_non_specialized_number[specialized_dim_numbers_raw_str]