[NewIR]Call _C_ops.xx in both dygraph and static mode (#56809)

0x45f · web-flow · commit 11a526ab2d38 · 2023-09-04T14:01:31.000+08:00
diff --git a/.gitignore b/.gitignore
@@ -72,7 +72,7 @@ tools/nvcc_lazy
 
 # This file is automatically generated.
 # TODO(zhiqiang) Move this file to build directory.
-paddle/fluid/pybind/eager_op_function.cc
+paddle/fluid/pybind/eager_op_function.*
 tools/nvcc_lazy
 paddle/phi/kernels/sparse/gpu/cutlass_generator/all_gemm_operations.h
 paddle/phi/kernels/sparse/gpu/cutlass_generator/configurations.h
diff --git a/paddle/fluid/eager/api/utils/global_utils.cc b/paddle/fluid/eager/api/utils/global_utils.cc
@@ -18,5 +18,7 @@
 namespace egr {
 
 Controller* Controller::controller_ = new Controller();
+thread_local std::shared_ptr<paddle::imperative::Tracer> Controller::tracer_ =
+    std::make_shared<paddle::imperative::Tracer>();
 
 }  // namespace egr
diff --git a/paddle/fluid/eager/api/utils/global_utils.h b/paddle/fluid/eager/api/utils/global_utils.h
@@ -145,8 +145,7 @@ class Controller {
  private:
   Controller() = default;
   static Controller* controller_;
-  std::shared_ptr<paddle::imperative::Tracer> tracer_{
-      new paddle::imperative::Tracer()};
+  static thread_local std::shared_ptr<paddle::imperative::Tracer> tracer_;
   std::unordered_map<std::string, std::vector<paddle::OpMetaInfo>>
       op_meta_info_map_;
   /* op_type : {{{grad_outputs}, {grad_inputs}, {input}, {output}, {attrs}},
diff --git a/paddle/fluid/eager/auto_code_generator/generator/CMakeLists.txt b/paddle/fluid/eager/auto_code_generator/generator/CMakeLists.txt
@@ -53,18 +53,25 @@ add_custom_target(
           ${nodes_h_path}
   VERBATIM)
 
-set(tmp_python_c_output_path
+set(tmp_python_c_source_path
     "${PADDLE_SOURCE_DIR}/paddle/fluid/pybind/eager_op_function.cc.tmp")
-set(python_c_output_path
+set(python_c_source_path
     "${PADDLE_SOURCE_DIR}/paddle/fluid/pybind/eager_op_function.cc")
+set(tmp_python_c_header_path
+    "${PADDLE_SOURCE_DIR}/paddle/fluid/pybind/eager_op_function.h.tmp")
+set(python_c_header_path
+    "${PADDLE_SOURCE_DIR}/paddle/fluid/pybind/eager_op_function.h")
 
 add_custom_target(
   eager_python_c_codegen
   COMMAND
     "${PYTHON_EXECUTABLE}"
     "${PADDLE_SOURCE_DIR}/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py"
     "--api_yaml_path=${api_yaml_path},${fwd_api_yaml_path}"
-    "--output_path=${tmp_python_c_output_path}"
-  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_python_c_output_path}
-          ${python_c_output_path}
+    "--source_path=${tmp_python_c_source_path}"
+    "--header_path=${tmp_python_c_header_path}"
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_python_c_source_path}
+          ${python_c_source_path}
+  COMMAND ${CMAKE_COMMAND} -E copy_if_different ${tmp_python_c_header_path}
+          ${python_c_header_path}
   VERBATIM)
diff --git a/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py b/paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py
@@ -87,7 +87,7 @@ def FindParsingFunctionFromAttributeType(atype):
 
 
 PYTHON_C_FUNCTION_TEMPLATE = """
-static PyObject * eager_api_{}(PyObject *self, PyObject *args, PyObject *kwargs) {{
+PyObject * eager_api_{}(PyObject *self, PyObject *args, PyObject *kwargs) {{
   {}
   PyThreadState *tstate = nullptr;
   try {{
@@ -173,6 +173,7 @@ def FindParsingFunctionFromAttributeType(atype):
 #include "paddle/fluid/pybind/eager.h"
 #include "paddle/fluid/eager/amp_utils.h"
 #include "paddle/fluid/eager/eager_amp_auto_cast.h"
+#include "paddle/fluid/pybind/eager_op_function.h"
 namespace paddle {{
 namespace pybind {{
 
@@ -253,6 +254,29 @@ def FindParsingFunctionFromAttributeType(atype):
 }}
 """
 
+PYTHON_C_H_TEMPLATE = """
+#pragma once
+
+#include <Python.h>
+
+// Avoid a problem with copysign defined in pyconfig.h on Windows.
+#ifdef copysign
+#undef copysign
+#endif
+
+namespace paddle {{
+namespace pybind {{
+
+{body}
+
+}} // namespace pybind
+}} // namespace paddle
+"""
+
+PYTHON_C_FUNCTION_DECLARE_TEMPLATE = """
+PyObject *eager_api_{name}(PyObject *self, PyObject *args, PyObject *kwargs);
+"""
+
 
 #####################
 # Generator Classes #
@@ -279,6 +303,7 @@ def __init__(self, forward_api_contents, namespace):
         # Generated Results
         self.python_c_function_str = ""
         self.python_c_function_reg_str = ""
+        self.python_c_funcion_declare_str = ""
 
     def CollectIsForwardOnly(self):
         forward_api_contents = self.forward_api_contents
@@ -428,6 +453,9 @@ def GeneratePythonCFunction(self):
             noamp_dygraph_function_str,
             return_str,
         )
+        self.python_c_funcion_declare_str = (
+            PYTHON_C_FUNCTION_DECLARE_TEMPLATE.format(name=forward_api_name)
+        )
 
         # Set prefix of forward_api_name to avoid conflicts
         prefix = self.namespace.strip("::")
@@ -483,6 +511,12 @@ def GeneratePythonCFunction(self):
                 return_str,
             )
 
+            python_c_funcion_declare_str = (
+                PYTHON_C_FUNCTION_DECLARE_TEMPLATE.format(
+                    name=inplaced_forward_api_name
+                )
+            )
+
             python_c_inplace_func_reg_str = (
                 PYTHON_C_FUNCTION_REG_TEMPLATE.format(
                     forward_api_name_prefix,
@@ -496,10 +530,14 @@ def GeneratePythonCFunction(self):
             # self.forward_api_name ending with '_' means it only has inplace api
             if self.forward_api_name[-1] == '_':
                 self.python_c_function_str = python_c_inplace_func_str
+                self.python_c_funcion_declare_str = python_c_funcion_declare_str
                 # Generate Python-C Function Registration
                 self.python_c_function_reg_str = python_c_inplace_func_reg_str
             else:
                 self.python_c_function_str += python_c_inplace_func_str
+                self.python_c_funcion_declare_str += (
+                    python_c_funcion_declare_str
+                )
                 # Generate Python-C Function Registration
                 self.python_c_function_reg_str += python_c_inplace_func_reg_str
 
@@ -541,6 +579,7 @@ def __init__(self, path):
         # Generated Result
         self.python_c_functions_str = ""
         self.python_c_functions_reg_str = ""
+        self.python_c_funcion_declare_str = ""
 
     def GeneratePythonCFunctions(self):
         namespace = self.namespace
@@ -559,6 +598,9 @@ def GeneratePythonCFunctions(self):
                 self.python_c_functions_reg_str += (
                     f_generator.python_c_function_reg_str
                 )
+                self.python_c_funcion_declare_str += (
+                    f_generator.python_c_funcion_declare_str
+                )
 
     def AttachNamespace(self):
         namespace = self.namespace
@@ -570,6 +612,11 @@ def AttachNamespace(self):
             self.python_c_functions_str = NAMESPACE_WRAPPER_TEMPLATE.format(
                 namespace, python_c_functions_str
             )
+            self.python_c_funcion_declare_str = (
+                NAMESPACE_WRAPPER_TEMPLATE.format(
+                    namespace, self.python_c_funcion_declare_str
+                )
+            )
 
     def run(self):
         # Infer namespace from yaml_path
@@ -593,7 +640,8 @@ def ParseArguments():
         description='Eager Code Generator Args Parser'
     )
     parser.add_argument('--api_yaml_path', type=str)
-    parser.add_argument('--output_path', type=str)
+    parser.add_argument('--source_path', type=str)
+    parser.add_argument('--header_path', type=str)
 
     args = parser.parse_args()
     return args
@@ -631,6 +679,7 @@ def GeneratePythonCFile(filepath, python_c_str):
 
     generated_python_c_functions = ""
     generated_python_c_registration = ""
+    generated_python_c_functions_header = ""
     for i in range(len(api_yaml_paths)):
         api_yaml_path = api_yaml_paths[i]
 
@@ -643,14 +692,22 @@ def GeneratePythonCFile(filepath, python_c_str):
         generated_python_c_registration += (
             py_c_generator.python_c_functions_reg_str
         )
+        generated_python_c_functions_header += (
+            py_c_generator.python_c_funcion_declare_str
+        )
 
     python_c_str = GeneratePythonCWrappers(
         generated_python_c_functions, generated_python_c_registration
     )
 
-    output_path = args.output_path
-    for path in [output_path]:
+    soucre_path = args.source_path
+    header_path = args.header_path
+    for path in [soucre_path, header_path]:
         if os.path.exists(path):
             os.remove(path)
 
-    GeneratePythonCFile(output_path, python_c_str)
+    GeneratePythonCFile(soucre_path, python_c_str)
+    GeneratePythonCFile(
+        header_path,
+        PYTHON_C_H_TEMPLATE.format(body=generated_python_c_functions_header),
+    )
diff --git a/paddle/fluid/imperative/tracer.cc b/paddle/fluid/imperative/tracer.cc
@@ -56,7 +56,7 @@ thread_local AmpLevel Tracer::amp_level_ = AmpLevel::O0;
 
 thread_local phi::DataType Tracer::amp_dtype_ = phi::DataType::FLOAT32;
 
-static std::shared_ptr<Tracer> g_current_tracer(nullptr);
+static thread_local std::shared_ptr<Tracer> g_current_tracer(nullptr);
 
 const std::shared_ptr<Tracer>& GetCurrentTracer() { return g_current_tracer; }
 
diff --git a/paddle/fluid/ir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/ir/dialect/op_generator/ops_api_gen.py
@@ -21,7 +21,9 @@
 #include <pybind11/pybind11.h>
 
 #include "paddle/fluid/pybind/static_op_function.h"
+#include "paddle/fluid/pybind/eager_op_function.h"
 #include "paddle/phi/core/enforce.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
 
 {body}
 
@@ -44,19 +46,63 @@
 
 FUNCTION_IMPL_TEMPLATE = """
 static PyObject *{name}(PyObject *self, PyObject *args, PyObject *kwargs) {{
+  if (egr::Controller::Instance().GetCurrentTracer() == nullptr) {{
+    VLOG(6) << "Call static_api_{name}";
+    return static_api_{name}(self, args, kwargs);
+  }} else {{
+    VLOG(6) << "Call eager_api_{name}";
+    return eager_api_{name}(self, args, kwargs);
+  }}
+}}"""
+
+NO_DY_FUNCTION_IMPL_TEMPLATE = """
+static PyObject *{name}(PyObject *self, PyObject *args, PyObject *kwargs) {{
+  VLOG(6) << "Call static_api_{name}";
   return static_api_{name}(self, args, kwargs);
 }}"""
 
 OPS_API_TEMPLATE = """
 {{"{name}", (PyCFunction)(void (*)(void)){name}, METH_VARARGS | METH_KEYWORDS, "C++ interface function for {name}."}},"""
 
+SPECIAL_STATIC_ONLY_APIS = [
+    'fetch',
+    'set_value_with_tensor',
+    'set_value_with_tensor_',
+    'fused_bn_add_activation_',
+    'fused_batch_norm_act_',
+    'add_n_',
+    'set_value',
+    'assign_value',
+    'set_value_',
+    'embedding_grad_sparse',
+    'add_n_with_kernel',
+    'print',
+    'send_v2',
+    'shadow_feed',
+    'recv_v2',
+    'rnn_',
+    'fused_scale_bias_relu_conv_bnstats',
+    'batch_norm_',
+    'c_allreduce_sum',
+    'c_embedding',
+    'c_identity',
+]
+
 
 class OpsAPIGen(CodeGen):
     def __init__(self) -> None:
         super().__init__()
 
     def _gen_one_function_impl(self, name):
-        return FUNCTION_IMPL_TEMPLATE.format(name=name)
+        if (
+            name.endswith('grad')
+            or name.endswith('grad_')
+            or name.endswith('xpu')
+            or name in SPECIAL_STATIC_ONLY_APIS
+        ):
+            return NO_DY_FUNCTION_IMPL_TEMPLATE.format(name=name)
+        else:
+            return FUNCTION_IMPL_TEMPLATE.format(name=name)
 
     def _gen_one_ops_api(self, name):
         return OPS_API_TEMPLATE.format(name=name)
diff --git a/paddle/fluid/pybind/.gitignore b/paddle/fluid/pybind/.gitignore
@@ -1,3 +1,3 @@
 pybind.h
-eager_op_function.cc
+eager_op_function.*
 eager_legacy_op_function.cc
diff --git a/python/paddle/_C_ops.py b/python/paddle/_C_ops.py
@@ -16,6 +16,14 @@
 
 __all__ = []
 
+UNIFIED_APIS = ['mean']
+
 for name in dir(core.eager.ops):
     globals()[name] = getattr(core.eager.ops, name)
     __all__.append(name)
+
+for name in dir(core.ir.ops):
+    if name in UNIFIED_APIS:
+        globals()[name] = getattr(core.ir.ops, name)
+        if name not in __all__:
+            __all__.append(name)
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -459,6 +459,7 @@
 from . import linalg  # noqa: F401
 from . import fft  # noqa: F401
 from . import signal  # noqa: F401
+from . import _ir_ops  # noqa: F401
 
 import paddle.text  # noqa: F401
 import paddle.vision  # noqa: F401
diff --git a/python/paddle/fluid/framework.py b/python/paddle/fluid/framework.py
diff --git a/python/paddle/framework/__init__.py b/python/paddle/framework/__init__.py
diff --git a/python/paddle/jit/dy2static/partial_program.py b/python/paddle/jit/dy2static/partial_program.py
diff --git a/python/paddle/tensor/stat.py b/python/paddle/tensor/stat.py
diff --git a/test/legacy_test/test_inference_api.py b/test/legacy_test/test_inference_api.py
