[SOT] merge PaddleSOT into Paddle (PaddlePaddle#57824)

PaddleSOT is a Bytecode level Implementation of Symbolic OpCode Translator For PaddlePaddle.

We originally developed in [PaddleSOT](https://github.com/PaddlePaddle/PaddleSOT), and to ensure consistency in Paddle versions, we are now merging PaddleSOT into Paddle.

Thanks to all the contributors of this project!

See more details in https://github.com/PaddlePaddle/PaddleSOT/graphs/contributors

---------

Co-authored-by: xiongkun <xiongkun03@baidu.com>
Co-authored-by: feifei-111 <2364819892@qq.com>
Co-authored-by: 0x45f <23097963+0x45f@users.noreply.github.com>
Co-authored-by: gouzil <66515297+gouzil@users.noreply.github.com>
Co-authored-by: 六个骨头 <46243324+zrr1999@users.noreply.github.com>
Co-authored-by: Aurelius84 <zhangliujie@baidu.com>
Co-authored-by: Wang Xin <xinwang614@gmail.com>
Co-authored-by: haozi <64006169+NotHaozi@users.noreply.github.com>
Co-authored-by: RedContritio <RedContritio@qq.com>
Co-authored-by: Sanbu <96160062+sanbuphy@users.noreply.github.com>
Co-authored-by: Difer <c7070655110@gmail.com>
Co-authored-by: cyberslack_lee <luhputu0815@gmail.com>
Co-authored-by: jjyaoao <jjyaoao@126.com>
Co-authored-by: PuQing <me@puqing.work>
Co-authored-by: Ran chongzhi <57489288+ranchongzhi@users.noreply.github.com>
Co-authored-by: Zhenghai Zhang <65210872+ccsuzzh@users.noreply.github.com>

.flake8

@@ -26,6 +26,9 @@ per-file-ignores =
     # These files need tabs for testing.
     test/dygraph_to_static/test_error.py:E101,W191
 
+    # Ignore compare with True in sot unittest
+    test/sot/test_dup_top.py:E712
+
     # temp ignore base directory
     python/paddle/base/*:
         E712,

paddle/scripts/paddle_build.sh

@@ -933,23 +933,21 @@ set -ex
 }
 
 function run_sot_test() {
-    PADDLE_SOT_ROOT=$1
-    PY_VERSION=$2
+    PY_VERSION=$1
     PYTHON_WITH_SPECIFY_VERSION=python$PY_VERSION
     PY_VERSION_NO_DOT=`echo $PY_VERSION | sed 's/\.//g'`
 
     export STRICT_MODE=1
     export COST_MODEL=False
     export MIN_GRAPH_SIZE=0
+    export SOT_LOG_LEVEL=0
 
     # Install PaddlePaddle
     $PYTHON_WITH_SPECIFY_VERSION -m pip install ${PADDLE_ROOT}/dist/paddlepaddle-0.0.0-cp${PY_VERSION_NO_DOT}-cp${PY_VERSION_NO_DOT}-linux_x86_64.whl
     # Install PaddleSOT
-    cd $PADDLE_SOT_ROOT
-    $PYTHON_WITH_SPECIFY_VERSION -m pip install -e .
+    cd $PADDLE_ROOT/test/sot/
 
     # Run unittest
-    cd tests
     failed_tests=()
 
     for file in ./test_*.py; do
@@ -4128,14 +4126,12 @@ function main() {
         ;;
       cicheck_sot)
         export WITH_SHARED_PHI=ON
-        PADDLE_SOT_ROOT=${PADDLE_ROOT}/sot
-        git clone https://github.com/PaddlePaddle/PaddleSOT.git ${PADDLE_SOT_ROOT}
         PYTHON_VERSIONS=(3.8 3.9 3.10 3.11)
         for PY_VERSION in ${PYTHON_VERSIONS[@]}; do
             ln -sf $(which python${PY_VERSION}) /usr/local/bin/python
             ln -sf $(which pip${PY_VERSION}) /usr/local/bin/pip
             run_setup ${PYTHON_ABI:-""} bdist_wheel ${parallel_number}
-            run_sot_test $PADDLE_SOT_ROOT $PY_VERSION
+            run_sot_test $PY_VERSION
             rm -rf ${PADDLE_ROOT}/build/CMakeCache.txt
         done
         ;;

python/paddle/jit/dy2static/program_translator.py

@@ -692,30 +692,22 @@ class SymbolicStaticFunction(StaticFunction):
     def __init__(self, function, input_spec=None, **kwargs):
         if input_spec is not None:
             warnings.warn(
-                "\nSymbolic Trace don't support input_spec arguments. It will Will not produce any effect.\n"
+                "\nSymbolic Trace don't support input_spec arguments. It will not produce any effect.\n"
                 "1. You can disable fallback mode by `paddle.jit.to_static(enable_fallback=False)` to switch to AST to static, then you can assign input spec.\n"
             )
         super().__init__(function, input_spec, **kwargs)
         self.last_call_input_spec = None
 
     def _perform_call(self, *args, **kwargs):
+        from ..sot import symbolic_translate
+
         args, kwargs = self._function_spec.unified_args_and_kwargs(args, kwargs)
         (
             input_args_with_spec,
             input_kwargs_with_spec,
         ) = self._function_spec.args_to_input_spec(args, kwargs)
         self.last_call_input_spec = input_args_with_spec
 
-        try:
-            from sot import symbolic_translate
-        except:
-            import os
-
-            os.system(
-                "pip install git+https://github.com/PaddlePaddle/PaddleSOT@develop"
-            )
-            from sot import symbolic_translate
-
         build_strategy = self._kwargs.get("build_strategy", None)
         backend = self._kwargs.get("backend", None)
         traced_fun = symbolic_translate(

python/paddle/jit/sot/__init__.py

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

python/paddle/jit/sot/infer_meta.py

@@ -0,0 +1,282 @@
+# 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 paddle
+from paddle.amp.auto_cast import amp_state
+from paddle.base.unique_name import UniqueNameGenerator
+from paddle.base.unique_name import guard as UniqueNameGuard
+from paddle.static import Program
+from paddle.utils import flatten, is_sequence
+
+from .utils import Cache, Singleton, map_if_extend, meta_str
+
+
+class MetaInfo:
+    def __init__(
+        self, shape, dtype, stop_gradient, name, persistable, type, place
+    ):
+        self.name = name
+        self.persistable = persistable
+        self.type = type
+        self.place = place
+        self.shape = shape
+        self.dtype = dtype
+        self.stop_gradient = stop_gradient
+
+    @staticmethod
+    def from_tensor(tensor):
+        # We always use float32 in simulation if AMP is enabled.
+        dtype = tensor.dtype
+        current_amp_state = amp_state()
+        if (
+            dtype == paddle.float16
+            and current_amp_state is not None
+            and current_amp_state["dtype"] == "float16"
+        ):
+            dtype = paddle.float32
+        return MetaInfo(
+            list(tensor.shape),
+            dtype,
+            tensor.stop_gradient,
+            tensor.name,
+            tensor.persistable,
+            tensor.type,
+            tensor.place,
+        )
+
+    def is_dynamic_shape(self):
+        """
+        if -1 in shape, return True
+        else: return False
+        """
+        return -1 in self.shape
+
+    def to_input_spec(self):
+        return paddle.static.InputSpec(
+            self.shape, dtype=self.dtype, stop_gradient=self.stop_gradient
+        )
+
+    def guard_str(self):
+        return f"({self.shape}, {self.dtype}, {self.stop_gradient})"
+
+    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))
+
+
+@Singleton
+class VariableCreator:
+    """
+    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_cache = {}
+        self.main_program = Program()
+        self.startup_program = Program()
+        self.var_name_generator = UniqueNameGenerator("infer_meta_variable_")
+
+    def gen_name(self, meta):
+        name = f"{meta.dtype}_{meta.stop_gradient}"
+        for l in meta.shape:
+            name += f"_{l}"
+        return name
+
+    def create_var(self, meta):
+        var = self.main_program.global_block().create_var(
+            shape=meta.shape,
+            dtype=meta.dtype,
+            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):
+        var_feature_name = self.gen_name(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
+        ):
+            args, kwargs = convert_meta_to_variable(
+                args
+            ), convert_meta_to_variable(kwargs)
+
+            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?
+                    out = getattr(args[0], func)(*args[1:], **kwargs)
+                else:
+                    out = func(*args, **kwargs)
+
+        return convert_variable_to_meta_info(out)
+
+
+def convert_meta_to_variable(args):
+    return map_if_extend(
+        args,
+        pred=lambda x: isinstance(x, MetaInfo),
+        true_fn=lambda x: VariableCreator().get_variable(x),
+        false_fn=lambda x: x,
+    )
+
+
+def convert_meta_to_input_spec(args):
+    return map_if_extend(
+        args,
+        pred=lambda x: isinstance(x, MetaInfo),
+        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.static.Variable),
+        true_fn=lambda x: MetaInfo.from_tensor(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, enable_fallback=False)
+
+    args_, kwargs_ = convert_meta_to_input_spec((args, kwargs))
+
+    (
+        concrete_program,
+        partial_program_layer,
+    ) = layer.forward.get_concrete_program(*args_, **kwargs_)
+
+    out = partial_program_layer._restore_out(
+        paddle.utils.flatten(
+            convert_variable_to_meta_info(concrete_program.outputs)
+        )
+    )
+    layer.forward.rollback()
+    return out
+
+
+@Singleton
+class SpecialInferMeta:
+    """
+    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
+
+
+@Singleton
+class InferMetaCache(Cache):
+    def key_fn(
+        self, func, *args, **kwargs
+    ):  # args & kwargs have transformed to MetaInfo
+        try:
+            retval = hash(
+                (
+                    func,
+                    tuple(flatten(args)),
+                    tuple(kwargs.keys()),
+                    tuple(flatten(kwargs)),
+                )
+            )
+        except Exception as e:
+            return None
+        return retval
+
+    def value_fn(self, func, *args, **kwargs):
+        return infer_meta(func, *args, **kwargs)
+
+
+@Singleton
+class LayerInferMetaCache(Cache):
+    def key_fn(self, layer, *args, **kwargs):
+        params = [
+            MetaInfo.from_tensor(x)
+            for x in layer.parameters(include_sublayers=True)
+        ]
+        try:
+            retval = hash(
+                (
+                    layer,
+                    tuple(params),
+                    tuple(flatten(args)),
+                    tuple(kwargs.keys()),
+                    tuple(flatten(kwargs)),
+                )
+            )
+        except Exception as e:
+            return None
+        return retval
+
+    def value_fn(self, layer, *args, **kwargs):
+        return infer_meta_for_layer(layer, *args, **kwargs)