chore: import upstream snapshot with attribution

python/paddle/jit/sot/__init__.py

@@ -0,0 +1,24 @@
+# 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 . import (
+    profiler as profiler,
+    psdb,  # noqa: F401
+)
+from .opcode_translator.breakpoint import (  # noqa: F401
+    BM,
+    add_breakpoint,
+    add_event,
+)
+from .translate import symbolic_translate  # noqa: F401

python/paddle/jit/sot/infer_meta.py

@@ -0,0 +1,656 @@
+# 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 copy
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Any, TypeVar
+
+import paddle
+from paddle.base.data_feeder import convert_dtype
+from paddle.base.framework import convert_nptype_to_datatype_or_vartype
+from paddle.base.unique_name import (
+    UniqueNameGenerator,
+    guard as UniqueNameGuard,
+)
+from paddle.distributed.auto_parallel.placement_type import (
+    get_shard_spec,
+    to_placements,
+)
+from paddle.distributed.auto_parallel.static.dist_input_spec import (
+    DistributedInputSpec,
+)
+from paddle.distributed.auto_parallel.static.utils import (
+    convert_to_dims_mapping,
+)
+from paddle.jit.dy2static.utils import (
+    ALREADY_D2S,
+    extract_tensor_dynamic_dims,
+    graph_tracing_guard,
+)
+from paddle.pir import is_fake_value
+from paddle.static import InputSpec
+from paddle.utils import flatten, is_sequence
+
+from .symbolic_shape.symbolic_value import SymbolicInt
+from .utils import (
+    Cache,
+    Singleton,
+    get_min_non_specialized_number,
+    map_if_extend,
+    meta_str,
+)
+from .utils.exceptions import BreakGraphError, NullMetaBreak
+
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+DynamicSymbolT = TypeVar("DynamicSymbolT")
+SOT_INFER_META_INNER_VAR = "___SOT_INFER_META_INNER_VAR"
+
+
+class DistInfo:
+    def __init__(self, mesh=None, dims_mapping=None, local_shape=None):
+        self.mesh = mesh
+        self.dims_mapping = dims_mapping
+        self.local_shape = local_shape
+
+    @staticmethod
+    def from_tensor(tensor: paddle.Tensor) -> DistInfo:
+        assert isinstance(tensor, paddle.Tensor) and tensor.is_dist(), (
+            f"Expect a Tensor, but got a {type(tensor)}."
+        )
+
+        mesh = tensor.process_mesh
+        sharding_specs = get_shard_spec(
+            mesh, tensor.placements, len(tensor.shape)
+        )
+        dims_mapping = convert_to_dims_mapping(sharding_specs, mesh)
+        local_shape = tensor._local_value().shape
+        return DistInfo(mesh, dims_mapping, local_shape)
+
+    @staticmethod
+    def from_value(value: paddle.pir.Value) -> DistInfo:
+        assert isinstance(value, paddle.pir.Value) and value.is_dist(), (
+            f"Expect a Value, but got a {type(value)}."
+        )
+        return DistInfo(
+            value.dist_attr().process_mesh,
+            value.dist_attr().dims_mapping,
+            value._local_shape,
+        )
+
+    def __deepcopy__(self, memo):
+        return DistInfo(
+            mesh=copy.deepcopy(self.mesh),
+            dims_mapping=copy.deepcopy(self.dims_mapping),
+            local_shape=copy.deepcopy(self.local_shape),
+        )
+
+    def __repr__(self) -> str:
+        return f"DistInfo(mesh={self.mesh}, dims_mapping={self.dims_mapping}, local_shape={self.local_shape})"
+
+
+class MetaInfoOrNull:
+    def __init__(self, meta: MetaInfo | None):
+        self.meta = meta
+
+    @staticmethod
+    def null():
+        return MetaInfoOrNull(None)
+
+    def is_null(self):
+        return self.meta is None
+
+    def unwrap_or_breakgraph(self):
+        if self.meta is None:
+            raise BreakGraphError(NullMetaBreak())
+        return self.meta
+
+    def unwrap_unsafe(self):
+        assert self.meta is not None, "MetaInfo is None"
+        return self.meta
+
+    def with_dynamic_axes(
+        self, name: str, dynamic_axes: list[int]
+    ) -> MetaInfoOrNull:
+        if self.meta is None:
+            return MetaInfoOrNull.null()
+        return self.meta.with_dynamic_axes(name, dynamic_axes).wrap()
+
+    def to_input_spec(self):
+        if self.meta is None:
+            return None
+        return self.meta.to_input_spec()
+
+    def guard_str(self):
+        if self.meta is None:
+            return "(Null)"
+        return self.meta.guard_str()
+
+    def __deepcopy__(self, memo):
+        if self.meta is None:
+            return MetaInfoOrNull(None)
+        return MetaInfoOrNull(copy.deepcopy(self.meta))
+
+    @staticmethod
+    def mix_axes(axes1: list[int], axes2: list[int]) -> list[int]:
+        return sorted(set(axes1 + axes2))
+
+    @staticmethod
+    def from_tensor(
+        tensor: paddle.Tensor, *, dynamic_axes: list[int] | None = None
+    ) -> MetaInfoOrNull:
+        if not tensor._is_dense_tensor_hold_allocation():
+            return MetaInfoOrNull.null()
+        assert isinstance(tensor, paddle.Tensor), (
+            "Expect a Tensor, but got a Value."
+        )
+
+        assert -1 not in tensor.shape, (
+            "Tensor shape should not contain -1, maybe you pass a Value to from_tensor"
+        )
+        user_specified_dynamic_axes = extract_tensor_dynamic_dims(tensor)
+        dynamic_axes = dynamic_axes or []
+        dynamic_axes = MetaInfoOrNull.mix_axes(
+            dynamic_axes, list(user_specified_dynamic_axes)
+        )
+        shape = [
+            SymbolicInt(dim) if i in dynamic_axes else dim
+            for i, dim in enumerate(tensor.shape)
+        ]
+        if tensor.is_dist():
+            dist_info = DistInfo.from_tensor(tensor)
+        else:
+            dist_info = None
+        return MetaInfo(
+            shape,
+            tensor.dtype,
+            tensor.stop_gradient,
+            tensor.name,
+            tensor.persistable,
+            tensor.type,
+            tensor.place,
+            None,
+            dist_info=dist_info,
+        ).wrap()
+
+    @staticmethod
+    def from_numpy(
+        nparray: npt.NDArray[Any], *, dynamic_axes: list[int] | None = None
+    ) -> MetaInfoOrNull:
+        dtype = convert_nptype_to_datatype_or_vartype(nparray.dtype)
+        dynamic_axes = dynamic_axes or []
+        shape = [
+            SymbolicInt() if i in dynamic_axes else dim
+            for i, dim in enumerate(nparray.shape)
+        ]
+        return MetaInfo(
+            shape,
+            dtype,
+            True,  # stop_gradient
+            None,
+            None,  # persistable
+            None,
+            None,
+            None,
+            dist_info=None,
+        ).wrap()
+
+    @staticmethod
+    def from_value(value) -> MetaInfoOrNull:
+        if is_fake_value(value):
+            return MetaInfoOrNull.null()
+        name = SOT_INFER_META_INNER_VAR
+        shape = [SymbolicInt() if dim == -1 else dim for dim in value.shape]
+        for dim in shape:
+            if isinstance(dim, int):
+                assert dim >= 0, (
+                    "Dimensions must be non-negative integers or SymbolicInt. "
+                    f"Encountered value {dim} in shape {shape}."
+                )
+
+        if isinstance(value, paddle.pir.Value) and value.is_dist():
+            dist_info = DistInfo.from_value(value)
+        else:
+            dist_info = None
+        return MetaInfo(
+            shape,
+            value.dtype,
+            value.stop_gradient,
+            name,
+            value.persistable,
+            None,  # type is not a unified attribute in dygraph and static mode.
+            None,  # We can't infer the right place in compile time.
+            None,  # there's no spec_name specified when from_value.
+            dist_info=dist_info,
+        ).wrap()
+
+    def __repr__(self):
+        if self.meta is None:
+            return "MetaInfoOrNull(None)"
+        return f"MetaInfoOrNull({self.meta})"
+
+    def __eq__(self, other):
+        if self.meta is None:
+            return other.meta is None
+        if other.meta is None:
+            return False
+        return self.meta == other.meta
+
+    def __hash__(self):
+        if self.meta is None:
+            return hash(None)
+        return hash(self.meta)
+
+
+class MetaInfo:
+    shape: list[int | SymbolicInt]
+
+    def __init__(
+        self,
+        shape,
+        dtype,
+        stop_gradient,
+        name,
+        persistable,
+        type,
+        place,
+        spec_name=None,
+        dist_info=None,
+    ):
+        assert -1 not in shape, (
+            "NOTE: Shape should not contain -1, consider convert it to SymbolicInt."
+        )
+        self.name = name
+        self.persistable = persistable
+        self.type = type
+        self.place = place
+        self.shape = shape
+        self.dtype = dtype
+        self.stop_gradient = stop_gradient
+        self.dist_info = dist_info
+        self.spec_name = spec_name
+
+    def wrap(self):
+        return MetaInfoOrNull(self)
+
+    def shape_with_special_symbol(
+        self, dynamic_symbol: DynamicSymbolT = -1
+    ) -> list[int | DynamicSymbolT]:
+        return [
+            dynamic_symbol if isinstance(dim, SymbolicInt) else dim
+            for dim in self.shape
+        ]
+
+    def with_dynamic_axes(self, name: str, dynamic_axes: list[int]) -> MetaInfo:
+        mixed_dynamic_axes = MetaInfoOrNull.mix_axes(
+            self.dynamic_axes, dynamic_axes
+        )
+        # NOTE(SigureMo): Make sure create a new shape list with dynamic axes.
+        # We will create a new shape list variable lazily in the future.
+        shape = [
+            (
+                SymbolicInt(dim)
+                if (
+                    i in mixed_dynamic_axes and not isinstance(dim, SymbolicInt)
+                )
+                else dim
+            )
+            for i, dim in enumerate(self.shape)
+        ]
+        return MetaInfo(
+            shape,
+            self.dtype,
+            self.stop_gradient,
+            self.name,
+            self.persistable,
+            self.type,
+            self.place,
+            spec_name=name,
+            dist_info=self.dist_info,
+        )
+
+    @property
+    def dynamic_axes(self):
+        return [
+            i
+            for i, dim in enumerate(self.shape)
+            if isinstance(dim, SymbolicInt)
+        ]
+
+    def is_inner_var(self):
+        return self.name == SOT_INFER_META_INNER_VAR
+
+    def is_dynamic_shape(self):
+        """
+        if SymbolicInt in shape, return True
+        else: return False
+        """
+        return len(self.dynamic_axes) > 0
+
+    def to_input_spec(self) -> DistributedInputSpec | ConstrainedInputSpec:
+        shape = self.shape_with_special_symbol(None)
+        if self.dist_info is not None:
+            placements = to_placements(
+                self.dist_info.dims_mapping, self.dist_info.mesh
+            )
+            return DistributedInputSpec(
+                shape,
+                dtype=self.dtype,
+                stop_gradient=self.stop_gradient,
+                mesh=self.dist_info.mesh,
+                placements=placements,
+                local_shape=self.dist_info.local_shape,
+            )
+        else:
+            return ConstrainedInputSpec(
+                self.dynamic_axes,
+                shape,
+                dtype=self.dtype,
+                name=self.spec_name,
+                stop_gradient=self.stop_gradient,
+            )
+
+    def guard_str(self):
+        shape = self.shape_with_special_symbol(SymbolicInt())
+        return f"({shape}, {self.dtype}, {self.stop_gradient})"
+
+    def __deepcopy__(self, memo):
+        return MetaInfo(
+            list(self.shape),
+            self.dtype,
+            self.stop_gradient,
+            self.name,
+            self.persistable,
+            self.type,
+            self.place,
+            self.spec_name,
+            dist_info=copy.deepcopy(self.dist_info),
+        )
+
+    def __repr__(self):
+        return meta_str(self.shape, self.dtype, self.stop_gradient)
+
+    def __eq__(self, meta):
+        return (
+            self.shape == meta.shape
+            and self.dtype == meta.dtype
+            and self.stop_gradient == meta.stop_gradient
+        )
+
+    def __hash__(self):
+        return hash((tuple(self.shape), self.dtype, self.stop_gradient))
+
+
+class VariableCreator(metaclass=Singleton):
+    """
+    We use the static graph Variable to infer the meta information of Tensor.
+    This singleton class is used to create Variable for infer meta.
+    """
+
+    def __init__(self):
+        self.var_name_generator = UniqueNameGenerator(SOT_INFER_META_INNER_VAR)
+        self.var_cache = {}
+        self.main_program = paddle.static.Program()
+        self.startup_program = paddle.static.Program()
+
+    def gen_name(self, meta_or_null: MetaInfoOrNull):
+        if meta_or_null.is_null():
+            return "null"
+        meta = meta_or_null.unwrap_unsafe()
+        name = f"{meta.dtype}_{meta.stop_gradient}_"
+        name += "_".join(map(str, meta.shape))
+        return name
+
+    def create_var(self, meta_or_null: MetaInfoOrNull):
+        if meta_or_null.is_null():
+            return None
+        meta = meta_or_null.unwrap_unsafe()
+        shape = meta.shape_with_special_symbol(-1)
+
+        with paddle.static.program_guard(
+            self.main_program, self.startup_program
+        ):
+            var = paddle.static.input.data(
+                name=self.gen_name(meta.wrap()),
+                shape=shape,
+                dtype=convert_dtype(meta.dtype),
+            )
+            var.stop_gradient = meta.stop_gradient
+
+            if meta.dist_info is not None:
+                mesh = meta.dist_info.mesh
+                placements = to_placements(meta.dist_info.dims_mapping, mesh)
+                var = paddle._pir_ops.shard_tensor(var, mesh, placements)
+                var.stop_gradient = meta.stop_gradient
+        assert not isinstance(var, paddle.Tensor), (
+            "Expect a Variable, but got a Tensor."
+        )
+        return var
+
+    def get_variable(self, meta: MetaInfoOrNull, without_cache=False):
+        var_feature_name = self.gen_name(meta)
+        if without_cache:
+            return self.create_var(meta)
+        if var_feature_name not in self.var_cache:
+            self.var_cache[var_feature_name] = self.create_var(meta)
+        return self.var_cache[var_feature_name]
+
+    def infer_meta(self, func, *args, **kwargs):
+        with (
+            paddle.base.framework._dygraph_guard(None),
+            UniqueNameGuard(self.var_name_generator),
+        ):
+            if func is paddle.distributed.shard_tensor:
+                args, kwargs = (
+                    convert_meta_to_variable(args, without_cache=True),
+                    convert_meta_to_variable(kwargs, without_cache=True),
+                )
+            else:
+                args, kwargs = (
+                    convert_meta_to_variable(args),
+                    convert_meta_to_variable(kwargs),
+                )
+
+            graph_tracing_context_manager = nullcontext()
+            with paddle.static.program_guard(
+                self.main_program, self.startup_program
+            ):
+                if isinstance(func, str):
+                    # TODO(Aurelius84): Is length of args always greater than 0?
+                    # Do we need add condition check here?
+                    func = getattr(args[0], func)
+                    args = args[1:]
+                if hasattr(func, ALREADY_D2S):
+                    graph_tracing_context_manager = graph_tracing_guard(
+                        self.main_program
+                    )
+                with graph_tracing_context_manager:
+                    out = func(*args, **kwargs)
+        return convert_variable_to_meta_info(out)
+
+
+def convert_meta_to_variable(args, without_cache=False):
+    return map_if_extend(
+        args,
+        pred=lambda x: isinstance(x, MetaInfoOrNull),
+        true_fn=lambda x: VariableCreator().get_variable(
+            x, without_cache=without_cache
+        ),
+        false_fn=lambda x: x,
+    )
+
+
+def convert_meta_to_input_spec(args):
+    return map_if_extend(
+        args,
+        pred=lambda x: isinstance(x, MetaInfoOrNull),
+        true_fn=lambda x: x.to_input_spec(),
+        # TODO(xiongkun): can x be tensor ?
+        false_fn=lambda x: (
+            paddle.static.InputSpec.from_tensor(x)
+            if isinstance(x, paddle.Tensor)
+            else x
+        ),
+    )
+
+
+def convert_variable_to_meta_info(args):
+    return map_if_extend(
+        args,
+        pred=lambda x: isinstance(x, paddle.pir.Value),
+        true_fn=lambda x: MetaInfoOrNull.from_value(x),
+        false_fn=lambda x: x,
+    )
+
+
+def infer_meta(func, *args, **kwargs):
+    fn = SpecialInferMeta().get_infermeta_fn(func)
+    if fn:
+        return fn(*args, **kwargs)
+    return VariableCreator().infer_meta(func, *args, **kwargs)
+
+
+def infer_meta_for_layer(layer, *args, **kwargs):
+    assert isinstance(layer, paddle.nn.Layer), (
+        f"Expect a Layer, but got {layer}."
+    )
+    layer = paddle.jit.to_static(layer, full_graph=True)
+
+    args_, kwargs_ = convert_meta_to_input_spec((args, kwargs))
+
+    (
+        concrete_program,
+        partial_program_layer,
+    ) = layer.forward.get_concrete_program(*args_, **kwargs_)
+
+    output_values = partial_program_layer._outputs.var_list
+
+    out = partial_program_layer._restore_out(
+        [
+            x
+            for x in paddle.utils.flatten(
+                convert_variable_to_meta_info(output_values)
+            )
+            if isinstance(x, MetaInfoOrNull)
+        ]
+    )
+    layer.forward.rollback()
+    return out
+
+
+def ast_infer_meta(static_function, *args, **kwargs):
+    args_, kwargs_ = convert_meta_to_input_spec((args, kwargs))
+
+    (
+        concrete_program,
+        partial_program_layer,
+    ) = static_function.get_concrete_program(*args_, **kwargs_)
+
+    out = partial_program_layer._restore_out(
+        [
+            x
+            for x in paddle.utils.flatten(
+                convert_variable_to_meta_info(concrete_program.outputs)
+            )
+            if isinstance(x, MetaInfoOrNull)
+        ]
+    )
+
+    return out
+
+
+class SpecialInferMeta(metaclass=Singleton):
+    """
+    There are some functions that cannot be inferred directly through static graph,
+    and need to be implemented manually. This class is used to implement infer meta
+    for these functions.
+    """
+
+    def __init__(self):
+        pass
+
+    def get_infermeta_fn(self, fn):
+        try:
+            funcname = fn.__name__
+            return getattr(self, f"infermeta_{funcname}")
+        except:
+            pass
+        return None
+
+    def infermeta_grad(
+        self,
+        outputs,
+        inputs,
+        grad_outputs=None,
+        retain_graph=None,
+        create_graph=False,
+        only_inputs=True,
+        allow_unused=False,
+        no_grad_vars=None,
+    ):
+        if not is_sequence(inputs):
+            inputs = [inputs]
+        return inputs
+
+
+class InferMetaCache(Cache, metaclass=Singleton):
+    def __init__(self):
+        super().__init__(copy=True)
+
+    def key_fn(
+        self, func, *args, **kwargs
+    ):  # args & kwargs have transformed to MetaInfo
+        return (
+            func,
+            tuple(flatten(args)),
+            tuple(kwargs.keys()),
+            tuple(flatten(kwargs)),
+        )
+
+    def value_fn(self, func, *args, **kwargs):
+        return infer_meta(func, *args, **kwargs)
+
+
+class LayerInferMetaCache(Cache, metaclass=Singleton):
+    def __init__(self):
+        super().__init__(copy=True)
+
+    def key_fn(self, layer, *args, **kwargs):
+        params = [
+            MetaInfoOrNull.from_tensor(x)
+            for x in layer.parameters(include_sublayers=True)
+        ]
+        return (
+            layer,
+            tuple(params),
+            tuple(flatten(args)),
+            tuple(kwargs.keys()),
+            tuple(flatten(kwargs)),
+        )
+
+    def value_fn(self, layer, *args, **kwargs):
+        return infer_meta_for_layer(layer, *args, **kwargs)
+
+
+class ConstrainedInputSpec(InputSpec):
+    def __init__(self, dynamic_axes: list[int], *args, **kwargs):
+        self.ranges: list[
+            tuple[int, int | None, int | None]
+        ] = []  # (idx of dim, min, max)
+        super().__init__(*args, **kwargs)
+        min_non_specialized_number = get_min_non_specialized_number()
+        for i in dynamic_axes:
+            self.ranges.append((i, min_non_specialized_number, None))

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

@@ -0,0 +1,18 @@
+# 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 .eval_frame_callback import eval_frame_callback  # noqa: F401
+from .skip_files import setup_skip_files
+
+setup_skip_files()

python/paddle/jit/sot/opcode_translator/breakpoint.py

@@ -0,0 +1,180 @@
+# 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 traceback
+from dataclasses import dataclass
+
+from ..opcode_translator.instruction_utils import instrs_info
+from ..utils import Singleton, log
+from .executor.opcode_executor import OpcodeExecutorBase
+
+# this file is a debug utils files for quick debug
+# >>> sot.add_breakpoint(file, line)
+# >>> sot.remove_breakpoint(file, line)
+
+
+@dataclass
+class Breakpoint:
+    file: str
+    line: int
+    co_name: str
+    offset: int
+
+    def __hash__(self):
+        return hash((self.file, self.line, self.co_name, self.offset))
+
+
+class BreakpointManager(metaclass=Singleton):
+    def __init__(self):
+        self.breakpoints = set()
+        self.executors = OpcodeExecutorBase.call_stack
+        self.activate = 0
+        self.record_event = []
+
+    def clear_event(self, event):
+        self.record_event.clear()
+
+    def add_event(self, event):
+        """
+        event in ['All' ,'FallbackError', 'BreakGraphError', 'InnerError']
+        """
+        self.record_event.append(event)
+
+    def add(self, file, line, co_name=None, offset=None):
+        log(1, f"add breakpoint at {file}:{line}\n")
+        self.breakpoints.add(Breakpoint(file, line, co_name, offset))
+
+    def addn(self, *lines):
+        """
+        called inside a executor. add a list of line number in current file.
+        """
+        if not isinstance(lines, (list, tuple)):
+            lines = [lines]
+        for line in lines:
+            file = self.cur_exe.vframe.code.co_filename
+            self.add(file, line)
+
+    def clear(self):
+        self.breakpoints.clear()
+
+    def hit(self, file, line, co_name, offset):
+        if Breakpoint(file, line, None, None) in self.breakpoints:
+            return True
+        if Breakpoint(file, line, co_name, offset) in self.breakpoints:
+            return True
+        return False
+
+    def locate(self, exe):
+        for i, _e in enumerate(self.executors):
+            if _e is exe:
+                self.activate = i
+                return
+        raise RuntimeError("Not found executor.")
+
+    def up(self):
+        if self.activate == 0:
+            return
+        self.activate -= 1
+        print("current function is: ", self.cur_exe.vframe.code.co_name)
+
+    def down(self):
+        if self.activate >= len(self.executors) - 1:
+            return
+        self.activate += 1
+        print("current function is: ", self.cur_exe.vframe.code.co_name)
+
+    def opcode(self, cur_exe=None):
+        if cur_exe is None:
+            cur_exe = self.cur_exe
+        instr = cur_exe._instructions[cur_exe.vframe.lasti - 1]
+        message = f"[Translate {cur_exe}] (line {cur_exe._current_line:>3}) {instr.opname:<12} {instr.argval}, stack is {cur_exe._stack}\n"
+        return message
+
+    def bt(self):
+        """
+        display all inline calls: backtrace.
+        """
+        for exe in self.executors:
+            lines, _ = inspect.getsourcelines(exe.vframe.code)
+            print(
+                "  "
+                + exe.vframe.code.co_filename
+                + f"({exe._current_line})"
+                + f"{exe.vframe.code.co_name}()"
+            )
+            print(f"-> {lines[0].strip()}")
+            print(f"-> {self.opcode(exe)}")
+        pass
+
+    def on_event(self, event):
+        if "All" in self.record_event or event in self.record_event:
+            print("event captured.")
+            self.activate = len(self.executors) - 1
+            breakpoint()  # noqa: T100
+
+    def _dis_source_code(self):
+        cur_exe = self.executors[self.activate]
+        lines, start_line = inspect.getsourcelines(cur_exe.vframe.code)
+        cur_line = cur_exe._current_line
+        lines[cur_line - start_line + 1 : cur_line - start_line + 1] = (
+            "  ^^^^^ HERE  \n"
+        )
+        print("\033[31mSource Code is: \033[0m")
+        print("".join(lines))
+
+    def dis(self, range=5):
+        """
+        display all instruction code and source code.
+        """
+        print("displaying debug info...")
+        cur_exe = self.cur_exe
+        print(self._dis_source_code())
+
+        print(f"\n{cur_exe.vframe.code}")
+        lasti = cur_exe.vframe.lasti
+        instr_str = instrs_info(
+            cur_exe._instructions, lasti - 1, range, want_str=True
+        )
+        print(instr_str)
+
+    @property
+    def cur_exe(self):
+        exe = self.executors[self.activate]
+        return exe
+
+    def sir(self):
+        """
+        display sir in a page.
+        """
+        print("displaying sir...")
+        self.cur_exe.print_sir()
+
+    def pe(self, e):
+        """
+        print exception.
+        """
+        lines = traceback.format_tb(e.__traceback__)
+        print("".join(lines))
+
+
+def add_breakpoint(file, line, co_name=None, offset=None):
+    BM.add(file, line, co_name, offset)
+
+
+def add_event(event):
+    BM.add_event(event)
+
+
+BM = BreakpointManager()

python/paddle/jit/sot/opcode_translator/custom_code.py

@@ -0,0 +1,25 @@
+# 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
+
+from typing import TYPE_CHECKING, NamedTuple
+
+if TYPE_CHECKING:
+    from types import CodeType
+
+
+class CustomCode(NamedTuple):
+    code: CodeType | None
+    disable_eval_frame: bool

python/paddle/jit/sot/opcode_translator/eval_frame_callback.py

@@ -0,0 +1,92 @@
+# 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 dis
+from typing import TYPE_CHECKING
+
+from ..profiler import EventGuard
+from ..utils import log_do, log_enabled, log_format
+from .executor.executor_cache import OpcodeExecutorCache
+
+if TYPE_CHECKING:
+    from .custom_code import CustomCode
+
+
+def print_locals(frame):
+    local_key = [
+        key for key in frame.f_locals.keys() if not key.startswith("__")
+    ]
+    print(
+        f"[eval_frame_callback] {frame.f_code.co_name} with locals {local_key}"
+    )
+    print(
+        f"[eval_frame_callback] {' ' * len(frame.f_code.co_name)} with cellvars + freevars:  {frame.f_code.co_cellvars + frame.f_code.co_freevars}"
+    )
+
+    def convert_obj(obj):
+        import paddle
+
+        if isinstance(obj, paddle.Tensor):
+            return "Tensor(" + str(obj.shape) + ")"
+        if isinstance(obj, list):
+            return [convert_obj(i) for i in obj]
+        return obj
+
+    for key in local_key:
+        print(
+            f"[eval_frame_callback] {' ' * len(frame.f_code.co_name)} {key} = {convert_obj(frame.f_locals[key])}"
+        )
+
+
+def eval_frame_callback(frame, **kwargs) -> CustomCode:
+    with EventGuard(
+        f"eval_frame_callback: {frame.f_code.co_name}", event_level=2
+    ):
+        # A quick way to check if the log is enabled
+        need_log = log_enabled(1)
+        if need_log:
+            log_format(
+                2,
+                "[eval_frame_callback] start to translate: {}\n",
+                frame.f_code,
+            )
+            log_do(4, lambda: print_locals(frame))
+
+            log_format(
+                3,
+                "[eval_frame_callback] OriginCode: {}\n",
+                frame.f_code.co_name,
+            )
+            log_do(3, lambda: dis.dis(frame.f_code))
+
+        custom_code = OpcodeExecutorCache()(frame, **kwargs)
+
+        if need_log:
+            if custom_code.code is None:
+                log_format(
+                    3,
+                    "[eval_frame_callback] NewCode (same as origin code): {}\n",
+                    frame.f_code.co_name,
+                )
+            else:
+                log_format(
+                    3,
+                    "[eval_frame_callback] NewCode: {}\n",
+                    custom_code.code.co_name,
+                )
+                log_do(3, lambda: dis.dis(custom_code.code))
+
+        return custom_code

python/paddle/jit/sot/opcode_translator/executor/__init__.py

@@ -0,0 +1,15 @@
+# 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 . import variable_dispatch  # noqa: F401

python/paddle/jit/sot/opcode_translator/executor/dispatch_functions.py

@@ -0,0 +1,71 @@
+# 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.
+
+# This file stores the customized function that will be called by the dispatch mechanism.
+
+from __future__ import annotations
+
+from ...utils import BreakGraphError, BreakGraphReasonBase, FallbackError
+
+
+def create_raise_break_graph_handler(reason: BreakGraphReasonBase):
+    def raise_break_graph_fn(*args, **kwarg):
+        raise BreakGraphError(reason)
+
+    return raise_break_graph_fn
+
+
+def raise_not_implement_fn(*args, **kwarg):
+    raise FallbackError("raise by raise_break_graph_fn.")
+
+
+# just a function for operator.in
+def operator_in(left, right):
+    return left in right
+
+
+def operator_not_in(left, right):
+    return left not in right
+
+
+def operator_exception_match(left, right):
+    pass
+
+
+def operator_BAD(left, right):
+    pass
+
+
+def operator_is_none(val):
+    pass
+
+
+def operator_is_not_none(val):
+    pass
+
+
+def tensor_dim(x):
+    pass
+
+
+def generator_send(x):
+    pass
+
+
+def place_get_device_id():
+    pass
+
+
+def place_get_device_type():
+    pass

python/paddle/jit/sot/opcode_translator/executor/dispatcher.py

@@ -0,0 +1,298 @@
+# 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 copy
+import inspect
+import operator
+from functools import cached_property, reduce
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from ...utils import InnerError, NameGenerator, hashable
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    T = TypeVar("T")
+    Args = tuple[T, ...]
+    Kwargs = dict[str, T]
+
+
+def format_type(type_: type[Any] | tuple[type[Any], ...]) -> str:
+    if not isinstance(type_, tuple):
+        type_ = (type_,)
+    return " | ".join([t.__name__ for t in type_])
+
+
+def format_param(param: Parameter) -> str:
+    kind = param.kind
+    if kind == inspect.Parameter.VAR_POSITIONAL:
+        return f"*{format_type(param.type)}"
+    elif kind == inspect.Parameter.VAR_KEYWORD:
+        return f"**{format_type(param.type)}"
+    else:
+        return format_type(param.type)
+
+
+def convert_annotation_to_type(type_str: str) -> tuple[type[Any], ...]:
+    """
+    Convert type annotation to runtime value. Because we are using :pep:`563`
+    to use the future annotation syntax, we cannot use `get_type_hints <https://docs.python.org/3.8/library/typing.html#typing.get_type_hints>`_
+    directly. Currently, only the builtins and variables namespaces are supported.
+
+    Returns:
+        tuple: The converted type.
+    """
+
+    import builtins
+
+    from . import variables
+
+    type_str = type_str.strip()
+    if type_str == "Any":
+        type_str = "object"
+
+    if "|" in type_str:
+        return reduce(
+            operator.add, map(convert_annotation_to_type, type_str.split("|"))
+        )
+
+    search_namespaces = [variables, builtins]
+    for namespace in search_namespaces:
+        if hasattr(namespace, type_str):
+            return (getattr(namespace, type_str),)
+    raise InnerError(f"Cannot find type {type_str} in {search_namespaces}")
+
+
+class Parameter:
+    name_gen = NameGenerator("param_")
+    annotation: str
+    name: str
+
+    def __init__(
+        self,
+        annotation: str,
+        *,
+        kind: inspect._ParameterKind = inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        name: str | None = None,
+        default: Any = inspect._empty,
+    ):
+        self.name = name if name is not None else Parameter.name_gen.next()
+        self.annotation = annotation
+        self.kind = kind
+        self.default = default
+
+    def to_parameter(self) -> inspect.Parameter:
+        return inspect.Parameter(
+            self.name,
+            kind=self.kind,
+            annotation=self.annotation,
+            default=copy.copy(self.default),
+        )
+
+    @cached_property
+    def type(self) -> tuple[type[Any], ...]:
+        return convert_annotation_to_type(self.annotation)
+
+    def match_arg(self, arg: Any) -> bool:
+        if self.kind == inspect.Parameter.VAR_POSITIONAL:
+            is_tuple = isinstance(arg, tuple)
+            return is_tuple and all(isinstance(a, self.type) for a in arg)
+        elif self.kind == inspect.Parameter.VAR_KEYWORD:
+            is_dict = isinstance(arg, dict)
+            return is_dict and all(
+                isinstance(a, self.type) for a in arg.values()
+            )
+        else:
+            return isinstance(arg, self.type)
+
+    @staticmethod
+    def from_str(annotation: str) -> Parameter:
+        return Parameter(annotation)
+
+    @staticmethod
+    def from_parameter(parameter: inspect.Parameter) -> Parameter:
+        if parameter.annotation != parameter.empty and not isinstance(
+            parameter.annotation, str
+        ):
+            raise InnerError(
+                f"Parameter {parameter} has annotation {parameter.annotation} "
+                "which is not a string. Please add `from __future__ import annotations` "
+                "to the top of your file."
+            )
+        annotation = (
+            parameter.annotation
+            if parameter.annotation != parameter.empty
+            else "Any"
+        )
+
+        return Parameter(
+            annotation,
+            kind=parameter.kind,
+            name=parameter.name,
+            default=parameter.default,
+        )
+
+    def __repr__(self) -> str:
+        default_repr = f"= {self.default!r}"
+        return f"Parameter({', '.join([self.annotation, default_repr])})"
+
+
+def optional(annotation: str, default: Any = None) -> Parameter:
+    return Parameter(annotation, default=default)
+
+
+class Pattern:
+    parameters: dict[str, Parameter]
+    signature: inspect.Signature
+
+    def __init__(
+        self,
+        *parameters: Parameter,
+    ):
+        self.parameters = {
+            parameter.name: parameter for parameter in parameters
+        }
+        self.signature = inspect.Signature(
+            [parameter.to_parameter() for parameter in self.parameters.values()]
+        )
+
+    def match_inputs(self, /, *args: Any, **kwargs: Any) -> bool:
+        """
+        Match the input parameters of the function.
+
+        Returns:
+            bool: Whether the input parameters match the pattern.
+        """
+        try:
+            bound_args = self.signature.bind(*args, **kwargs)
+        except TypeError:
+            return False
+        for arg_name, arg_value in bound_args.arguments.items():
+            if arg_name not in self.parameters:
+                continue
+            if not self.parameters[arg_name].match_arg(arg_value):
+                return False
+        return True
+
+    def __repr__(self) -> str:
+        types_repr = ", ".join(
+            [format_param(param) for param in self.parameters.values()]
+        )
+        return f"Pattern({types_repr})"
+
+
+class Dispatcher:
+    """
+    Used for pattern registration and distribution.
+
+    For more design ideas, refer to the `Builtin dispatcher <https://github.com/PaddlePaddle/PaddleSOT/blob/develop/docs/design/builtin-dispatcher.md>`_ for details.
+
+    Examples:
+
+        >>> def builtin_add(a: int, b: int) -> int: ...
+        >>> Dispatcher.register(builtin_add, ("int", "int"), lambda a, b: a + b)
+        >>> handler = Dispatcher.dispatch(builtin_add, 1, 2)
+        >>> handler(1, 2)
+        3
+    """
+
+    handlers: dict[
+        Callable[..., Any], list[tuple[Pattern, Callable[..., Any]]]
+    ] = {}
+    graph: Any = None
+
+    @classmethod
+    def register(
+        cls,
+        fn: Callable[..., Any],
+        parameters: tuple[str | Parameter, ...],
+        handler: Callable[..., Any],
+    ):
+        """
+        Registering function signature.
+
+        Args:
+            fn: The function to be registered.
+            parameters: The parameters of the function to be registered.
+            handler: The handler function.
+        """
+        _parameters = tuple(
+            (
+                Parameter.from_str(parameter)
+                if isinstance(parameter, str)
+                else parameter
+            )
+            for parameter in parameters
+        )
+        if fn not in cls.handlers:
+            cls.handlers[fn] = []
+        cls.handlers[fn].append((Pattern(*_parameters), handler))
+
+    @classmethod
+    def register_decorator(cls, fn: Callable[..., Any]):
+        """
+        Decorator mode of register, Used to register some complex functions.
+
+        Args:
+            fn: The function to be registered.
+
+        Examples:
+            >>> def builtin_add(a: int, b: int) -> int: ...
+            >>> @Dispatcher.register_decorator(builtin_add)
+            ... def builtin_add_dispatcher(a: int, b: int) -> int:
+            ...     return a + b
+            >>> handler = Dispatcher.dispatch(builtin_add, 1, 2)
+            >>> handler(1, 2)
+            3
+        """
+
+        def decorator(handler: Callable[..., Any]):
+            signature = inspect.signature(handler)
+            parameters = tuple(
+                Parameter.from_parameter(parameter)
+                for parameter in signature.parameters.values()
+            )
+            cls.register(fn, parameters, handler)
+
+        return decorator
+
+    @classmethod
+    def call(cls, fn, *args, **kwargs):
+        func = cls.dispatch(fn, *args, **kwargs)
+        if func is None:
+            raise InnerError(
+                f"Cannot find handler for {fn} with args {args} and kwargs {kwargs}"
+            )
+        return func(*args, **kwargs)
+
+    @classmethod
+    def dispatch(
+        cls, fn: Callable[..., Any], *args: Any, **kwargs: Any
+    ) -> Callable[..., Any] | None:
+        """
+        Find the matching handler from the registered functions.
+
+        Args:
+            fn: The function to be dispatched.
+            args: The args of the function.
+            kwargs: The kwargs of the function.
+        """
+        if not hashable(fn) or fn not in cls.handlers:
+            return None
+        for pattern, handler in cls.handlers[fn]:
+            if pattern.match_inputs(*args, **kwargs):
+                return handler
+        return None

python/paddle/jit/sot/opcode_translator/executor/exception_stack.py

@@ -0,0 +1,129 @@
+# 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.
+
+from __future__ import annotations
+
+import dataclasses
+from typing import TYPE_CHECKING
+
+from ...utils import InnerError
+from .variables import ConstantVariable, ExceptionVariable
+
+if TYPE_CHECKING:
+    from .function_graph import FunctionGraph
+
+
+@dataclasses.dataclass
+class ExceptionStack:
+    # This data structure manages exceptions as in CPython, primarily handling
+    # the __context__ attribute of SotCapturedException.
+
+    _exception_stack: list[ExceptionVariable] = dataclasses.field(
+        default_factory=list
+    )
+    _current_exception: ExceptionVariable | None = dataclasses.field(
+        default=None
+    )
+
+    def clear_current_exception(self):
+        self._current_exception = None
+
+    def set_current_exception(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
+        self._set_context_and_break_context_reference_cycle(val, graph)
+        self._current_exception = val
+
+    def move_current_exception_to_stack(self):
+        self.push(self.get_current_exception())
+        self.clear_current_exception()
+
+    def get_current_exception(self):
+        if self._current_exception is None:
+            raise InnerError("Current exception should not be None")
+        return self._current_exception
+
+    def _set_context_and_break_context_reference_cycle(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
+        # set Exception.__context__
+        self._set_context_recursive(val, len(self._exception_stack) - 1)
+        self._break_context_reference_cycle(val, graph)
+
+    def _set_context_recursive(self, val: ExceptionVariable, prev_idx: int):
+        # Recursively sets the __context__ attribute for ExceptionVariable objects
+        # in self._exception_stack. Ensures that __context__ is properly linked
+        # to the previous exception in the stack.
+        if (ctx := val.__context__) and not isinstance(ctx, ConstantVariable):
+            return val
+        if (
+            len(self._exception_stack) + prev_idx > 0
+        ):  # Prevent invalid negative indexing
+            prev = self._exception_stack[prev_idx]
+            self._set_context_recursive(prev, prev_idx - 1)
+            val.setattr("__context__", prev)
+        return val
+
+    def _break_context_reference_cycle(
+        self, val: ExceptionVariable, graph: FunctionGraph
+    ):
+        # Detects and breaks cycles in exception __context__ chains using Floyd's algorithm,
+        # following CPython's implementation.
+
+        fast = slow = val
+        slow_update_toggle = False
+        while True:
+            context = fast.__context__
+            if isinstance(
+                context, ConstantVariable
+            ):  # End of the chain; no context set
+                break
+
+            if context is val:
+                # The chain loops back to the original exception; break the cycle.
+                fast.setattr(
+                    "__context__", ConstantVariable.wrap_literal(None, graph)
+                )
+                break
+
+            fast = context
+            if fast is slow:
+                # Cycle detected; all exceptions on the path have been visited and checked.
+                break
+
+            if slow_update_toggle:
+                slow = slow.__context__
+            slow_update_toggle = not slow_update_toggle
+
+    def pop(self) -> ExceptionVariable:
+        return self._exception_stack.pop()
+
+    def push(self, val: ExceptionVariable) -> None:
+        self._exception_stack.append(val)
+
+    def empty(self) -> bool:
+        return len(self._exception_stack) == 0
+
+    def __len__(self):
+        return len(self._exception_stack)
+
+    def __repr__(self):
+        return f"ExceptionStack({self._exception_stack})"
+
+    def __getitem__(self, idx: int) -> ExceptionVariable:
+        return self._exception_stack[idx]
+
+    def cleanup(self) -> None:
+        self._exception_stack[:] = []
+        self._current_exception = None

python/paddle/jit/sot/opcode_translator/executor/executor_cache.py

@@ -0,0 +1,483 @@
+# 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 gc
+import traceback
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+import paddle
+
+from ...profiler import EventGuard, event_register
+from ...psdb import NO_FALLBACK_CODES
+from ...utils import (
+    ENV_SOT_ALLOW_DYNAMIC_SHAPE,
+    ENV_SOT_ENABLE_COMPILE_TIME_LIMIT,
+    ENV_SOT_ENABLE_GUARD_TREE,
+    ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
+    ENV_SOT_UNSAFE_CACHE_FASTPATH,
+    BreakGraphError,
+    CompileCountInfo,
+    ConditionalFallbackError,
+    FallbackError,
+    InfoCollector,
+    InnerError,
+    Singleton,
+    SotCapturedException,
+    is_strict_mode,
+    log,
+    log_do,
+    log_once,
+)
+from ..custom_code import CustomCode
+from .function_graph import FunctionGraph
+from .guard import Guard
+from .opcode_executor import OpcodeExecutor, OpcodeExecutorBase
+from .virtual_frame import VirtualFrame
+
+if TYPE_CHECKING:
+    import types
+
+GuardedFunction = tuple[CustomCode, Guard]
+GuardedFunctions = list[GuardedFunction]
+GuardChain = list[paddle.framework.core.GuardNodeBase]
+GuardChainList = list[GuardChain]
+
+dummy_guard: Guard = lambda frame: True
+dummy_guard.expr = "lambda frame: True"
+dummy_guard.inlined_expr = "lambda frame: True"
+if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+    dummy_guard.mirror_guard = lambda frame: True
+
+
+class OpcodeExecutorCache(metaclass=Singleton):
+    """
+    A singleton class that implements a cache for translated instructions.
+    This cache is used to store previously translated instructions along with their corresponding guard functions.
+
+    Attributes:
+        cache (dict): A dictionary that maps code objects to tuples of a cache getter function and a list of guarded functions.
+        translate_count (int): The count of how many instructions have been translated. It is used to test whether the cache hits.
+    """
+
+    MAX_CACHE_SIZE = 20
+    MAX_COMPILE_TIME_PER_CODE = 40
+    MAX_COMPILE_TIME_TOTAL = 15 * 60
+    CACHE_HIT_FASTPATH_THRESHOLD = 32
+    cache: dict[
+        types.CodeType, tuple[GuardedFunctions, paddle.framework.core.GuardTree]
+    ]
+    translate_count: int
+    code_symbolic_inputs: dict[types.CodeType, dict[str, None | dict[int, int]]]
+    compile_time_stats: dict[types.CodeType, float]
+    consecutive_cache_hit_count: defaultdict[types.CodeType, int]
+
+    def __init__(self):
+        self.cache = {}
+        self.translate_count = 0
+        self.code_symbolic_inputs = {}
+        self.compile_time_stats = {}
+        self.consecutive_cache_hit_count = defaultdict(int)
+
+    def get_symbolic_inputs(
+        self, code: types.CodeType
+    ) -> dict[str, dict[int, int] | None]:
+        self.code_symbolic_inputs.setdefault(code, {})
+        return self.code_symbolic_inputs[code]
+
+    def clear(self):
+        """
+        Clears the cache and resets the translate count.
+        """
+        self.cache.clear()
+        self.translate_count = 0
+        self.code_symbolic_inputs.clear()
+        self.compile_time_stats.clear()
+
+    def dump_state(self):
+        return {
+            "cache": self.cache,
+            "translate_count": self.translate_count,
+            "code_symbolic_inputs": self.code_symbolic_inputs,
+            "compile_time_stats": self.compile_time_stats,
+        }
+
+    def load_state(self, state):
+        self.cache = state["cache"]
+        self.translate_count = state["translate_count"]
+        self.code_symbolic_inputs = state["code_symbolic_inputs"]
+        self.compile_time_stats = state["compile_time_stats"]
+
+    def __call__(self, frame: types.FrameType, **kwargs) -> CustomCode:
+        code: types.CodeType = frame.f_code
+        if code not in self.cache:
+            log(2, f"[Cache] Firstly call {code}\n")
+            new_custom_code, guard_fn, guard_chain = self.translate(
+                frame, **kwargs
+            )
+            assert guard_fn is not None
+            assert guard_chain is not None
+            self.cache[code] = (
+                [(new_custom_code, guard_fn)],
+                paddle.framework.core.GuardTree([guard_chain]),
+            )
+            return new_custom_code
+        guarded_fns, guard_tree = self.cache[code]
+        compile_time_for_code = self.compile_time_stats.get(code, 0)
+        compile_time_total = sum(self.compile_time_stats.values())
+        return self.lookup(
+            frame,
+            guarded_fns,
+            guard_tree,
+            compile_time_for_code,
+            compile_time_total,
+            **kwargs,
+        )
+
+    def is_fastpath_threshold_reached(self, code):
+        # Returns True if the number of consecutive cache hits for the given code
+        # exceeds the UNSAFE_CACHE_FASTPATH threshold.
+        return (
+            self.consecutive_cache_hit_count.get(code, 0)
+            >= self.CACHE_HIT_FASTPATH_THRESHOLD
+        )
+
+    @event_register("lookup")
+    def lookup(
+        self,
+        frame: types.FrameType,
+        guarded_fns: GuardedFunctions,
+        guard_tree: paddle.framework.core.GuardTree,
+        compile_time_for_code: float,
+        compile_time_total: float,
+        **kwargs,
+    ) -> CustomCode:
+        """
+        Looks up the cache for a matching code object and returns a custom code object if a matching guard function is found, otherwise None.
+
+        Args:
+            frame (types.FrameType): The frame whose code object needs to be looked up in the cache.
+            guarded_fns (GuardedFunctions): The list of guarded functions associated with the code object.
+
+        Returns:
+            CustomCode: The custom code object if a matching guard function is found, otherwise None.
+        """
+        code: types.CodeType = frame.f_code
+
+        if len(guarded_fns) >= self.MAX_CACHE_SIZE:
+            log(2, "[Cache] Exceed max cache size, skip it\n")
+            return CustomCode(None, False)
+
+        enable_strict_guard = ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get()
+        enable_guard_tree = ENV_SOT_ENABLE_GUARD_TREE.get()
+        enable_unsafe_cache_fastpath = ENV_SOT_UNSAFE_CACHE_FASTPATH.get()
+        enable_compile_time_limit = ENV_SOT_ENABLE_COMPILE_TIME_LIMIT.get()
+
+        if enable_unsafe_cache_fastpath and (
+            self.is_fastpath_threshold_reached(code)
+        ):
+            # NOTE: In inference scenarios, cache misses are generally rare, so we can enable this unsafe short path.
+            log(
+                2,
+                "[Cache] The CACHE_HIT_FASTPATH_THRESHOLD has been reached, so fast path is now enabled\n",
+            )
+            return guarded_fns[0][0]
+
+        cache_index = None
+        if enable_strict_guard or enable_guard_tree:
+            log(4, f"[Cache] Guard tree: \n{guard_tree.stringify()}")
+            cache_index = guard_tree.lookup(frame)
+
+        if not enable_strict_guard and enable_guard_tree:
+            if cache_index is not None:
+                # TODO(zrr1999): add a mapping between custom_code and cache_index
+                return guarded_fns[cache_index][0]
+            else:
+                log(2, "[Cache] all guards missed (guard tree mode)\n")
+                if (
+                    enable_compile_time_limit
+                    and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
+                ):
+                    log(
+                        2,
+                        "[Cache] Exceed max compile time per code, skip it\n",
+                    )
+                    return CustomCode(None, False)
+                if (
+                    enable_compile_time_limit
+                    and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
+                ):
+                    log_once(
+                        f"[SOT] Current total compile time is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
+                    )
+                    log(
+                        2,
+                        "[Cache] Exceed max compile time total, skip it\n",
+                    )
+                    return CustomCode(None, False)
+                new_custom_code, guard_fn, guard_chain = self.translate(
+                    frame, **kwargs
+                )
+                if guard_fn is not None:
+                    assert guard_chain is not None
+                    guarded_fns.append((new_custom_code, guard_fn))
+                    guard_tree.add_guard_chain(guard_chain)
+                return new_custom_code
+
+        for index, (custom_code, guard_fn) in enumerate(guarded_fns):
+            if enable_strict_guard:
+                mirror_guard_error = None
+                try:
+                    with EventGuard("try mirror guard"):
+                        mirror_guard_result = guard_fn.mirror_guard(frame)
+                except Exception as e:
+                    log(2, f"[Cache] Mirror guard error: {e}\n")
+                    mirror_guard_error = e
+
+            try:
+                with EventGuard("try guard"):
+                    guard_result = guard_fn(frame)
+                if enable_strict_guard and (not enable_unsafe_cache_fastpath):
+                    assert mirror_guard_result == guard_result, (
+                        "faster guard result is not equal to guard result, "
+                        f"guard_expr: {getattr(guard_fn, 'expr', 'None')} \n"
+                        f"faster_guard_expr: {getattr(guard_fn.mirror_guard, 'expr', 'None')},"
+                    )
+                if guard_result:
+                    log(
+                        2,
+                        f"[Cache] Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                    )
+                    if not enable_unsafe_cache_fastpath:
+                        # TODO(zrr1999): cache_index should be equal to index when enable_strict_guard.
+                        assert cache_index is None or index == cache_index, (
+                            f"cache_index({cache_index}) is not equal to index({index})"
+                        )
+
+                    if enable_unsafe_cache_fastpath:
+                        if index == 0:
+                            self.consecutive_cache_hit_count[code] += 1
+                        else:
+                            # Move the current hit to the front
+                            # Note: Be cautious when modifying the order of elements in a list during iteration,
+                            # as it can lead to unexpected behavior.
+                            guarded_fns[:] = [
+                                guarded_fns[index],
+                                *guarded_fns[:index],
+                                *guarded_fns[index + 1 :],
+                            ]
+                            self.consecutive_cache_hit_count[code] = 0
+
+                    return custom_code
+                else:
+                    log_do(
+                        4,
+                        self.analyse_guard_global_object(guard_fn),
+                    )
+                    log(
+                        2,
+                        f"[Cache] Cache miss, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                    )
+                    log_do(
+                        2,
+                        self.analyse_guard_error(guard_fn, frame),
+                    )
+            except Exception as e:
+                log(2, f"[Cache] Guard function error: {e}\n")
+                log(
+                    2,
+                    f"[Cache] Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                )
+                log_do(
+                    2,
+                    self.analyse_guard_error(guard_fn, frame),
+                )
+                if enable_strict_guard and (not enable_unsafe_cache_fastpath):
+                    assert type(e) == type(mirror_guard_error) and str(
+                        e
+                    ) == str(mirror_guard_error), (
+                        "mirror guard error is not equal to guard error, "
+                        f"guard_error: {e} \n"
+                        f"mirror_guard_error: {mirror_guard_error},"
+                    )
+
+        log(2, "[Cache] all guards missed\n")
+        if (
+            enable_compile_time_limit
+            and compile_time_for_code >= self.MAX_COMPILE_TIME_PER_CODE
+        ):
+            log(2, "[Cache] Exceed max compile time per code, skip it\n")
+            return CustomCode(None, False)
+        if (
+            enable_compile_time_limit
+            and compile_time_total >= self.MAX_COMPILE_TIME_TOTAL
+        ):
+            log_once(
+                f"[SOT] Current compile time total is {compile_time_total}, exceed max compile time total {self.MAX_COMPILE_TIME_TOTAL}, fallback new function to dygraph"
+            )
+            log(
+                2,
+                "[Cache] Exceed max compile time total, skip it\n",
+            )
+            return CustomCode(None, False)
+        new_custom_code, guard_fn, guard_chain = self.translate(frame, **kwargs)
+        if guard_fn is not None:
+            assert guard_chain is not None
+            guarded_fns.append((new_custom_code, guard_fn))
+            guard_tree.add_guard_chain(guard_chain)
+        return new_custom_code
+
+    def before_translate_hook(self, frame: types.FrameType):
+        if not ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
+            return
+
+    def translate(
+        self, frame: types.FrameType, **kwargs
+    ) -> tuple[CustomCode, Guard | None, GuardChain | None]:
+        """
+        Translates the given frame's code object and returns the cache getter function and a guarded function for the translated code object.
+
+        Args:
+            frame (types.FrameType): The frame whose code object needs to be translated.
+
+        Returns:
+            tuple[CustomCode, Guard]: The cache getter function and a guarded function for the translated code object.
+        """
+        self.before_translate_hook(frame)
+        self.translate_count += 1
+        custom_new_code, guard_fn, guard_chain = start_translate(
+            frame, **kwargs
+        )
+        return custom_new_code, guard_fn, guard_chain
+
+    def analyse_guard_global_object(self, guard_fn):
+        def inner():
+            for key in guard_fn.__globals__.keys():
+                if key.startswith("__object"):
+                    print(
+                        f"[Cache] meet global object: {key} : {guard_fn.__globals__[key]}",
+                    )
+
+        return inner
+
+    def analyse_guard_error(self, guard_fn, frame):
+        def inner():
+            guard_expr = guard_fn.inlined_expr
+            lambda_head = "lambda frame: "
+            guard_expr = guard_expr.replace(lambda_head, "")
+            guards = guard_expr.split(" and ")
+            for guard_str in guards:
+                guard = eval(lambda_head + guard_str, guard_fn.__globals__)
+                result = False
+                try:
+                    result = guard(frame)
+                except Exception as e:
+                    print(
+                        f"[Cache] Error occurred when checking guard {guard_str}: {e}"
+                    )
+                    return
+                if result is False:
+                    print(f"[Cache] missed at {guard_str}")
+                    return
+            print("[Cache] missed guard not found.")
+
+        return inner
+
+
+def start_translate(
+    frame: types.FrameType,
+    **kwargs,
+) -> tuple[CustomCode, Guard | None, GuardChain | None]:
+    """
+    Starts the translation process for the given frame and returns the translated code object, its guard function and its guard tree node, or None if translation fails.
+
+    Args:
+        frame: The frame to be translated.
+
+    Returns:
+        tuple[CustomCode, Guard | None, GuardChain | None]: The translated code object, its guard function and its guard tree node, or None if translation fails.
+    """
+    simulator = None
+    graph = FunctionGraph(frame.f_code, frame.f_globals, **kwargs)
+    try:
+        vframe = VirtualFrame.from_real_frame(frame, graph)
+        simulator = OpcodeExecutor(vframe, graph)
+        simulator.check_code_simulatable()
+        InfoCollector().attach(CompileCountInfo, frame.f_code)
+
+        new_custom_code, guard_fn = simulator.transform(frame)
+        if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+            assert guard_fn(frame)
+            assert guard_fn.mirror_guard(frame)
+
+        if not simulator._graph.need_cache:
+            return (
+                CustomCode(None, True),
+                None,
+                None,
+            )
+        guard_chain = simulator.guard_chain
+        if len(guard_chain) == 0:
+            guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
+        return new_custom_code, guard_fn, guard_chain
+    # TODO(0x45f): handle BreakGraphError to trigger fallback
+    except BreakGraphError as e:
+        raise RuntimeError(
+            f"Found BreakGraphError raised, it should not be catch at start_translate!\n{e}"
+        )
+    except FallbackError as e:
+        if frame.f_code in NO_FALLBACK_CODES:
+            raise InnerError(
+                f"{frame.f_code.co_name} should not fallback, but got '{e}'"
+            )
+        if is_strict_mode():
+            raise
+        log(
+            2,
+            f"Unsupported Frame is {frame.f_code}, error message is: \n"
+            + "".join(traceback.format_exception(type(e), e, e.__traceback__)),
+        )
+        dummy_guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
+        guard, guard_chain = dummy_guard, dummy_guard_chain
+
+        if isinstance(e, ConditionalFallbackError):
+            # Guard global variables only
+            graph.input_variables.clear()
+            guard = graph.guard_fn
+            guard_chain = graph.guard_chain
+
+        return (
+            CustomCode(None, e.disable_eval_frame),
+            guard,
+            guard_chain,
+        )
+    except SotCapturedException as e:
+        log(
+            1,
+            "Note: This fallback may be triggered by user code, or it could result from an internal "
+            "SOT exception being incorrectly captured. Please investigate carefully.\n",
+        )
+        if is_strict_mode():
+            raise
+        dummy_guard_chain: GuardChain = [paddle.framework.core.DummyGuardNode()]
+        return (CustomCode(None, True), dummy_guard, dummy_guard_chain)
+    except Exception as e:
+        raise InnerError(OpcodeExecutorBase.error_message_summary(e)) from e
+    finally:
+        if simulator is not None:
+            simulator.cleanup()
+        del simulator
+        gc.collect()

python/paddle/jit/sot/opcode_translator/executor/guard.py

@@ -0,0 +1,327 @@
+# 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 types
+import weakref
+from collections.abc import Callable
+from functools import cached_property
+from typing import TYPE_CHECKING, Any, TypeVar
+
+import paddle
+
+from ...profiler import EventGuard
+from ...utils import (
+    ENV_SOT_ENABLE_FASTER_GUARD,
+    ENV_SOT_ENABLE_STRICT_GUARD_CHECK,
+    current_symbol_registry,
+    log,
+    log_do,
+)
+
+Guard = Callable[[types.FrameType], bool]
+
+if TYPE_CHECKING:
+    from .variables import VariableBase
+
+    GuardBase = paddle.framework.core.GuardBase
+    CheckGuardInputT = TypeVar("CheckGuardInputT", bound=VariableBase)
+
+# NOTE(SigureMo): [How to write Stringified Guard?]
+# 1. we should capture free variables manually, the string cannot capture free
+#    variables automatically.
+# 2. Be aware that the comparison logic before and after stringify may be different.
+# 3. we should compute as much as possible at "compile time" and encode the
+#    computation in the Guard string, rather than passing it to runtime to minimize
+#    runtime overhead.
+
+
+class StringifiedExpression:
+    """
+    Used to store string based expressions for generating Guard.
+    """
+
+    def __init__(
+        self,
+        expr_template: str,
+        sub_exprs: list[StringifiedExpression],
+        free_vars: dict[str, Any],
+    ):
+        self.expr_template = expr_template
+        expr = self.expr_template.format(
+            *[sub_expr.symbol for sub_expr in sub_exprs]
+        )
+        self.registered_expr = expr
+        self.symbol = current_symbol_registry().request_symbol(expr)
+        self.sub_exprs = sub_exprs
+        self.free_vars = free_vars
+
+    @cached_property
+    def inlined_expr(self):
+        return self.expr_template.format(
+            *[sub_expr.inlined_expr for sub_expr in self.sub_exprs]
+        )
+
+    def gen_expr(self):
+        def gen_expr_fn():
+            return self.expr_template.format(
+                *[sub_expr.gen_expr() for sub_expr in self.sub_exprs]
+            )
+
+        return current_symbol_registry().gen_expr(
+            self.registered_expr, gen_expr_fn
+        )
+
+    def __hash__(self):
+        if self.free_vars:
+            return hash((self.inlined_expr, id(self)))
+        else:
+            return hash(self.inlined_expr)
+
+
+class FasterStringifiedExpression(StringifiedExpression):
+    def __init__(
+        self,
+        expr_template: str,
+        faster_guard: GuardBase,
+        sub_exprs: list[StringifiedExpression],
+        free_vars: dict[str, Any],
+    ):
+        self.faster_guard = faster_guard
+        if ENV_SOT_ENABLE_FASTER_GUARD.get():
+            if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+                self.py_guard_expr_template = original_expr_template = (
+                    expr_template
+                )
+            else:
+                original_expr_template = expr_template
+            expr_template, free_vars = gen_faster_guard_expr_template(
+                faster_guard, sub_exprs, free_vars
+            )
+            log(
+                3,
+                f"[FasterGuard] transform {original_expr_template} to {expr_template}\n",
+            )
+
+        super().__init__(expr_template, sub_exprs, free_vars)
+
+    def gen_mirror_guard(
+        self, enable_faster_gurad: bool
+    ) -> StringifiedExpression:
+        if not enable_faster_gurad:
+            # gen faster_guard_expr
+            expr_template, expr_free_vars = gen_faster_guard_expr_template(
+                self.faster_guard,
+                self.sub_exprs,
+                self.free_vars,
+            )
+            return StringifiedExpression(
+                expr_template, self.sub_exprs, expr_free_vars
+            )
+        # gen pyGuard_expr
+        return StringifiedExpression(
+            self.py_guard_expr_template, self.sub_exprs, self.free_vars
+        )
+
+
+def gen_faster_guard_expr_template(
+    faster_guard: GuardBase,
+    sub_exprs: list[StringifiedExpression],
+    free_vars: dict[str, Any],
+) -> tuple[str, dict[str, Any]]:
+    guard_cls_name = faster_guard.__class__.__name__
+    guard_name = f"{guard_cls_name}_{id(faster_guard)}"
+    expr_template = guard_name + "(" + ", ".join(["{}"] * len(sub_exprs)) + ")"
+    free_vars = union_free_vars(free_vars, {guard_name: faster_guard.check})
+    return expr_template, free_vars
+
+
+def union_free_vars(*free_vars: dict[str, Any]):
+    return {k: v for d in free_vars for k, v in d.items()}
+
+
+def make_guard(stringified_guards: list[StringifiedExpression]) -> Guard:
+    """
+    Make a guard from a list of StringifiedExpression.
+
+    For more design ideas, refer to the `Stringified guard <https://github.com/PaddlePaddle/PaddleSOT/blob/develop/docs/design/stringify-guard.md>`_ for details.
+
+    Args:
+        stringified_guards: a list of StringifiedExpression.
+    """
+    with EventGuard("make_guard"):
+        num_guards = len(stringified_guards)
+        if not num_guards:
+            guard = lambda frame: True
+            guard.expr = "lambda frame: True"
+            guard.original_guard = guard
+            if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+                guard.mirror_guard = lambda frame: True
+            return guard
+
+        free_vars = union_free_vars(
+            *(expr.free_vars for expr in stringified_guards)
+        )
+        inlined_guard_expr = "lambda frame: " + " and ".join(
+            [expr.inlined_expr for expr in stringified_guards]
+        )
+        guard_expr: str = "lambda frame: " + " and ".join(
+            [expr.gen_expr() for expr in stringified_guards]
+        )
+
+        guard = eval(guard_expr, free_vars)
+
+        log(3, f"[Guard] {inlined_guard_expr}\n")
+        guard.inlined_expr = inlined_guard_expr
+        guard.expr = guard_expr
+
+        def check_guard_callable(guard: GuardBase):
+            assert callable(guard), "guard must be callable."
+
+        if ENV_SOT_ENABLE_STRICT_GUARD_CHECK.get():
+            mirror_guard_expr_list: list[str] = []
+            mirror_guard_temp_free_vars: dict[str, Any] = {}
+            enable_faster_gurad = ENV_SOT_ENABLE_FASTER_GUARD.get()
+            for expr in stringified_guards:
+                if isinstance(expr, FasterStringifiedExpression):
+                    expr = expr.gen_mirror_guard(enable_faster_gurad)
+                mirror_guard_expr_list.append(expr.inlined_expr)
+                mirror_guard_temp_free_vars.update(expr.free_vars)
+            mirror_guard_expr = "lambda frame: " + " and ".join(
+                mirror_guard_expr_list
+            )
+            mirror_guard_free_vars = union_free_vars(
+                mirror_guard_temp_free_vars
+            )
+            guard.mirror_guard = eval(mirror_guard_expr, mirror_guard_free_vars)
+            guard.mirror_guard.expr = mirror_guard_expr
+            check_guard_callable(guard.mirror_guard)
+
+        check_guard_callable(guard)
+
+        return guard
+
+
+def support_weak_ref(obj):
+    if isinstance(obj, types.FunctionType):
+        return True
+    return False
+
+
+# TODO(zrr1999): unify check_guard and check_faster_guard
+def check_guard(
+    fn: Callable[[CheckGuardInputT], list[StringifiedExpression]],
+) -> Callable[[CheckGuardInputT], list[StringifiedExpression]]:
+    def wrapper(self: CheckGuardInputT) -> list[StringifiedExpression]:
+        assert self.tracker.is_traceable(), (
+            "Cannot make guard from a non-tracable guard variable."
+        )
+
+        def guard_log():
+            frame_value_tracer = self.tracker.trace_value_from_frame()
+            print(
+                f"[Guard] guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.registered_expr}"
+            )
+
+        log_do(4, guard_log)
+        return fn(self)
+
+    return wrapper
+
+
+def check_faster_guard(
+    fn: Callable[[CheckGuardInputT], list[paddle.framework.core.GuardNodeBase]],
+) -> Callable[[CheckGuardInputT], list[paddle.framework.core.GuardNodeBase]]:
+    def wrapper(
+        self: CheckGuardInputT,
+    ) -> list[paddle.framework.core.GuardNodeBase]:
+        assert self.tracker.is_traceable(), (
+            "Cannot make guard from a non-tracable guard variable."
+        )
+
+        def guard_log():
+            frame_value_tracer = self.tracker.trace_value_from_frame()
+            print(
+                f"[Guard Tree] guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.registered_expr}"
+            )
+
+        log_do(4, guard_log)
+        return fn(self)
+
+    return wrapper
+
+
+@check_guard
+def object_equal_stringified_guard(self) -> list[StringifiedExpression]:
+    frame_value_tracer = self.tracker.trace_value_from_frame()
+
+    obj_free_var_name = f"__{self.id}"
+    weak_ref_obj = self.get_py_value()
+    if support_weak_ref(weak_ref_obj):
+        weak_ref_obj = weakref.ref(self.get_py_value())
+        return [
+            FasterStringifiedExpression(
+                f"{obj_free_var_name}() is not None and {{}} == {obj_free_var_name}()",
+                paddle.framework.core.WeakRefMatchGuard(self.get_py_value()),
+                [frame_value_tracer],
+                union_free_vars(
+                    frame_value_tracer.free_vars,
+                    {obj_free_var_name: weak_ref_obj},
+                ),
+            )
+        ]
+    return [
+        FasterStringifiedExpression(
+            f"{{}} == {obj_free_var_name}",
+            paddle.framework.core.ValueMatchGuard(weak_ref_obj),
+            [frame_value_tracer],
+            union_free_vars(
+                frame_value_tracer.free_vars,
+                {obj_free_var_name: self.get_py_value()},
+            ),
+        )
+    ]
+
+
+@check_faster_guard
+def object_equal_faster_guard(
+    self,
+) -> list[paddle.framework.core.GuardNodeBase]:
+    expr_node = self.tracker.guard_tree_expr_node()
+
+    weak_ref_obj = self.get_py_value()
+    if support_weak_ref(weak_ref_obj):
+        weak_ref_obj = weakref.ref(self.get_py_value())
+        return [
+            paddle.framework.core.GuardNode(
+                paddle.framework.core.WeakRefMatchGuard(self.get_py_value()),
+                [expr_node],
+            )
+        ]
+    return [
+        paddle.framework.core.GuardNode(
+            paddle.framework.core.ValueMatchGuard(weak_ref_obj),
+            [expr_node],
+        )
+    ]
+
+
+def stringify_pyobject(obj: object) -> tuple[str, dict[str, Any]]:
+    if isinstance(obj, paddle.core.VarDesc.VarType):
+        return f"paddle.core.VarDesc.VarType({obj.value})", {"paddle": paddle}
+    elif isinstance(obj, paddle.core.DataType):
+        return f"paddle.core.DataType({obj.value})", {"paddle": paddle}
+    # For builtin values
+    return f"{obj!r}", {}

python/paddle/jit/sot/opcode_translator/executor/instr_flag.py

@@ -0,0 +1,72 @@
+# 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.
+
+# flags for instructions
+from enum import Enum
+
+
+class FORMAT_VALUE_FLAG:
+    FVC_MASK = 0x3
+    FVC_NONE = 0x0
+    FVC_STR = 0x1
+    FVC_REPR = 0x2
+    FVC_ASCII = 0x3
+    FVS_MASK = 0x4
+    FVS_HAVE_SPEC = 0x4
+
+
+class CONVERT_VALUE_FLAG:
+    CV_STR = 1
+    CV_REPR = 2
+    CV_ASCII = 3
+
+
+# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_utils.h#L63-L68
+class MAKE_FUNCTION_FLAG:
+    MF_HAS_ANNOTATE = 0x10
+    MF_HAS_CLOSURE = 0x08
+    MF_HAS_ANNOTATION = 0x04
+    MF_HAS_KWDEFAULTS = 0x02
+    MF_HAS_DEFAULTS = 0x01
+
+
+class CALL_FUNCTION_EX_FLAG:
+    CFE_HAS_KWARGS = 0x01
+
+
+# see https://github.com/python/cpython/blob/3.12/Python/intrinsics.c#L211-L225
+class IntrinsicsUnaryFunctions(Enum):
+    INTRINSIC_1_INVALID = 0
+    INTRINSIC_PRINT = 1  # no support, only non-interactive mode
+    INTRINSIC_IMPORT_STAR = 2  # no support, `from module import *`
+    INTRINSIC_STOPITERATION_ERROR = 3  # no support, generator or coroutine
+    INTRINSIC_ASYNC_GEN_WRAP = 4  # no support, async
+    INTRINSIC_UNARY_POSITIVE = 5
+    INTRINSIC_LIST_TO_TUPLE = 6
+    INTRINSIC_TYPEVAR = 7  # no support, PEP 695
+    INTRINSIC_PARAMSPEC = 8  # no support, PEP 695
+    INTRINSIC_TYPEVARTUPLE = 9  # no support, PEP 695
+    INTRINSIC_SUBSCRIPT_GENERIC = 10  # no support, PEP 695
+    INTRINSIC_TYPEALIAS = 11  # no support, PEP 695
+
+
+# https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_utils.h#L70-L76
+# All are attributes of 'builtins'
+LOAD_COMMON_CONSTANT_FLAG = (
+    "AssertionError",
+    "NotImplementedError",
+    "tuple",
+    "all",
+    "any",
+)

python/paddle/jit/sot/opcode_translator/executor/mutable_data.py

@@ -0,0 +1,306 @@
+# 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
+
+from collections.abc import Callable
+from typing import Any, Concatenate, Generic, TypeAlias, TypeVar
+
+from typing_extensions import ParamSpec
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+MutableDataT = TypeVar("MutableDataT", bound="MutableData")
+DataGetter: TypeAlias = Callable[[MutableDataT, Any], Any]
+
+InnerMutableDataT = TypeVar(
+    "InnerMutableDataT", bound="dict[str, Any] | list[Any]"
+)
+
+
+class Mutation:
+    ABBR: str
+
+
+class MutationSet(Mutation):
+    """
+    Setting a value.
+    This mutation is used for MutableDictLikeData and MutableListLikeData.
+    """
+
+    ABBR = "S"
+
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+
+    def __repr__(self):
+        return f"MutationSet({self.key}, {self.value})"
+
+
+class MutationDel(Mutation):
+    """
+    Deleting a value.
+    This mutation is used for MutableDictLikeData and MutableListLikeData.
+    """
+
+    ABBR = "D"
+
+    def __init__(self, key):
+        self.key = key
+
+    def __repr__(self):
+        return f"MutationDel({self.key})"
+
+
+class MutationNew(Mutation):
+    """
+    Adding a new value.
+    This mutation is only used for MutableDictLikeData.
+    """
+
+    ABBR = "N"
+
+    def __init__(self, key, value):
+        self.key = key
+        self.value = value
+
+    def __repr__(self):
+        return f"MutationNew({self.key}, {self.value})"
+
+
+class MutationInsert(Mutation):
+    """
+    Inserting a value.
+    This mutation is only used for MutableListLikeData.
+    """
+
+    ABBR = "I"
+
+    def __init__(self, index, value):
+        self.index = index
+        self.value = value
+
+    def __repr__(self):
+        return f"MutationInsert({self.index}, {self.value})"
+
+
+class MutationPermutate(Mutation):
+    """
+    Permutating all the values.
+    This mutation is only used for MutableListLikeData.
+    """
+
+    ABBR = "P"
+
+    def __init__(self, permutation):
+        self.permutation = permutation
+
+    def __repr__(self):
+        return f"MutationPermutate({self.permutation})"
+
+
+def record_mutation(
+    mutation_fn: Callable[Concatenate[MutableDataT, P], Mutation],
+) -> Callable[Concatenate[MutableDataT, P], None]:
+    def wrapper(self, *args: P.args, **kwargs: P.kwargs):
+        mutation = mutation_fn(self, *args, **kwargs)
+        self.records.append(mutation)
+
+    return wrapper
+
+
+class MutableData(Generic[InnerMutableDataT]):
+    """
+    An intermediate data structure between data and variable, it records all the mutations.
+    """
+
+    read_cache: InnerMutableDataT
+
+    class Empty:
+        def __repr__(self):
+            return "Empty()"
+
+    def __init__(self, data: Any, getter: DataGetter):
+        self.original_data = data
+        self.getter = getter
+        self.records: list[Mutation] = []
+
+    def is_empty(self, value):
+        return isinstance(value, MutableData.Empty)
+
+    @property
+    def version(self):
+        return len(self.records)
+
+    @property
+    def has_changed(self):
+        return self.version != 0
+
+    def check_changed(self, key: Any) -> bool:
+        raise NotImplementedError
+
+    def rollback(self, version: int):
+        assert version <= self.version
+        self.records[:] = self.records[:version]
+
+    def get(self, key):
+        raise NotImplementedError
+
+    def set(self, key, value):
+        raise NotImplementedError
+
+    def apply(self, mutation: Mutation, write_cache: InnerMutableDataT):
+        raise NotImplementedError
+
+    def reproduce(self, version: int | None = None) -> InnerMutableDataT:
+        if version is None:
+            version = self.version
+        write_cache = self.read_cache.copy()
+        for mutation in self.records[:version]:
+            self.apply(mutation, write_cache)
+        return write_cache
+
+    def __repr__(self) -> str:
+        records_abbrs = "".join([mutation.ABBR for mutation in self.records])
+        return f"{self.__class__.__name__}({records_abbrs})"
+
+
+class MutableDictLikeData(MutableData["dict[str, Any]"]):
+    def __init__(self, data: Any, getter: DataGetter):
+        super().__init__(data, getter)
+        self.read_cache = {}
+
+    def clear_read_cache(self):
+        self.read_cache.clear()
+
+    def check_changed(self, key: Any) -> bool:
+        if not self.has_changed:
+            return False
+        for mutation in self.records:
+            if (
+                isinstance(mutation, (MutationNew, MutationDel, MutationSet))
+                and mutation.key == key
+            ):
+                return True
+        return False
+
+    def get(self, key: Any):
+        # TODO(SigureMo): Optimize performance of this.
+        write_cache = self.reproduce(self.version)
+        if key not in write_cache:
+            self.read_cache[key] = self.getter(self, key)
+        return self.reproduce(self.version)[key]
+
+    def get_all(self):
+        original_keys = list(self.original_data.keys())
+        for mutation in self.records:
+            if isinstance(mutation, MutationNew):
+                original_keys.append(mutation.key)
+            elif isinstance(mutation, MutationDel):
+                original_keys.remove(mutation.key)
+        return {key: self.get(key) for key in original_keys}
+
+    @record_mutation
+    def set(self, key: Any, value: Any) -> Mutation:
+        is_new = False
+        if self.is_empty(self.get(key)):
+            is_new = True
+        return (
+            MutationSet(key, value) if not is_new else MutationNew(key, value)
+        )
+
+    @record_mutation
+    def delete(self, key):
+        return MutationDel(key)
+
+    def apply(self, mutation: Mutation, write_cache: dict[str, Any]):
+        if isinstance(mutation, MutationNew):
+            write_cache[mutation.key] = mutation.value
+        elif isinstance(mutation, MutationSet):
+            write_cache[mutation.key] = mutation.value
+        elif isinstance(mutation, MutationDel):
+            write_cache[mutation.key] = MutableData.Empty()
+        else:
+            raise ValueError(f"Unknown mutation type {mutation}")
+
+    def reproduce(self, version: int | None = None):
+        if version is None:
+            version = self.version
+        write_cache = self.read_cache.copy()
+        for mutation in self.records[:version]:
+            self.apply(mutation, write_cache)
+        return write_cache
+
+
+class MutableListLikeData(MutableData["list[Any]"]):
+    def __init__(self, data: Any, getter: DataGetter):
+        super().__init__(data, getter)
+        self.read_cache = [
+            self.getter(self, idx) for idx in range(len(self.original_data))
+        ]
+
+    def clear_read_cache(self):
+        self.read_cache[:] = []
+
+    def check_changed(self, key: Any) -> bool:
+        return self.has_changed
+
+    @property
+    def length(self):
+        return len(self.reproduce())
+
+    def get(self, key):
+        write_cache = self.reproduce(self.version)
+        return write_cache[key]
+
+    def get_all(self) -> list[Any]:
+        items = self.reproduce(self.version)
+        return items
+
+    @record_mutation
+    def set(self, key: int, value: Any):
+        return MutationSet(self._regularize_index(key), value)
+
+    @record_mutation
+    def delete(self, key: int):
+        return MutationDel(self._regularize_index(key))
+
+    @record_mutation
+    def insert(self, index: int, value: Any):
+        return MutationInsert(self._regularize_index(index), value)
+
+    @record_mutation
+    def permutate(self, permutation: list[int]):
+        return MutationPermutate(permutation)
+
+    def _regularize_index(self, index: int):
+        if index < 0:
+            index += self.length
+        return index
+
+    def apply(self, mutation: Mutation, write_cache: list[Any]):
+        if isinstance(mutation, MutationSet):
+            write_cache[mutation.key] = mutation.value
+        elif isinstance(mutation, MutationDel):
+            write_cache[:] = (
+                write_cache[: mutation.key] + write_cache[mutation.key + 1 :]
+            )
+        elif isinstance(mutation, MutationInsert):
+            write_cache.insert(mutation.index, mutation.value)
+        elif isinstance(mutation, MutationPermutate):
+            write_cache[:] = [write_cache[i] for i in mutation.permutation]
+        else:
+            raise ValueError(f"Unknown mutation type {mutation}")

python/paddle/jit/sot/opcode_translator/executor/opcode_inline_executor.py

@@ -0,0 +1,228 @@
+# 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 inspect
+import sys
+from typing import TYPE_CHECKING
+
+from ...utils import (
+    BreakGraphError,
+    DataDependencyControlFlowBreak,
+    FallbackError,
+    UnsupportedIteratorBreak,
+)
+from ...utils.exceptions import SotCapturedStopIteration
+from ..instruction_utils import Instruction
+from .dispatch_functions import generator_send
+from .opcode_executor import OpcodeExecutorBase, Stop
+from .tracker import DanglingTracker
+from .variables import (
+    BuiltinVariable,
+    ConstantVariable,
+    GeneratorVariable,
+    IterVariable,
+    ObjectVariable,
+    UserDefinedIterVariable,
+    VariableBase,
+)
+
+if TYPE_CHECKING:
+    from .function_graph import FunctionGraph
+    from .virtual_frame import VirtualFrame
+
+
+def inline_for_iter_impl(exe: OpcodeExecutorBase, instr: Instruction):
+    iterator = exe.stack.top
+    assert isinstance(iterator, IterVariable)
+
+    exe._graph.add_global_guarded_variable(iterator)
+
+    # simply get next
+    if not isinstance(iterator, UserDefinedIterVariable):
+        try:
+            exe.stack.push(iterator.next())
+        except SotCapturedStopIteration:
+            exe.stack.pop()
+            assert isinstance(instr.jump_to, Instruction)
+            exe.vframe.lasti = exe.indexof(instr.jump_to)
+            if sys.version_info >= (3, 12):
+                assert exe._instructions[exe.vframe.lasti].opname == "END_FOR"
+                skip_n_instrs = 2 if sys.version_info >= (3, 13) else 1
+                exe.vframe.lasti += skip_n_instrs
+
+    else:
+        exe._graph.remove_global_guarded_variable(iterator)
+        raise BreakGraphError(
+            UnsupportedIteratorBreak(
+                reason_str=f"Found {iterator.__class__.__name__} as iterator."
+            )
+        )
+
+
+class OpcodeInlineExecutor(OpcodeExecutorBase):
+    """
+    A class that represents an executor for inlined opcode operations.
+
+    Args:
+        fn_variable: The function variable.
+
+    """
+
+    def __init__(
+        self,
+        vframe: VirtualFrame,
+        code_var: VariableBase,
+        graph: FunctionGraph,
+    ):
+        super().__init__(vframe, graph)
+        self.return_value: VariableBase | None = None
+        self._code_var = code_var
+        self._name = "InlineFn"
+
+    def inline_call(self) -> VariableBase:
+        """
+        Execute the inline call of the function.
+        """
+        self._graph.add_global_guarded_variable(self._code_var)
+        self.run()
+        assert self.return_value is not None
+        return self.return_value
+
+    def RETURN_VALUE(self, instr: Instruction):
+        assert len(self.stack) == 1, (
+            f"Stack must have one element, but get {len(self.stack)} elements."
+        )
+        self.return_value = self.stack.pop()
+        return Stop(state="Return")
+
+    def RETURN_CONST(self, instr: Instruction):
+        self.return_value = self.vframe.consts[instr.arg]
+        return Stop(state="Return")
+
+    def _break_graph_when_if(self, result, instr: Instruction):
+        """
+        Helper method to raise a BreakGraphError when breaking the graph in a jump operation.
+
+        Args:
+            result: The result of the operation.
+            instr (Instruction): The jump instruction.
+        """
+
+        raise BreakGraphError(DataDependencyControlFlowBreak())
+
+    def FOR_ITER(self, instr: Instruction):
+        return inline_for_iter_impl(self, instr)
+
+
+class OpcodeInlineGeneratorExecutor(OpcodeExecutorBase):
+    def __init__(
+        self,
+        vframe: VirtualFrame,
+        code_var: VariableBase,
+        graph: FunctionGraph,
+    ):
+        super().__init__(vframe, graph)
+        self.return_value: VariableBase | None = None
+        self._code_var = code_var
+        self._name = "InlineGen"
+
+    def inline_call(self) -> VariableBase:
+        self._graph.add_global_guarded_variable(self._code_var)
+        self.run()
+        assert self.return_value is not None
+        return self.return_value
+
+    def RETURN_GENERATOR(self, instr: Instruction):
+        vframe = self.vframe
+        code_var = self._code_var
+        # NOTE: we set the real tracker in calling function
+        self.return_value = GeneratorVariable(
+            code_var, vframe, self._graph, DanglingTracker()
+        )
+        return Stop(state="Return")
+
+    def SEND(self, instr: Instruction):
+        assert len(self.stack) >= 2
+        recv = self.stack.pop()
+        source_obj = self.stack.top
+        if not isinstance(source_obj, IterVariable):
+            raise FallbackError(
+                "Yield from for non-generator object is not supported."
+            )
+        self.stack.push(
+            BuiltinVariable(generator_send, self._graph, DanglingTracker())(
+                source_obj, recv
+            )
+        )
+
+    def END_SEND(self, instr: Instruction):
+        value = self.stack.pop()
+        receiver = self.stack.pop()  # pop the receiver
+        self.stack.push(value)
+
+    def GEN_START(self, instr: Instruction):
+        tos = self.stack.pop()
+        assert isinstance(tos, ConstantVariable)
+        assert tos.value is None
+
+    def YIELD_VALUE(self, instr: Instruction):
+        assert len(self.stack) >= 1
+        self.return_value = self.stack.pop()
+        return Stop(state="Yield")
+
+    def GET_YIELD_FROM_ITER(self, instr: Instruction):
+        source_obj = self.stack.top
+        if isinstance(source_obj, ObjectVariable) and inspect.iscoroutine(
+            source_obj.value
+        ):
+            raise FallbackError(
+                "Get yield from iter for coroutine object is not supported."
+            )
+        if isinstance(source_obj, GeneratorVariable):
+            return
+        source_obj = self.stack.pop()
+        iter_variable = BuiltinVariable(iter, self._graph, DanglingTracker())(
+            source_obj
+        )
+        self.stack.push(iter_variable)
+
+    def YIELD_FROM(self, instr: Instruction):
+        recv = self.stack.pop()
+        source_obj = self.stack.top
+        if not isinstance(source_obj, IterVariable):
+            raise FallbackError(
+                "Yield from for non-generator object is not supported."
+            )
+        self.return_value = BuiltinVariable(
+            generator_send, self._graph, DanglingTracker()
+        )(source_obj, recv)
+        assert self.vframe.lasti > 0
+        self.vframe.lasti -= 1
+        return Stop(state="Yield")
+
+    def FOR_ITER(self, instr: Instruction):
+        return inline_for_iter_impl(self, instr)
+
+    def RETURN_VALUE(self, instr: Instruction):
+        assert len(self.stack) == 1, (
+            f"Stack must have one element, but get {len(self.stack)} elements."
+        )
+        self.return_value = self.stack.pop()
+        return Stop(state="Return")
+
+    def RETURN_CONST(self, instr: Instruction):
+        self.return_value = self.vframe.consts[instr.arg]
+        return Stop(state="Return")

python/paddle/jit/sot/opcode_translator/executor/side_effects.py

@@ -0,0 +1,237 @@
+# 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
+
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import TypeAlias
+
+    from .mutable_data import DataGetter, MutableData
+    from .pycode_generator import PyCodeGen
+    from .variables import VariableBase
+
+    IdGetter: TypeAlias = Callable[[Any], int]
+    MutableDataT = TypeVar("MutableDataT", bound=MutableData)
+
+
+class SideEffectsState(NamedTuple):
+    data_id_to_proxy: dict[int, MutableData]
+    proxy_variables: list[VariableBase]
+    mutable_variables: list[VariableBase]
+    proxy_versions: list[int]
+    mutable_attrs: list[dict[str, Any]]
+
+
+class SideEffects:
+    def __init__(self):
+        self.data_id_to_proxy: dict[int, MutableData] = {}
+        self.proxy_variables: list[VariableBase] = []
+        self.mutable_variables: list[VariableBase] = []
+
+    def record_proxy_variable(self, variable: VariableBase):
+        if variable not in self.proxy_variables:
+            self.proxy_variables.append(variable)
+
+    def record_mutable_variable(self, variable: VariableBase):
+        if variable not in self.mutable_variables:
+            self.mutable_variables.append(variable)
+
+    def get_proxy(
+        self,
+        proxy_type: type[MutableDataT],
+        data: Any,
+        getter: DataGetter,
+        id_getter: IdGetter = id,
+    ) -> MutableDataT:
+        data_id = id_getter(data)
+        if data_id not in self.data_id_to_proxy:
+            self.data_id_to_proxy[data_id] = proxy_type(data, getter)
+        return self.data_id_to_proxy[data_id]  # type: ignore
+
+    def get_state(self):
+        return SideEffectsState(
+            self.data_id_to_proxy.copy(),
+            self.proxy_variables.copy(),
+            self.mutable_variables.copy(),
+            [proxy.version for proxy in self.data_id_to_proxy.values()],
+            [
+                {attr: getattr(var, attr)}
+                for var in self.mutable_variables
+                for attr in var.mutable_attrs
+            ],
+        )
+
+    def restore_state(self, state: SideEffectsState):
+        self.data_id_to_proxy = state.data_id_to_proxy
+        self.proxy_variables = state.proxy_variables
+        self.mutable_variables = state.mutable_variables
+        # NOTE(SigureMo): We can use the `strict=True` option in zip after
+        # Python 3.10.
+        assert len(self.data_id_to_proxy.values()) == len(
+            state.proxy_versions
+        ), "proxy_versions length not match"
+        assert sum(
+            len(var.mutable_attrs) for var in self.mutable_variables
+        ) == len(state.mutable_attrs), "mutable_attrs length not match"
+
+        for proxy, version in zip(
+            self.data_id_to_proxy.values(), state.proxy_versions
+        ):
+            proxy.rollback(version)
+
+        for (variable, attr), attr_dict in zip(
+            (
+                (var, attr)
+                for var in self.mutable_variables
+                for attr in var.mutable_attrs
+            ),
+            (attr_dict for attr_dict in state.mutable_attrs),
+        ):
+            setattr(variable, attr, attr_dict[attr])
+
+
+class SideEffectRestorer:
+    def pre_gen(self, codegen: PyCodeGen):
+        raise NotImplementedError
+
+    def post_gen(self, codegen: PyCodeGen):
+        raise NotImplementedError
+
+
+class DictSideEffectRestorer(SideEffectRestorer):
+    """
+    old_dict.clear()
+    old_dict.update(new_dict)
+    """
+
+    def __init__(self, var: VariableBase):
+        super().__init__()
+        self.var = var
+
+    def pre_gen(self, codegen: PyCodeGen):
+        # Reference to the original dict.
+        # load old_dict.update and new_dict to stack.
+        self.var.reconstruct(codegen)
+        codegen.gen_load_method("update")
+        # Generate dict by each key-value pair.
+        self.var.reconstruct(codegen, use_tracker=False)
+        # load old_dict.clear to stack.
+        self.var.reconstruct(codegen)
+        codegen.gen_load_method("clear")
+
+    def post_gen(self, codegen: PyCodeGen):
+        # Call methods to apply side effects.
+        codegen.gen_call_method(0)  # call clear
+        codegen.gen_pop_top()
+        codegen.gen_call_method(1)  # call update
+        codegen.gen_pop_top()
+
+
+class ListSideEffectRestorer(SideEffectRestorer):
+    """
+    old_list[:] = new_list
+    """
+
+    def __init__(self, var: VariableBase):
+        super().__init__()
+        self.var = var
+
+    def pre_gen(self, codegen: PyCodeGen):
+        # Reference to the original list.
+        # load new_list to stack.
+        self.var.reconstruct(codegen, use_tracker=False)
+        # load old_list[:] to stack.
+        self.var.reconstruct(codegen)
+        codegen.gen_load_const(None)
+        codegen.gen_load_const(None)
+        codegen.gen_build_slice(2)
+
+    def post_gen(self, codegen: PyCodeGen):
+        # Call STORE_SUBSCR to apply side effects.
+        codegen.gen_store_subscr()
+
+
+class GlobalSetSideEffectRestorer(SideEffectRestorer):
+    """
+    global_var = new_value
+    """
+
+    def __init__(self, name: str, var: VariableBase):
+        super().__init__()
+        self.name = name
+        self.var = var
+
+    def pre_gen(self, codegen: PyCodeGen):
+        self.var.reconstruct(codegen)
+
+    def post_gen(self, codegen: PyCodeGen):
+        codegen.gen_store_global(self.name)
+
+
+class GlobalDelSideEffectRestorer(SideEffectRestorer):
+    """
+    del global_var
+    """
+
+    def __init__(self, name: str):
+        super().__init__()
+        self.name = name
+
+    def pre_gen(self, codegen: PyCodeGen):
+        # do nothing
+        ...
+
+    def post_gen(self, codegen: PyCodeGen):
+        codegen.gen_delete_global(self.name)
+
+
+class ObjSetSideEffectRestorer(SideEffectRestorer):
+    """
+    obj.attr = new_value
+    """
+
+    def __init__(self, obj: VariableBase, name: str, var: VariableBase):
+        super().__init__()
+        self.obj = obj
+        self.name = name
+        self.var = var
+
+    def pre_gen(self, codegen: PyCodeGen):
+        # value
+        self.var.reconstruct(codegen)
+        # obj
+        self.obj.reconstruct(codegen)
+
+    def post_gen(self, codegen: PyCodeGen):
+        codegen.gen_store_attr(self.name)
+
+
+class ObjDelSideEffectRestorer(SideEffectRestorer):
+    """
+    del obj.attr
+    """
+
+    def __init__(self, obj: VariableBase, name: str):
+        super().__init__()
+        self.obj = obj
+        self.name = name
+
+    def pre_gen(self, codegen: PyCodeGen):
+        self.obj.reconstruct(codegen)
+
+    def post_gen(self, codegen: PyCodeGen):
+        codegen.gen_delete_attr(self.name)

python/paddle/jit/sot/opcode_translator/executor/tracker.py

@@ -0,0 +1,619 @@
+# 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 sys
+from itertools import chain
+from typing import TYPE_CHECKING
+
+import paddle
+
+from ...utils import InnerError, NameGenerator
+from .guard import StringifiedExpression, stringify_pyobject, union_free_vars
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from ...utils.magic_methods import BinaryOp, UnaryOp
+    from .pycode_generator import PyCodeGen
+    from .variables import FunctionVariable, VariableBase
+
+
+class Tracker:
+    """
+    Tracker is a base class responsible for tracking variables or objects in Python code.
+    It is used to identify how a variable is derived from the initial state of the frame.
+
+    Args:
+        inputs: The list of variables to be tracked.
+
+    Note:
+        It serves as an abstract class and should not be instantiated directly.
+    """
+
+    inputs: Sequence[VariableBase]
+    name_generator = NameGenerator("tracker_")
+
+    def __init__(self, inputs: Sequence[VariableBase], changed: bool = False):
+        self.inputs = inputs
+        self.changed = changed
+        self.id = Tracker.name_generator.next()
+
+    def gen_instructions(self, codegen: PyCodeGen) -> None:
+        """
+        Generate instructions based on the tracked variables.
+
+        Args:
+            codegen (PyCodeGen): An instance of PyCodeGen to generate instructions.
+        """
+        raise NotImplementedError
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        raise NotImplementedError(
+            f"{self.__class__.__name__} has no guard_tree_expr_node"
+        )
+
+    # TODO(xiongkun): trace_value_from_frame is not a good name, it should be more related to guard but not traceable.
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        """
+        Trace the value of the tracked variables from the frame. It used for generating the guard.
+
+        Returns:
+            The value of the tracked variables.
+        """
+        raise NotImplementedError
+
+    def is_traceable(self) -> bool:
+        """
+        Determine if all the tracked variables can be traced from the frame.
+
+        Returns:
+            bool: True if all tracked variables are traceable, False otherwise.
+        """
+        if self.changed:
+            return False
+        for input in self.inputs:
+            if not input.tracker.is_traceable():
+                return False
+        return True
+
+    def need_guard(self) -> bool:
+        return self.is_traceable()
+
+
+class DummyTracker(Tracker):
+    """
+    DummyTracker is a subclass of Tracker that specifically tracks variables cannot be reproduced from the frame.
+    It is mostly generated by complex operations (instructions).
+
+    Args:
+        inputs (list[VariableBase]): The input variables associated with the generated variables.
+    """
+
+    def __init__(self, inputs: Sequence[VariableBase]):
+        super().__init__(inputs)
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        raise InnerError("DummyTracker has no instructions")
+
+    def trace_value_from_frame(self):
+        raise InnerError("DummyTracker can't trace value from frame")
+
+    def is_traceable(self):
+        return False
+
+    def __repr__(self) -> str:
+        return f"DummyTracker(num_inputs={len(self.inputs)})"
+
+    def need_guard(self) -> bool:
+        return False
+
+
+class SymbolicOperationTracker(Tracker):
+    """
+    SymbolicOperationTracker is a subclass of Tracker that specifically tracks variables cannot be reproduced from the frame.
+    It is mostly generated by complex operations of symbolic variables.
+
+    Args:
+        inputs (list[VariableBase]): The input variables associated with the generated variables.
+    """
+
+    def __init__(self, inputs: Sequence[VariableBase], op: UnaryOp | BinaryOp):
+        super().__init__(inputs)
+        self.op = op
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        raise InnerError("SymbolicOperationTracker has no instructions")
+
+    def trace_value_from_frame(self):
+        raise InnerError(
+            "SymbolicOperationTracker can't trace value from frame"
+        )
+
+    def __repr__(self) -> str:
+        return f"SymbolicOperationTracker(num_inputs={len(self.inputs)})"
+
+    def is_traceable(self):
+        return False
+
+
+class DanglingTracker(Tracker):
+    """
+    DanglingTracker is a subclass of Tracker that specifically tracks variables that are not in the frame.
+    Variables whose tracker is DanglingTracker should not be placed on the stack, except for NullVariable.
+    DanglingTracker is often used in conjunction with BuiltinVariable to reuse the dispatch mechanism.
+
+    Examples:
+        >>> import operator
+        >>> from sot.opcode_translator.executor.variables import (
+        ...     BuiltinVariable,
+        ...     ConstantVariable,
+        ... )
+        >>> a = ConstantVariable.wrap_literal(1, None)
+        >>> b = ConstantVariable.wrap_literal(2, None)
+        >>> c = BuiltinVariable(operator.add, None, DanglingTracker())(a, b)
+        >>> c.value
+        3
+    """
+
+    def __init__(self):
+        super().__init__([])
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        raise InnerError("DanglingTracker has no instructions")
+
+    def trace_value_from_frame(self):
+        raise InnerError("DanglingTracker can't trace value from frame")
+
+    def is_traceable(self):
+        return False
+
+    def __repr__(self) -> str:
+        return "DanglingTracker()"
+
+
+class LocalTracker(Tracker):
+    """
+    LocalTracker is a subclass of Tracker that specifically tracks variables from f_locals of frame.
+
+    Args:
+        name (str): The name of the variable in f_locals to be tracked.
+    """
+
+    def __init__(self, name: str):
+        super().__init__([])
+        self.name = name
+
+    def gen_instructions(self, codegen: PyCodeGen) -> None:
+        codegen.gen_load_fast(self.name)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.LocalVarExprNode(self.name)
+
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        return StringifiedExpression(f"frame.f_locals['{self.name}']", [], {})
+
+    def __repr__(self) -> str:
+        return f"LocalTracker(name={self.name})"
+
+
+class CellTracker(LocalTracker):
+    def gen_instructions(self, codegen: PyCodeGen):
+        codegen.gen_load_deref(self.name)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.LocalVarExprNode(self.name)
+
+    def trace_value_from_frame(self):
+        return StringifiedExpression(f"frame.f_locals['{self.name}']", [], {})
+
+    def __repr__(self) -> str:
+        return f"CellTracker(name={self.name})"
+
+
+class GlobalTracker(Tracker):
+    """
+    GlobalTracker is a subclass of Tracker that specifically tracks variables from f_globals of frame.
+
+    Args:
+        name (str): The name of the variable in f_globals to be tracked.
+    """
+
+    def __init__(self, name: str):
+        super().__init__([])
+        self.name = name
+
+    def gen_instructions(self, codegen: PyCodeGen) -> None:
+        codegen.gen_load_global(self.name, push_null=False)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.GlobalVarExprNode(self.name)
+
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        return StringifiedExpression(f"frame.f_globals['{self.name}']", [], {})
+
+    def __repr__(self) -> str:
+        return f"GlobalTracker(name={self.name})"
+
+
+class BuiltinTracker(Tracker):
+    """
+    BuiltinTracker is a subclass of Tracker that specifically tracks variables from f_builtins of frame.
+
+    Args:
+        name (str): The name of the variable in f_builtins to be tracked.
+    """
+
+    def __init__(self, name: str):
+        super().__init__([])
+        self.name = name
+
+    def gen_instructions(self, codegen: PyCodeGen) -> None:
+        codegen.gen_load_global(self.name, push_null=False)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.ConstantExprNode(
+            getattr(builtins, self.name)
+        )
+
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        return StringifiedExpression(
+            f"builtins.__dict__['{self.name}']", [], {"builtins": builtins}
+        )
+
+    def __repr__(self) -> str:
+        return f"BuiltinTracker(name={self.name})"
+
+
+class ConstTracker(Tracker):
+    """
+    ConstTracker is a subclass of Tracker that specifically tracks a constant value.
+
+    Args:
+        value (Any): The value of the constant.
+    """
+
+    def __init__(self, value):
+        super().__init__([])
+        self.value = value
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        codegen.gen_load_const(self.value)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.ConstantExprNode(self.value)
+
+    def trace_value_from_frame(self):
+        value_str, value_free_vars = stringify_pyobject(self.value)
+        return StringifiedExpression(
+            value_str, [], union_free_vars(value_free_vars)
+        )
+
+    def __repr__(self) -> str:
+        return f"ConstTracker(value={self.value})"
+
+    def need_guard(self) -> bool:
+        return False
+
+
+class GetAttrTracker(Tracker):
+    """
+    GetAttrTracker is a subclass of Tracker that specifically tracks the attribute access of an variable.
+
+    Args:
+        obj (VariableBase): The object whose attribute is to be tracked.
+        attr (str): The attribute to be tracked.
+    """
+
+    def __init__(self, obj: VariableBase, attr: str, changed: bool = False):
+        super().__init__([obj], changed)
+        self.obj = obj
+        self.attr = attr
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        self.obj.tracker.gen_instructions(codegen)
+        codegen.gen_load_attr(self.attr)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        obj_tracer = self.obj.tracker.guard_tree_expr_node()
+        return paddle.framework.core.AttributeExprNode(
+            obj_tracer,
+            self.attr,
+        )
+
+    def trace_value_from_frame(self):
+        obj_tracer = self.obj.tracker.trace_value_from_frame()
+        if self.attr.isidentifier():
+            expr = f"{{}}.{self.attr}"
+        else:
+            expr = f"getattr({{}}, '{self.attr}')"
+        return StringifiedExpression(
+            expr,
+            [obj_tracer],
+            union_free_vars(obj_tracer.free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return f"GetAttrTracker(attr={self.attr})"
+
+    def need_guard(self) -> bool:
+        return self.is_traceable() and self.obj.tracker.need_guard()
+
+
+class GetItemTracker(Tracker):
+    """
+    GetItemTracker is a subclass of Tracker that specifically tracks item access of a container variable.
+
+    It generates instructions and traces the item value from the frame.
+
+    Args:
+        container_var (VariableBase): The container object whose item is to be tracked.
+        key: The key/index of the item to be tracked.
+    """
+
+    def __init__(self, container_var: VariableBase, key: object, changed=False):
+        super().__init__([container_var], changed)
+        self.container = container_var
+        self.key = key
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        self.container.tracker.gen_instructions(codegen)
+        if isinstance(self.key, slice):
+            codegen.gen_load_const(self.key.start)
+            codegen.gen_load_const(self.key.stop)
+            codegen.gen_load_const(self.key.step)
+            codegen.gen_build_slice(3)
+        else:
+            codegen.gen_load_const(self.key)
+        codegen.gen_subscribe()
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        container_tracer = self.container.tracker.guard_tree_expr_node()
+        return paddle.framework.core.ItemExprNode(
+            container_tracer,
+            paddle.framework.core.ConstantExprNode(self.key),
+        )
+
+    def trace_value_from_frame(self):
+        container_tracer = self.container.tracker.trace_value_from_frame()
+        key_string, key_free_vars = stringify_pyobject(self.key)
+        return StringifiedExpression(
+            f"{{}}[{key_string}]",
+            [container_tracer],
+            union_free_vars(container_tracer.free_vars, key_free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return f"GetItemTracker(key={self.key!r})"
+
+    def need_guard(self) -> bool:
+        return self.is_traceable() and self.container.tracker.need_guard()
+
+
+class GetIterTracker(Tracker):
+    """
+    GetIterTracker is a subclass of Tracker that specifically tracks iteration of a variable.
+
+    It generates instructions and traces the iterator from the frame.
+
+    Args:
+        iter_source (VariableBase): The source variable to be iterated.
+    """
+
+    def __init__(self, iter_source: VariableBase):
+        super().__init__([iter_source])
+        self.iter_source = iter_source
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        self.iter_source.tracker.gen_instructions(codegen)
+        codegen.add_instr("GET_ITER")
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        # TODO(zrr1999): implement IterExprNode
+        raise NotImplementedError("IterExprNode is not implemented")
+
+    def trace_value_from_frame(self):
+        iter_source_tracer = self.iter_source.tracker.trace_value_from_frame()
+        return StringifiedExpression(
+            "iter({})",
+            [iter_source_tracer],
+            union_free_vars(iter_source_tracer.free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return "GetIterTracker()"
+
+
+class CreateLayerTracker(Tracker):
+    def __init__(self, layer_class, args, kwargs):
+        super().__init__([layer_class, *list(args), *list(kwargs.values())])
+        self.layer_class = layer_class
+        self.args = args
+        self.kwargs = kwargs
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        if sys.version_info >= (3, 11):
+            codegen.gen_push_null()
+
+        self.layer_class.reconstruct(codegen)
+        for variable in self.args:
+            variable.reconstruct(codegen)
+
+        if len(self.kwargs) == 0:
+            codegen.gen_call_function(argc=len(self.args))
+        else:
+            codegen.gen_build_tuple(len(self.args))
+            for k, v in self.kwargs.items():
+                codegen.gen_load_const(k)
+                v.reconstruct(codegen)
+            codegen.gen_build_map(len(self.kwargs))
+            codegen.gen_call_function_ex(has_kwargs=True)
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        # TODO(zrr1999): implement LayerExprNode.guard_tree_expr_node
+        raise NotImplementedError("LayerExprNode is not implemented")
+
+    def trace_value_from_frame(self):
+        class_tracer = self.layer_class.tracker.trace_value_from_frame()
+        arg_tracers = [
+            arg.tracker.trace_value_from_frame() for arg in self.args
+        ]
+        kwarg_tracers_dict = {
+            k: v.tracker.trace_value_from_frame()
+            for k, v in self.kwargs.items()
+        }
+        kwarg_tracers = list(kwarg_tracers_dict.values())
+
+        expr = "{}("
+        expr += ", ".join(["{}"] * len(arg_tracers))
+        if len(arg_tracers) and len(kwarg_tracers) > 0:
+            expr += ", "
+        expr += ", ".join(f"{k}={{}}" for k in kwarg_tracers_dict.keys())
+        expr += ")"
+
+        return StringifiedExpression(
+            expr,
+            [class_tracer, *arg_tracers, *kwarg_tracers],
+            union_free_vars(
+                *(
+                    tracer.free_vars
+                    for tracer in chain(
+                        [class_tracer], arg_tracers, kwarg_tracers
+                    )
+                )
+            ),
+        )
+
+    def __repr__(self) -> str:
+        return f"CreateLayerTracker(Layer={self.layer_class}, args={self.args}, kwargs={self.kwargs})"
+
+
+class FunctionClosureTracker(Tracker):
+    """
+    A tracker class that represents a function closure variable.
+
+    Args:
+        fn: The FunctionVariable object.
+        idx: The index of the closure variable.
+
+    """
+
+    def __init__(self, fn: FunctionVariable, idx: int):
+        super().__init__([fn])
+        self.fn = fn
+        self.idx = idx
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        """
+        Generate bytecode instructions to trace the value of the function closure variable.
+
+        Args:
+            codegen: The PyCodeGen object used to generate bytecode.
+
+        """
+        self.fn.tracker.gen_instructions(codegen)
+        codegen.gen_load_attr("__closure__")
+        codegen.gen_load_const(self.idx)
+        codegen.gen_subscribe()
+        codegen.gen_load_attr("cell_contents")
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        fn_tracer = self.fn.tracker.guard_tree_expr_node()
+        return paddle.framework.core.AttributeExprNode(
+            paddle.framework.core.ItemExprNode(
+                paddle.framework.core.AttributeExprNode(
+                    fn_tracer,
+                    "__closure__",
+                ),
+                paddle.framework.core.ConstantExprNode(self.idx),
+            ),
+            "cell_contents",
+        )
+
+    def trace_value_from_frame(self):
+        """
+        Trace the value of the function closure variable from the frame.
+
+        Returns:
+            The traced value of the function closure variable.
+
+        """
+        fn_tracer = self.fn.tracker.trace_value_from_frame()
+        return StringifiedExpression(
+            f"{{}}.__closure__[{self.idx}].cell_contents",
+            [fn_tracer],
+            union_free_vars(fn_tracer.free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return f"FunctionClosureTracker(fn={self.fn}, idx={self.idx})"
+
+
+class FunctionGlobalTracker(Tracker):
+    """
+    A tracker class that represents a function global variable.
+
+    Args:
+        fn: FunctionVariable object.
+        name: The name of the global variable.
+
+    """
+
+    def __init__(self, fn: FunctionVariable, name: str):
+        super().__init__([fn])
+        self.fn = fn
+        self.name = name
+
+    def gen_instructions(self, codegen: PyCodeGen):
+        """
+        Generate bytecode instructions in order to put the variables at the top of the stack.
+
+        Args:
+            codegen: The PyCodeGen object used to generate bytecode.
+
+        """
+        self.fn.tracker.gen_instructions(codegen)
+        codegen.gen_load_attr("__globals__")
+        codegen.gen_load_const(self.name)
+        codegen.gen_subscribe()
+
+    def guard_tree_expr_node(self) -> paddle.framework.core.ExprNodeBase:
+        fn_tracer = self.fn.tracker.guard_tree_expr_node()
+        return paddle.framework.core.ItemExprNode(
+            paddle.framework.core.AttributeExprNode(
+                fn_tracer,
+                "__globals__",
+            ),
+            paddle.framework.core.ConstantExprNode(self.name),
+        )
+
+    def trace_value_from_frame(self) -> StringifiedExpression:
+        """
+        Trace the value of the function global variable from the frame.
+
+        Returns:
+            StringifiedExpression: The traced value of the function global variable.
+
+        """
+        fn_tracer = self.fn.tracker.trace_value_from_frame()
+        return StringifiedExpression(
+            f"{{}}.__globals__['{self.name}']",
+            [fn_tracer],
+            union_free_vars(fn_tracer.free_vars),
+        )
+
+    def __repr__(self) -> str:
+        return f"FunctionGlobalTracker(fn={self.fn}, name={self.name})"

python/paddle/jit/sot/opcode_translator/executor/variable_stack.py

@@ -0,0 +1,217 @@
+# 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
+
+from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    ValidateValueFunc = Callable[[Any], None]
+
+
+StackDataT = TypeVar("StackDataT")
+
+
+class VariableStack(Generic[StackDataT]):
+    """
+    A stack class for storing variables.
+
+    Examples:
+        >>> var1, var2, var3, var4 = range(1, 5)
+        >>> stack = VariableStack()
+        >>> stack.push(var1)
+        >>> stack.push(var3)
+        >>> stack.insert(1, var2)
+        >>> stack
+        [1, 2, 3]
+        >>> stack.pop()
+        3
+        >>> stack.pop_n(2)
+        [1, 2]
+        >>> stack.push(var1)
+        >>> stack.push(var2)
+        >>> stack.push(var3)
+        >>> stack
+        [1, 2, 3]
+        >>> stack.top
+        3
+        >>> stack.peek[1]
+        3
+        >>> stack.peek[:1]
+        [3]
+        >>> stack.peek[:2]
+        [2, 3]
+        >>> stack.peek[1] = var4
+        >>> stack
+        [1, 2, 4]
+
+    """
+
+    class VariablePeeker:
+        @overload
+        def __getitem__(self, index: int) -> StackDataT: ...
+
+        @overload
+        def __getitem__(self, index: slice) -> list[StackDataT]: ...
+
+        @overload
+        def __call__(self, index: int = 1) -> StackDataT: ...
+
+        @overload
+        def __call__(self, index: slice) -> list[StackDataT]: ...
+
+        def __init__(
+            self, data: list[StackDataT], validate_value_func: ValidateValueFunc
+        ):
+            self._data = data
+            self.validate_value_func = validate_value_func
+
+        def __getitem__(
+            self, index: int | slice
+        ) -> StackDataT | list[StackDataT]:
+            if isinstance(index, int):
+                assert 0 < index <= len(self._data)
+                return self._data[-index]
+            if isinstance(index, slice):
+                assert index.start is None and index.step is None, (
+                    "slice which has start or step not supported"
+                )
+                assert 0 < index.stop <= len(self._data)
+                return self._data[-index.stop :]
+            raise NotImplementedError(f"index type {type(index)} not supported")
+
+        def __setitem__(self, index: int, value: Any):
+            assert isinstance(index, int), (
+                f"index type {type(index)} not supported"
+            )
+            assert 0 < index <= len(self._data), (
+                f"index should be in [1, {len(self._data)}], but get {index}"
+            )
+            self.validate_value_func(value)
+            self._data[-index] = value
+
+        def __call__(
+            self, index: int | slice = 1
+        ) -> StackDataT | list[StackDataT]:
+            return self[index]
+
+    def __init__(
+        self,
+        data: list[StackDataT] | None = None,
+        *,
+        validate_value_func: ValidateValueFunc | None = None,
+    ):
+        if data is None:
+            data = []
+        else:
+            data = data.copy()
+        self.validate_value_func = (
+            (lambda _: None)
+            if validate_value_func is None
+            else validate_value_func
+        )
+        self._data = data
+        self._peeker = VariableStack.VariablePeeker(
+            self._data, self.validate_value_func
+        )
+
+    def copy(self):
+        return VariableStack(
+            self._data, validate_value_func=self.validate_value_func
+        )
+
+    def push(self, val: StackDataT):
+        """
+        Pushes a variable onto the stack.
+
+        Args:
+            val: The variable to be pushed.
+
+        """
+        self.validate_value_func(val)
+        self._data.append(val)
+
+    def insert(self, index: int, val: StackDataT):
+        """
+        Inserts a variable onto the stack.
+
+        Args:
+            index: The index at which the variable is to be inserted, the top of the stack is at index 0.
+            val: The variable to be inserted.
+
+        """
+        assert 0 <= index <= len(self), (
+            f"index should be in [0, {len(self)}], but get {index}"
+        )
+        self.validate_value_func(val)
+        self._data.insert(len(self) - index, val)
+
+    def pop(self) -> StackDataT:
+        """
+        Pops the top value from the stack.
+
+        Returns:
+            The popped value.
+
+        """
+        assert len(self) > 0, "stack is empty"
+        return self._data.pop()
+
+    def pop_n(self, n: int) -> list[StackDataT]:
+        """
+        Pops the top n values from the stack.
+
+        Args:
+            n: The number of values to pop.
+
+        Returns:
+            A list of the popped values.
+
+        """
+        assert len(self) >= n >= 0, (
+            f"n should be in [0, {len(self)}], but get {n}"
+        )
+        if n == 0:
+            return []
+        retval = self._data[-n:]
+        self._data[-n:] = []
+        return retval
+
+    @property
+    def peek(self) -> VariablePeeker:
+        return self._peeker
+
+    @property
+    def top(self) -> StackDataT:
+        assert len(self) > 0, "stack is empty"
+        return self.peek[1]
+
+    @top.setter
+    def top(self, value):
+        assert len(self) > 0, "stack is empty"
+        self.peek[1] = value
+
+    def __contains__(self, value):
+        return value in self._data
+
+    def __iter__(self):
+        return iter(self._data)
+
+    def __len__(self) -> int:
+        return len(self._data)
+
+    def __repr__(self) -> str:
+        return str(self._data)

python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py

@@ -0,0 +1,80 @@
+# 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 .base import (  # noqa: F401
+    VariableBase,
+    VariableFactory,
+    find_traceable_vars,
+    map_variables,
+)
+from .basic import (  # noqa: F401
+    CellVariable,
+    ConstantVariable,
+    DataClassInstanceVariable,
+    DataVariable,
+    DygraphTracerVariable,
+    EnumVariable,
+    ExceptionVariable,
+    FunctionGlobalVariable,
+    GlobalVariable,
+    InterpolationVariable,
+    ModuleVariable,
+    NullVariable,
+    NumPyArrayVariable,
+    NumPyNumberVariable,
+    NumPyVariable,
+    ObjectVariable,
+    ParameterVariable,
+    PlaceVariable,
+    SliceVariable,
+    SuperVariable,
+    SymbolicVariable,
+    TemplateVariable,
+    TensorVariable,
+)
+from .callable import (  # noqa: F401
+    BuiltinVariable,
+    CallableVariable,
+    ClassVariable,
+    ContainerLayerVariable,
+    DataClassVariable,
+    FunctionVariable,
+    LayerVariable,
+    MethodVariable,
+    NumPyApiVariable,
+    PaddleApiVariable,
+    PaddleLayerVariable,
+    PartialVariable,
+    UserCodeVariable,
+    UserDefinedFunctionVariable,
+    UserDefinedGeneratorFunctionVariable,
+    UserDefinedLayerVariable,
+)
+from .container import (  # noqa: F401
+    ContainerVariable,
+    DictVariable,
+    ListVariable,
+    RangeVariable,
+    SizeVariable,
+    TupleVariable,
+)
+from .iter import (  # noqa: F401
+    EnumerateVariable,
+    GeneratorVariable,
+    IterVariable,
+    MapVariable,
+    SequenceIterVariable,
+    UserDefinedIterVariable,
+    ZipVariable,
+)

python/paddle/jit/sot/opcode_translator/executor/variables/base.py

@@ -0,0 +1,731 @@
+# 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 inspect
+import operator
+from contextlib import contextmanager
+from dataclasses import fields
+from functools import cached_property
+from queue import Queue
+from typing import TYPE_CHECKING, Any
+
+import paddle
+from paddle.jit.dy2static.utils import (
+    dataclass_from_dict,
+)
+
+from ....profiler import event_register
+from ....utils import (
+    NameGenerator,
+    get_unbound_method,
+    log,
+)
+from ....utils.exceptions import FallbackError, HasNoAttributeError
+from ..dispatcher import Dispatcher
+from ..guard import (
+    FasterStringifiedExpression,
+    StringifiedExpression,
+    check_faster_guard,
+    check_guard,
+    union_free_vars,
+)
+from ..mutable_data import MutableDictLikeData
+from ..tracker import (
+    BuiltinTracker,
+    ConstTracker,
+    DummyTracker,
+    GetAttrTracker,
+    GetItemTracker,
+    GetIterTracker,
+    GlobalTracker,
+    LocalTracker,
+    Tracker,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import TypeAlias
+
+    from ..function_graph import FunctionGraph
+    from ..pycode_generator import PyCodeGen
+
+    # Each variable object should implement a method called `from_value`,
+    # which should adhere to the FromValueFunc signature.
+    FromValueFunc: TypeAlias = Callable[
+        [Any, FunctionGraph, Tracker], "VariableBase | None"
+    ]
+
+
+@event_register("find_traceable_vars")
+def find_traceable_vars(
+    root_vars: list[VariableBase],
+) -> list[VariableBase]:
+    """
+    This function is used to find all traceable variables in the given list of variables.
+
+    Args:
+        root_vars (list[VariableBase]): A list of root variables from which the ordering starts.
+
+    Returns:
+        list[VariableBase]: A list of variables that are traceable.
+    """
+    results: list[VariableBase] = []
+    visited: set[VariableBase] = set()
+    queue: Queue[VariableBase] = Queue()
+
+    for root in root_vars:
+        queue.put(root)
+
+    while not queue.empty():
+        var = queue.get()
+        if var in visited:
+            continue
+
+        visited.add(var)
+        if var.tracker.need_guard():
+            results.append(var)
+            continue
+
+        # Pruning traceable variable, if the variable is traceable, we don't need to
+        # trace its inputs.
+        inputs = var.get_inputs()
+
+        for var in inputs:
+            if var not in visited and var not in queue.queue:
+                queue.put(var)
+
+    return results
+
+
+def map_variables(
+    map_func,
+    variables: list[VariableBase],
+    *,
+    restore_variable=False,
+) -> list[VariableBase]:
+    """
+    This function maps the given map_func to the given list of variables in a recursive manner.
+    Args:
+        map_func (Callable[[VariableBase], Any]): The function to be mapped to each variable.
+        variables (list[VariableBase]): A list of variables to which the map_func is to be applied.
+
+    Returns:
+        tuple: The result of applying the map_func to the variables.
+    """
+    from .basic import DataClassInstanceVariable, SliceVariable
+    from .container import ContainerVariable
+
+    def _map_container_variable(variable: VariableBase | object):
+        if not isinstance(variable, ContainerVariable):
+            return variable
+        new_container = paddle.utils.map_structure(
+            _map_variable, variable.get_wrapped_items()
+        )
+        if not restore_variable:
+            return new_container
+        return VariableFactory.from_value(
+            new_container,
+            variable.graph,
+            DummyTracker(paddle.utils.flatten(new_container)),
+        )
+
+    def _map_slice_variable(variable: VariableBase | object):
+        if not isinstance(variable, SliceVariable):
+            return variable
+        new_slice = slice(
+            map_func(variable.getattr("start")),
+            map_func(variable.getattr("stop")),
+            map_func(variable.getattr("step")),
+        )
+        if not restore_variable:
+            return new_slice
+        return VariableFactory.from_value(
+            new_slice,
+            variable.graph,
+            DummyTracker([new_slice.start, new_slice.stop, new_slice.step]),
+        )
+
+    def _map_dataclass_variable(variable: VariableBase | object):
+        if not isinstance(variable, DataClassInstanceVariable):
+            return variable
+        new_dataclass = dataclass_from_dict(
+            variable.get_py_type(),
+            {
+                fd.name: _map_variable(variable.getattr(fd.name))
+                for fd in fields(variable.get_py_type())
+            },
+        )
+        if not restore_variable:
+            return new_dataclass
+        return VariableFactory.from_value(
+            new_dataclass,
+            variable.graph,
+            DummyTracker(
+                [
+                    variable.getattr(fd.name)
+                    for fd in fields(variable.get_py_type())
+                ]
+            ),
+        )
+
+    def _map_variable(variable: VariableBase | object):
+        variable = _map_container_variable(variable)
+        variable = _map_slice_variable(variable)
+        variable = _map_dataclass_variable(variable)
+        return map_func(variable)
+
+    return paddle.utils.map_structure(_map_variable, variables)
+
+
+class VariableFactory:
+    """
+    A factory class for creating variables from arbitrary values.
+
+    This class provides a set of registration and factory methods for creating variables
+    of different types based on the type of the input value.
+
+    """
+
+    registered_funcs: dict[str, list[str]] = {"default": []}
+    mapping_str_func: dict[str, FromValueFunc] = {}
+
+    @staticmethod
+    def default_from_value(value, graph, tracker):
+        """
+        A default factory function that creates an ObjectVariable from the given value.
+
+        Args:
+            value: The input value.
+            graph: The FunctionGraph object that this variable is associated with.
+            tracker: The Tracker object that tracks the information of this variable.
+
+        Returns:
+            ObjectVariable: A new ObjectVariable representing the input value.
+        """
+        from .basic import ObjectVariable
+
+        return ObjectVariable(value, graph, tracker)
+
+    @staticmethod
+    def register_from_value(*, successor: str | None = None):
+        """
+        A decorator function that registers a function for creating a Variable from a value.
+
+        Args:
+            successor (str | None, optional): The name of the successor function that will be called after this function when creating a Variable. If None, the function is added to a default list of functions.
+
+        Returns:
+            The _register_from_value decorator function, which takes the function to be registered as an argument.
+        """
+        registered_funcs = VariableFactory.registered_funcs
+        mapping_str_func = VariableFactory.mapping_str_func
+
+        def _register_from_value(func: FromValueFunc):
+            """
+            Function to register a function for creating a Variable from a value
+            """
+            # Get the name of the function
+            name = func.__qualname__.split(".")[0]
+            # Map the name of the function to the function
+            mapping_str_func[name] = func
+            if successor is None:
+                registered_funcs["default"].append(
+                    name
+                )  # If successor is None, add the function to the "default" list
+            elif successor not in registered_funcs.keys():
+                registered_funcs[successor] = [
+                    name
+                ]  # If the successor is not in the registered_funcs dictionary, set the value to a list containing only name
+            else:
+                registered_funcs[successor].append(
+                    name
+                )  # If the successor is in the registered_funcs dictionary, append name to the existing list of functions for that successor
+
+        log(
+            4, VariableFactory.registered_funcs
+        )  # Print the registered_funcs dictionary if the logging level is at least 4
+        return _register_from_value
+
+    @staticmethod
+    def from_value(
+        value: Any,
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ) -> VariableBase:
+        """
+        Create a new variable object from the given value.
+
+        This method searches through the registered from_value functions to find one
+        that can create a variable object from the given value. If no matching function
+        is found, the default_from_value function is used.
+
+        Args:
+            value (Any): The input value.
+            graph (FunctionGraph): The FunctionGraph object that this variable is associated with.
+            tracker (Tracker): The Tracker object that tracks the information of this variable.
+
+        Returns:
+            VariableBase: A new variable object representing the input value.
+        """
+        registered_funcs = VariableFactory.registered_funcs
+
+        def _find_var(key: str = "default") -> VariableBase | None:
+            for name in registered_funcs[key]:
+                if name in registered_funcs.keys():
+                    # If the function name is a key in the registered_funcs dictionary, recursively find a Variable using that function
+                    var = _find_var(name)
+                    if var is not None:
+                        return var
+                # Get the function corresponding to the name from the mapping_str_func dictionary
+                func = VariableFactory.mapping_str_func[name]
+                var = func(
+                    value, graph, tracker
+                )  # Call the function to create a Variable from the value
+                if var is not None:
+                    return var
+
+        var = _find_var()
+        if var is None:
+            var = VariableFactory.default_from_value(
+                value, graph, tracker
+            )  # If a Variable could not be found using the registered functions, use the default function to create a new Variable
+        return var
+
+
+def infer_debug_name_from_tracker(tracker: Tracker) -> str | None:
+    res = None
+    if isinstance(tracker, (LocalTracker, GlobalTracker, BuiltinTracker)):
+        res = f"{tracker.name}"
+    elif isinstance(tracker, ConstTracker):
+        res = f"{tracker.value}"
+    elif isinstance(tracker, GetItemTracker) and tracker.container.debug_name:
+        res = f"{tracker.container.debug_name}[{tracker.key}]"
+    elif isinstance(tracker, GetAttrTracker) and tracker.obj.debug_name:
+        res = f"{tracker.obj.debug_name}.{tracker.attr}"
+    return res
+
+
+class VariableBase:
+    """
+    VariableBase is a basic concept and each symbols in VM stack is regarded as
+    an Variable Object in symbolic tracing process.
+
+    There are two key data structures during Python runtime:
+    PyFrameObject, which provides the instance for function logical lock usage,
+    and PyCodeObject, which provides the bytecode for the corresponding function.
+    With these data, the Python virtual machine executes the bytecode sequentially on a stack to complete function logic.
+
+    Args:
+        tracker(Tracker): The Tracker object that tracks the information of this variable.
+
+    Note:
+        We should push an object of a subclass of VariableBase instead of an object of VariableBase onto the VM stack.
+        It serves as an abstract class and should not be instantiated directly.
+    """
+
+    tracker: Tracker  # An attribute to store the Tracker object associated with the variable
+    value: Any
+    name_generator = NameGenerator(
+        "object_"
+    )  # A class-level attribute to generate names for new variables
+    mutable_attrs = []
+
+    def __init__(self, graph: FunctionGraph, tracker: Tracker):
+        self.graph = graph
+        self.tracker = tracker
+        self.id = VariableBase.name_generator.next()
+        self.debug_name = infer_debug_name_from_tracker(tracker)
+
+    @property
+    def main_info(self) -> dict[str, Any]:
+        """
+        Property method to return a dictionary of main information about the variable
+
+        Returns:
+            main_info: Main information of the variable.
+        """
+        return {}
+
+    @property
+    def debug_info(self) -> dict[str, Any]:
+        """
+        Property method to return a dictionary of debug information about the variable
+        """
+        info = {
+            "id": self.id,
+        }
+        if self.debug_name:
+            info["debug_name"] = self.debug_name
+        return info
+
+    def __hash__(self):
+        return hash(self.id)
+
+    @check_faster_guard
+    def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
+        expr_node = self.tracker.guard_tree_expr_node()
+        return [
+            paddle.framework.core.GuardNode(
+                paddle.framework.core.ValueMatchGuard(self.get_py_value()),
+                [expr_node],
+            )
+        ]
+
+    @check_guard
+    def make_stringified_guard(self) -> list[StringifiedExpression]:
+        """
+        Create a StringifiedExpression object that represents a guard expression for this variable.
+
+        Returns:
+            StringifiedExpression: An object that contains the guard expression and the free variables used in the expression.
+        """
+
+        # Get a ValueTracer object from the Tracker object associated with the variable
+        frame_value_tracer = self.tracker.trace_value_from_frame()
+        return [
+            FasterStringifiedExpression(
+                f"id(type({{0}})) == {id(self.get_py_type())} and {{0}} == {self.get_py_value()!r}",
+                paddle.framework.core.ValueMatchGuard(self.get_py_value()),
+                [frame_value_tracer],
+                union_free_vars(frame_value_tracer.free_vars),
+            )
+        ]
+
+    def get_py_value(self, allow_tensor=False) -> Any:
+        """
+        Abstract method to get the value of the variable
+        """
+        raise NotImplementedError
+
+    def get_py_type(self):
+        """
+        Method to get the type of the variable's value
+        """
+        return type(self.get_py_value())
+
+    def is_none(self) -> bool:
+        """
+        Method to check if the variable's value is None
+        """
+        return self.get_py_value() is None
+
+    def reconstruct(
+        self,
+        codegen: PyCodeGen,
+        *,
+        use_tracker: bool = True,
+        add_to_global_guarded_vars: bool = True,
+    ):
+        if self.tracker.is_traceable() and use_tracker:
+            self.tracker.gen_instructions(codegen)
+        else:
+            if add_to_global_guarded_vars:
+                self.graph.add_global_guarded_variable(self)
+            self._reconstruct(codegen)
+
+    def _reconstruct(self, codegen: PyCodeGen) -> None:
+        """
+        Abstract method to construct an opcode and append it into codegen.instructions
+        """
+        raise FallbackError(
+            f'{self.__class__.__name__} does not implement "_reconstruct" method'
+        )
+
+    def flatten_inner_vars(self) -> list[VariableBase]:
+        """
+        Recursively flatten the items in this container variable to a list of Variable objects.
+
+        Returns:
+            list[VariableBase]: Flattened items of a container variable.
+        """
+        return [self]
+
+    def get_inputs(self) -> list[VariableBase]:
+        """
+        This method is used to get the inputs for the current variable.
+
+        Returns:
+            list[VariableBase]: Inputs for the current variable.
+        """
+        return self.tracker.inputs
+
+    def get_traceable_inputs(self) -> list[VariableBase]:
+        """
+        This method is used to get the traceable inputs for the current variable.
+
+        Returns:
+            list[VariableBase]: Traceable inputs for the current variable.
+        """
+        return list(
+            filter(lambda x: x.tracker.is_traceable(), self.tracker.inputs)
+        )
+
+    def call_function(self, /, *args, **kwargs):
+        pass
+
+    @cached_property
+    def attr_proxy(self):
+        return self.graph.side_effects.get_proxy(
+            MutableDictLikeData, self.get_py_value(), self.attr_proxy_getter
+        )
+
+    def attr_proxy_getter(self, proxy: MutableDictLikeData, name: str):
+        if not hasattr(proxy.original_data, name):  # can't true.
+            return MutableDictLikeData.Empty()
+
+        attr = getattr(proxy.original_data, name)
+        if inspect.ismethod(attr) or (
+            hasattr(attr, "__self__")
+            and inspect.ismethoddescriptor(
+                getattr(attr.__self__.__class__, name, None)
+            )
+        ):
+            from .callable import MethodVariable
+
+            fn = None
+            instance = self
+            if inspect.ismethoddescriptor(
+                getattr(attr.__self__.__class__, name, None)
+            ):
+                class_var = VariableFactory.from_value(
+                    self.get_py_type(),
+                    self.graph,
+                    GetAttrTracker(self, "__class__"),
+                )
+                fn = VariableFactory.from_value(
+                    getattr(attr.__self__.__class__, name),
+                    self.graph,
+                    GetAttrTracker(class_var, name),
+                )
+            if not hasattr(self.get_py_type(), name):
+                instance = None
+            return MethodVariable.wrap_method(
+                value=attr,
+                instance=instance,
+                fn=fn,
+                graph=self.graph,
+                tracker=GetAttrTracker(self, name),
+            )
+
+        return VariableFactory.from_value(
+            attr, self.graph, tracker=GetAttrTracker(self, name)
+        )
+
+    def hasattr(self, name: str):
+        from .basic import ConstantVariable
+
+        try:
+            self.getattr(name)
+            return ConstantVariable(
+                True, graph=self.graph, tracker=DummyTracker([self])
+            )
+        except HasNoAttributeError:
+            # NOTE(SigureMo): Only the HasNoAttributeError is raised, we can
+            # ensure that the attribute does not exist. Otherwise, we should
+            # raise the error.
+            return ConstantVariable(
+                False, graph=self.graph, tracker=DummyTracker([self])
+            )
+
+    def getattr(self, name: str, default=None):
+        result = self.attr_proxy.get(name)
+        if isinstance(result, MutableDictLikeData.Empty):
+            if default is not None:
+                assert isinstance(default, VariableBase)
+                return default
+            raise HasNoAttributeError(
+                f"{self.__class__.__name__} {self} has no attribute {name}"
+            )
+        return result
+
+    def setattr(self, key: str, value):
+        from .basic import ConstantVariable
+
+        self.attr_proxy.set(key, value)
+        self.graph.side_effects.record_proxy_variable(self)
+        return ConstantVariable.wrap_literal(None, self.graph)
+
+    def delattr(self, key: str):
+        from .basic import ConstantVariable
+
+        self.attr_proxy.delete(key)
+        self.graph.side_effects.record_proxy_variable(self)
+        return ConstantVariable.wrap_literal(None, self.graph)
+
+    def __setitem__(self, key, value):
+        return self.setitem(key, value)
+
+    def setitem(self, key, value):
+        raise FallbackError(f"{self} is not support setitem.")
+
+    def __repr__(self):
+        info = self.main_info | self.debug_info
+        info_str = ", ".join([f"{value}" for value in info.values()])
+        return f"{self.__class__.__name__}({info_str})"
+
+    def __str__(self):
+        return self.__repr__()
+
+    def __getitem__(self, idx):
+        return Dispatcher.call(operator.getitem, self, idx)
+
+    def getitem(self, item):
+        class_var = VariableFactory.from_value(
+            self.get_py_value().__class__,
+            self.graph,
+            GetAttrTracker(self, '__class__'),
+        )
+        fn_var = VariableFactory.from_value(
+            get_unbound_method(self.get_py_value(), '__getitem__'),
+            self.graph,
+            GetAttrTracker(class_var, '__getitem__'),
+        )
+        self.graph.add_global_guarded_variable(item)
+        item = item.get_py_value()
+        output = fn_var(self, item)
+        return output
+
+    def __call__(self, /, *args, **kwargs):
+        """
+        Call the object represented by this variable with the given arguments.
+
+        Args:
+            *args: Positional arguments to pass to the object's __call__ method.
+            **kwargs: Keyword arguments to pass to the object's __call__ method.
+
+        Returns:
+            VariableBase: A new variable representing the result of calling the object's __call__ method.
+        """
+        from .callable import BuiltinVariable, UserDefinedFunctionVariable
+
+        class_var = VariableFactory.from_value(
+            self.get_py_value().__class__,
+            self.graph,
+            GetAttrTracker(self, '__class__'),
+        )
+        assert class_var is not None
+        # if __call__ is a method, we should add self to arguments.
+        if inspect.ismethod(self.get_py_value().__call__):
+            args = (self, *args)
+        unbound_method = get_unbound_method(self.get_py_value(), '__call__')
+        if hasattr(unbound_method, "__code__"):
+            fn_var = UserDefinedFunctionVariable(
+                unbound_method,
+                self.graph,
+                GetAttrTracker(class_var, '__call__'),
+            )
+        else:
+            fn_var = BuiltinVariable(
+                self.value,
+                self.graph,
+                GetAttrTracker(class_var, '__call__'),
+            )
+        output = fn_var(*args, **kwargs)
+        return output
+
+    def get_iter(self):
+        from . import (
+            BuiltinVariable,
+            ConstantVariable,
+            SequenceIterVariable,
+            UserDefinedFunctionVariable,
+            UserDefinedIterVariable,
+        )
+
+        if not hasattr(self.value, "__iter__"):
+            return UserDefinedIterVariable(
+                self, self.graph, GetIterTracker(self)
+            )
+        iter_name_var = ConstantVariable.wrap_literal("__iter__", self.graph)
+        iter_method = BuiltinVariable(
+            getattr, graph=self.graph, tracker=DummyTracker([self])
+        )(self, iter_name_var)
+        # If the target object is a builtin object like list_iterator, the iter_method's fn will be a ObjectVariable instead of UserDefinedFunctionVariable.
+        if not isinstance(iter_method.fn, UserDefinedFunctionVariable):
+            return UserDefinedIterVariable(
+                self, self.graph, GetIterTracker(self)
+            )
+        iter_result = iter_method()
+
+        if not isinstance(iter_result, SequenceIterVariable):
+            return UserDefinedIterVariable(
+                self, self.graph, GetIterTracker(self)
+            )
+
+        return iter_result
+
+    @VariableFactory.register_from_value()
+    def from_value(
+        value: Any,
+        graph: FunctionGraph | None,
+        tracker: Tracker,
+    ) -> VariableBase | None:
+        """
+        Create a new variable from a given value, or return None if the value cannot be converted to a variable.
+        Args:
+            value (Any): The value to create a variable from.
+            graph (FunctionGraph | None): The graph in which the variable will be used.
+            tracker (Tracker): The variable tracker to put the new variable in if created.
+
+        Returns:
+            VariableBase | None: A new variable if one can be created from the given value, or None if the value cannot be converted to a variable.
+        """
+        if isinstance(value, VariableBase):
+            return value
+        return None
+
+
+@contextmanager
+def signature_clear_guard(fn, name):
+    if not hasattr(fn, name):
+        yield
+    else:
+        saved_attr = getattr(fn, name)
+        delattr(fn, name)
+        yield
+        setattr(fn, name, saved_attr)
+
+
+def fn_bind_inputs(
+    fn: Callable[..., Any],
+    graph: FunctionGraph,
+    *args: Any,
+    **kwargs: Any,
+):
+    # temparay clear the fn.__signature__ to avoid signature check error
+    with (
+        signature_clear_guard(fn, "__signature__"),
+        signature_clear_guard(fn, "__wrapped__"),
+    ):
+        sig = inspect.signature(fn)
+        bound_args = sig.bind(*args, **kwargs)
+    bound_args.apply_defaults()
+    parameters = {}
+    for name, value in bound_args.arguments.items():
+        assert name in sig.parameters
+        # Convert varargs and kwargs to Variable
+        if sig.parameters[name].kind == inspect.Parameter.VAR_POSITIONAL:
+            tracker = DummyTracker(value)
+        elif sig.parameters[name].kind == inspect.Parameter.VAR_KEYWORD:
+            tracker = DummyTracker(list(value.values()))
+        # Convert default args to Variable
+        elif not isinstance(value, VariableBase):
+            tracker = ConstTracker(value)
+        else:
+            tracker = value.tracker
+        value = VariableFactory.from_value(value, graph, tracker)
+        parameters[name] = value
+    return parameters

python/paddle/jit/sot/opcode_translator/executor/variables/iter.py

@@ -0,0 +1,460 @@
+# 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 types
+from typing import TYPE_CHECKING, Any
+
+from paddle._typing import unreached
+
+from ....profiler import EventGuard
+from ....utils import do_until_stop_iteration
+from ....utils.exceptions import (
+    BreakGraphError,
+    BreakGraphInlineCallBreak,
+    FallbackError,
+    FallbackInlineCallBreak,
+    OtherInlineCallBreak,
+    SotCapturedExceptionFactory,
+    SotCapturedStopIteration,
+    SotErrorBase,
+    UnsupportedOperationBreak,
+)
+from ..guard import check_faster_guard
+from ..tracker import ConstTracker, DanglingTracker, DummyTracker
+from .base import (
+    VariableBase,
+    VariableFactory,
+)
+from .basic import ConstantVariable
+from .callable import BuiltinVariable
+from .container import TupleVariable
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    import paddle
+
+    from ..function_graph import FunctionGraph
+    from ..pycode_generator import PyCodeGen
+    from ..tracker import Tracker
+    from ..virtual_frame import VirtualFrame
+
+
+class IterVariable(VariableBase):
+    """
+    This Variable (include subclasses) should be generated only when simulate GET_ITER opcode
+    """
+
+    def __init__(self, graph: FunctionGraph, tracker: Tracker):
+        super().__init__(graph, tracker)
+
+    def next(self):
+        raise NotImplementedError(f"Can not simulate `next` for {type(self)}")
+
+    def to_list(self):
+        raise NotImplementedError(
+            f"Can not simulate `to_list` for {type(self)}"
+        )
+
+    def send(self, value: VariableBase):
+        return self.next()
+
+    def get_iter(self):
+        return self
+
+
+class SequenceIterVariable(IterVariable):
+    """
+    The basic SequenceIterVariable wraps iterators which can be simulated by call getitem
+    Currently includes: List | Tuple | Dict (keys) | Range | Tensor | nn.LayerList
+
+    these interfaces is needed:
+    - next
+    - to_list
+    - has_side_effect
+    - _reconstruct
+    """
+
+    mutable_attrs = ["idx"]
+
+    def __init__(
+        self,
+        held: VariableBase | list[VariableBase],
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        if not isinstance(held, list):
+            held = [held]
+        super().__init__(graph, tracker)
+        self.holds = held
+        self.idx = 0
+        self.graph.side_effects.record_mutable_variable(self)
+
+    @check_faster_guard
+    def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
+        return [
+            guard for held in self.holds for guard in held.make_faster_guard()
+        ]
+
+    def make_stringified_guard(self):
+        return [
+            guard
+            for held in self.holds
+            for guard in held.make_stringified_guard()
+        ]
+
+    def next(self):
+        held = self.holds[0]
+        if self.idx < len(held):
+            val = held[self.idx]
+            self.idx += 1
+            return val
+        else:
+            raise SotCapturedExceptionFactory.create(StopIteration())
+
+    def to_list(self) -> list:
+        if self.has_side_effect():
+            raise FallbackError("Can not convert an used iterator into list")
+        held = self.holds[0]
+        self.idx = len(held)
+        retval = []
+        for i in range(len(held)):
+            retval.append(held[i])
+        return retval
+
+    def has_side_effect(self) -> bool:
+        return self.idx != 0
+
+    def _reconstruct(self, codegen: PyCodeGen):
+        if self.has_side_effect():
+            super()._reconstruct(codegen)
+        else:
+            self.holds[0].reconstruct(codegen)
+            codegen.gen_get_iter()
+
+    @property
+    def main_info(self) -> dict[str, Any]:
+        return {
+            "idx": self.idx,
+        }
+
+    def flatten_inner_vars(self) -> list[VariableBase]:
+        held = self.holds
+        return [
+            inner_var for obj in held for inner_var in obj.flatten_inner_vars()
+        ]
+
+
+class EnumerateVariable(SequenceIterVariable):
+    """
+    EnumerateVariable holds a SequenceIterVariable and return additional index
+    """
+
+    def __init__(
+        self, val_iterator: IterVariable, graph: FunctionGraph, tracker: Tracker
+    ):
+        super().__init__(val_iterator, graph, tracker)
+
+    def next(self):
+        val = self.holds[0].next()
+        idx_var = ConstantVariable(self.idx, self.graph, ConstTracker(self.idx))
+        self.idx += 1
+        return TupleVariable(
+            (idx_var, val), self.graph, DummyTracker([idx_var, val])
+        )
+
+    def to_list(self):
+        values = self.holds[0].to_list()
+        idx = [
+            ConstantVariable(i, self.graph, ConstTracker(i))
+            for i in range(len(values))
+        ]
+        return list(zip(idx, values))
+
+    def has_side_effect(self) -> bool:
+        return self.holds[0].has_side_effect()
+
+    def _reconstruct(self, codegen: PyCodeGen):
+        if self.has_side_effect():
+            super()._reconstruct(codegen)
+        else:
+            codegen.gen_load_global("enumerate", push_null=True)
+            self.holds[0].reconstruct(codegen)
+            codegen.gen_call_function(1)
+
+    @staticmethod
+    def from_iterator(value, graph: FunctionGraph | None, tracker: Tracker):
+        iter_variable = value.get_iter()
+        if isinstance(iter_variable, UserDefinedIterVariable):
+            return UserDefinedIterVariable(value, graph, tracker)
+        else:
+            return EnumerateVariable(iter_variable, graph, tracker)
+
+
+class ZipVariable(SequenceIterVariable):
+    """
+    ZipVariable holds a list of SequenceIterVariable
+    """
+
+    def __init__(
+        self, iters: list[IterVariable], graph: FunctionGraph, tracker: Tracker
+    ):
+        super().__init__(iters, graph, tracker)
+
+    def next(self):
+        # can not use <listcomp> here, because it will raise a RuntimeError("StopIteration")
+        # but we want a StopIteration Exception
+        values = []
+        for iter_var in self.holds:
+            next_var = iter_var.next()
+            values.append(next_var)
+
+        return VariableFactory.from_value(
+            tuple(values), self.graph, DummyTracker(values)
+        )
+
+    def to_list(self):
+        lists = [iter_vars.to_list() for iter_vars in self.holds]
+        min_len = min(len(l) for l in lists)
+        result = []
+        for i in range(min_len):
+            result.append(
+                VariableFactory.from_value(
+                    tuple(l[i] for l in lists),
+                    self.graph,
+                    DummyTracker(list(self.holds)),
+                )
+            )
+        return result
+
+    def has_side_effect(self) -> bool:
+        return any(iter_var.has_side_effect() for iter_var in self.holds)
+
+    def _reconstruct(self, codegen: PyCodeGen):
+        if self.has_side_effect():
+            super()._reconstruct(codegen)
+        else:
+            codegen.gen_load_global("zip", push_null=True)
+            for iter_var in self.holds:
+                iter_var.reconstruct(codegen)
+            codegen.gen_call_function(len(self.holds))
+
+    @staticmethod
+    def from_iterator(
+        value: Sequence[VariableBase],
+        graph: FunctionGraph | None,
+        tracker: Tracker,
+    ):
+        assert isinstance(value, (list, tuple))
+        zip_targets = []
+
+        for variable in value:
+            iter_variable = variable.get_iter()
+            if isinstance(iter_variable, UserDefinedIterVariable):
+                return UserDefinedIterVariable(value, graph, tracker)
+            zip_targets.append(iter_variable)
+
+        return ZipVariable(zip_targets, graph, tracker)
+
+
+class MapVariable(SequenceIterVariable):
+    """
+    MapVariable holds a SequenceIterVariable and return a Iterable Variable after map function
+    """
+
+    def __init__(self, fn, iters: list[IterVariable], graph, tracker):
+        super().__init__(iters, graph, tracker)
+        self.fn = fn
+
+    def next(self):
+        return self.fn(*[iter_var.next() for iter_var in self.holds])
+
+    def to_list(self) -> list:
+        lists = [iter_var.to_list() for iter_var in self.holds]
+        min_len = min(len(l) for l in lists)
+        result = []
+        for i in range(min_len):
+            result.append(self.fn(*(l[i] for l in lists)))
+        return result
+
+    def has_side_effect(self) -> bool:
+        return any(iter_var.has_side_effect() for iter_var in self.holds)
+
+    def _reconstruct(self, codegen: PyCodeGen):
+        if self.has_side_effect():
+            super()._reconstruct(codegen)
+        else:
+            codegen.gen_load_global("map", push_null=True)
+            self.fn.reconstruct(codegen)
+            for iter_var in self.holds:
+                iter_var.reconstruct(codegen)
+            codegen.gen_call_function(len(self.holds) + 1)
+
+    @staticmethod
+    def from_iterator(
+        fn,
+        value: Sequence[VariableBase],
+        graph: FunctionGraph | None,
+        tracker: Tracker,
+    ):
+        map_targets = []
+
+        for variable in value:
+            iter_variable = variable.get_iter()
+            if isinstance(iter_variable, UserDefinedIterVariable):
+                return UserDefinedIterVariable(value, graph, tracker)
+            map_targets.append(iter_variable)
+
+        return MapVariable(fn, map_targets, graph, tracker)
+
+
+class GeneratorVariable(IterVariable):
+    def __init__(
+        self,
+        code_var: VariableBase,
+        vframe: VirtualFrame,
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        self.code_var = code_var
+        self.vframe = vframe
+        self.shared_stack = []
+        super().__init__(graph, tracker)
+
+    def send(self, /, value: VariableBase):
+        from ..opcode_inline_executor import OpcodeInlineGeneratorExecutor
+
+        checkpoint = self.graph.save_memo()
+        frame_state = self.vframe.get_state()
+        try:
+            inline_gen_executor = OpcodeInlineGeneratorExecutor(
+                self.vframe, self.code_var, self.graph
+            )
+            self.vframe.stack.push(value)
+            with EventGuard(
+                f"Inline Gen Call: {inline_gen_executor.vframe.code.co_name}, file {inline_gen_executor.vframe.code.co_filename}, line {int(inline_gen_executor.vframe.code.co_firstlineno)}"
+            ):
+                output: VariableBase = inline_gen_executor.inline_call()
+                if inline_gen_executor.stop_state == "Return":
+                    raise SotCapturedExceptionFactory.create(StopIteration())
+        except SotCapturedStopIteration:
+            raise
+        except SotErrorBase as error:
+            self.graph.restore_memo(checkpoint)
+            self.vframe.restore_state(frame_state)
+            filename = self.code_var.value.co_filename
+            lineno = self.code_var.value.co_firstlineno
+            code_name = self.code_var.value.co_name
+            location_info = f'File "{filename}", line {lineno}, in {code_name}'
+
+            exception_class = OtherInlineCallBreak
+            if isinstance(error, BreakGraphError):
+                exception_class = BreakGraphInlineCallBreak
+            elif isinstance(error, FallbackError):
+                exception_class = FallbackInlineCallBreak
+
+            raise BreakGraphError(
+                exception_class(
+                    f"{location_info} encountered breakgraph error caused by\n    {error}"
+                )
+            )
+
+        return output
+
+    def getattr(self, name: str, default=None):
+        from ..dispatch_functions import generator_send
+
+        known_generator_attrs = {"send"}
+        if name not in known_generator_attrs:
+            raise BreakGraphError(
+                UnsupportedOperationBreak(
+                    reason_str=f"Get attribute {name} from generator is not supported."
+                )
+            )
+        if name == "send":
+            return BuiltinVariable(
+                generator_send, self.graph, DanglingTracker()
+            ).bind_dangling_fn(self, "send")
+        unreached()
+
+    def get_py_value(self, allow_tensor=False):
+        raise BreakGraphError(
+            UnsupportedOperationBreak(
+                reason_str="Get real value from generator is not supported."
+            )
+        )
+
+    def get_py_type(self):
+        return types.GeneratorType
+
+    def next(self):
+        return self.send(ConstantVariable.wrap_literal(None, self.graph))
+
+    def to_list(self):
+        return do_until_stop_iteration(lambda: self.next())
+
+    @property
+    def main_info(self) -> dict[str, Any]:
+        return {
+            "co_name": self.code_var.value.co_name,
+        }
+
+    # @VariableFactory.register_from_value()
+    # def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+    #     if inspect.isgenerator(value):
+    #         return GeneratorVariable()
+    #     return None
+
+
+# what UserDefinedIterVariable holds doesn't matter, because use user defined iterator will trigger break graph
+class UserDefinedIterVariable(IterVariable):
+    def __init__(
+        self,
+        held: VariableBase | list[VariableBase],
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        if not isinstance(held, list):
+            held = [held]
+        self.holds = held
+        super().__init__(graph, tracker)
+
+    def to_list(self):
+        raise BreakGraphError(
+            UnsupportedOperationBreak(
+                reason_str="Break graph when iterating user defined iterator"
+            )
+        )
+
+    def next(self):
+        raise BreakGraphError(
+            UnsupportedOperationBreak(
+                reason_str="Break graph when iterating user defined iterator"
+            )
+        )
+
+    @check_faster_guard
+    def make_faster_guard(self) -> list[paddle.framework.core.GuardNodeBase]:
+        return [
+            guard for held in self.holds for guard in held.make_faster_guard()
+        ]
+
+    def make_stringified_guard(self):
+        return [
+            guard
+            for held in self.holds
+            for guard in held.make_stringified_guard()
+        ]

python/paddle/jit/sot/opcode_translator/executor/virtual_frame.py

@@ -0,0 +1,238 @@
+# 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.
+
+from __future__ import annotations
+
+import builtins
+import re
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+from ...utils import log
+from .tracker import (
+    BuiltinTracker,
+    CellTracker,
+    ConstTracker,
+    DanglingTracker,
+    FunctionClosureTracker,
+    LocalTracker,
+)
+from .variable_stack import VariableStack
+from .variables.base import VariableBase, VariableFactory, fn_bind_inputs
+from .variables.basic import (
+    CellVariable,
+    FunctionGlobalVariable,
+    GlobalVariable,
+    NullVariable,
+)
+
+if TYPE_CHECKING:
+    import types
+    from typing import TypeAlias
+
+    from ..instruction_utils import Instruction
+    from .function_graph import FunctionGraph
+    from .variables.callable import FunctionVariable
+
+    # The type to represent the (*args, **kwargs) pack in the call.
+    CallArgsPack: TypeAlias = tuple[tuple[Any, ...], dict[str, Any]]
+
+
+def validate_value(value):
+    assert isinstance(value, VariableBase), (
+        f"value: {value}, type should be VariableBase(or derived), but get {type(value)}"
+    )
+    assert not isinstance(value.tracker, DanglingTracker) or isinstance(
+        value, (NullVariable, CellVariable)
+    ), f"dangling variable {value} should not be pushed into stack."
+
+
+@dataclass
+class BlockStackItem:
+    # `PyTryBlock` in CPython source code
+    type: str
+    inst: Instruction
+    handler: Instruction
+    level: int
+
+
+class VirtualFrameState(NamedTuple):
+    locals: dict[str, VariableBase]
+    builtins: dict[str, VariableBase]
+    cells: dict[str, VariableBase]
+    lasti: int
+    stack_data: list[VariableBase]
+    block_stack: list[BlockStackItem]
+
+
+class VirtualFrame:
+    code: types.CodeType
+    locals: dict[str, Any]  # TODO: should we use DictVariable instead of dict?
+    globals: GlobalVariable
+    builtins: dict[str, Any]
+    consts: list[Any]
+    cells: dict[str, Any]
+    lasti: int
+    stack: VariableStack
+    block_stack: list[BlockStackItem]
+
+    def __init__(self, code: types.CodeType):
+        self.code = code
+        self.locals = {}
+        self.globals = None  # type: ignore
+        self.builtins = {}
+        self.cells = {}
+        self.lasti = 0
+        self.consts = []
+        self.stack = VariableStack(validate_value_func=validate_value)
+        self.block_stack: list[BlockStackItem] = []
+
+    @staticmethod
+    def from_real_frame(frame: types.FrameType, graph: FunctionGraph):
+        code = frame.f_code
+        locals = frame.f_locals
+        vframe = VirtualFrame(code)
+
+        # convert locals
+        free_or_cell_vars = code.co_cellvars + code.co_freevars
+        for name, value in locals.items():
+            tracker = (
+                CellTracker(name)
+                if name in free_or_cell_vars
+                else LocalTracker(name)
+            )
+            vframe.locals[name] = VariableFactory.from_value(
+                value, graph, tracker
+            )
+
+        for name in free_or_cell_vars:
+            # create a cell for each variable.
+            vframe.cells[name] = CellVariable()  # put in cells.
+            if name in vframe.locals:
+                vframe.cells[name].set_value(vframe.locals[name])
+
+        # convert globals
+        vframe.globals = GlobalVariable(
+            frame.f_globals,
+            graph,
+            DanglingTracker(),
+        )
+
+        # convert builtins
+        for name, value in builtins.__dict__.items():
+            vframe.builtins[name] = VariableFactory.from_value(
+                value, graph, BuiltinTracker(name)
+            )
+        # Temporarily use the builtins from the graph to avoid the conversion overhead.
+        graph.builtins = vframe.builtins
+
+        # prepare consts
+        for value in code.co_consts:
+            vframe.consts.append(
+                VariableFactory.from_value(value, graph, ConstTracker(value))
+            )
+        return vframe
+
+    @staticmethod
+    def from_inline_call(
+        code: types.CodeType,
+        fn_var: FunctionVariable,
+        fn_value: types.FunctionType,
+        graph: FunctionGraph,
+        call_args_pack: CallArgsPack,
+    ):
+        call_args, call_kwargs = call_args_pack
+        vframe = VirtualFrame(code)
+        vframe.globals = FunctionGlobalVariable(
+            fn_var,
+            fn_value.__globals__,
+            graph,
+            DanglingTracker(),
+        )
+
+        # convert builtins
+        # NOTE(SigureMo): inline call should inherit the builtins from the caller to reduce the conversion overhead.
+        vframe.builtins = graph.builtins
+
+        # prepare consts
+        for value in code.co_consts:
+            vframe.consts.append(
+                VariableFactory.from_value(value, graph, ConstTracker(value))
+            )
+
+        # convert locals
+        vframe.locals.update(
+            fn_bind_inputs(fn_value, graph, *call_args, **call_kwargs)
+        )
+
+        log(
+            5,
+            f"[INLINE CALL] {code.co_name} with locals: ",
+            vframe.locals,
+        )
+
+        # handle implicit variables in comprehensions
+        vframe.handle_comps(fn_value)
+
+        # convert closures
+        closure = fn_var.get_py_value().__closure__
+        for name in code.co_cellvars + code.co_freevars:
+            # create a cell for each variable.
+            vframe.cells[name] = CellVariable()  # put in cells.
+            if name in vframe.locals:
+                vframe.cells[name].set_value(vframe.locals[name])
+
+        if closure is None:
+            return vframe
+        assert len(closure) == len(code.co_freevars)
+        for idx, (name, cell) in enumerate(zip(code.co_freevars, closure)):
+            value = cell.cell_contents
+            value = VariableFactory.from_value(
+                value, graph, FunctionClosureTracker(fn_var, idx)
+            )
+            # wrapped by a CellVariable
+            if not isinstance(value, CellVariable):
+                value = CellVariable(value)
+            vframe.cells[name] = value
+        return vframe
+
+    def handle_comps(self, fn_value):
+        is_comp = any(
+            x in fn_value.__name__
+            for x in ['<listcomp>', '<dictcomp>', '<setcomp>', '<genexpr>']
+        )
+        if not is_comp:
+            return
+        pattern = r'implicit\d+'
+        for name in list(self.locals.keys()):
+            if re.match(pattern, name):
+                self.locals[name.replace('implicit', '.')] = self.locals[name]
+
+    def get_state(self):
+        return VirtualFrameState(
+            locals=self.locals.copy(),
+            builtins=self.builtins.copy(),
+            cells=self.cells.copy(),
+            lasti=self.lasti,
+            stack_data=list(self.stack._data),
+            block_stack=self.block_stack.copy(),
+        )
+
+    def restore_state(self, state: VirtualFrameState):
+        self.locals = state.locals
+        self.builtins = state.builtins
+        self.cells = state.cells
+        self.lasti = state.lasti
+        self.stack._data = state.stack_data
+        self.block_stack = state.block_stack

python/paddle/jit/sot/opcode_translator/instruction_utils/__init__.py

@@ -0,0 +1,35 @@
+# 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 .instruction_pass import apply_instr_pass  # noqa: F401
+from .instruction_utils import (  # noqa: F401
+    Instruction,
+    Space,
+    calc_offset_from_bytecode_offset,
+    calc_stack_effect,
+    convert_instruction,
+    gen_instr,
+    get_instruction_size,
+    get_instructions,
+    instrs_info,
+    modify_extended_args,
+    modify_instrs,
+    modify_vars,
+    relocate_jump_target,
+    replace_instr,
+    reset_offset,
+)
+from .opcode_analysis import (  # noqa: F401
+    analysis_used_names,
+)

python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_pass.py

@@ -0,0 +1,334 @@
+# 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 dis
+import sys
+from typing import TYPE_CHECKING
+
+from paddle.jit.sot.utils import log, log_do
+
+from ...utils import InnerError
+from .instruction_utils import instrs_info
+from .opcode_info import TO_FUSED_INSTS
+from .stack_analyse import StackAnalyser
+
+if TYPE_CHECKING:
+    from .instruction_utils import Instruction
+
+
+def apply_instr_pass(instrs: list[Instruction], code_options):
+    log(4, f"[Opcode Pass]: Original New Code {code_options['co_name']}:\n")
+    log_do(4, lambda: print(instrs_info(instrs)))
+    supported_passes = [
+        remove_load_store_pass,
+        remove_duplicate_resume,
+        check_precall_followed_by_call,
+    ]
+
+    if sys.version_info >= (3, 12):
+        supported_passes.append(check_for_iter_jump_to)
+
+    if sys.version_info >= (3, 13):
+        supported_passes.append(fuse_double_super_instrs)
+
+    for instr_pass in supported_passes:
+        instr_pass(instrs, code_options)
+
+    log(
+        4,
+        f"[Opcode Pass]: New Code After Opcode Pass {code_options['co_name']}:\n",
+    )
+    log_do(4, lambda: print(instrs_info(instrs)))
+
+
+def find_stored_once_local_vars(instrs: list[Instruction], code_options):
+    """
+    find out the local var names which is only stored once
+    """
+    stored_vars = {}
+
+    # The input vars are considered as stored at the beginning
+    input_names = code_options['co_varnames'][: code_options['co_argcount']]
+
+    for name in input_names:
+        stored_vars[name] = 1
+
+    for instr in instrs:
+        if instr.opname == "STORE_FAST":
+            if instr.argval in stored_vars:
+                stored_vars[instr.argval] += 1
+            else:
+                stored_vars[instr.argval] = 1
+
+    stored_once = {name for name, count in stored_vars.items() if count == 1}
+    return stored_once
+
+
+def find_loaded_once_local_vars(instrs: list[Instruction], code_options):
+    """
+    find out the local var names which is only stored once
+    """
+    loaded_vars = {}
+    for instr in instrs:
+        if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
+            if instr.argval in loaded_vars:
+                loaded_vars[instr.argval] += 1
+            else:
+                loaded_vars[instr.argval] = 1
+
+    loaded_once = {name for name, count in loaded_vars.items() if count == 1}
+    return loaded_once
+
+
+def find_related_local_opcodes(instrs: list[Instruction], code_options):
+    """
+    Find opcode pairs consisting of LOAD_FAST, LOAD_FAST_BORROW, STORE_FAST, and LOAD_FAST_CHECK.
+    """
+    stack = []
+    opcode_pairs = []
+    for instr in instrs:
+        if instr.opname in ["LOAD_FAST", "LOAD_FAST_BORROW", "LOAD_FAST_CHECK"]:
+            stack.append(instr)
+        elif instr.opname == "STORE_FAST":
+            if len(stack) > 0 and stack[-1] is not None:
+                opcode_pairs.append((stack[-1], instr))
+            stack.pop()
+        elif "ROT" in instr.opname or "DUP" in instr.opname:
+            return []
+        else:
+            try:
+                pop_n, push_n = StackAnalyser().stack_effect(instr)
+                if pop_n == 0:
+                    stack.extend([None] * push_n)
+                else:
+                    stack = stack[:-pop_n] + [None] * push_n
+            except AttributeError:
+                break
+
+    return opcode_pairs
+
+
+def remove_load_store_pass(instrs: list[Instruction], code_options):
+    """
+    This question is extremely complex, so we just simplify it as
+    'remove renames which is between var names who only stored once'
+    and we only consider the local vars.
+    """
+
+    def stored_from(load_instr, instrs):
+        idx = instrs.index(load_instr) - 1
+        while idx >= 0:
+            instr = instrs[idx]
+            if (
+                instr.opname == "STORE_FAST"
+                and instr.argval == load_instr.argval
+            ):
+                return instr
+            idx -= 1
+        return None
+
+    def code_exist(opname, argval, instrs):
+        for instr in instrs:
+            if instr.opname == opname and instr.argval == argval:
+                return True
+        return False
+
+    # remove rename and load store
+    jump_target = {
+        instr.jump_to for instr in instrs if instr.jump_to is not None
+    }
+
+    modified = True
+    while modified:
+        modified = False
+        stored_once = find_stored_once_local_vars(instrs, code_options)
+
+        # find out all LOAD_FAST -> STORE_FAST pair
+        opcode_pairs = find_related_local_opcodes(instrs, code_options)
+
+        for load_a, store_b in opcode_pairs:
+            if load_a in jump_target or store_b in jump_target:
+                continue
+            a_name = load_a.argval
+            b_name = store_b.argval
+
+            # if these two names are only stored once
+            # it means these two name only have one value all the time
+            # so we can just rename them, to delete some codes
+            if a_name in stored_once and b_name in stored_once:
+                instrs.remove(load_a)
+                instrs.remove(store_b)
+                if a_name != b_name:
+                    for instr in instrs:
+                        if (
+                            instr.opname
+                            in (
+                                "LOAD_FAST_CHECK",
+                                "LOAD_FAST",
+                                "LOAD_FAST_BORROW",
+                                "STORE_FAST",
+                            )
+                            and instr.argval == b_name
+                        ):
+                            instr.argval = a_name
+                            instr.arg = load_a.arg
+                modified = True
+
+            # if
+            #       LOAD A
+            #       STORE B
+            # A or B is not stored only once (maybe it is input)
+            # we give a more general way to simplify the codes
+            #
+            # if A will not be loaded again after (6)STORE B, it means we can move (6)STORE B ahead to (1)STORE A
+            # TIP: there is no more STORE A between (1) and (5)
+            #  (1)      STORE A             ->          STORE B
+            #           ...                             ...
+            #  (2)      LOAD A              ->          LOAD B
+            #           ...
+            #  (3)      LOAD B              ->          not support
+            #           ...
+            #  (4)      STORE B             ->          not support
+            #           ...                             ...
+            #  (5)      LOAD A              ->          ---- (rm)
+            #  (6)      STORE B                         ---- (rm)
+            #           ...
+            #  (7)      STORE B
+            #  (8)      LOAD A
+            # so we can rename the rest LOAD A below as LOAD B
+            #
+            # What changed:
+            #   1. if (4) exist, B changed:
+            #       (1) ~ (4), (6) ~
+            #   2. if (4) not exist, B changed:
+            #       (1), (6)
+            #   3. A changed:
+            #       (1) ~
+            #
+            # To do this transform, we should make sure
+            #   1. (4) is not exist in (1) ~ (5): it is too complex
+            #   2. (3) is not exist in (1) ~ (5): load B in the range that B value is changed
+            #   3. (7) (8) is not exist in (6)~: load A in range that A value is changed, if we load B instead, but B also changed
+            #       we can simplify this as "no more LOAD A after (6)"
+            else:
+                last_store_a = stored_from(load_a, instrs)
+                if last_store_a is None:
+                    # if last store a just not exist, we can not do this transform
+                    continue
+
+                last_store_idx = instrs.index(last_store_a)
+                code_range = instrs[last_store_idx : instrs.index(store_b)]
+                if (
+                    not code_exist("STORE_FAST", b_name, code_range)
+                    and not code_exist("LOAD_FAST_CHECK", b_name, code_range)
+                    and not code_exist("LOAD_FAST", b_name, code_range)
+                    and not code_exist("LOAD_FAST_BORROW", b_name, code_range)
+                    and not code_exist(
+                        "LOAD_FAST_CHECK",
+                        a_name,
+                        instrs[instrs.index(store_b) :],
+                    )
+                    and not code_exist(
+                        "LOAD_FAST", a_name, instrs[instrs.index(store_b) :]
+                    )
+                    and not code_exist(
+                        "LOAD_FAST_BORROW",
+                        a_name,
+                        instrs[instrs.index(store_b) :],
+                    )
+                ):
+                    last_store_a.argval = b_name
+                    last_store_a.arg = store_b.arg
+                    instrs.remove(load_a)
+                    instrs.remove(store_b)
+                    for instr in instrs[last_store_idx:]:
+                        if (
+                            instr.opname
+                            in (
+                                "LOAD_FAST_CHECK",
+                                "LOAD_FAST",
+                                "LOAD_FAST_BORROW",
+                                "STORE_FAST",
+                            )
+                            and instr.argval == a_name
+                        ):
+                            instr.argval = b_name
+                            instr.arg = store_b.arg
+
+
+def remove_duplicate_resume(instrs: list[Instruction], code_options):
+    resumes = list(filter(lambda instr: instr.opname == "RESUME", instrs))
+    if not resumes:
+        return
+    for resume in resumes[1:]:
+        instrs.remove(resume)
+
+
+def check_precall_followed_by_call(instrs: list[Instruction], code_options):
+    """
+    PRECALL should be followed by CALL, otherwise it will cause a segmentation fault
+    """
+    for instr, next_instr in zip(instrs[:-1], instrs[1:]):
+        if instr.opname == "PRECALL" and next_instr.opname != "CALL":
+            raise InnerError(
+                f"PRECALL is not followed by CALL in {code_options['co_name']}"
+            )
+
+
+def check_for_iter_jump_to(instrs: list[Instruction], code_options):
+    """
+    Check if the `jump_to` of FOR_ITER is END_FOR, in Python3.12+
+    """
+    for instr in instrs:
+        if instr.opname == "FOR_ITER":
+            assert instr.jump_to is not None
+            if instr.jump_to.opname != "END_FOR":
+                raise InnerError("FOR_ITER jump_to is not END_FOR")
+
+
+def fuse_double_super_instrs(instrs: list[Instruction], code_options):
+    """
+    Fuse two consecutive LOAD_FAST or STORE_FAST instructions into one.
+    """
+    co_varnames = code_options['co_varnames']
+
+    def able_to_merge(idx: int):
+        return (
+            idx > 0
+            and (instrs[idx - 1].opname, instrs[idx].opname)
+            in TO_FUSED_INSTS.keys()
+            and not instrs[idx].is_jump_target
+            and not instrs[idx - 1].is_jump_target
+            and co_varnames.index(instrs[idx - 1].argval) < 16
+            and co_varnames.index(instrs[idx].argval) < 16
+        )
+
+    def merge_two_op(prev_instr: Instruction, instr: Instruction):
+        merge_key = (instrs[idx - 1].opname, instrs[idx].opname)
+        prev_instr.opname = TO_FUSED_INSTS[merge_key]
+        prev_instr.opcode = dis.opmap[prev_instr.opname]
+        prev_instr.is_generated = True
+        prev_instr.argval = (prev_instr.argval, instr.argval)
+        instrs.remove(instr)
+
+    idx = 0
+    # We must manually control the indices, so we cannot use a for loop.
+    while idx < len(instrs):
+        if able_to_merge(idx):
+            merge_two_op(instrs[idx - 1], instrs[idx])
+            continue
+
+        idx += 1

python/paddle/jit/sot/opcode_translator/instruction_utils/instruction_utils.py

@@ -0,0 +1,586 @@
+# 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 dataclasses
+import dis
+import sys
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+from ...utils import InnerError
+from .opcode_info import (
+    ABS_JUMP,
+    ALL_JUMP,
+    FUSED_INSTS,
+    PYOPCODE_CACHE_SIZE,
+    REL_BWD_JUMP,
+    REL_JUMP,
+)
+
+if TYPE_CHECKING:
+    import types
+
+
+@dataclasses.dataclass
+class Instruction:
+    opcode: int
+    opname: str
+    arg: int | None
+    argval: Any
+    offset: int | None = None
+    starts_line: int | None = None
+    is_jump_target: bool = False
+    jump_to: Instruction | None = None
+    is_generated: bool = True
+
+    # for analysis EXTENDED_ARG
+    first_ex_arg: Instruction | None = None
+    ex_arg_for: Instruction | None = None
+
+    # used in modify_extended_args
+    def __hash__(self):
+        return id(self)
+
+    def __eq__(self, instr):
+        return id(self) == id(instr)
+
+
+def get_instruction_size(instr: Instruction) -> int:
+    cache_size = 0
+    if sys.version_info >= (3, 11):
+        cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
+    return 2 * (cache_size + 1)
+
+
+def gen_instr(name, arg=None, argval=None, gened=True, jump_to=None):
+    return Instruction(
+        opcode=dis.opmap[name],
+        opname=name,
+        arg=arg,
+        argval=argval,
+        is_generated=gened,
+        jump_to=jump_to,
+    )
+
+
+def convert_instruction(instr: dis.Instruction) -> Instruction:
+    """
+    Converts a disassembled instruction to a customized Instruction object.
+
+    Args:
+        instr (dis.Instruction): The disassembled instruction.
+
+    Returns:
+        Instruction: A customized Instruction object.
+    """
+    return Instruction(
+        instr.opcode,
+        instr.opname,
+        instr.arg,
+        instr.argval,
+        instr.offset,
+        instr.line_number if sys.version_info >= (3, 13) else instr.starts_line,
+        instr.is_jump_target,
+        jump_to=None,
+        is_generated=False,
+    )
+
+
+def replace_jump_target(
+    instrs: list[Instruction],
+    replacements: dict[Instruction, Instruction],
+) -> None:
+    """Replace jump targets based on the replacements dictionary.
+
+    Args:
+        instrs (list[Instruction]): The list of instructions to modify.
+        replacements (dict[Instruction, Instruction]): Mapping from old jump targets to new ones.
+    """
+    for instr in instrs:
+        if instr.jump_to in replacements:
+            instr.jump_to = replacements[instr.jump_to]
+
+
+def expand_super_instrs(instructions: list[Instruction]) -> list[Instruction]:
+    expanded_instrs = []
+    replacements = {}
+
+    def copy_instruction(
+        instr, opname, argval, arg, is_jump_target, is_generated
+    ):
+        return Instruction(
+            opcode=dis.opmap[opname],
+            opname=opname,
+            arg=arg,
+            argval=argval,
+            is_jump_target=is_jump_target,
+            is_generated=is_generated,
+            jump_to=instr.jump_to,
+        )
+
+    for instr in instructions:
+        if instr.opname in FUSED_INSTS:
+            instr1 = copy_instruction(
+                instr,
+                FUSED_INSTS[instr.opname][0],
+                instr.argval[0],
+                instr.arg >> 4,
+                instr.is_jump_target,
+                True,
+            )
+            instr2 = copy_instruction(
+                instr,
+                FUSED_INSTS[instr.opname][1],
+                instr.argval[1],
+                instr.arg & 15,
+                False,
+                False,
+            )
+            replacements[instr] = instr1
+            expanded_instrs.append(instr1)
+            expanded_instrs.append(instr2)
+        # If the LOAD_ATTR opcode will lead to load_method in 3.13+, we manually split it into two instructions,
+        # to avoid Uncontrollable specialization that changes the behavior of LOAD_ATTR,
+        # which can lead to incorrect results when the current graph is smaller than the MIN_GRAPH_SIZE
+        elif (
+            sys.version_info >= (3, 13)
+            and instr.opname == "LOAD_ATTR"
+            and instr.arg & 1
+        ):
+            instr1 = copy_instruction(
+                instr,
+                "LOAD_ATTR",
+                instr.argval,
+                instr.arg & ~1,
+                instr.is_jump_target,
+                True,
+            )
+            instr2 = Instruction(
+                dis.opmap["PUSH_NULL"],
+                "PUSH_NULL",
+                None,
+                None,
+                is_generated=True,
+            )
+            replacements[instr] = instr1
+            expanded_instrs.append(instr1)
+            expanded_instrs.append(instr2)
+        else:
+            expanded_instrs.append(instr)
+
+    replace_jump_target(expanded_instrs, replacements)
+    return expanded_instrs
+
+
+def replace_load_fast_borrow_with_strong_ref(
+    instructions: list[Instruction],
+) -> list[Instruction]:
+    """
+    Patch LOAD_FAST_BORROW to LOAD_FAST for Python 3.14+.
+
+    LOAD_FAST_BORROW loads a value using a borrowing reference and does not
+    increment the reference count. In some cases this can cause subsequent
+    STORE_FAST or other operations to retain a reference that becomes invalid
+    once the borrowed value is released, leading to incorrect behavior or
+    crashes when the variable is accessed later.
+    To avoid these issues, we replace LOAD_FAST_BORROW with LOAD_FAST here.
+    """
+    replacements = {}
+    expanded_instrs = []
+    for instr in instructions:
+        if instr.opname == "LOAD_FAST_BORROW":
+            instr1 = Instruction(
+                dis.opmap["LOAD_FAST"],
+                "LOAD_FAST",
+                instr.arg,
+                instr.argval,
+                is_generated=instr.is_generated,
+                is_jump_target=instr.is_jump_target,
+                jump_to=instr.jump_to,
+            )
+            replacements[instr] = instr1
+            expanded_instrs.append(instr1)
+        else:
+            expanded_instrs.append(instr)
+
+    replace_jump_target(expanded_instrs, replacements)
+    return expanded_instrs
+
+
+def get_instructions(code: types.CodeType) -> list[Instruction]:
+    """
+    Returns parsed instructions from the given code object and exclude
+    any opcodes that contain `EXTENDED_ARG`.
+
+    Args:
+        code (types.CodeType): The code object to extract instructions from.
+
+    Returns:
+        list[Instruction]: A list of Instruction objects representing the
+            bytecode instructions in the code object.
+    """
+    # instrs do not contain EXTENDED_ARG
+    instrs = list(map(convert_instruction, dis.get_instructions(code)))
+    for instr in instrs:
+        if instr.opname in ALL_JUMP:
+            origin_jump_target = calc_offset_from_bytecode_offset(
+                instr.argval, instrs
+            )
+            jump_offset = origin_jump_target
+
+            while instrs[jump_offset].opname == "EXTENDED_ARG":
+                jump_offset += 1
+
+            if origin_jump_target != jump_offset:
+                # copy infos from EXTENDED_ARG to other opcode
+
+                if instrs[origin_jump_target].is_jump_target:
+                    instrs[jump_offset].is_jump_target = instrs[
+                        origin_jump_target
+                    ].is_jump_target
+                if instrs[origin_jump_target].starts_line:
+                    instrs[jump_offset].starts_line = instrs[
+                        origin_jump_target
+                    ].starts_line
+
+            instr.jump_to = instrs[jump_offset]
+
+    # if the origin opcode contains EXTENDED_ARG, it should be like:
+    #     >>  EXTENDED_ARG 1
+    #         XX 388    <-  256 + 132
+    # filter all EXTENDED_ARG here
+    instrs = [x for x in instrs if x.opname != "EXTENDED_ARG"]
+    prepare_passes = [expand_super_instrs]
+    if sys.version_info >= (3, 14):
+        prepare_passes.append(replace_load_fast_borrow_with_strong_ref)
+
+    for pass_fn in prepare_passes:
+        instrs = pass_fn(instrs)
+    return instrs
+
+
+def modify_instrs(instructions: list[Instruction]) -> None:
+    """
+    Modifies the given list of instructions. It contains three steps:
+
+    1. reset offset
+    2. relocate jump target
+    3. add EXTENDED_ARG instruction if needed
+
+    Args:
+        instructions (list): The list of Instruction objects representing bytecode instructions.
+
+    Returns:
+        None
+    """
+    modify_completed = False
+    while not modify_completed:
+        reset_offset(instructions)
+        has_inverted_jump = relocate_jump_target(instructions)
+        modify_completed = (
+            modify_extended_args(instructions) and not has_inverted_jump
+        )
+
+
+def reset_offset(instructions: list[Instruction]) -> None:
+    """
+    Resets the offset for each instruction in the list.
+
+    Args:
+        instructions (list): The list of Instruction objects representing bytecode instructions.
+
+    Returns:
+        None
+    """
+    from ..executor.pycode_generator import get_instruction_size
+
+    if sys.version_info >= (3, 11):
+        current_offset = 0
+        for instr in instructions:
+            instr.offset = current_offset
+            current_offset += get_instruction_size(instr)
+        return
+    for idx, instr in enumerate(instructions):
+        instr.offset = idx * 2
+
+
+def correct_jump_direction(
+    instr: Instruction, arg: int
+) -> tuple[Instruction, bool]:
+    """
+    Corrects the jump direction of the given instruction.
+    NOTE(zrr1999): In Python 3.11, JUMP_ABSOLUTE is removed, so python generates JUMP_FORWARD or JUMP_BACKWARD instead,
+    but in for loop breakgraph, we reuse JUMP_BACKWARD to jump forward, so we need to change it to JUMP_FORWARD.
+
+    Args:
+        instr (Instruction): The instruction to be corrected.
+        invert_jump (bool): Whether to invert the jump direction.
+    """
+    if instr.opname in ABS_JUMP:
+        instr.arg = arg
+        return instr, False
+    elif instr.opname in REL_JUMP:
+        if arg < 0:
+            if instr.opname in REL_BWD_JUMP:
+                forward_op_name = instr.opname.replace("BACKWARD", "FORWARD")
+                if forward_op_name not in dis.opmap:
+                    raise InnerError(f"Unknown jump type {instr.opname}")
+                instr.opname = forward_op_name
+                instr.opcode = dis.opmap[forward_op_name]
+            else:  # instr.opname in REL_FWD_JUMP
+                backward_op_name = instr.opname.replace("FORWARD", "BACKWARD")
+                if backward_op_name not in dis.opmap:
+                    raise InnerError(f"Unknown jump type {instr.opname}")
+                instr.opname = backward_op_name
+                instr.opcode = dis.opmap[backward_op_name]
+            instr.arg = -arg
+            invert_jump = True
+        else:
+            instr.arg = arg
+            invert_jump = False
+        return instr, invert_jump
+    else:
+        raise ValueError(f"unknown jump type: {instr.opname}")
+
+
+def relocate_jump_target(instructions: list[Instruction]) -> bool:
+    """
+    If a jump instruction is found, this function will adjust the jump targets based on the presence of EXTENDED_ARG instructions.
+    If an EXTENDED_ARG instruction exists for the jump target, use its offset as the new target.
+
+    Args:
+        instructions (list): The list of Instruction objects representing bytecode instructions.
+
+    Returns:
+        bool: True if the jump direction is inverted, False otherwise.
+    """
+    has_inverted_jump = False
+    extended_arg = []
+    for instr in instructions:
+        if instr.opname == "EXTENDED_ARG":
+            extended_arg.append(instr)
+            continue
+
+        if instr.opname in ALL_JUMP:
+            assert instr.jump_to is not None
+            assert instr.offset is not None
+            # if jump target has extended_arg, should jump to the first extended_arg opcode
+            jump_target = (
+                instr.jump_to.offset
+                if instr.jump_to.first_ex_arg is None
+                else instr.jump_to.first_ex_arg.offset
+            )
+            assert jump_target is not None
+
+            if instr.opname in ABS_JUMP:
+                new_arg = jump_target
+            else:  # instr.opname in REL_JUMP
+                cache_size = PYOPCODE_CACHE_SIZE.get(instr.opname, 0)
+                new_arg = jump_target - (2 * cache_size) - instr.offset - 2
+                if instr.opname in REL_BWD_JUMP:
+                    new_arg = -new_arg
+
+            new_arg //= 2
+            _, invert_jump = correct_jump_direction(instr, new_arg)
+            has_inverted_jump = has_inverted_jump or invert_jump
+            assert instr.arg is not None
+            if extended_arg:
+                instr.arg &= 0xFF
+                new_arg = new_arg >> 8
+                for ex in reversed(extended_arg):
+                    ex.arg = new_arg & 0xFF
+                    new_arg = new_arg >> 8
+
+                # need more extended_args instr
+                # set arg in the first extended_arg
+                if new_arg > 0:
+                    extended_arg[0].arg += new_arg << 8
+        extended_arg.clear()
+    return has_inverted_jump
+
+
+def modify_extended_args(instructions: list[Instruction]) -> bool:
+    """
+    This function replaces any instruction with an argument greater than or equal to 256 with one or more EXTENDED_ARG instructions.
+
+    Args:
+        instructions (list): The list of Instruction objects representing bytecode instructions.
+
+    Returns:
+        bool: True if the modification is completed, False otherwise.
+    """
+
+    modify_completed = True
+    extend_args_record = {}
+    for instr in instructions:
+        if instr.arg and instr.arg >= 256:  # more than one byte
+            _instrs = [
+                instr
+            ]  # replace instr with _instrs later (it is a set of instrs), all operations will be recorded in extend_args_record
+            val = instr.arg
+            instr.arg = val & 0xFF
+            val = val >> 8
+            while val > 0:
+                _instrs.append(gen_instr("EXTENDED_ARG", arg=val & 0xFF))
+                val = val >> 8
+
+            extend_args_record.update({instr: list(reversed(_instrs))})
+
+    if extend_args_record:
+        # if new EXTENDED_ARG inserted, we need update offset and jump target
+        modify_completed = False
+
+        def bind_ex_arg_with_instr(ex_arg, instr):
+            # move opcode info to EXTENDED_ARG
+            ex_arg.starts_line = instr.starts_line
+            instr.starts_line = None
+            ex_arg.is_jump_target = instr.is_jump_target
+            instr.is_jump_target = False
+
+            if instr.ex_arg_for is not None:
+                # instr is also an ex_arg for another instr
+                instr.ex_arg_for.first_ex_arg = ex_arg
+                ex_arg.ex_arg_for = instr.ex_arg_for
+                instr.ex_arg_for = None
+            else:
+                instr.first_ex_arg = ex_arg
+                ex_arg.ex_arg_for = instr
+
+        for key, val in extend_args_record.items():
+            bind_ex_arg_with_instr(val[0], key)
+            replace_instr(instructions, instr=key, new_instr=val)
+
+    return modify_completed
+
+
+def modify_vars(instructions: list[Instruction], code_options):
+    co_varnames = code_options['co_varnames']
+    co_freevars = code_options['co_freevars']
+    for instrs in instructions:
+        if instrs.opname in [
+            'LOAD_FAST',
+            'LOAD_FAST_BORROW',
+            'LOAD_FAST_CHECK',
+            'STORE_FAST',
+            'DELETE_FAST',
+        ]:
+            assert instrs.argval in co_varnames, (
+                f"`{instrs.argval}` not in {co_varnames}"
+            )
+            instrs.arg = co_varnames.index(instrs.argval)
+        elif instrs.opname == "LOAD_DEREF" or instrs.opname == "STORE_DEREF":
+            if sys.version_info >= (3, 11):
+                namemap = co_varnames + co_freevars
+                assert instrs.argval in namemap, (
+                    f"`{instrs.argval}` not in {namemap}"
+                )
+                instrs.arg = namemap.index(instrs.argval)
+        elif instrs.opname in FUSED_INSTS.keys():
+            assert instrs.argval[0] in co_varnames, (
+                f"`{instrs.argval[0]}` not in {co_varnames}"
+            )
+            assert instrs.argval[1] in co_varnames, (
+                f"`{instrs.argval[1]}` not in {co_varnames}"
+            )
+            instrs.arg = (
+                co_varnames.index(instrs.argval[0]) << 4
+            ) + co_varnames.index(instrs.argval[1])
+
+
+def calc_offset_from_bytecode_offset(
+    bytecode_offset: int,
+    instructions: list[dis.Instruction] | list[Instruction],
+) -> int:
+    """
+    Calculate the index from bytecode offset, because it have 2 bytes per instruction (for Python <= 3.10).
+
+    Args:
+        bytecode_offset (int): The bytecode offset of the instruction.
+
+    Returns:
+        int: The index of the instruction in the instruction list.
+    """
+
+    if sys.version_info >= (3, 11):
+        instruction_offsets = [x.offset for x in instructions]
+        return instruction_offsets.index(bytecode_offset)
+    return bytecode_offset // 2
+
+
+def replace_instr(instructions, instr, new_instr):
+    idx = instructions.index(instr)
+    instructions[idx : idx + 1] = new_instr
+
+
+def instrs_info(instrs, mark=None, range=None, want_str=True):
+    ret = []
+    start = -1
+    end = 1000000
+    if mark is not None and range is not None:
+        start = mark - range
+        end = mark + range + 1
+    for idx, instr in enumerate(instrs):
+        if idx < start or idx >= end:
+            continue
+        if instr.starts_line is not None:
+            ret.append("")
+        ret.append(
+            "{line:<8s}{is_jump_target:>2s}{offset:>4d} {opname:<30s}{arg:<4s}{argval:<40s}{mark}".format(
+                line=str(instr.starts_line) if instr.starts_line else "",
+                is_jump_target=">>" if instr.is_jump_target else "  ",
+                offset=(
+                    instr.offset if instr.offset or instr.offset == 0 else -1
+                ),
+                opname=instr.opname,
+                arg=str(instr.arg) if instr.arg is not None else "",
+                argval=f"({instr.argval})" if instr.argval else "",
+                mark="",
+            )
+        )
+        if idx == mark:
+            ret[-1] = "\033[31m" + ret[-1] + "\033[0m"
+    if want_str:
+        return "\n".join(ret)
+    return ret
+
+
+def calc_stack_effect(instr: Instruction, *, jump: bool | None = None) -> int:
+    """
+    Gets the stack effect of the given instruction. In Python 3.11, the stack effect of `CALL` is -1,
+    refer to https://github.com/python/cpython/blob/3.11/Python/compile.c#L1123-L1124.
+
+    Args:
+        instr: The instruction.
+
+    Returns:
+        The stack effect of the instruction.
+
+    """
+    if sys.version_info[:2] == (3, 11):
+        if instr.opname == "PRECALL":
+            return 0
+        elif instr.opname == "CALL":
+            # NOTE(zrr1999): push_n = 1, pop_n = oparg + 2, stack_effect = push_n - pop_n = -oparg-1
+            assert instr.arg is not None
+            return -instr.arg - 1
+    return dis.stack_effect(instr.opcode, instr.arg, jump=jump)
+
+
+class Space(Enum):
+    locals = 1
+    globals = 2
+    cells = 3
+    not_found = 4

python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_analysis.py

@@ -0,0 +1,161 @@
+# 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 dataclasses
+from typing import TYPE_CHECKING
+
+from paddle.jit.utils import OrderedSet
+
+from .opcode_info import (
+    ALL_JUMP,
+    HAS_FREE,
+    HAS_LOCAL,
+    UNCONDITIONAL_JUMP,
+)
+
+if TYPE_CHECKING:
+    from .instruction_utils import Instruction
+
+
+@dataclasses.dataclass
+class NameRecorder:
+    reads: OrderedSet[str]
+    writes: OrderedSet[str]
+
+    def __or__(self, other):
+        reads = self.reads | other.reads
+        writes = self.writes | other.writes
+        return NameRecorder(reads, writes)
+
+
+def is_read_opcode(opname):
+    if opname in [
+        "LOAD_FAST",
+        "LOAD_FAST_CHECK",
+        "LOAD_DEREF",
+        "LOAD_NAME",
+        "LOAD_GLOBAL",
+        "LOAD_CLOSURE",
+        "LOAD_FAST_BORROW",
+    ]:
+        return True
+    if opname in (
+        "DELETE_FAST",
+        "DELETE_DEREF",
+        "DELETE_NAME",
+        "DELETE_GLOBAL",
+    ):
+        return True
+    return False
+
+
+def is_write_opcode(opname):
+    if opname in ["STORE_FAST", "STORE_NAME", "STORE_DEREF", "STORE_GLOBAL"]:
+        return True
+    if opname in (
+        "DELETE_FAST",
+        "DELETE_DEREF",
+        "DELETE_NAME",
+        "DELETE_GLOBAL",
+    ):
+        return True
+    return False
+
+
+def analysis_used_names(
+    instructions: list[Instruction],
+    current_instr_idx: int,
+    stop_instr_idx: int | None = None,
+) -> tuple[OrderedSet[str], OrderedSet[str]]:
+    """
+    Analyze the inputs of the instructions from current_instr_idx to stop_instr_idx.
+
+    Args:
+        instructions (list[Instruction]): The instructions to analyze.
+        current_instr_idx (int): The index of the current instruction.
+        stop_instr_idx (int | None, optional): The index of the instruction to stop. Defaults to None.
+            If None, the analysis will stop at the end of the instructions.
+
+    Returns:
+        State: The analysis result.
+    """
+    name_recorder = NameRecorder(OrderedSet(), OrderedSet())
+
+    # start idx and writes names can decide the analysis result below
+    # so, just check the pair of (idx, writes), to skip repeat simulation
+    # (writes can decide if a name should be add to reads)
+    # one idx can has multi writes for whom is not subset with each other
+    # if A is subset of B, we just record A, simulate A might add more reads
+    visited_states = {}
+
+    def check_and_update_visited_states(idx, writes):
+        writes = set(writes)
+
+        if idx in visited_states:
+            history = visited_states[idx]
+            for record in history:
+                if record.issubset(writes):
+                    return True
+                elif writes.issubset(record):
+                    history.remove(record)
+                    history.append(writes)
+                    return False
+        else:
+            visited_states[idx] = [writes]
+
+        return False
+
+    def fork(
+        name_recorder: NameRecorder, start: int, jump: bool, jump_target: int
+    ) -> NameRecorder:
+        new_start = start + 1 if not jump else jump_target
+        new_state = NameRecorder(
+            OrderedSet(name_recorder.reads),
+            OrderedSet(name_recorder.writes),
+        )
+        return walk(new_state, new_start)
+
+    def walk(name_recorder: NameRecorder, start: int) -> NameRecorder:
+        end = len(instructions) if stop_instr_idx is None else stop_instr_idx
+        for i in range(start, end):
+            if check_and_update_visited_states(i, name_recorder.writes):
+                return name_recorder
+
+            instr = instructions[i]
+            if instr.opname in HAS_LOCAL | HAS_FREE:
+                if is_read_opcode(instr.opname) and instr.argval not in (
+                    name_recorder.writes
+                ):
+                    name_recorder.reads.add(instr.argval)
+                elif is_write_opcode(instr.opname):
+                    name_recorder.writes.add(instr.argval)
+            elif instr.opname in ALL_JUMP:
+                assert instr.jump_to is not None
+                target_idx = instructions.index(instr.jump_to)
+                # Fork to two branches, jump or not
+                jump_branch = fork(name_recorder, i, True, target_idx)
+                not_jump_branch = (
+                    fork(name_recorder, i, False, target_idx)
+                    if instr.opname not in UNCONDITIONAL_JUMP
+                    else NameRecorder(OrderedSet(), OrderedSet())
+                )
+                return jump_branch | not_jump_branch
+            elif instr.opname == "RETURN_VALUE":
+                return name_recorder
+        return name_recorder
+
+    name_recorder = walk(name_recorder, current_instr_idx)
+    return name_recorder.reads, name_recorder.writes

python/paddle/jit/sot/opcode_translator/instruction_utils/opcode_info.py

@@ -0,0 +1,167 @@
+# 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 opcode
+import sys
+from enum import Enum
+
+REL_JUMP = {opcode.opname[x] for x in opcode.hasjrel}
+REL_BWD_JUMP = {opname for opname in REL_JUMP if "BACKWARD" in opname}
+REL_FWD_JUMP = REL_JUMP - REL_BWD_JUMP
+ABS_JUMP = {opcode.opname[x] for x in opcode.hasjabs}
+HAS_LOCAL = {opcode.opname[x] for x in opcode.haslocal}
+HAS_FREE = {opcode.opname[x] for x in opcode.hasfree}
+NEED_TO_BOOL = {"UNARY_NOT", "POP_JUMP_IF_FALSE", "POP_JUMP_IF_TRUE"}
+ALL_JUMP = REL_JUMP | ABS_JUMP
+UNCONDITIONAL_JUMP = {"JUMP_ABSOLUTE", "JUMP_FORWARD"}
+if sys.version_info >= (3, 11):
+    UNCONDITIONAL_JUMP.add("JUMP_BACKWARD")
+RETURN = {"RETURN_VALUE"}
+if (3, 12) <= sys.version_info < (3, 14):
+    RETURN.add("RETURN_CONST")
+
+
+class JumpDirection(Enum):
+    FORWARD = "FORWARD"
+    BACKWARD = "BACKWARD"
+
+
+class PopJumpCond(Enum):
+    FALSE = "FALSE"
+    TRUE = "TRUE"
+    NONE = "NONE"
+    NOT_NONE = "NOT_NONE"
+
+
+def _get_pyopcode_cache_size() -> dict[str, int]:
+    if sys.version_info >= (3, 11) and sys.version_info < (3, 12):
+        # Cache for some opcodes, it's for Python 3.11
+        # https://github.com/python/cpython/blob/3.11/Include/internal/pycore_opcode.h#L41-L53
+        return {
+            "BINARY_SUBSCR": 4,
+            "STORE_SUBSCR": 1,
+            "UNPACK_SEQUENCE": 1,
+            "STORE_ATTR": 4,
+            "LOAD_ATTR": 4,
+            "COMPARE_OP": 2,
+            "LOAD_GLOBAL": 5,
+            "BINARY_OP": 1,
+            "LOAD_METHOD": 10,
+            "PRECALL": 1,
+            "CALL": 4,
+        }
+    elif sys.version_info >= (3, 12) and sys.version_info < (3, 13):
+        # Cache for some opcodes, it's for Python 3.12
+        # https://github.com/python/cpython/blob/3.12/Include/internal/pycore_opcode.h#L34-L47
+        return {
+            "BINARY_SUBSCR": 1,
+            "STORE_SUBSCR": 1,
+            "UNPACK_SEQUENCE": 1,
+            "FOR_ITER": 1,
+            "STORE_ATTR": 4,
+            "LOAD_ATTR": 9,
+            "COMPARE_OP": 1,
+            "LOAD_GLOBAL": 4,
+            "BINARY_OP": 1,
+            "SEND": 1,
+            "LOAD_SUPER_ATTR": 1,
+            "CALL": 3,
+        }
+    elif sys.version_info >= (3, 13) and sys.version_info < (3, 14):
+        # Cache for some opcodes, it's for Python 3.13
+        # https://github.com/python/cpython/blob/3.13/Include/internal/pycore_opcode_metadata.h#L1598-L1618
+        return {
+            "JUMP_BACKWARD": 1,
+            "TO_BOOL": 3,
+            "BINARY_SUBSCR": 1,
+            "STORE_SUBSCR": 1,
+            "SEND": 1,
+            "UNPACK_SEQUENCE": 1,
+            "STORE_ATTR": 4,
+            "LOAD_GLOBAL": 4,
+            "LOAD_SUPER_ATTR": 1,
+            "LOAD_ATTR": 9,
+            "COMPARE_OP": 1,
+            "CONTAINS_OP": 1,
+            "POP_JUMP_IF_TRUE": 1,
+            "POP_JUMP_IF_FALSE": 1,
+            "POP_JUMP_IF_NONE": 1,
+            "POP_JUMP_IF_NOT_NONE": 1,
+            "FOR_ITER": 1,
+            "CALL": 3,
+            "BINARY_OP": 1,
+        }
+    elif sys.version_info >= (3, 14) and sys.version_info < (3, 15):
+        # Cache for some opcodes, it's for Python 3.14
+        # https://github.com/python/cpython/blob/3.14/Include/internal/pycore_opcode_metadata.h#L1764-L1784
+        return {
+            "TO_BOOL": 3,
+            "STORE_SUBSCR": 1,
+            "SEND": 1,
+            "UNPACK_SEQUENCE": 1,
+            "STORE_ATTR": 4,
+            "LOAD_GLOBAL": 4,
+            "LOAD_SUPER_ATTR": 1,
+            "LOAD_ATTR": 9,
+            "COMPARE_OP": 1,
+            "CONTAINS_OP": 1,
+            "JUMP_BACKWARD": 1,
+            "POP_JUMP_IF_TRUE": 1,
+            "POP_JUMP_IF_FALSE": 1,
+            "POP_JUMP_IF_NONE": 1,
+            "POP_JUMP_IF_NOT_NONE": 1,
+            "FOR_ITER": 1,
+            "CALL": 3,
+            "CALL_KW": 3,
+            "BINARY_OP": 5,
+        }
+    elif sys.version_info >= (3, 15):
+        raise NotImplementedError(
+            f"Need to supplement cache operation code, for Python {sys.version_info}"
+        )
+    else:
+        return {}
+
+
+PYOPCODE_CACHE_SIZE = _get_pyopcode_cache_size()
+
+
+class ExceptionHandler:
+    opcode = 257
+    opname = "EXCEPT_HANDLER"
+
+
+def _get_binary_op_arg_map() -> dict[str, int]:
+    if sys.version_info < (3, 11):
+        return {}
+    res = {}
+    for i, op in enumerate(opcode._nb_ops):
+        res[op[0]] = i
+    return res
+
+
+BINARY_OP_ARG_MAP: dict[str, int] = _get_binary_op_arg_map()
+
+FUSED_INSTS: dict[str, tuple[str, str]] = {
+    "LOAD_FAST_LOAD_FAST": ("LOAD_FAST", "LOAD_FAST"),
+    "LOAD_FAST_BORROW_LOAD_FAST_BORROW": (
+        "LOAD_FAST_BORROW",
+        "LOAD_FAST_BORROW",
+    ),
+    "STORE_FAST_STORE_FAST": ("STORE_FAST", "STORE_FAST"),
+    "STORE_FAST_LOAD_FAST": ("STORE_FAST", "LOAD_FAST"),
+}
+
+TO_FUSED_INSTS = {v: k for k, v in FUSED_INSTS.items()}

python/paddle/jit/sot/opcode_translator/instruction_utils/stack_analyse.py

@@ -0,0 +1,93 @@
+# 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 ...utils import Singleton
+
+
+class StackAnalyser(metaclass=Singleton):
+    def stack_effect(self, instr):
+        if "BINARY" in instr.opname or "INPLACE" in instr.opname:
+            return 2, 1
+        elif "UNARY" in instr.opname:
+            return 1, 1
+        return getattr(self, instr.opname)(instr)
+
+    def LOAD_GLOBAL(self, instr):
+        return 0, 1
+
+    def LOAD_CONST(self, instr):
+        return 0, 1
+
+    def LOAD_FAST(self, instr):
+        return 0, 1
+
+    def LOAD_ATTR(self, instr):
+        return 1, 1
+
+    def LOAD_METHOD(self, instr):
+        return 1, 2
+
+    def STORE_FAST(self, instr):
+        return 1, 0
+
+    def BUILD_TUPLE(self, instr):
+        return instr.arg, 1
+
+    def BUILD_LIST(self, instr):
+        return instr.arg, 1
+
+    def BUILD_SLICE(self, instr):
+        if instr.arg == 3:
+            return 3, 1
+        else:
+            return 2, 1
+
+    def UNPACK_SEQUENCE(self, instr):
+        return 1, instr.arg
+
+    def CALL_FUNCTION(self, instr):
+        return instr.arg + 1, 1
+
+    def DUP_TOP(self, instr):
+        return 0, 1
+
+    def DUP_TOP_TWO(self, instr):
+        return 0, 2
+
+    def ROT_N(self, instr):
+        return 0, 0
+
+    def ROT_TWO(self, instr):
+        return 0, 0
+
+    def ROT_THREE(self, instr):
+        return 0, 0
+
+    def ROT_FOUR(self, instr):
+        return 0, 0
+
+    def GET_ITER(self, instr):
+        return 1, 1
+
+    def POP_TOP(self, instr):
+        return 1, 0
+
+    def PUSH_NULL(self, instr):
+        return 0, 1
+
+    def NOP(self, instr):
+        return 0, 0
+
+    def EXTENDED_ARG(self, instr):
+        return 0, 0

python/paddle/jit/sot/opcode_translator/skip_files.py

@@ -0,0 +1,167 @@
+# 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 _collections_abc
+import _weakrefset
+import abc
+import codecs
+import collections
+import contextlib
+import copy
+import copyreg
+import dataclasses
+import enum
+import functools
+import genericpath
+import importlib
+import inspect
+import linecache
+import logging
+import multiprocessing
+import operator
+import os
+import posixpath
+import random
+import re
+import selectors
+import signal
+import sys
+import tempfile
+import threading
+import tokenize
+import traceback
+import types
+import typing
+import unittest
+import uuid
+import warnings
+import weakref
+
+import google.protobuf
+import numpy
+import setuptools
+
+import paddle
+
+NEED_SKIP_THIRD_PARTY_MODULES = {
+    abc,
+    collections,
+    contextlib,
+    copy,
+    copyreg,
+    dataclasses,
+    enum,
+    functools,
+    google.protobuf,
+    importlib,
+    inspect,
+    linecache,
+    logging,
+    multiprocessing,
+    numpy,
+    operator,
+    os,
+    posixpath,
+    random,
+    re,
+    selectors,
+    signal,
+    tempfile,
+    threading,
+    tokenize,
+    traceback,
+    types,
+    typing,
+    unittest,
+    weakref,
+    _collections_abc,
+    _weakrefset,
+    codecs,
+    uuid,
+    setuptools,
+    warnings,
+    genericpath,
+}
+
+if sys.version_info < (3, 11):
+    import sre_compile
+    import sre_parse
+
+    NEED_SKIP_THIRD_PARTY_MODULES.add(sre_compile)
+    NEED_SKIP_THIRD_PARTY_MODULES.add(sre_parse)
+
+if sys.version_info < (3, 12):
+    import distutils
+
+    NEED_SKIP_THIRD_PARTY_MODULES.add(distutils)
+
+
+def _extend_skip_modules_if_exists(module_name: str):
+    """Extend skip modules set if the module exists."""
+    try:
+        module = importlib.import_module(module_name)
+        NEED_SKIP_THIRD_PARTY_MODULES.add(module)
+    except ImportError:
+        pass
+
+
+# Some modules may not be installed in the environment, so we try to import them.
+_extend_skip_modules_if_exists("coverage")
+_extend_skip_modules_if_exists("colorama")
+
+
+def _strip_init_py(s):
+    return re.sub(r"__init__.py$", "", s)
+
+
+def _module_dir(m: types.ModuleType):
+    module_file = getattr(m, "__file__", None)
+    if module_file is None:
+        return None
+    return _strip_init_py(module_file)
+
+
+skip_file_names = {
+    path
+    for path in [_module_dir(m) for m in NEED_SKIP_THIRD_PARTY_MODULES]
+    if path is not None
+}
+
+
+sot_path = os.path.dirname(__file__).rpartition(os.sep)[0] + os.sep
+paddle_path = sys.modules["paddle"].__file__.rpartition(os.sep)[0] + os.sep
+
+skip_file_names.add(sot_path)
+skip_file_names.add(paddle_path)
+skip_file_names.add(
+    "<frozen importlib",
+)
+skip_file_names.add("<__array_function__ internals>")
+
+skip_file_name_re = re.compile(
+    f"^({'|'.join(map(re.escape, skip_file_names))})"
+)
+
+no_skip_code = {paddle.nn.Sequential.forward.__code__}
+
+with_graph_codes = (
+    paddle.nn.Layer.__call__.__code__,
+    paddle.nn.Layer._dygraph_call_func.__code__,
+)
+
+
+def setup_skip_files():
+    paddle.framework.core.eval_frame_skip_file_prefix(tuple(skip_file_names))
+    paddle.framework.core.eval_frame_no_skip_codes(tuple(no_skip_code))
+    paddle.framework.core.sot_setup_codes_with_graph(with_graph_codes)

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

@@ -0,0 +1,21 @@
+# 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 paddle.jit.profiler import (
+    EventGuard as EventGuard,
+    SotProfiler as SotProfiler,
+    event_register as event_register,
+)
+
+from .kernel_stats import SotStepProfilerGuard as SotStepProfilerGuard

python/paddle/jit/sot/profiler/kernel_stats.py

@@ -0,0 +1,236 @@
+# 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 re
+from dataclasses import dataclass
+from enum import Enum
+
+from paddle import profiler
+from paddle.base.core import tracer_event_type_to_string
+
+from ..utils.envs import ENV_ENABLE_SOT_STEP_PROFILER
+
+
+class EventVisitor:
+    def visit(self, event_node):
+        event_type_name = tracer_event_type_to_string(event_node.type)
+        visit_method_name = f"visit_{event_type_name}"
+        if not hasattr(self, visit_method_name):
+            self.generic_visit(event_node)
+            return
+        getattr(self, visit_method_name)(event_node)
+
+    def generic_visit(self, event_node):
+        for child in event_node.children_node:
+            self.visit(child)
+
+    def __call__(self, events):
+        for event_node in events.values():
+            self.visit(event_node)
+
+
+class KernelRunMode(Enum):
+    Dygraph = 1
+    Static = 2
+
+
+def safe_divide(a, b):
+    # Avoid division by zero
+    return a / b if b != 0 else 0
+
+
+@dataclass
+class KernelInfo:
+    name: str
+    run_mode: KernelRunMode
+    duration: float
+    cuda_kernels: list[CudaKernelInfo]
+
+
+@dataclass
+class CudaKernelInfo:
+    name: str
+    duration: float
+
+
+# TODO(SigureMo): Split into multiple files by visitor type
+class KernelStatsVisitor(EventVisitor):
+    SKIP_KERNEL_NAMES = {"full", "full_int_array", "shadow_feed"}
+    KERNEL_NAME_REGEX = re.compile("(?P<kernel_name>.+) kernel launch")
+
+    def __init__(self):
+        self.kernels = []
+
+    def get_kernel_name(self, event_name):
+        if match_obj := self.KERNEL_NAME_REGEX.match(event_name):
+            return match_obj.group("kernel_name")
+        raise ValueError(f"Unexpected event name: {event_name}")
+
+    def calc_kernel_count(self, mode):
+        return len(
+            [
+                kernel
+                for kernel in self.kernels
+                if (
+                    kernel.run_mode == mode
+                    and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
+                )
+            ]
+        )
+
+    def calc_kernel_duration(self, mode):
+        return sum(
+            [
+                kernel.duration
+                for kernel in self.kernels
+                if (
+                    kernel.run_mode == mode
+                    and kernel.name not in KernelStatsVisitor.SKIP_KERNEL_NAMES
+                )
+            ]
+        )
+
+    def find_all_cuda_kernels(self, host_event):
+        # TODO(SigureMo): Find a better way to find all CUDA kernels
+        return [
+            CudaKernelInfo(
+                device_event.name, device_event.end_ns - device_event.start_ns
+            )
+            for runtime_event in host_event.runtime_node
+            for device_event in runtime_event.device_node
+        ]
+
+    def visit_DygraphKernelLaunch(self, event_node):
+        duration = event_node.end_ns - event_node.start_ns
+        kernel_name = self.get_kernel_name(event_node.name)
+        all_cuda_kernels = self.find_all_cuda_kernels(event_node)
+        self.kernels.append(
+            KernelInfo(
+                kernel_name, KernelRunMode.Dygraph, duration, all_cuda_kernels
+            )
+        )
+        self.generic_visit(event_node)
+
+    def visit_StaticKernelLaunch(self, event_node):
+        duration = event_node.end_ns - event_node.start_ns
+        kernel_name = self.get_kernel_name(event_node.name)
+        all_cuda_kernels = self.find_all_cuda_kernels(event_node)
+        self.kernels.append(
+            KernelInfo(
+                kernel_name, KernelRunMode.Static, duration, all_cuda_kernels
+            )
+        )
+        self.generic_visit(event_node)
+
+    def summary(self) -> str:
+        static_kernel_duration = self.calc_kernel_duration(KernelRunMode.Static)
+        dygraph_kernel_duration = self.calc_kernel_duration(
+            KernelRunMode.Dygraph
+        )
+        static_kernel_count = self.calc_kernel_count(KernelRunMode.Static)
+        dygraph_kernel_count = self.calc_kernel_count(KernelRunMode.Dygraph)
+
+        percentage_static_kernel_count = safe_divide(
+            static_kernel_count, static_kernel_count + dygraph_kernel_count
+        )
+        percentage_static_kernel_duration = safe_divide(
+            static_kernel_duration,
+            static_kernel_duration + dygraph_kernel_duration,
+        )
+        step_summary = ""
+        step_summary += f"dygraph kernel count: {dygraph_kernel_count}\n"
+        step_summary += f"static kernel count: {static_kernel_count}\n"
+        step_summary += f"percentage static kernel count: {percentage_static_kernel_count:.2%}\n"
+
+        step_summary += f"dygraph kernel duration: {dygraph_kernel_duration / 1000000:.2f} ms\n"
+        step_summary += f"static kernel duration: {static_kernel_duration / 1000000:.2f} ms\n"
+        step_summary += f"percentage static kernel duration: {percentage_static_kernel_duration:.2%}\n"
+        return step_summary
+
+
+class SotStepProfilerGuard:
+    EXPORT_CHROME_TRACING_PATH = "./sot-chrome-tracing/"
+    STEP_CNT = 0
+    LAST_INFO_SUMMARY = None
+
+    def __init__(self, enable_kernel_stats=True, enable_chrome_tracing=False):
+        self.enable_kernel_stats = enable_kernel_stats
+        self.enable_chrome_tracing = enable_chrome_tracing
+        self.started = False
+        self.record_event = None
+        self.summary = None
+
+    def _kernel_stats(self, prof) -> str:
+        kernel_stats_visitor = KernelStatsVisitor()
+        kernel_stats_visitor(prof.profiler_result.get_data())
+        return kernel_stats_visitor.summary()
+
+    def collect_step_info_summary(self, prof) -> str:
+        summary = ""
+        if self.enable_kernel_stats:
+            summary += self._kernel_stats(prof)
+        if self.enable_chrome_tracing:
+            # If you want to export chrome tracing, you can enable this flag
+            # and view the tracing result in https://ui.perfetto.dev/#!/viewer
+            profiler.export_chrome_tracing(
+                SotStepProfilerGuard.EXPORT_CHROME_TRACING_PATH,
+                f"step_{SotStepProfilerGuard.STEP_CNT:03d}",
+            )(prof)
+        return summary
+
+    def on_trace_ready(self, prof):
+        summary = self.collect_step_info_summary(prof)
+        if SotStepProfilerGuard.STEP_CNT == 0:
+            coldstart_title = f"SOT step profiler info summary (ColdStart, step#{SotStepProfilerGuard.STEP_CNT}):"
+            coldstart_report = f"{coldstart_title}\n{summary}"
+            print(coldstart_report)
+            self.summary = coldstart_report
+        else:
+            warmup_title = f"SOT step profiler info summary (Warmup, step#{SotStepProfilerGuard.STEP_CNT}):"
+            warmup_report = f"{warmup_title}\n{summary}"
+            if SotStepProfilerGuard.LAST_INFO_SUMMARY is None:
+                atexit.register(
+                    lambda: print(SotStepProfilerGuard.LAST_INFO_SUMMARY)
+                )
+            SotStepProfilerGuard.LAST_INFO_SUMMARY = warmup_report
+            self.summary = warmup_report
+
+    def start(self):
+        if ENV_ENABLE_SOT_STEP_PROFILER.get():
+            self.profiler = profiler.Profiler(
+                targets=[
+                    profiler.ProfilerTarget.CPU,
+                    profiler.ProfilerTarget.GPU,
+                ],
+                on_trace_ready=self.on_trace_ready,
+            )
+            self.profiler.start()
+            self.started = True
+
+    def stop(self):
+        if self.started:
+            assert self.profiler is not None
+            self.profiler.stop()
+            self.profiler = None  # Avoid to hold the profiler instance
+        SotStepProfilerGuard.STEP_CNT += 1
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()

python/paddle/jit/sot/psdb.py

@@ -0,0 +1,68 @@
+# 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
+from typing import TYPE_CHECKING, TypeVar
+
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from types import CodeType
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+NO_BREAKGRAPH_CODES: set[CodeType] = set()
+NO_FALLBACK_CODES: set[CodeType] = set()
+
+
+def assert_true(input: bool):
+    assert input
+
+
+def print(*args, **kwargs):
+    builtins.print("[Dygraph]", *args, **kwargs)
+
+
+def breakpoint():
+    import paddle
+
+    old = paddle.framework.core.set_eval_frame(None)
+    builtins.breakpoint()  # noqa: T100
+    paddle.framework.core.set_eval_frame(old)
+
+
+def check_no_breakgraph(fn: Callable[P, T]) -> Callable[P, T]:
+    NO_BREAKGRAPH_CODES.add(fn.__code__)
+    return fn
+
+
+def breakgraph():
+    pass
+
+
+def check_no_fallback(fn: Callable[P, T]) -> Callable[P, T]:
+    NO_FALLBACK_CODES.add(fn.__code__)
+    return fn
+
+
+def fallback(recursive=False):
+    pass
+
+
+def in_sot():
+    return False

python/paddle/jit/sot/symbolic/builder.py

@@ -0,0 +1,192 @@
+# 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
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+from ..utils import log
+from .compile_cache import CompileSIRCache
+from .statement_ir import (
+    ApiStatement,
+    ASTStatement,
+    CallStatement,
+    LayerStatement,
+    MethodStatement,
+    ParametersHolder,
+    StatementContext,
+    StatementIR,
+    StatementIRFactory,
+    Symbol,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle.static import InputSpec
+
+
+class StatementIRBuilder:
+    """
+    A class to build a StatementIR.
+    """
+
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        """
+        Reset the context.
+        """
+
+        self.statement_factory = StatementIRFactory()
+        self._current_statement_ctxs = []
+        self._current_sir: StatementIR = self.statement_factory.create()
+
+    @property
+    def current_sir(self) -> StatementIR:
+        """
+        Get the current SIR.
+        """
+        return self._current_sir
+
+    def replace_current_sir(self, sir: StatementIR):
+        """
+        Replace the current SIR with a new SIR.
+        """
+        self._current_sir = sir
+        self.statement_factory.update(sir)
+
+    def call_SIR(self, sirname, inputs, outputs, stacks):
+        """
+        Call a SIR, which is a subgraph.
+        """
+
+        stmt = CallStatement(
+            sirname, inputs, outputs, list(self._current_statement_ctxs), stacks
+        )
+        self.current_sir.add_statement(stmt)
+
+    def call_API(self, api, inputs, outputs, stacks):
+        """
+        Call a paddle api.
+        """
+        assert callable(api), "call_API must receive a paddle api."
+        stmt = ApiStatement(
+            api, inputs, outputs, list(self._current_statement_ctxs), stacks
+        )
+        self.current_sir.add_statement(stmt)
+
+    def call_METHOD(self, method_name, inputs, outputs, stacks):
+        """
+        Call a method of a api. The API here can be python or Paddle
+        """
+        assert isinstance(method_name, str), (
+            "call_METHOD must method api name. string."
+        )
+        assert isinstance(inputs[0][0], Symbol), (
+            "call_METHOD first argument must be Symbol Variable."
+        )
+        stmt = MethodStatement(
+            method_name,
+            inputs,
+            outputs,
+            list(self._current_statement_ctxs),
+            stacks,
+        )
+        self.current_sir.add_statement(stmt)
+
+    def call_LAYER(self, layer, inputs, outputs, stacks):
+        """
+        Call a layer of a api.
+        """
+        stmt = LayerStatement(
+            layer, inputs, outputs, list(self._current_statement_ctxs), stacks
+        )
+        self.current_sir.add_statement(stmt)
+
+    def call_AST(self, static_function, inputs, outputs, stacks):
+        stmt = ASTStatement(
+            static_function,
+            inputs,
+            outputs,
+            list(self._current_statement_ctxs),
+            stacks,
+        )
+        self.current_sir.add_statement(stmt)
+
+    def get_sir(self, name: str):
+        """
+        Get a SIR from statement_factory.
+
+        Args:
+            name (str): the name of SIR.
+
+        Returns:
+            StatementIR: the SIR.
+        """
+        return self.statement_factory[name]
+
+    @contextmanager
+    def attach_statement_context_guard(self, ctx: StatementContext):
+        """
+        Attach a statement context to the current SIR.
+        """
+        self._current_statement_ctxs.append(ctx)
+        try:
+            yield
+        finally:
+            self._current_statement_ctxs.pop()
+
+    def finalize(self, ret_vals):
+        current_sir: StatementIR = self.current_sir
+        current_sir.inputs, current_sir.params = current_sir.analyse_inputs()
+        current_sir.outputs = ret_vals
+        log(2, "start subgraph compile and execution.\n")
+        log(2, current_sir, "\n")
+        return current_sir
+
+    def compile_do_nothing(self) -> Callable[..., Any]:
+        """
+        Return a dummy function, which will return an empty list.
+
+        Args:
+            ret_vals (list[Symbol]): the return values of the function.
+        """
+
+        class DummyFunc:
+            def __call__(*args, **kwargs):
+                return []
+
+            def graph_size(self):
+                return 0
+
+        return DummyFunc()
+
+    def compile_fn(
+        self,
+        sir_name: str,
+        parameters_holder: ParametersHolder,
+        input_spec: tuple[InputSpec | None, ...],
+        **kwargs,
+    ):
+        """
+        start compile and return the python function, which must can be to_static without errors.
+        """
+        static_func = CompileSIRCache()(
+            self, sir_name, parameters_holder, input_spec, **kwargs
+        )
+
+        return static_func

python/paddle/jit/sot/symbolic/compile_cache.py

@@ -0,0 +1,322 @@
+# 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 inspect
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.jit.profiler import EventGuard, event_register
+
+from ..infer_meta import convert_meta_to_input_spec
+from ..utils import (
+    ENV_SOT_EXPORT,
+    Cache,
+    InfoCollector,
+    NewSymbolHitRateInfo,
+    Singleton,
+    SIRToCodeMap,
+    StepInfoManager,
+    SubGraphInfo,
+    SubGraphRelationInfo,
+    log_do,
+)
+from .export import export
+from .interpreter import compile_sir
+
+if TYPE_CHECKING:
+    from paddle.static import InputSpec, Program
+
+    from .builder import StatementIRBuilder
+    from .statement_ir import ParametersHolder
+
+
+def trace_back_frames():
+    frame = inspect.currentframe()
+    while frame.f_back is not None:
+        frame = frame.f_back
+        code = frame.f_code
+        paddle.framework.core.sot_set_with_graph(code)
+
+
+def clear_eager_tensor_name(output_tensors):
+    for output_tensor in output_tensors:
+        output_tensor.name = ""
+
+
+def _is_builtin_op(op):
+    dialect_name, opname = op.name().split(".")
+    return dialect_name == "builtin"
+
+
+def _is_computation_op(op):
+    return not _is_builtin_op(op) and op.name() not in ["pd_op.data"]
+
+
+class UniqueIdGenerator:
+    def __init__(self):
+        self._id = 0
+
+    def generate(self):
+        self._id += 1
+        return self._id
+
+    def __call__(self):
+        return self.generate()
+
+
+class TensorIdAllocator(metaclass=Singleton):
+    TENSOR_ID_ATTR = "__tensor_id__"
+
+    def __init__(self):
+        self._id_generator = UniqueIdGenerator()
+
+    def allocate(self, tensor):
+        if not hasattr(tensor, self.TENSOR_ID_ATTR):
+            setattr(tensor, self.TENSOR_ID_ATTR, self._id_generator())
+        return getattr(tensor, self.TENSOR_ID_ATTR)
+
+
+class FallbackWrapper:
+    """
+    Used to store and call static graph methods generated by paddle.jit.to_static
+    """
+
+    def __init__(self, compiled_fn, SIR, is_training: bool):
+        self.compiled_fn = compiled_fn
+        self.partial_program = None
+        self.concrete_program = None
+        self.SIR = SIR  # for debug
+        self.is_training = is_training
+        self.exported = False
+        self.is_first_call = True
+
+    def graph_size(self):
+        if self.partial_program is None:
+            input_spec = convert_meta_to_input_spec(
+                tuple(
+                    self.SIR.symbol_meta_map[symbol]
+                    for symbol in self.SIR.inputs
+                )
+            )
+            (
+                self.concrete_program,
+                self.partial_program,
+            ) = self.compiled_fn.get_concrete_program(input_spec)
+            self.partial_program.training = self.is_training
+        global_block_ops = self.concrete_program.main_program.global_block().ops
+        non_builtin_ops = list(filter(_is_computation_op, global_block_ops))
+        return len(non_builtin_ops)
+
+    def collect_new_symbol_hit_rate(self, inputs, outputs):
+        if not InfoCollector().need_collect(NewSymbolHitRateInfo):
+            return
+        input_tensor_ids = []
+        output_tensor_ids = []
+        assert len(inputs) == 1
+        assert isinstance(inputs[0], tuple)
+        for i, arg in enumerate(inputs[0]):
+            assert isinstance(arg, paddle.Tensor), f"Expect Tensor, got {arg}"
+            tensor_id = TensorIdAllocator().allocate(arg)
+            input_tensor_ids.append(tensor_id)
+
+        for i, out in enumerate(outputs):
+            assert isinstance(out, paddle.Tensor)
+            tensor_id = TensorIdAllocator().allocate(out)
+            output_tensor_ids.append(tensor_id)
+
+        InfoCollector().attach(
+            NewSymbolHitRateInfo, input_tensor_ids, output_tensor_ids
+        )
+
+    def collect_subgraph_relation(self, inputs, outputs, partial_program_layer):
+        if not InfoCollector().need_collect(SubGraphRelationInfo):
+            return
+        input_shape_infos = []
+        output_shape_infos = []
+        forward_input_values = partial_program_layer.program.program_attr['fx']
+        forward_output_values = partial_program_layer.program.program_attr['fo']
+        assert len(inputs) == 1
+        assert isinstance(inputs[0], tuple)
+        assert len(inputs[0]) == len(forward_input_values)
+        assert len(outputs) == len(forward_output_values)
+        for i, arg in enumerate(inputs[0]):
+            assert isinstance(arg, paddle.Tensor), f"Expect Tensor, got {arg}"
+            tensor_id = TensorIdAllocator().allocate(arg)
+            input_ir_shape = forward_input_values[i].shape
+            input_real_shape = arg.shape
+            input_shape_info = SubGraphRelationInfo.ConcreteShapeInfo(
+                tensor_id, input_ir_shape, input_real_shape
+            )
+            input_shape_infos.append(input_shape_info)
+
+        for i, out in enumerate(outputs):
+            assert isinstance(out, paddle.Tensor)
+            tensor_id = TensorIdAllocator().allocate(out)
+            output_ir_shape = forward_output_values[
+                partial_program_layer._outputs.quick_index_map[i]
+            ].shape
+            output_real_shape = out.shape
+            output_shape_info = SubGraphRelationInfo.ConcreteShapeInfo(
+                tensor_id, output_ir_shape, output_real_shape
+            )
+            output_shape_infos.append(output_shape_info)
+
+        InfoCollector().attach(
+            SubGraphRelationInfo,
+            self.SIR.name,
+            input_shape_infos,
+            output_shape_infos,
+            self.is_first_call,
+            self.graph_size(),
+        )
+
+    def collect_subgraph_info(self, program: Program):
+        if not InfoCollector().need_collect(SubGraphInfo):
+            return
+
+        InfoCollector().attach(
+            SubGraphInfo,
+            str(program),
+            self.graph_size(),
+            self.SIR.name,
+        )
+
+    def update_compile_time_info(self, SIR, partial_program_layer):
+        if not self.is_first_call:
+            return
+        from ..opcode_translator.executor.executor_cache import (
+            OpcodeExecutorCache,
+        )
+
+        code = SIRToCodeMap().get(SIR)
+        assert code is not None, f"Cannot find code for SIR: {SIR}"
+
+        OpcodeExecutorCache().compile_time_stats.setdefault(code, 0)
+        OpcodeExecutorCache().compile_time_stats[code] += (
+            partial_program_layer._compile_time_counter.get_total_time()
+        )
+
+    @event_register(
+        lambda self, *args, **kwargs: f"FallbackWrapper: {self.SIR.name}"
+    )
+    def __call__(self, *args, **kwargs):
+        if StepInfoManager().need_back_trace:
+            trace_back_frames()
+
+        log_do(
+            2,
+            lambda: print("[FallbackWrapper] start run SIR: \n", self.SIR),
+        )
+        log_do(
+            4,
+            lambda: print(
+                self.compiled_fn.get_concrete_program(*args, **kwargs)[
+                    1
+                ].train_program
+            ),
+        )
+        if self.partial_program is None:
+            with EventGuard("FallbackWrapper: get_concrete_program"):
+                (
+                    self.concrete_program,
+                    self.partial_program,
+                ) = self.compiled_fn.get_concrete_program(*args, **kwargs)
+                self.partial_program.training = self.is_training
+        outputs = self.partial_program.sot_call(*args, **kwargs)
+
+        clear_eager_tensor_name(outputs)
+        log_do(
+            4,
+            lambda: print("[CompileCache] run sir forward success."),
+        )
+        self.collect_new_symbol_hit_rate(args, outputs)
+        self.collect_subgraph_relation(args, outputs, self.partial_program)
+        self.collect_subgraph_info(self.concrete_program.main_program)
+        self.update_compile_time_info(self.SIR, self.partial_program)
+        if ENV_SOT_EXPORT.get() != "" and not self.exported:
+            export(self.SIR, ENV_SOT_EXPORT.get())
+            self.exported = True
+
+        self.is_first_call = False
+        return outputs
+
+
+class CompileSIRCache(Cache, metaclass=Singleton):
+    """
+    Cache the compiled function of SIR
+    """
+
+    def __init__(self):
+        super().__init__(weak=False)
+
+    def key_fn(
+        self,
+        builder: StatementIRBuilder,
+        sir_name: str,
+        parameters_holder: ParametersHolder,
+        input_spec: tuple[InputSpec | None, ...],
+        **kwargs,
+    ):
+        """
+        generate a hash key for a SIR
+
+        Args:
+            context: The context to compile
+            sir_name: The name of the sir to compile
+            build_strategy: The build strategy to compile
+
+        Returns:
+            The hash key of the SIR
+        """
+        sir = builder.get_sir(sir_name)
+        # NOTE(dev): Is str(sir) a heavy operation ?
+        hash_key = hash(
+            (str(sir), *input_spec, id(parameters_holder), kwargs['training'])
+        )
+        return hash_key
+
+    def value_fn(
+        self,
+        builder: StatementIRBuilder,
+        sir_name: str,
+        parameters_holder: ParametersHolder,
+        input_spec: tuple[InputSpec | None, ...],
+        **kwargs,
+    ):
+        """
+        Generate static graph function
+
+        Args:
+            context: The context to compile
+            sir_name: The name of the sir to compile
+            build_strategy: The build strategy to compile
+
+        Returns:
+            The static graph function
+        """
+        build_strategy = kwargs.get("build_strategy", None)
+        backend = kwargs.get("backend", None)
+        return FallbackWrapper(
+            paddle.jit.to_static(
+                compile_sir(builder, sir_name, parameters_holder),
+                input_spec=[input_spec],
+                build_strategy=build_strategy,
+                backend=backend,
+                full_graph=True,
+            ),
+            builder.get_sir(sir_name),
+            is_training=kwargs['training'],
+        )

python/paddle/jit/sot/symbolic/export.py

@@ -0,0 +1,387 @@
+# 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.
+
+import os
+from itertools import chain
+
+import paddle
+from paddle.utils import flatten
+
+from ..utils import ConstTypes, ExportError, NameGenerator, get_api_fullname
+from .statement_ir import Symbol
+
+
+class PyStatement:
+    tab = " " * 4
+
+    def __init__(self, *lines):
+        self.sub_statement = []
+        self.lines = lines
+
+    def get_lines(self, prefix=""):
+        lines = [prefix + line for line in self.lines]
+        for statement in self.sub_statement:
+            lines.extend(statement.get_lines(self.tab + prefix))
+        return lines
+
+    def add_sub(self, *lines):
+        sub = PyStatement(*lines)
+        self.sub_statement.append(sub)
+        return sub
+
+    def __str__(self):
+        return "\n".join(self.get_lines())
+
+
+class NameGener:
+    def __init__(self, SIR):
+        self.SIR = SIR
+        self.name_map = {}
+        self.param_name_generator = NameGenerator("self.parameter_")
+        self.non_param_name_generator = NameGenerator("var_")
+
+    def __call__(self, var):
+        return self.get_str(var)
+
+    def get_str(self, var):
+        if isinstance(var, list):
+            return self.get_list_str(var)
+        elif isinstance(var, tuple):
+            return self.get_tuple_str(var)
+        elif isinstance(var, dict):
+            return self.get_dict_str(var)
+        elif isinstance(var, set):
+            return self.get_set_str(var)
+        else:
+            return self.get_obj_str(var)
+
+    def get_list_str(self, list_):
+        return "[{}]".format(", ".join(self.get_str(var) for var in list_))
+
+    def get_tuple_str(self, tuple_):
+        return "({},)".format(", ".join(self.get_str(var) for var in tuple_))
+
+    def get_dict_str(self, dict_):
+        return "{{{},}}".format(
+            ", ".join(
+                f"{self.get_str(k)}: {self.get_str(v)}"
+                for k, v in dict_.items()
+            )
+        )
+
+    def get_set_str(self, set_):
+        return "{{{},}}".format(", ".join(self.get_str(var) for var in set_))
+
+    def get_obj_str(self, var):
+        if isinstance(var, Symbol):
+            if var not in self.name_map:
+                self.register_symbol(var)
+            return self.name_map[var]
+
+        elif isinstance(var, str):
+            return f"'{var}'"
+        else:
+            return str(var)
+
+    def register_symbol(self, symbol):
+        if symbol in self.SIR.param_symbol:
+            name = self.param_name_generator.next()
+        else:
+            name = self.non_param_name_generator.next()
+        self.name_map[symbol] = name
+
+
+class PyFileGen:
+    def __init__(self, SIR):
+        self.SIR = SIR
+        self.roots = []
+
+        self.name_gener = NameGener(self.SIR)
+
+        self.SIR_sig = "||".join(
+            f"{stmt.type}:{stmt.name}" for stmt in SIR.statements
+        )
+
+    def new_root(self, *args):
+        stmt = PyStatement(*args)
+        self.roots.append(stmt)
+        return stmt
+
+    def roots_to_string(self):
+        lines = []
+        for root in self.roots:
+            lines.extend(root.get_lines())
+        return "\n".join(lines)
+
+    def gen_py_codes(self):
+        self.check_exportable()
+        self.create_header()
+        self.new_root("\n")
+        self.create_layer()
+        self.new_root("\n")
+        self.create_inputs()
+        self.new_root("\n")
+        self.create_test()
+        self.new_root("\n")
+        self.create_tail()
+        return self.roots_to_string()
+
+    def is_exportable_type(self, value):
+        if (
+            isinstance(value, (ConstTypes, Symbol, paddle.dtype))
+            or value is Ellipsis  # NOINT
+        ):
+            return True
+        if isinstance(value, slice):
+            return (
+                self.is_exportable_type(value.start)
+                and self.is_exportable_type(value.stop)
+                and self.is_exportable_type(value.step)
+            )
+        return False
+
+    def check_exportable(self):
+        for stmt in self.SIR.statements:
+            for inp in flatten(stmt.inputs):
+                if not self.is_exportable_type(inp):
+                    raise ExportError(
+                        f"Not support create python file with input: {inp}"
+                    )
+
+    def create_header(self):
+        self.new_root(
+            f"# {self.SIR_sig}",
+            "import paddle",
+            "import unittest",
+            "import numpy as np",
+        )
+
+    def create_layer(self):
+        layer_class = self.new_root("class LayerCase(paddle.nn.Layer):")
+
+        init_fn = layer_class.add_sub("def __init__(self):")
+        init_fn.add_sub("super().__init__()")
+
+        for param in self.SIR.param_symbol:
+            meta = self.SIR.symbol_meta_map[param].unwrap_unsafe()
+            init_fn.add_sub(
+                f"{self.name_gener(param)} = self.create_parameter(",
+                f"   shape={meta.shape},",
+                f"   dtype={meta.dtype},",
+                ")",
+            )
+
+        for stmt in self.SIR.statements:
+            if stmt.type == "layer":
+                layer = stmt.layer()
+                init_fn.add_sub(self.init_sub_layer(layer))
+
+        forward_definition = ["def forward(", "    self,"]
+
+        for inp in self.SIR.inputs:
+            if inp in self.SIR.non_param_symbol:
+                meta = self.SIR.symbol_meta_map[inp]
+                forward_definition.append(
+                    f"    {self.name_gener(inp)},    # {meta}"
+                )
+        forward_definition.append("):")
+
+        forward_fn = layer_class.add_sub(*forward_definition)
+
+        for stmt in self.SIR.statements:
+            forward_fn.add_sub(*self.create_stmt_line(stmt))
+
+        forward_fn.add_sub(
+            "return {}".format(
+                ", ".join(self.name_gener(out) for out in self.SIR.outputs)
+            )
+        )
+
+    def create_inputs(self):
+        create_paddle_inputs = self.new_root("def create_paddle_inputs():")
+        self.new_root("\n")
+        create_numpy_inputs = self.new_root("def create_numpy_inputs():")
+
+        paddle_inputs = ["inputs = ("]
+        numpy_inputs = ["inputs = ("]
+
+        for inp in self.SIR.inputs:
+            if inp in self.SIR.non_param_symbol:
+                meta = self.SIR.symbol_meta_map[inp.name].unwrap_unsafe()
+                shape_str = "[1]" if len(meta.shape) == 0 else str(meta.shape)
+                if meta.dtype in (
+                    paddle.int8,
+                    paddle.int16,
+                    paddle.int32,
+                    paddle.int64,
+                ):
+                    paddle_inputs.append(
+                        f"    paddle.randint(low=0, high=10, shape={shape_str}, dtype={meta.dtype}),"
+                    )
+                    numpy_inputs.append(
+                        "    np.random.randint(low=0, high=10, size={}, dtype='{}'),".format(
+                            shape_str, str(meta.dtype).replace('paddle.', '')
+                        )
+                    )
+                elif meta.dtype is paddle.bool:
+                    paddle_inputs.append(
+                        f"    paddle.randint(low=0, high=2, shape={shape_str}, dtype=paddle.int32).cast(paddle.bool),"
+                    )
+                    numpy_inputs.append(
+                        f"    np.random.randint(low=0, high=2, size={shape_str}, dtype='int').astype('bool'),"
+                    )
+                else:
+                    paddle_inputs.append(
+                        f"    paddle.rand(shape={shape_str}, dtype={meta.dtype}),"
+                    )
+                    numpy_inputs.append(
+                        "    np.random.random(size={}).astype('{}'),".format(
+                            shape_str, str(meta.dtype).replace('paddle.', '')
+                        )
+                    )
+
+        paddle_inputs.append(")")
+        paddle_inputs.append("return inputs")
+        numpy_inputs.append(")")
+        numpy_inputs.append("return inputs")
+
+        create_paddle_inputs.add_sub(*paddle_inputs)
+        create_numpy_inputs.add_sub(*numpy_inputs)
+
+    def create_test(self):
+        test_class = self.new_root("class TestLayer(unittest.TestCase):")
+
+        setup = test_class.add_sub("def setUp(self):")
+        setup.add_sub("self.inputs = create_paddle_inputs()")
+        setup.add_sub("self.net = LayerCase()")
+
+        train = test_class.add_sub(
+            "def train(self, net, to_static, with_prim=False, with_cinn=False):"
+        )
+        train.add_sub(
+            "if to_static:",
+            "    paddle.base.core._set_prim_all_enabled(with_prim)",
+            "    if with_cinn:",
+            '        assert with_prim, "with_cinn=True but with_prim=False is unsupported"',
+            '        net = paddle.jit.to_static(net, backend="CINN", full_graph=True)',
+            "    else:",
+            "        net = paddle.jit.to_static(net, backend=None, full_graph=True)",
+            "paddle.seed(123)",
+            "outs = net(*self.inputs)",
+            "return outs",
+        )
+
+        test_ast_cinn_static = test_class.add_sub(
+            "def test_ast_prim_cinn(self):"
+        )
+        test_ast_cinn_static.add_sub(
+            "st_out = self.train(self.net, to_static=True)",
+            "cinn_out = self.train(self.net, to_static=True, with_prim=True, with_cinn=True)",
+            "for st, cinn in zip(paddle.utils.flatten(st_out), paddle.utils.flatten(cinn_out)):",
+            "    np.testing.assert_allclose(st.numpy(), cinn.numpy(), atol=1e-8)",
+        )
+
+    def create_tail(self):
+        self.new_root(
+            "if __name__ == '__main__':",
+            "    unittest.main()",
+        )
+
+    def init_sub_layer(self, layer, layer_name):
+        # TODO @wuzhanfei need more efficient way to create a sub layer
+        # now, we just close call_Layer behavior
+        raise ExportError("Not support create sub layer now.")
+
+    def create_input_string(self, args, kwargs):
+        return ", ".join(
+            chain(
+                (self.name_gener(arg) for arg in args),
+                (f"{k}={self.name_gener(v)}" for k, v in kwargs.items()),
+            )
+        )
+
+    def create_unpack_output_string(self, outputs):
+        path = ["out"]
+        result = []
+
+        def search(outputs, path, result):
+            if isinstance(outputs, (list, tuple)):
+                search_sequence(outputs, path, result)
+            elif isinstance(outputs, dict):
+                search_dict(outputs, path, result)
+            elif isinstance(outputs, Symbol):
+                result.append(self.name_gener(outputs) + " = " + "".join(path))
+
+        def search_sequence(outputs, path, result):
+            for idx, out in enumerate(outputs):
+                path.append(f"[{idx}]")
+                search(out, path, result)
+                path.pop()
+
+        def search_dict(outputs, path, result):
+            for k, out in outputs.items():
+                path.append(f"[{k}]")
+                search(out, path, result)
+                path.pop()
+
+        search(outputs, path, result)
+        return result
+
+    def create_stmt_line(self, stmt):
+        return getattr(self, "create_" + stmt.type + "_stmt")(stmt)
+
+    def create_api_stmt(self, stmt):
+        args, kwargs = stmt.inputs
+        input_str = self.create_input_string(args, kwargs)
+        api = stmt.api
+        api_str = get_api_fullname(api)
+        if api_str is None:
+            raise ExportError(f"Can not find module of {api}")
+        if isinstance(stmt.outputs, Symbol):
+            return [f"{self.name_gener(stmt.outputs)} = {api_str}({input_str})"]
+        else:
+            compute_code = f"out = {api_str}({input_str})"
+            unpack_codes = self.create_unpack_output_string(stmt.outputs)
+            return [compute_code, *unpack_codes]
+
+    def create_method_stmt(self, stmt):
+        args, kwargs = stmt.inputs
+        input_str = self.create_input_string(args[1:], kwargs)
+        method_str = self.name_gener(args[0]) + "." + stmt.method
+        if isinstance(stmt.outputs, Symbol):
+            return [
+                f"{self.name_gener(stmt.outputs)} = {method_str}({input_str})"
+            ]
+        else:
+            compute_code = f"out = {method_str}({input_str})"
+            unpack_codes = self.create_unpack_output_string(stmt.outputs)
+            return [compute_code, *unpack_codes]
+
+
+def export(SIR, path):
+    try:
+        pygen = PyFileGen(SIR)
+        string = pygen.gen_py_codes()
+    except ExportError as e:
+        print(f"[SOT] Export {SIR.name} Failed:", e)
+        return
+
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+    with open(os.path.join(path, f"{SIR.name}.py"), "w") as f:
+        f.write(string)
+        print(
+            f"[SOT] Export {SIR.name} Success with size {len(SIR.statements)}"
+        )

python/paddle/jit/sot/symbolic/interpreter.py

@@ -0,0 +1,220 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+import paddle
+from paddle.jit.dy2static.utils import compose_guards
+from paddle.utils import to_sequence
+
+from ..utils import (
+    InnerError,
+    log_do,
+    map_if,
+    map_if_extend,
+)
+from .statement_ir import (
+    ParametersHolder,
+    StatementContext,
+    StatementContextRegistry,
+    Symbol,
+)
+
+if TYPE_CHECKING:
+    from .builder import StatementIRBuilder
+    from .statement_ir import Statement, StatementIR
+
+
+def replace_symbol(
+    values: list[Symbol] | list[object], state: dict[str, Symbol]
+):
+    """
+    Replaces Symbol objects with their corresponding values.
+
+    Args:
+        values: A list of values that may contain Symbol objects.
+        state: A dict mapping Symbol names to their corresponding values.
+
+    Returns:
+        A new list with Symbol objects replaced by their corresponding values in the state dict.
+    """
+    # deal with list / map etc.
+    values = map_if_extend(
+        values,
+        pred=lambda x: isinstance(x, Symbol),
+        true_fn=lambda x: state[x.name],
+        false_fn=lambda x: x,
+    )
+    return values
+
+
+def _append_opstack_between(start, end, stack):
+    # The range is [start, end)
+    for op in for_each_ops_between(start, end):
+        op.callstack = stack
+
+
+def for_each_ops_between(start, end):
+    # [start, end)
+    program = paddle.static.default_main_program()
+    ops = program.global_block().ops[start:end]
+    yield from ops
+
+
+def opnum_in_program():
+    program = paddle.static.default_main_program()
+    return len(program.global_block().ops)
+
+
+class Interpreter:
+    """
+    Interpreter is used to interpret and execute SIR.
+    """
+
+    def __init__(self, builder: StatementIRBuilder):
+        self._builder = builder
+
+    def get_sir(self, name: str) -> StatementIR:
+        """
+        Returns the StatementIR object by given name.
+
+        Args:
+            name: The name of the StatementIR.
+
+        Returns:
+            The StatementIR object with the given name.
+        """
+        return self._builder.get_sir(name)
+
+    def run_sir(self, name: str, state: dict[str, Symbol]):
+        """
+        Runs the StatementIR with the given name using the provided state.
+
+        Args:
+            name: The name of the given StatementIR to run.
+            state: A dict mapping Symbol names to their corresponding values.
+
+        Returns:
+            A list of the Symbol of the StatementIR after execution.
+        """
+
+        def _set(v, s):
+            state[s.name] = v
+
+        SIR = self.get_sir(name)
+        for stmt in SIR.statements:
+            stmt: Statement
+            before_stmt_opnum = opnum_in_program()
+            inputs = replace_symbol(stmt.inputs, state)
+
+            with create_context_guard(stmt.contexts)():
+                outs = getattr(self, stmt.type)(stmt, inputs)
+
+            if len(to_sequence(outs)) != len(to_sequence(stmt.outputs)):
+                raise InnerError("Number output mismatch, some error happen.")
+
+            log_do(
+                3,
+                lambda: _append_opstack_between(
+                    before_stmt_opnum, opnum_in_program() + 1, stmt.stmt_stack
+                ),
+            )
+
+            map_if(
+                outs,
+                stmt.outputs,
+                pred=lambda v, s: isinstance(s, Symbol),
+                true_fn=lambda v, s: _set(v, s),
+                false_fn=lambda v, s: None,
+            )
+        # fetch outputs
+        return replace_symbol(SIR.outputs, state)
+
+    def api(self, stmt, inputs):
+        args, kwargs = inputs
+        return stmt.api(*args, **kwargs)
+
+    def method(self, stmt, inputs):
+        args, kwargs = inputs
+        var = args[0]
+        return getattr(var, stmt.method)(*args[1:], **kwargs)
+
+    def layer(self, stmt, inputs):
+        args, kwargs = inputs
+        layer = stmt.layer()
+        assert layer is not None, "SIR bound layer is None."
+        return layer(*args, **kwargs)
+
+    def AST(self, stmt, inputs):
+        args, kwargs = inputs
+        return stmt.converted_func(*args, **kwargs)
+
+
+def compile_sir(
+    builder: StatementIRBuilder, name: str, parameters_holder: ParametersHolder
+):
+    """
+    Compile a SIR to a new function
+
+    Args:
+        context: The context to compile
+        name: The name of the sir to compile
+
+    """
+
+    @paddle.jit.not_to_static
+    def wrapper(args):
+        """
+        This function will be decorated by paddle.to_static.
+        so the args is variables, not eager tensors.
+        """
+        interpreter = Interpreter(builder)
+        SIR = interpreter.get_sir(name)
+        state = prepare_state(SIR, args, parameters_holder)
+        return interpreter.run_sir(name, state)
+
+    return wrapper
+
+
+def prepare_state(
+    SIR: StatementIR, inputs, parameters_holder: ParametersHolder
+):
+    state = {}
+    # bind inputs
+    assert len(SIR.inputs) == len(inputs), "Inputs length mismatch."
+    for sir_inp, inp in zip(SIR.inputs, inputs):
+        state[sir_inp.name] = inp
+
+    for sir_param in SIR.params:
+        state[sir_param.name] = paddle.base.dygraph.base._convert_into_variable(
+            parameters_holder.get(sir_param.name)
+        )
+
+    return state
+
+
+def create_context_guard(contexts: list[StatementContext]):
+    guards = list(
+        map(
+            lambda ctx: (
+                lambda: StatementContextRegistry.get_context_guard(type(ctx))(
+                    ctx
+                )
+            ),
+            contexts,
+        )
+    )
+    return compose_guards(*guards)

python/paddle/jit/sot/symbolic/statement_ir.py

@@ -0,0 +1,419 @@
+# 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 functools
+import weakref
+from typing import TYPE_CHECKING, Any, TypeVar
+from weakref import WeakValueDictionary
+
+import paddle
+from paddle.jit.dy2static.utils import parameters_persistent_mode_is_enabled
+from paddle.jit.utils import OrderedSet
+from paddle.utils import flatten, map_structure
+
+from ..utils import (
+    InnerError,
+    NameGenerator,
+    Singleton,
+    flatten_extend,
+    get_api_fullname,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from contextlib import AbstractContextManager
+
+
+_StatementContextT = TypeVar("_StatementContextT", bound="StatementContext")
+
+
+class ParametersHolder:
+    def __init__(self):
+        self._params = WeakValueDictionary[
+            str, paddle.base.framework.EagerParamBase
+        ]()
+
+    def set(self, name, param):
+        self._params[name] = param
+
+    def get(self, name):
+        if (param := self._params.get(name)) is None:
+            raise InnerError(
+                f"Parameter '{name}' not found in ParametersHolder."
+            )
+        return param
+
+    def copy(self):
+        new_holder = ParametersHolder()
+        new_holder._params = self._params.copy()
+        return new_holder
+
+
+class Reference:  # to unify weak_ref and strong_ref
+    def __init__(self, value, is_weak):
+        self.is_weak = is_weak
+        if is_weak is True:
+            self.ref = weakref.ref(value)
+        else:
+            self.ref = value
+
+    def __call__(self):
+        if self.is_weak is True:
+            return self.ref()
+        else:
+            return self.ref
+
+
+class Symbol:
+    """
+    Symbol is used to distinguish a string and a `math variable`.
+    """
+
+    def __init__(self, name: str):
+        self.name = name
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return str(self)
+
+    def __eq__(self, other):
+        if isinstance(other, str):
+            return self.name == other
+        return self.name == other.name
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def __deepcopy__(self, memo=None):
+        return Symbol(self.name)
+
+
+class StatementContext: ...
+
+
+class StatementContextRegistry:
+    _ctx_map: dict[
+        type[Any],
+        Callable[[Any], AbstractContextManager[None]],
+    ] = {}
+
+    @classmethod
+    def register_context_guard(
+        cls,
+        ctx_cls: type[_StatementContextT],
+        handler: Callable[[_StatementContextT], AbstractContextManager[None]],
+    ):
+        """
+        Register a context handler for the given context.
+        """
+        if ctx_cls in cls._ctx_map:
+            raise ValueError(f"Context {ctx_cls} is already registered.")
+        cls._ctx_map[ctx_cls] = handler
+
+    @classmethod
+    def register_context(
+        cls,
+        handler: Callable[[_StatementContextT], AbstractContextManager[None]],
+    ):
+        def decorator(ctx_cls: type[_StatementContextT]):
+            cls.register_context_guard(ctx_cls, handler)
+            return ctx_cls
+
+        return decorator
+
+    @classmethod
+    def get_context_guard(
+        cls,
+        ctx_cls: type[_StatementContextT],
+    ) -> Callable[[_StatementContextT], AbstractContextManager[None]]:
+        """
+        Get the context handler for the given context.
+        """
+        if ctx_cls not in cls._ctx_map:
+            raise ValueError(f"Context {ctx_cls} is not registered.")
+        return cls._ctx_map[ctx_cls]
+
+
+class Statement:
+    """
+    Statement is used to represent a sentence of code for building the neural network model,
+    which has four types: "call", "api", "method", and "layer".
+
+    Note:
+        Statement temporarily does not support control flow.
+    """
+
+    def __init__(
+        self,
+        type: str,
+        name: str,
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        assert type in ["call", "api", "method", "layer", "AST"]
+        self.name = name
+        self.inputs = inputs  # (list of Symbols, dict of Symbols)
+        self.outputs = outputs  # list of Symbol | PythonObj
+        self.contexts = contexts  # list of StatementContext
+        self.stmt_stack = (
+            stacks  # a list of string to record the source code callstack.
+        )
+        self.type = type
+
+    def __str__(self):
+        return "{} || {} = {} ({}) ".format(
+            self.type + " " * (10 - len(self.type)),
+            self.to_string(self.outputs),
+            self.name,
+            self.to_string(self.inputs),
+        )
+
+    def __repr__(self):
+        return self.__str__()
+
+    @staticmethod
+    def to_string(inps):
+        return ", ".join(repr(x) for x in flatten(inps))
+
+
+class CallStatement(Statement):
+    def __init__(
+        self,
+        name: str,
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        super().__init__("call", name, inputs, outputs, contexts, stacks)
+        self.sir_name = name
+
+
+class ApiStatement(Statement):
+    def __init__(
+        self,
+        api: Callable,
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        fullname = get_api_fullname(api)
+        if fullname is None:
+            fullname = "paddle." + api.__name__
+        super().__init__("api", fullname, inputs, outputs, contexts, stacks)
+        self.api = api
+
+
+class MethodStatement(Statement):
+    def __init__(
+        self,
+        name: str,
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        super().__init__("method", name, inputs, outputs, contexts, stacks)
+        self.method = name
+
+
+class LayerStatement(Statement):
+    def __init__(
+        self,
+        layer: Reference,  # Reference of paddle.nn.Layer
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        if isinstance(layer, Reference):
+            name = layer().__class__.__name__
+        else:
+            name = layer.__class__.__name__
+        super().__init__(
+            "layer",
+            name,
+            inputs,
+            outputs,
+            contexts,
+            stacks,
+        )
+        self.layer = layer
+
+
+class ASTStatement(Statement):
+    def __init__(
+        self,
+        static_function,
+        inputs: list[Symbol],
+        outputs: list[Symbol],
+        contexts: list[StatementContext],
+        stacks: list[str],
+    ):
+        # this dygraph_function always has attr __code__, which is checked before
+        dygraph_func = static_function.dygraph_function
+        super().__init__(
+            "AST",
+            dygraph_func.__code__.co_name,
+            inputs,
+            outputs,
+            contexts,
+            stacks,
+        )
+        converted_func = paddle.jit.dy2static.convert_to_static(dygraph_func)
+        func_self = getattr(dygraph_func, '__self__', None)
+        if func_self is not None:
+            converted_func = functools.partial(converted_func, func_self)
+        self.converted_func = converted_func
+
+
+class StatementIR:
+    """
+    StatementIR is the carrier that records the code for building the neural network model.It is
+    a representation of a purely computational structure, and does not care about specific values.
+    The function converted from StatementIR can ensure that it can be turned into a static state.
+    In this way, we can reuse the original `to_static` function to realize the execution of the static graph.
+
+    Note:
+        Don't create by yourself, just use the StatementIRFactory.create()
+    """
+
+    def __init__(self, name: str):
+        self.name = name
+        self.inputs: list[Symbol] = []
+        self.params: list[Symbol] = []
+        self.outputs: list[Symbol] = []
+        self.statements: list[Statement] = []
+
+        self.symbol_meta_map = {}
+        self.param_symbol = set()
+        self.non_param_symbol = set()
+
+    @property
+    def input_with_params(self):
+        return self.inputs + self.params
+
+    def __len__(self):
+        return len(self.statements)
+
+    def __deepcopy__(self, memo=None):
+        new_sir = StatementIR(self.name)
+        new_sir.inputs = list(self.inputs)
+        new_sir.params = list(self.params)
+        new_sir.outputs = list(self.outputs)
+        new_sir.statements = list(self.statements)
+        new_sir.symbol_meta_map = dict(self.symbol_meta_map.items())
+        new_sir.param_symbol = set(self.param_symbol)
+        new_sir.non_param_symbol = set(self.non_param_symbol)
+        return new_sir
+
+    def set_parameter_info(self, params, non_params):
+        self.param_symbol.update(params)
+        self.non_param_symbol.update(non_params)
+
+    def set_symbol_meta_map(self, meta_map):
+        # if the meta of a input symbol inplace changed, we should get the origin meta as input of SIR
+        meta_map.update(self.symbol_meta_map)
+        self.symbol_meta_map = meta_map
+
+    def add_input(self, input):
+        self.inputs.append(input)
+
+    def add_output(self, output):
+        self.outputs.append(output)
+
+    def add_statement(self, statement):
+        assert isinstance(statement, Statement)
+        self.statements.append(statement)
+
+    def analyse_inputs(self):
+        used_symbols = OrderedSet()
+        generated_symbols = OrderedSet()
+        for stmt in self.statements:
+            for inp in flatten_extend(stmt.inputs):
+                if isinstance(inp, Symbol) and inp not in generated_symbols:
+                    used_symbols.add(inp)
+            for out in flatten_extend(stmt.outputs):
+                if isinstance(out, Symbol):
+                    generated_symbols.add(out)
+
+        used_symbols = sorted(used_symbols, key=lambda x: x.name)
+        if not parameters_persistent_mode_is_enabled():
+            return used_symbols, []
+        input_symbols = [
+            symbol for symbol in used_symbols if symbol not in self.param_symbol
+        ]
+        param_symbols = [
+            symbol for symbol in used_symbols if symbol in self.param_symbol
+        ]
+        return input_symbols, param_symbols
+
+    def __str__(self):
+        strs = []
+        strs.append(f"StatementIR: {self.name}")
+        strs.append(f"  inputs: {map_structure(lambda x: x.name, self.inputs)}")
+        strs.append(f"  params: {map_structure(lambda x: x.name, self.params)}")
+        strs.append(
+            f"  outputs: {map_structure(lambda x: x.name, self.outputs)}"
+        )
+        strs.append("  statements: ")
+        for stmt in self.statements:
+            strs.append(f"    {stmt}")
+        return "\n".join(strs)
+
+    def __repr__(self):
+        return self.__str__()
+
+
+class StatementIRFactory(metaclass=Singleton):
+    """
+    It is used to create a StatementIR.
+    """
+
+    def __init__(self):
+        self.cache = {}
+        self.name_generator = NameGenerator("SIR_")
+
+    def __getitem__(self, key):
+        return self.cache[key]
+
+    def create(self, input_name=None):
+        if input_name:
+            name = input_name
+        else:
+            name = self.name_generator.next()
+
+        sir = StatementIR(name)
+        self.cache[name] = sir
+        return sir
+
+    def update(self, stmt_ir):
+        name = stmt_ir.name
+        self.cache[name] = stmt_ir
+
+    def clear(self):
+        want_clear = [
+            key
+            for key in self.cache.keys()
+            if self.name_generator.match_name(key)
+        ]
+        for key in want_clear:
+            del self.cache[key]

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

@@ -0,0 +1,13 @@
+# 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.

python/paddle/jit/sot/symbolic_shape/constraints.py

@@ -0,0 +1,262 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import paddle
+
+from ..utils.exceptions import InnerError
+
+if TYPE_CHECKING:
+    from ..opcode_translator.executor.guard import StringifiedExpression
+
+
+class ConstraintNode:
+    def __init__(self, inputs: list[ConstraintNode]):
+        self.inputs = inputs
+
+    def create_guard_expr(
+        self, extern_vars: dict[str, StringifiedExpression]
+    ) -> StringifiedExpression:
+        raise NotImplementedError
+
+    def create_guard_node(
+        self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
+    ) -> paddle.framework.core.ExprNodeBase:
+        raise NotImplementedError(
+            f"{self.__class__.__name__}.create_guard_node is not implemented"
+        )
+
+
+class LeafConstraintNode(ConstraintNode):
+    def __init__(self):
+        super().__init__([])
+
+
+class UnaryConstraintNode(ConstraintNode):
+    READABLE_SYMBOL: str
+
+    def __init__(self, input: ConstraintNode):
+        super().__init__([input])
+        self.input = input
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self.input})"
+
+    def create_guard_expr(
+        self, extern_vars: dict[str, StringifiedExpression]
+    ) -> StringifiedExpression:
+        from ..opcode_translator.executor.guard import (
+            StringifiedExpression,
+            union_free_vars,
+        )
+
+        input = self.input.create_guard_expr(extern_vars)
+        return StringifiedExpression(
+            f"{self.READABLE_SYMBOL}({{}})",
+            [input],
+            union_free_vars(input.free_vars),
+        )
+
+    def create_guard_node(
+        self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
+    ) -> paddle.framework.core.ExprNodeBase:
+        input = self.input.create_guard_node(extern_vars)
+        return paddle.framework.core.UnaryExprNode(input, self.READABLE_SYMBOL)
+
+
+class BinaryConstraintNode(ConstraintNode):
+    READABLE_SYMBOL: str
+
+    def __init__(self, lhs: ConstraintNode, rhs: ConstraintNode):
+        super().__init__([lhs, rhs])
+        self.lhs = lhs
+        self.rhs = rhs
+
+    def create_guard_expr(
+        self, extern_vars: dict[str, StringifiedExpression]
+    ) -> StringifiedExpression:
+        from ..opcode_translator.executor.guard import (
+            StringifiedExpression,
+            union_free_vars,
+        )
+
+        lhs = self.lhs.create_guard_expr(extern_vars)
+        rhs = self.rhs.create_guard_expr(extern_vars)
+        return StringifiedExpression(
+            f"({{}} {self.READABLE_SYMBOL} {{}})",
+            [lhs, rhs],
+            union_free_vars(lhs.free_vars, rhs.free_vars),
+        )
+
+    def create_guard_node(
+        self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
+    ) -> paddle.framework.core.ExprNodeBase:
+        lhs = self.lhs.create_guard_node(extern_vars)
+        rhs = self.rhs.create_guard_node(extern_vars)
+        return paddle.framework.core.BinaryExprNode(
+            lhs, rhs, self.READABLE_SYMBOL
+        )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.lhs}, {self.rhs})"
+
+
+class ConstantConstraintNode(LeafConstraintNode):
+    def __init__(self, value):
+        super().__init__()
+        self.value = value
+
+    def create_guard_expr(
+        self, extern_vars: dict[str, StringifiedExpression]
+    ) -> StringifiedExpression:
+        from ..opcode_translator.executor.guard import (
+            StringifiedExpression,
+        )
+
+        return StringifiedExpression(f"{self.value!r}", [], {})
+
+    def create_guard_node(
+        self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
+    ) -> paddle.framework.core.ExprNodeBase:
+        return paddle.framework.core.ConstantExprNode(self.value)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.value})"
+
+
+class SymbolicConstraintNode(LeafConstraintNode):
+    def __init__(self, name: str):
+        super().__init__()
+        self.name = name
+
+    def create_guard_expr(
+        self, extern_vars: dict[str, StringifiedExpression]
+    ) -> StringifiedExpression:
+        from ..opcode_translator.executor.guard import (
+            StringifiedExpression,
+            union_free_vars,
+        )
+
+        if self.name not in extern_vars:
+            raise InnerError(
+                f"Symbolic variable {self.name} not found in extern_vars."
+            )
+        return StringifiedExpression(
+            "{}",
+            [extern_vars[self.name]],
+            union_free_vars(extern_vars[self.name].free_vars),
+        )
+
+    def create_guard_node(
+        self, extern_vars: dict[str, paddle.framework.core.ExprNodeBase]
+    ) -> paddle.framework.core.ExprNodeBase:
+        if self.name not in extern_vars:
+            raise InnerError(
+                f"Symbolic variable {self.name} not found in extern_vars."
+            )
+        return extern_vars[self.name]
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({self.name})"
+
+
+class NegativeConstraintNode(UnaryConstraintNode):
+    READABLE_SYMBOL = "-"
+
+
+class BitwiseNotConstraintNode(UnaryConstraintNode):
+    READABLE_SYMBOL = "~"
+
+
+class AddConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "+"
+
+
+class SubConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "-"
+
+
+class MulConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "*"
+
+
+class TrueDivConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "/"
+
+
+class FloorDivConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "//"
+
+
+class ModConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "%"
+
+
+class PowConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "**"
+
+
+class BitwiseLShiftConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "<<"
+
+
+class BitwiseRShiftConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = ">>"
+
+
+class BitwiseAndConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "&"
+
+
+class BitwiseOrConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "|"
+
+
+class BitwiseXorConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "^"
+
+
+class LogicalToBoolConstraintNode(UnaryConstraintNode):
+    READABLE_SYMBOL = "bool"
+
+
+class LogicalNotConstraintNode(UnaryConstraintNode):
+    READABLE_SYMBOL = "not"
+
+
+class EqualConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "=="
+
+
+class NotEqualConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "!="
+
+
+class LessThanConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "<"
+
+
+class LessEqualConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = "<="
+
+
+class GreaterThanConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = ">"
+
+
+class GreaterEqualConstraintNode(BinaryConstraintNode):
+    READABLE_SYMBOL = ">="

python/paddle/jit/sot/symbolic_shape/operators.py

@@ -0,0 +1,104 @@
+# 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.
+
+from __future__ import annotations
+
+import operator
+from typing import TYPE_CHECKING
+
+import paddle
+
+if TYPE_CHECKING:
+    from ..utils.magic_methods import BinaryOp, UnaryOp
+
+
+def symbolic_to_bool(x):
+    # Unified api for python number and paddle Tensor
+    return x != 0
+
+
+def symbolic_not(x):
+    return x == 0
+
+
+def symbolic_truediv(x, y):
+    # NOTE(SigureMo): In Paddle, the truediv maybe has precision issue.
+    # For example, paddle.tensor(168) / 7, in Python it should be 24.0,
+    # but in Paddle it will construct a Scale OP, which will calculate
+    # as 168 * (1 / 7) = 24.00000191, which may cause some unexpected
+    # bugs. So we cast the tensor and scalar both to float64 to avoid
+    # this issue.
+    is_need_cast_tensor = lambda v: (
+        isinstance(v, paddle.pir.Value) and v.dtype is not paddle.float64
+    )
+    cast_tensor_if_needed = lambda v: (
+        v.cast(paddle.float64) if is_need_cast_tensor(v) else v
+    )
+    cast_scalar_if_needed = lambda v: (
+        paddle.full([], v, dtype=paddle.float64)
+        if isinstance(v, (int, float))
+        else v
+    )
+    cast_if_needed = lambda v: cast_tensor_if_needed(cast_scalar_if_needed(v))
+    has_tensor_need_cast = is_need_cast_tensor(x) or is_need_cast_tensor(y)
+    if not has_tensor_need_cast:
+        return operator.truediv(x, y)
+    x = cast_if_needed(x)
+    y = cast_if_needed(y)
+    return operator.truediv(x, y)
+
+
+# All symbolic operations need unified for python number and paddle Tensor
+SYMBOLIC_UNARY_MATH_OPS: list[UnaryOp] = [
+    # Basic
+    operator.neg,
+    # Bitwise
+    operator.invert,
+]
+SYMBOLIC_BINARY_MATH_OPS: list[BinaryOp] = [
+    # Basic
+    operator.add,
+    operator.sub,
+    operator.mul,
+    symbolic_truediv,
+    operator.floordiv,
+    operator.pow,
+    operator.mod,
+    # Bitwise
+    operator.lshift,
+    operator.rshift,
+    operator.and_,
+    operator.or_,
+    operator.xor,
+]
+SYMBOLIC_UNARY_LOGICAL_OPS: list[UnaryOp] = [
+    symbolic_to_bool,
+    symbolic_not,
+]
+SYMBOLIC_BINARY_LOGICAL_OPS: list[BinaryOp] = [
+    operator.eq,
+    operator.ne,
+    operator.lt,
+    operator.le,
+    operator.gt,
+    operator.ge,
+]
+SYMBOLIC_MATH_OPS = SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_BINARY_MATH_OPS
+SYMBOLIC_MATH_OPS = SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_BINARY_MATH_OPS
+SYMBOLIC_UNARY_OPS: list[UnaryOp] = (
+    SYMBOLIC_UNARY_MATH_OPS + SYMBOLIC_UNARY_LOGICAL_OPS
+)
+SYMBOLIC_BINARY_OPS: list[BinaryOp] = (
+    SYMBOLIC_BINARY_MATH_OPS + SYMBOLIC_BINARY_LOGICAL_OPS
+)

python/paddle/jit/sot/symbolic_shape/symbolic_value.py

@@ -0,0 +1,57 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+_T = TypeVar("_T", "int", "float", "bool")
+
+
+class SymbolicValue(Generic[_T]):
+    example_value: _T | None
+
+    def __init__(self, example_value=None):
+        self.example_value = example_value
+
+    def __repr__(self) -> str:
+        if self.is_backed():
+            return f"{self.__class__.__name__}({self.example_value})"
+        return f"{self.__class__.__name__}()"
+
+    def get_static_type(self) -> type[_T]:
+        raise NotImplementedError("get_py_type is not implemented.")
+
+    def is_backed(self):
+        return self.example_value is not None
+
+    def get_example_value(self) -> _T:
+        if self.example_value is None:
+            raise ValueError(f"{self} is not backed by a value.")
+        return self.example_value
+
+
+class SymbolicBool(SymbolicValue):
+    def get_static_type(self) -> type[bool]:
+        return bool
+
+
+class SymbolicInt(SymbolicValue):
+    def get_static_type(self) -> type[int]:
+        return int
+
+
+class SymbolicFloat(SymbolicValue):
+    def get_static_type(self) -> type[float]:
+        return float

python/paddle/jit/sot/translate.py

@@ -0,0 +1,118 @@
+# 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
+
+from typing import TYPE_CHECKING, TypeVar
+
+from typing_extensions import ParamSpec
+
+import paddle
+
+from .opcode_translator import eval_frame_callback
+from .profiler import SotStepProfilerGuard
+from .utils import (
+    InfoCollector,
+    StepInfoManager,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+def symbolic_translate(fn: Callable[P, R], **kwargs) -> Callable[P, R]:
+    """
+    This function is the entry point of PaddleSOT. It sets eval_frame_callback before input
+    function to achieve Opcode-level translation. The translation process depends on the
+    simulation execution, in which information will be collected, especially the network
+    code. After the simulation execution is completed, the network code will be compiled
+    into a static graph Program to improve performance.
+
+    Args:
+        fn: The input function.
+
+    Returns:
+        Callable, The wrapped function.
+
+    Examples:
+        >>> # doctest: +SKIP("Could not get source code of function foo."")
+        >>> import paddle
+        >>> import numpy as np
+        >>> from sot.translate import symbolic_translate
+        >>> def foo(cond: paddle.Tensor, x: paddle.Tensor):
+        ...     x += 1
+        ...     if cond:
+        ...         x += 1
+        ...     else:
+        ...         x -= 1
+        ...     return x
+        >>> symbolic_translate_foo = symbolic_translate(foo)
+        >>> # For the true branch, the output is 2.
+        >>> cond = paddle.to_tensor(True)
+        >>> x = paddle.to_tensor(0)
+        >>> dygraph_out = foo(cond, x)
+        >>> symbolic_translate_out = symbolic_translate_foo(cond, x)
+        >>> dygraph_out
+        Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
+        2)
+        >>> symbolic_translate_out
+        Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
+        2)
+        >>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
+        >>> # For the false branch, the output is 0.
+        >>> cond = paddle.to_tensor(False)
+        >>> dygraph_out = foo(cond, x)
+        >>> symbolic_translate_out = symbolic_translate_foo(cond, x)
+        >>> dygraph_out
+        Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
+        0)
+        >>> symbolic_translate_out
+        Tensor(shape=[], dtype=int64, place=Place(cpu), stop_gradient=True,
+        0)
+        >>> np.testing.assert_allclose(dygraph_out.numpy(), symbolic_translate_out.numpy())
+
+    """
+
+    if not paddle.framework.use_pir_api():
+        raise RuntimeError(
+            "SOT is only supported when running in PIR mode. Please set the environment variable "
+            "FLAGS_enable_pir_api=1 to enable it."
+        )
+
+    kwargs.setdefault('training', True)
+
+    def callback(frame):
+        return eval_frame_callback(frame, **kwargs)
+
+    def impl(*args: P.args, **kwargs: P.kwargs) -> R:
+        assert hasattr(fn, "__code__"), (
+            "Target function doesn't have code for simulating."
+        )
+        with StepInfoManager().step_guard(fn.__code__), SotStepProfilerGuard():
+            InfoCollector().clear_step_info()
+            paddle.framework.core.set_eval_frame(callback)
+            try:
+                outs = fn(*args, **kwargs)
+            except Exception as e:
+                raise e
+            finally:
+                paddle.framework.core.set_eval_frame(None)
+
+            InfoCollector().print_step_report()
+            return outs
+
+    return impl

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]