Refactored python-level trace_op to call through _C_ops instead of Tracer::TraceOp, under eager_mode (#38338)

* Replaced core.ops with _C_ops * Refactored python-level trace_op to call through _C_ops instead of Tracer::TraceOp, under eager_mode * Modified trace_op interface * Refactored trace_op logic for eager mode * Added Eager Dygraph support for OpTest * Fixed ci issues * Fixed CI failures * Fixed Coverage CI Issues * Fixed XPU CI Issues

Author: jim19930609

paddle/fluid/eager/auto_code_generator/eager_generator.cc

@@ -37,6 +37,8 @@ std::unordered_map<std::string, std::vector<std::string>>
  core_ops_returns_info = {};
 std::unordered_map<std::string, std::vector<std::string>> core_ops_args_info =
  {};
+std::unordered_map<std::string, std::vector<std::string>>
+ core_ops_args_type_info = {};
 
 /* --- Static maps to handle corner cases --- */
 static std::unordered_map<std::string, paddle::framework::AttributeMap>
@@ -1225,10 +1227,16 @@ static std::pair<std::string, std::string> GenerateForwardFunctionContents(
  */
  VLOG(6) << "Generating Dygraph Forward Function";
 
- std::string generated_function_body = "";
+ const char* FORWARD_FUNCTION_TEMPLATE =
+ " VLOG(3) << \"Running Eager Forward Op: %s\";\n";
+ std::string generated_function_body =
+ paddle::string::Sprintf(FORWARD_FUNCTION_TEMPLATE, op_type);
+
  std::string dygraph_function_args_str = "";
  core_ops_args_info[op_type] = {};
+ core_ops_args_type_info[op_type] = {};
  core_ops_args_info[op_type].resize(in_vars.size());
+ core_ops_args_type_info[op_type].resize(in_vars.size());
 
  /* ------ Dygraph forward function generation ------ */
  generated_function_body += " // Dygraph Forward Pass\n";
@@ -1246,10 +1254,14 @@ static std::pair<std::string, std::string> GenerateForwardFunctionContents(
  "const std::vector<egr::EagerTensor>& %s";
  input_args_str_list[input_position] =
  paddle::string::Sprintf(FWD_INS_ARG_TEMPLATE, input_name);
+
+ core_ops_args_type_info[op_type][input_position] = "list";
  } else {
  const char* FWD_INS_ARG_TEMPLATE = "const egr::EagerTensor& %s";
  input_args_str_list[input_position] =
  paddle::string::Sprintf(FWD_INS_ARG_TEMPLATE, input_name);
+
+ core_ops_args_type_info[op_type][input_position] = "tensor";
  }
  core_ops_args_info[op_type][input_position] = input_name;
 
@@ -1318,11 +1330,14 @@ static std::pair<std::string, std::string> GenerateForwardFunctionContents(
  paddle::string::Sprintf(FWD_NUM_ARG_TEMPLATE, output_var_name);
  dygraph_function_args_str += arg_str;
 
+ core_ops_args_type_info[op_type].push_back("list");
  } else {
  const char* FWD_NUM_ARG_TEMPLATE = ", egr::EagerTensor* %s";
  std::string arg_str =
  paddle::string::Sprintf(FWD_NUM_ARG_TEMPLATE, output_var_name);
  dygraph_function_args_str += arg_str;
+
+ core_ops_args_type_info[op_type].push_back("tensor");
  }
  const char* FWD_OUTS_CONTENT_TEMPLATE =
  "{ \"%s\", egr::EagerUtils::TrySyncToVars(%s) },";
@@ -1344,6 +1359,7 @@ static std::pair<std::string, std::string> GenerateForwardFunctionContents(
  outs_contents_str += paddle::string::Sprintf(FWD_OUTS_CONTENT_TEMPLATE,
  output_name, outnum);
  core_ops_args_info[op_type].push_back(outnum);
+ core_ops_args_type_info[op_type].push_back("int");
  } else {
  const char* FWD_OUTS_CONTENT_TEMPLATE =
  "{ \"%s\", "
@@ -1811,6 +1827,11 @@ static std::string GenerateGradNodeCCContents(
  }
  */
 
+ const char* EAGER_LOG_TEMPLATE =
+ " VLOG(3) << \"Running Eager Backward Node: GradNode%s\";\n";
+ std::string generated_grad_function_body =
+ paddle::string::Sprintf(EAGER_LOG_TEMPLATE, fwd_op_type);
+
  // This is a Copy
  auto op_base_infos = bwd_info.GetOpBaseInfos();
 
@@ -1829,7 +1850,6 @@ static std::string GenerateGradNodeCCContents(
  op_base_infos.emplace_back(std::move(op_base_info));
  }
 
- std::string generated_grad_function_body = "";
  size_t outs_size = 0;
  for (size_t i = 0; i < op_base_infos.size(); i++) {
  const auto& op_base_info = op_base_infos[i];
@@ -2030,6 +2050,9 @@ static std::string GenerateDygraphHFileIncludes() {
  dygraph_forward_api_includes_str +=
  "extern std::unordered_map<std::string, std::vector<std::string>> "
  "core_ops_args_info;\n";
+ dygraph_forward_api_includes_str +=
+ "extern std::unordered_map<std::string, std::vector<std::string>> "
+ "core_ops_args_type_info;\n";
  dygraph_forward_api_includes_str +=
  "extern std::unordered_map<std::string, std::vector<std::string>> "
  "core_ops_returns_info;\n\n";
@@ -2126,16 +2149,20 @@ static std::string GenerateCoreOpsReturnsInfo() {
  "std::unordered_map<std::string, std::vector<std::string>> "
  "core_ops_args_info = { %s };\n"
  "std::unordered_map<std::string, std::vector<std::string>> "
+ "core_ops_args_type_info = { %s };\n"
+ "std::unordered_map<std::string, std::vector<std::string>> "
  "core_ops_returns_info = { %s };\n";
 
  std::string core_ops_args_info_init_str =
  ConvertCoreOpsInfosToString(core_ops_args_info);
+ std::string core_ops_args_type_info_init_str =
+ ConvertCoreOpsInfosToString(core_ops_args_type_info);
  std::string core_ops_returns_info_init_str =
  ConvertCoreOpsInfosToString(core_ops_returns_info);
 
  std::string core_ops_info_str = paddle::string::Sprintf(
  Core_Ops_Returns_MAP_TEMPLATE, core_ops_args_info_init_str,
- core_ops_returns_info_init_str);
+ core_ops_args_type_info_init_str, core_ops_returns_info_init_str);
 
  return core_ops_info_str;
 }

paddle/fluid/pybind/eager_method.cc

@@ -121,6 +121,30 @@ static PyObject* eager_tensor_method__copy_to(EagerTensorObject* self,
  EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* eager_tensor_method_reconstruct_from_(EagerTensorObject* self,
+ PyObject* args,
+ PyObject* kwargs) {
+ EAGER_SYNC_TRY
+ egr::EagerTensor src_tensor =
+ CastPyArg2EagerTensor(PyTuple_GET_ITEM(args, 0), 0);
+ bool blocking = CastPyArg2AttrBoolean(PyTuple_GET_ITEM(args, 1), 1);
+ std::string orig_name = self->eager_tensor.name();
+ VLOG(6) << "Start Reconstructing Tensor from" << src_tensor.name() << " to "
+ << orig_name;
+ self->eager_tensor.copy_(src_tensor, blocking);
+ // Steal Tensor from src tensor
+ self->eager_tensor.set_tensor(src_tensor.Tensor());
+
+ // Recover source name
+ self->eager_tensor.set_name(orig_name);
+
+ VLOG(6) << "Finished Reconstructing Tensor from" << src_tensor.name()
+ << " to " << self->eager_tensor.name();
+ Py_INCREF(Py_None);
+ return Py_None;
+ EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* eager_tensor_method_copy_(EagerTensorObject* self,
  PyObject* args, PyObject* kwargs) {
  EAGER_SYNC_TRY
@@ -387,6 +411,9 @@ PyMethodDef variable_methods[] = {
  METH_VARARGS | METH_KEYWORDS, NULL},
  {"copy_", (PyCFunction)(void (*)(void))eager_tensor_method_copy_,
  METH_VARARGS | METH_KEYWORDS, NULL},
+ {"reconstruct_from_",
+ (PyCFunction)(void (*)(void))eager_tensor_method_reconstruct_from_,
+ METH_VARARGS | METH_KEYWORDS, NULL},
  {"retain_grads", (PyCFunction)(void (*)(void))eager_tensor_retain_grads,
  METH_VARARGS | METH_KEYWORDS, NULL},
  {"_clear_gradient",

paddle/fluid/pybind/eager_op_function_generator.cc

@@ -313,6 +313,21 @@ static std::string GenerateCoreOpsInfoMap() {
  " }\n"
  "}\n"
  "\n"
+ "static PyObject * eager_get_core_ops_args_type_info(PyObject *self) {\n"
+ " PyThreadState *tstate = nullptr;\n"
+ " try\n"
+ " {\n"
+ " return ToPyObject(core_ops_args_type_info);\n"
+ " }\n"
+ " catch(...) {\n"
+ " if (tstate) {\n"
+ " PyEval_RestoreThread(tstate);\n"
+ " }\n"
+ " ThrowExceptionToPython(std::current_exception());\n"
+ " return nullptr;\n"
+ " }\n"
+ "}\n"
+ "\n"
  "static PyObject * eager_get_core_ops_returns_info(PyObject *self) {\n"
  " PyThreadState *tstate = nullptr;\n"
  " try\n"
@@ -399,6 +414,10 @@ int main(int argc, char* argv[]) {
  "{\"get_core_ops_args_info\", "
  "(PyCFunction)(void(*)(void))eager_get_core_ops_args_info, METH_NOARGS, "
  "\"C++ interface function for eager_get_core_ops_args_info.\"},\n"
+ "{\"get_core_ops_args_type_info\", "
+ "(PyCFunction)(void(*)(void))eager_get_core_ops_args_type_info, "
+ "METH_NOARGS, "
+ "\"C++ interface function for eager_get_core_ops_args_type_info.\"},\n"
  " {\"get_core_ops_returns_info\", "
  "(PyCFunction)(void(*)(void))eager_get_core_ops_returns_info, "
  "METH_NOARGS, \"C++ interface function for "

python/paddle/fluid/dygraph/tracer.py

@@ -19,6 +19,7 @@
 from collections import defaultdict
 from paddle.fluid import core
 from paddle.fluid import framework
+from paddle import _C_ops
 
 
 class Tracer(core.Tracer):
@@ -46,9 +47,87 @@ def trace_op(self,
  attrs,
  stop_gradient=False,
  inplace_map=None):
- self.trace(type, inputs, outputs, attrs,
- framework._current_expected_place(), self._has_grad and
- not stop_gradient, inplace_map if inplace_map else {})
+ if framework._in_eager_mode():
+ # inputs : {"sum": [tensor], ...}
+ # outputs : {"sum": [tensor], ...}
+
+ function_ptr = _C_ops.__dict__[type]
+
+ core_ops_args_info = _C_ops.get_core_ops_args_info()
+ core_ops_args_type_info = _C_ops.get_core_ops_args_type_info()
+ core_ops_returns_info = _C_ops.get_core_ops_returns_info()
+
+ op_args = core_ops_args_info[type]
+ op_args_type = core_ops_args_type_info[type]
+ op_returns = core_ops_returns_info[type]
+
+ arg_list = []
+ for i in range(len(op_args)):
+ arg_name = op_args[i]
+ arg_type = op_args_type[i]
+ if arg_name in inputs.keys():
+ arg_to_append = inputs[arg_name]
+ elif arg_name in outputs.keys():
+ arg_to_append = outputs[arg_name]
+ else:
+ if "Num" in arg_name:
+ # Remove "Num" suffix to get out_name
+ out_name = arg_name[:-3]
+ assert out_name in outputs.keys()
+ num_outs = len(outputs[out_name])
+ arg_to_append = num_outs
+ else:
+ arg_to_append = None
+
+ if arg_to_append is None:
+ arg_list.append(arg_to_append)
+ elif arg_type == "tensor":
+ if isinstance(arg_to_append, list):
+ arg_list.append(arg_to_append[0])
+ else:
+ arg_list.append(arg_to_append)
+ elif arg_type == "list":
+ assert isinstance(arg_to_append, list)
+ arg_list.append(arg_to_append)
+ else:
+ assert arg_type == "int"
+ assert isinstance(arg_to_append, int)
+ arg_list.append(arg_to_append)
+
+ attrs_list = []
+ for k, v in attrs.items():
+ attrs_list.append(k)
+ attrs_list.append(v)
+ returns = function_ptr(*arg_list, *attrs_list)
+
+ if isinstance(returns, tuple):
+ for i in range(len(op_returns)):
+ retname = op_returns[i]
+ if retname in outputs.keys():
+ # Replaced outputs by function returns
+ if isinstance(returns[i], list):
+ for j in range(len(returns[i])):
+ outputs[retname][j].reconstruct_from_(
+ returns[i][j], False)
+ else:
+ outputs[retname][0].reconstruct_from_(returns[i],
+ False)
+ elif isinstance(returns, list):
+ assert len(outputs.keys()) == 1
+ key = list(outputs.keys())[0]
+ for j in range(len(returns)):
+ outputs[key][j].reconstruct_from_(returns[j], False)
+ else:
+ assert len(outputs.keys()) == 1
+ key = list(outputs.keys())[0]
+ if isinstance(outputs[key], list):
+ outputs[key][0].reconstruct_from_(returns, False)
+ else:
+ outputs[key].reconstruct_from_(returns, False)
+ else:
+ self.trace(type, inputs, outputs, attrs,
+ framework._current_expected_place(), self._has_grad and
+ not stop_gradient, inplace_map if inplace_map else {})
 
  def train_mode(self):
  self._train_mode = True

python/paddle/fluid/tests/unittests/CMakeLists.txt

@@ -104,6 +104,10 @@ foreach(TEST_OP ${MIXED_DIST_TEST_OPS})
  list(REMOVE_ITEM TEST_OPS ${TEST_OP})
 endforeach()
 
+if(ON_INFER)
+ LIST(REMOVE_ITEM TEST_OPS test_eager_trace_op)
+endif()
+
 if(NOT WITH_GPU)
  LIST(REMOVE_ITEM TEST_OPS test_fused_feedforward_op)
  LIST(REMOVE_ITEM TEST_OPS test_fused_attention_op)