chore: import upstream snapshot with attribution

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