[SOT] Add new internal API paddle.jit.marker.unified to mark an API as unified in dynamic and static mode and support custom op (#72466)

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: SigureMo

python/paddle/jit/__init__.py

@@ -13,16 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from . import marker as marker
 from .api import (
     ignore_module,
     json_to_pdmodel,  # noqa: F401
     load,
-    not_to_static,
     save,
     to_static,
 )
 from .dy2static.logging_utils import set_code_level, set_verbosity
 from .dy2static.program_translator import enable_to_static
+from .marker import (
+    not_to_static,
+)
 from .translated_layer import TranslatedLayer
 
 __all__ = [

python/paddle/jit/api.py

@@ -69,7 +69,7 @@
 )
 
 from .dy2static import logging_utils
-from .dy2static.convert_call_func import ConversionOptions, add_ignore_module
+from .dy2static.convert_call_func import add_ignore_module
 from .dy2static.program_translator import (
     ASTStaticFunction,
     ProgramTranslator,
@@ -346,69 +346,6 @@ def decorated(python_func):
     return decorated
 
 
-class _NotToStaticDecorator(Protocol):
-    @overload
-    def __call__(
-        self, func: Callable[_InputT, _RetT]
-    ) -> Callable[_InputT, _RetT]: ...
-
-    @overload
-    def __call__(self, func: None = ...) -> _NotToStaticDecorator: ...
-
-
-@overload
-def not_to_static(
-    func: Callable[_InputT, _RetT],
-) -> Callable[_InputT, _RetT]: ...
-
-
-@overload
-def not_to_static(func: None = ...) -> _NotToStaticDecorator: ...
-
-
-def not_to_static(func=None):
-    """
-    A Decorator to suppresses the convention of a function.
-
-    Args:
-        func(callable): The function to decorate.
-
-    Returns:
-        callable: A function which won't be converted in Dynamic-to-Static.
-
-    Examples:
-        .. code-block:: python
-
-            >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
-            >>> import paddle
-
-            >>> @paddle.jit.not_to_static
-            ... def func_not_to_static(x):
-            ...     res = x - 1
-            ...     return res
-
-            >>> @paddle.jit.to_static
-            ... def func(x):
-            ...     if paddle.mean(x) < 0:
-            ...         out = func_not_to_static(x)
-            ...     else:
-            ...         out = x + 1
-            ...     return out
-            ...
-            >>> x = paddle.ones([1, 2], dtype='float32')
-            >>> out = func(x)
-            >>> print(out)
-            Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
-            [[2., 2.]])
-    """
-    if func is None:
-        return not_to_static
-
-    options = ConversionOptions(not_convert=True)
-    options.attach(func)
-    return func
-
-
 class _SaveLoadConfig:
     def __init__(self):
         self._output_spec = None

python/paddle/jit/dy2static/convert_call_func.py

@@ -37,12 +37,16 @@
 )
 from .logging_utils import TranslatorLogger
 from .program_translator import (
-    CONVERSION_OPTIONS,
     StaticFunction,
     convert_to_static,
     unwrap_decorators,
 )
-from .utils import is_builtin, is_paddle_func, patch_method_guard
+from .utils import (
+    TransformOptions,
+    is_builtin,
+    is_paddle_func,
+    patch_method_guard,
+)
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -53,31 +57,6 @@
 translator_logger = TranslatorLogger()
 
 
-class ConversionOptions:
-    """
-    A container for conversion flags of a function in dynamic-to-static.
-
-    Attributes:
-        not_convert(bool): An attribute indicates that the function won't be converted in dynamic-to-static.
-
-    NOTE(liym27): More attributes and methods can be added in this class.
-    """
-
-    def __init__(self, not_convert=False):
-        self.not_convert = not_convert
-
-    def attach(self, func):
-        if inspect.ismethod(func):
-            func = func.__func__
-
-        if inspect.isfunction(func):
-            setattr(func, CONVERSION_OPTIONS, self)
-        else:
-            translator_logger.warn(
-                f"Only support @not_to_static to type(function) or type(method), but received {type(func)}"
-            )
-
-
 def builtin_modules():
     """
     Return builtin modules.
@@ -259,8 +238,9 @@ def convert_call(func):
     # in this case, unwraps it into a raw method or function.
     _, func = unwrap_decorators(func)
 
-    options = getattr(func, CONVERSION_OPTIONS, None)
-    if options is not None and options.not_convert:
+    if not TransformOptions.check_fn_need_transform(
+        func, TransformOptions.ToStaticMode.AST
+    ):
         translator_logger.log(
             2,
             "%s is not converted when it is decorated by 'paddle.jit.not_to_static'.",

python/paddle/jit/dy2static/program_translator.py

@@ -61,6 +61,7 @@
 from .utils import (
     ALREADY_D2S,
     NO_SHAPE_VAR_TYPE,
+    TransformOptions,
     ast_to_func,
     backend_guard,
     cuda_pinned_tensors_move_to_excepted_place,
@@ -86,8 +87,6 @@
 # Once exceeding the threshold, we will raise warning to users to make sure the conversion is as expected.
 MAX_TRACED_PROGRAM_COUNT = 10
 
-CONVERSION_OPTIONS = "__jst_not_to_static"
-
 
 def synchronized(func):
     func.__lock__ = threading.Lock()
@@ -252,12 +251,13 @@ def convert_to_static(function):
     if getattr(function, ALREADY_D2S, None):
         return function
 
-    # Return directly if decorated with @not_to_static and DO NOT Cache it
-    options = getattr(function, CONVERSION_OPTIONS, None)
+    # Return directly if decorated with @jit.marker.unified and DO NOT Cache it
     # or ignore paddle api
-    need_skip = (options is not None and options.not_convert) or is_paddle_func(
-        function
-    )
+    need_skip = (
+        not TransformOptions.check_fn_need_transform(
+            function, TransformOptions.ToStaticMode.AST
+        )
+    ) or is_paddle_func(function)
     if need_skip:
         return function.__func__ if inspect.ismethod(function) else function
 

python/paddle/jit/dy2static/utils.py

@@ -29,7 +29,7 @@
 import types
 import warnings
 from contextlib import contextmanager
-from enum import Enum, auto
+from enum import Enum, Flag, auto
 from importlib.machinery import SourceFileLoader
 from typing import TYPE_CHECKING, Any
 
@@ -115,6 +115,44 @@ def is_phi(self):
         return self == Backend.PHI
 
 
+class TransformOptions:
+
+    class ToStaticMode(Flag):
+        SOT = auto()
+        AST = auto()
+
+        @classmethod
+        def Nil(cls):
+            return cls(0)
+
+    TRANSFORM_OPTIONS_ATTR_NAME = "___jit_transform_options___"
+
+    def __init__(self, skip_transform_mode: ToStaticMode = ToStaticMode.Nil()):
+        self.skip_transform_mode = skip_transform_mode
+
+    def attach(self, fn):
+        if inspect.ismethod(fn):
+            fn = fn.__func__
+
+        if inspect.isfunction(fn):
+            setattr(fn, TransformOptions.TRANSFORM_OPTIONS_ATTR_NAME, self)
+        else:
+            warnings.warn(
+                f"Only support @jit.marker.unified to type(function) or type(method), but received {type(fn)}"
+            )
+
+    def need_transform(self, mode: ToStaticMode):
+        return not (self.skip_transform_mode & mode)
+
+    @staticmethod
+    def check_fn_need_transform(fn, mode: ToStaticMode):
+        if not hasattr(fn, TransformOptions.TRANSFORM_OPTIONS_ATTR_NAME):
+            return True
+        return getattr(
+            fn, TransformOptions.TRANSFORM_OPTIONS_ATTR_NAME
+        ).need_transform(mode)
+
+
 class TimeCounter:
     def __init__(self):
         self._time_history: list[float] = []

python/paddle/jit/marker.py

@@ -0,0 +1,125 @@
+# 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 (
+    Callable,
+    Protocol,
+    TypeVar,
+    overload,
+)
+
+from typing_extensions import (
+    ParamSpec,
+)
+
+from .dy2static.utils import (
+    TransformOptions,
+)
+
+_RetT = TypeVar("_RetT")
+_InputT = ParamSpec("_InputT")
+
+
+class _NotToStaticDecorator(Protocol):
+    @overload
+    def __call__(
+        self, func: Callable[_InputT, _RetT]
+    ) -> Callable[_InputT, _RetT]: ...
+
+    @overload
+    def __call__(self, func: None = ...) -> _NotToStaticDecorator: ...
+
+
+@overload
+def not_to_static(
+    func: Callable[_InputT, _RetT],
+) -> Callable[_InputT, _RetT]: ...
+
+
+@overload
+def not_to_static(func: None = ...) -> _NotToStaticDecorator: ...
+
+
+# Legacy decorator only for AST
+def not_to_static(func=None):
+    """
+    A Decorator to suppresses the convention of a function.
+
+    Args:
+        func(callable): The function to decorate.
+
+    Returns:
+        callable: A function which won't be converted in Dynamic-to-Static.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +SKIP('`paddle.jit.to_static` can not run in xdoctest')
+            >>> import paddle
+
+            >>> @paddle.jit.not_to_static
+            ... def func_not_to_static(x):
+            ...     res = x - 1
+            ...     return res
+
+            >>> @paddle.jit.to_static
+            ... def func(x):
+            ...     if paddle.mean(x) < 0:
+            ...         out = func_not_to_static(x)
+            ...     else:
+            ...         out = x + 1
+            ...     return out
+            ...
+            >>> x = paddle.ones([1, 2], dtype='float32')
+            >>> out = func(x)
+            >>> print(out)
+            Tensor(shape=[1, 2], dtype=float32, place=Place(cpu), stop_gradient=True,
+            [[2., 2.]])
+    """
+    return unified(func, for_sot=False, for_ast=True)
+
+
+def unified(
+    fn: Callable[_InputT, _RetT] | None = None,
+    *,
+    for_sot: bool = True,
+    for_ast: bool = True,
+) -> Callable[_InputT, _RetT]:
+    """
+    Mark a function already unified in dygraph and static mode. So
+    that it won't be transformed again in SOT or AST mode.
+
+    Args:
+        fn(callable): The function to decorate.
+        for_sot(bool): Whether to mark the function as unified in SOT mode.
+        for_ast(bool): Whether to mark the function as unified in AST mode.
+    """
+
+    def _mark_as_unified(fn, *, for_sot: bool, for_ast: bool):
+        mode = TransformOptions.ToStaticMode.Nil()
+        if for_sot:
+            mode |= TransformOptions.ToStaticMode.SOT
+        if for_ast:
+            mode |= TransformOptions.ToStaticMode.AST
+        options = TransformOptions(
+            skip_transform_mode=mode,
+        )
+        options.attach(fn)
+        return fn
+
+    if fn is None:
+        return lambda fn: _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)
+    return _mark_as_unified(fn, for_sot=for_sot, for_ast=for_ast)

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

@@ -54,9 +54,9 @@
     NameGenerator,
     SIRToCodeMap,
     SotUndefinedVar,
+    already_unified_in_dynamic_and_static_graph,
     inner_error_default_handler,
     is_inplace_api,
-    is_paddle_api,
     log,
     log_do,
     map_if,
@@ -549,7 +549,7 @@ def call_paddle_api(
         Args:
             func: paddle api
         """
-        assert is_paddle_api(func)
+        assert already_unified_in_dynamic_and_static_graph(func)
         log(3, f"call paddle.api : {func.__name__}", "\n")
 
         def message_handler(*args, **kwargs):