[Dy2St] Optimize ConstructAttrMapForRunProgram performance (#73682)

Author: SigureMo

paddle/fluid/pybind/eager_custom_python_api.h

@@ -96,8 +96,7 @@ static PyObject *eager_api_run_program(PyObject *self,
     }
     framework::AttributeMap attrs;
     VLOG(6) << "Start PIR ConstructAttrMapFromPyArgs";
-    ConstructAttrMapForRunProgram(
-        "run_program", args, 4, PyTuple_GET_SIZE(args), attrs);
+    ConstructAttrMapForRunProgram("run_program", args, 4, attrs);
 
     VLOG(6) << "Finish Pir ConstructAttrMapFromPyArgs";
     tstate = PyEval_SaveThread();

paddle/fluid/pybind/op_function_common.cc

@@ -846,17 +846,18 @@ void CastPyArg2AttrScalars(PyObject* obj,
 std::vector<std::string> CastPyArg2Strings(PyObject* obj,
                                            const std::string& op_type,
                                            ssize_t arg_pos) {
-  std::vector<std::string> value;
+  std::vector<std::string_view> views;
   if (PyList_Check(obj)) {
     Py_ssize_t len = PyList_Size(obj);
+    views.reserve(len);
     PyObject* item = nullptr;
     for (Py_ssize_t i = 0; i < len; i++) {
       item = PyList_GetItem(obj, i);
       if (PyObject_CheckString(item)) {
         Py_ssize_t size = 0;
         const char* data = nullptr;
         data = PyUnicode_AsUTF8AndSize(item, &size);
-        value.emplace_back(std::string(data, (size_t)size));  // NOLINT
+        views.emplace_back(std::string_view(data, (size_t)size));  // NOLINT
       } else {
         PADDLE_THROW(common::errors::InvalidType(
             "%s(): argument (position %d) must be "
@@ -869,14 +870,15 @@ std::vector<std::string> CastPyArg2Strings(PyObject* obj,
     }
   } else if (PyTuple_Check(obj)) {
     Py_ssize_t len = PyTuple_Size(obj);
+    views.reserve(len);
     PyObject* item = nullptr;
     for (Py_ssize_t i = 0; i < len; i++) {
       item = PyTuple_GetItem(obj, i);
       if (PyObject_CheckString(item)) {
         Py_ssize_t size = 0;
         const char* data = nullptr;
         data = PyUnicode_AsUTF8AndSize(item, &size);
-        value.emplace_back(std::string(data, (size_t)size));  // NOLINT
+        views.emplace_back(std::string_view(data, (size_t)size));  // NOLINT
       } else {
         PADDLE_THROW(common::errors::InvalidType(
             "%s(): argument (position %d) must be "
@@ -895,7 +897,11 @@ std::vector<std::string> CastPyArg2Strings(PyObject* obj,
         arg_pos + 1,
         ((PyTypeObject*)obj->ob_type)->tp_name));  // NOLINT
   }
-
+  std::vector<std::string> value;
+  value.reserve(views.size());
+  for (const auto& view : views) {
+    value.emplace_back(view);
+  }
   return value;
 }
 
@@ -1207,16 +1213,15 @@ void ConstructAttrMapForLegacyRunProgram(
 void ConstructAttrMapForRunProgram(
     const std::string& op_type,
     PyObject* args,
-    ssize_t attr_start,
-    ssize_t attr_end,
+    ssize_t arg_pos,
     paddle::framework::AttributeMap& attrs) {  // NOLINT
-  PADDLE_ENFORCE_EQ((attr_end - attr_start) % 2,
-                    0,
-                    common::errors::InvalidArgument(
-                        "The number of arguments for attributes should be even "
-                        "but attr_start = %d, attr_end = %d.",
-                        attr_start,
-                        attr_end));
+  PyObject* attrs_dict = PyTuple_GET_ITEM(args, arg_pos);
+  if (!PyDict_Check(attrs_dict)) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "%s(): argument must be dict, but got %s",
+        op_type,
+        reinterpret_cast<PyTypeObject*>(attrs_dict->ob_type)->tp_name));
+  }
 
   using CastFuncType = void (*)(PyObject*,
                                 paddle::framework::AttributeMap&,
@@ -1246,34 +1251,29 @@ void ConstructAttrMapForRunProgram(
       {"cuda_graph_dispatch_key", CastPyArg2AttrLong},
   };
 
-  PyObject* obj = nullptr;
-  for (ssize_t arg_pos = attr_start; arg_pos < attr_end; arg_pos += 2) {
-    VLOG(3) << "Start Process " << arg_pos;
+  PyObject *key, *value;
+  Py_ssize_t pos = 0;
+  while (PyDict_Next(attrs_dict, &pos, &key, &value)) {
     Py_ssize_t key_len = 0;
     const char* key_ptr = nullptr;
-    obj = PyTuple_GET_ITEM(args, arg_pos);
-    if (PyObject_CheckString(obj)) {
-      key_ptr = PyUnicode_AsUTF8AndSize(obj, &key_len);
+    if (PyObject_CheckString(key)) {
+      key_ptr = PyUnicode_AsUTF8AndSize(key, &key_len);
     } else {
       PADDLE_THROW(common::errors::InvalidArgument(
-          "%s(): argument (position %d) must be str, but got %s",
+          "%s(): dict key must be str, but got %s",
           op_type,
-          arg_pos,
-          ((PyTypeObject*)obj->ob_type)->tp_name));  // NOLINT
+          reinterpret_cast<PyTypeObject*>(key->ob_type)->tp_name));
     }
     std::string_view key_view(key_ptr, static_cast<size_t>(key_len));
-    VLOG(3) << "Start Process " << key_view;
-    obj = PyTuple_GET_ITEM(args, arg_pos + 1);
     auto it = kAttrFuncMap.find(std::string(key_view));
     if (it != kAttrFuncMap.end()) {
-      // Call Cast function
-      it->second(obj, attrs, std::string(key_view), op_type, arg_pos);
+      it->second(value, attrs, std::string(key_view), op_type, 0);
     } else {
       PADDLE_THROW(common::errors::InvalidArgument(
           "Attribute key %.*s is not recognized for operator %s.",
           static_cast<int>(key_view.size()),
           key_view.data(),
-          op_type.c_str()));  // NOLINT
+          op_type.c_str()));
     }
   }
 }

paddle/fluid/pybind/op_function_common.h

@@ -218,8 +218,7 @@ void ConstructAttrMapForLegacyRunProgram(
 void ConstructAttrMapForRunProgram(
     const std::string& op_type,
     PyObject* args,
-    ssize_t attr_start,
-    ssize_t attr_end,
+    ssize_t arg_pos,
     paddle::framework::AttributeMap& attrs);  // NOLINT
 
 unsigned long GetUnsignedLongFromArgs(  // NOLINT

python/paddle/jit/dy2static/pir_partial_program.py

@@ -731,9 +731,9 @@ def __init__(
     @staticmethod
     def run_impl(partial_program_layer, inputs, parameters, outputs, attrs):
         _C_ops.run_program(
-            inputs,
-            parameters,
-            outputs,
+            PartialProgramLayer._valid_vars(inputs),
+            PartialProgramLayer._valid_vars(parameters),
+            PartialProgramLayer._valid_vars(outputs),
             partial_program_layer._create_scope_vec(
                 cache_key=(
                     PartialProgramLayer._calc_scope_cache_key(
@@ -745,23 +745,21 @@ def run_impl(partial_program_layer, inputs, parameters, outputs, attrs):
                 ),
                 use_scope_cache=True,
             ),
-            *PartialProgramLayer._dict_attributes_to_op_fn_attrs(attrs),
+            attrs,
         )
 
     def __call__(self, inputs):
         """
         Execute static graph by Interpreter and Return dynamic Tensors.
         """
         attrs = self._prepare_attributes(in_sot_mode=False)
-        inputs = self._valid_vars(self._prepare_inputs(inputs))
-        parameters = self._valid_vars(self._params)
+        inputs = self._prepare_inputs(inputs)
         out_vars = self._prepare_outputs()
-        outputs = self._valid_vars(out_vars)
 
         self.call_run_impl_with_hook(
             inputs,
-            parameters,
-            outputs,
+            self._params,
+            out_vars,
             attrs,
         )
 
@@ -773,15 +771,12 @@ def sot_call(self, inputs):
         In sot, inputs and outputs of partial program only contain tensors, so we can skip some step to speed up
         """
         attrs = self._prepare_attributes(in_sot_mode=True)
-        inputs = self._valid_vars(inputs)
-        parameters = self._valid_vars(self._params)
         out_vars = self._prepare_outputs()
-        outputs = self._valid_vars(out_vars)
 
         self.call_run_impl_with_hook(
             inputs,
-            parameters,
-            outputs,
+            self._params,
+            out_vars,
             attrs,
         )
         return self._outputs.quick_restore(out_vars)
@@ -1198,27 +1193,15 @@ def _append_backward(
         return whole_program
 
     def _prepare_attributes(self, in_sot_mode=False):
-        attrs = {
+        return {
             'forward_program': self.program.forward_program,
             'backward_program': self.program.backward_program,
             'is_test': not self.training,
             'program_id': self.program_id,
             'in_sot_mode': in_sot_mode,
             'cuda_graph_state': CUDAGraphState.DISABLE,  # default value for not use cuda graph
             'cuda_graph_dispatch_key': 0,  # default value for not use cuda graph
-        }
-        attrs |= self.program.program_attr.items()
-        return attrs
-
-    @staticmethod
-    def _dict_attributes_to_op_fn_attrs(attrs):
-        op_fn_attrs = []
-        for k, v in attrs.items():
-            if k == "cuda_graph_state":
-                v = int(v)
-            op_fn_attrs.append(k)
-            op_fn_attrs.append(v)
-        return op_fn_attrs
+        } | self.program.program_attr
 
     def _prepare_inputs(self, inputs):
         """
@@ -1360,7 +1343,8 @@ def _check_params_all_inited(self, main_program):
                 )
             param_and_buffer_names_set.add(var.name)
 
-    def _valid_vars(self, vars):
+    @staticmethod
+    def _valid_vars(vars):
         return vars if vars else None