[CustomOP Inplace] Automap inplace dtype and shape, prepare for vector<Tensor> output (#52214)

* [CustomOP Inplace] Automap inplace dtype and shape, prepare for vector<Tensor> output * delete dtype,shape func of multi_inplace op * [CustomOP Inplace] Automap inplace dtype and shape, support vector<Tensor> output * [CustomOP Inplace] Auto-generate python API for inplace vector<Tensor> output

Author: jiahy0825

paddle/fluid/pybind/eager_functions.cc

@@ -22,6 +22,7 @@ typedef SSIZE_T ssize_t;
 #endif
 
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "paddle/fluid/eager/accumulation/accumulation_node.h"
@@ -50,6 +51,7 @@ typedef SSIZE_T ssize_t;
 #include "paddle/phi/core/dense_tensor.h"
 #include "paddle/phi/core/sparse_coo_tensor.h"
 #include "paddle/phi/core/sparse_csr_tensor.h"
+#include "paddle/utils/string/string_helper.h"
 #include "pybind11/numpy.h"
 #include "pybind11/pybind11.h"
 
@@ -489,6 +491,49 @@ static PyObject* eager_api_jit_function_call(PyObject* self,
  EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* eager_api__get_custom_operator_inplace_reverse_idx(
+ PyObject* self, PyObject* args, PyObject* kwargs) {
+ EAGER_TRY
+ std::string op_type = CastPyArg2AttrString(PyTuple_GET_ITEM(args, 0), 0);
+ auto meta_info_map = egr::Controller::Instance().GetOpMetaInfoMap();
+ PADDLE_ENFORCE_NE(meta_info_map.find(op_type),
+ meta_info_map.end(),
+ paddle::platform::errors::NotFound(
+ "Can't find %s in Eager OpMetaInfoMap which should be "
+ "created by LoadOpMetaInfoAndRegisterOp, please make "
+ "sure you registered your op first and try again. ",
+ op_type));
+
+ const auto& inputs =
+ paddle::OpMetaInfoHelper::GetInputs(meta_info_map.at(op_type)[0]);
+ const auto& outputs =
+ paddle::OpMetaInfoHelper::GetOutputs(meta_info_map.at(op_type)[0]);
+ const auto& inplace_map =
+ paddle::OpMetaInfoHelper::GetInplaceMap(meta_info_map.at(op_type)[0]);
+ VLOG(7) << "Custom operator " << op_type
+ << " get InplaceMap for python, inplace map size = "
+ << inplace_map.size();
+
+ std::unordered_map<int, int> inplace_idx_map;
+ for (size_t in_idx = 0; in_idx < inputs.size(); ++in_idx) {
+ auto& input = inputs[in_idx];
+ if (inplace_map.find(input) == inplace_map.end()) {
+ continue;
+ }
+ auto out_iter = find(outputs.begin(), outputs.end(), inplace_map.at(input));
+ PADDLE_ENFORCE(
+ out_iter != outputs.end(),
+ phi::errors::NotFound("Can't find the mapped value of %s, please check "
+ "the input of `Inplace` again and make "
+ "sure you registered your op accurately. ",
+ input));
+ inplace_idx_map[distance(outputs.begin(), out_iter)] = in_idx;
+ }
+
+ return ToPyObject(inplace_idx_map);
+ EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* eager_api_run_custom_op(PyObject* self,
  PyObject* args,
  PyObject* kwargs) {
@@ -1133,6 +1178,11 @@ PyMethodDef variable_functions[] = {
  (PyCFunction)(void (*)(void))eager_api_run_partial_grad,
  METH_VARARGS | METH_KEYWORDS,
  NULL},
+ {"_get_custom_operator_inplace_map",
+ (PyCFunction)(void (*)(
+ void))eager_api__get_custom_operator_inplace_reverse_idx,
+ METH_VARARGS | METH_KEYWORDS,
+ NULL},
  {"_run_custom_op",
  (PyCFunction)(void (*)(void))eager_api_run_custom_op,
  METH_VARARGS | METH_KEYWORDS,

paddle/fluid/pybind/eager_utils.cc

@@ -810,6 +810,27 @@ PyObject* ToPyObject(const void* value) {
  platform::errors::Fatal("ToPyObject do not support void* with value."));
 }
 
+PyObject* ToPyObject(const std::unordered_map<int, int>& value) {
+ PyObject* dict = PyDict_New();
+ for (const auto& map_iter : value) {
+ // Convert Key
+ PyObject* key = ToPyObject(map_iter.first);
+ // Convert Value
+ PyObject* value = ToPyObject(map_iter.second);
+
+ if (!key || !value) {
+ PADDLE_THROW(
+ platform::errors::Fatal("Unable to convert int to PyObject"));
+ }
+
+ if (PyDict_SetItem(dict, key, value) != 0) {
+ PADDLE_THROW(
+ platform::errors::Fatal("Unable to set key:value for py_dict"));
+ }
+ }
+ return dict;
+}
+
 PyObject* ToPyObject(
  const std::unordered_map<std::string, std::vector<std::string>>& value) {
  PyObject* dict = PyDict_New();

paddle/fluid/pybind/eager_utils.h

@@ -110,6 +110,7 @@ PyObject* ToPyObject(const phi::SelectedRows* value);
 PyObject* ToPyObject(const paddle::framework::proto::VarType::Type& dtype);
 PyObject* ToPyObject(const paddle::framework::proto::VarType& type);
 PyObject* ToPyObject(const void* value);
+PyObject* ToPyObject(const std::unordered_map<int, int>& value);
 PyObject* ToPyObject(
  const std::unordered_map<std::string, std::vector<std::string>>& value);
 PyObject* ToPyObject(const paddle::framework::Vocab& value);

python/paddle/utils/cpp_extension/extension_utils.py

@@ -950,12 +950,12 @@ def parse_op_info(op_name):
  op_proto = OpProtoHolder.instance().get_op_proto(op_name)
 
  in_names = [x.name for x in op_proto.inputs]
- out_names = [x.name for x in op_proto.outputs]
  attr_names = [
  x.name for x in op_proto.attrs if x.name not in DEFAULT_OP_ATTR_NAMES
  ]
+ out_names = [x.name for x in op_proto.outputs]
 
- return in_names, out_names, attr_names
+ return in_names, attr_names, out_names
 
 
 def _import_module_from_library(module_name, build_directory, verbose=False):
@@ -1038,28 +1038,58 @@ def remove_if_exit(filepath):
  return custom_module
 
 
+def _gen_output_content(in_names, out_names, inplace_reverse_idx):
+ # ' ' * tab space * tab number
+ indent = ' ' * 4 * 2
+ dynamic_content = ""
+ static_content = ""
+ for out_idx, out_name in enumerate(out_names):
+ in_idx = -1
+ if out_idx in inplace_reverse_idx:
+ in_idx = inplace_reverse_idx[out_idx]
+ if in_idx != -1 and "@VECTOR" in in_names[in_idx]:
+ lower_in_names = in_names[in_idx].split("@")[0].lower()
+ dynamic_content += f"""
+{indent}outs['{out_name}'] = [core.eager.Tensor() for _ in range(len({lower_in_names}))]
+{indent}ctx.add_outputs(outs['{out_name}'])"""
+ static_content += f"""
+{indent}outs['{out_name}'] = [helper.create_variable(dtype='float32') for _ in range(len({lower_in_names}))]"""
+ else:
+ dynamic_content += f"""
+{indent}outs['{out_name}'] = core.eager.Tensor()
+{indent}ctx.add_outputs(outs['{out_name}'])"""
+ static_content += f"""
+{indent}outs['{out_name}'] = helper.create_variable(dtype='float32')"""
+
+ return dynamic_content, static_content
+
+
 def _custom_api_content(op_name):
  (
- params_str,
- ins_str,
- attrs_str,
- outs_str,
+ params_list,
+ ins_map,
+ attrs_map,
+ outs_list,
  in_names,
- attrs_names,
+ attr_names,
+ out_names,
+ inplace_reverse_idx,
  ) = _get_api_inputs_str(op_name)
- lower_in_names = [p.split("@")[0].lower() for p in in_names]
+ dynamic_content, static_content = _gen_output_content(
+ in_names, out_names, inplace_reverse_idx
+ )
+ lower_in_list = [p.split("@")[0].lower() for p in in_names]
  API_TEMPLATE = textwrap.dedent(
  """
  import paddle.fluid.core as core
  from paddle.fluid.core import VarBase, CustomOpKernelContext
  from paddle.fluid.framework import _dygraph_tracer, in_dygraph_mode
  from paddle.fluid.layer_helper import LayerHelper
 
- def {op_name}({inputs}):
+ def {op_name}({params_list}):
  # prepare inputs and outputs
- attrs = {attrs}
  outs = {{}}
- out_names = {out_names}
+ outs_list = {outs_list}
 
  # The output variable's dtype use default value 'float32',
  # and the actual dtype of output variable will be inferred in runtime.
@@ -1069,23 +1099,19 @@ def {op_name}({inputs}):
  ctx.add_inputs(i)
  for j in {attr_names}:
  ctx.add_attr(j)
- for out_name in out_names:
- outs[out_name] = core.eager.Tensor()
- ctx.add_outputs(outs[out_name])
+ {dynamic_content}
  core.eager._run_custom_op(ctx, "{op_name}", True)
  else:
  ins = {{}}
- for key, value in dict({ins}).items():
+ for key, value in dict({ins_map}).items():
  # handle optional inputs
  if value is not None:
  ins[key] = value
  helper = LayerHelper("{op_name}", **locals())
- for out_name in out_names:
- outs[out_name] = helper.create_variable(dtype='float32')
-
- helper.append_op(type="{op_name}", inputs=ins, outputs=outs, attrs=attrs)
+ {static_content}
+ helper.append_op(type="{op_name}", inputs=ins, outputs=outs, attrs={attrs_map})
 
- res = [outs[out_name] for out_name in out_names]
+ res = [outs[out_name] for out_name in outs_list]
 
  return res[0] if len(res)==1 else res
  """
@@ -1094,13 +1120,15 @@ def {op_name}({inputs}):
  # generate python api file
  api_content = API_TEMPLATE.format(
  op_name=op_name,
- inputs=params_str,
- ins=ins_str,
- attrs=attrs_str,
+ params_list=params_list,
+ ins_map=ins_map,
+ attrs_map=attrs_map,
  # "[x, y, z]""
- in_names="[" + ",".join(lower_in_names) + "]",
- attr_names="[" + ",".join(attrs_names) + "]",
- out_names=outs_str,
+ in_names="[" + ",".join(lower_in_list) + "]",
+ attr_names="[" + ",".join(attr_names) + "]",
+ outs_list=outs_list,
+ dynamic_content=dynamic_content,
+ static_content=static_content,
  )
 
  return api_content
@@ -1132,30 +1160,42 @@ def _get_api_inputs_str(op_name):
  """
  Returns string of api parameters and inputs dict.
  """
- in_names, out_names, attr_names = parse_op_info(op_name)
+ in_names, attr_names, out_names = parse_op_info(op_name)
  # e.g: x, y, z
  param_names = in_names + attr_names
  # NOTE(chenweihang): we add suffix `@VECTOR` for std::vector<Tensor> input,
  # but the string contains `@` cannot used as argument name, so we split
  # input name by `@`, and only use first substr as argument
- params_str = ','.join([p.split("@")[0].lower() for p in param_names])
+ params_list = ','.join([p.split("@")[0].lower() for p in param_names])
  # e.g: {'X': x, 'Y': y, 'Z': z}
- ins_str = "{%s}" % ','.join(
+ ins_map = "{%s}" % ','.join(
  [
  "'{}' : {}".format(in_name, in_name.split("@")[0].lower())
  for in_name in in_names
  ]
  )
  # e.g: {'num': n}
- attrs_str = "{%s}" % ",".join(
+ attrs_map = "{%s}" % ",".join(
  [
  "'{}' : {}".format(attr_name, attr_name.split("@")[0].lower())
  for attr_name in attr_names
  ]
  )
  # e.g: ['Out', 'Index']
- outs_str = "[%s]" % ','.join(["'{}'".format(name) for name in out_names])
- return params_str, ins_str, attrs_str, outs_str, in_names, attr_names
+ outs_list = "[%s]" % ','.join(["'{}'".format(name) for name in out_names])
+
+ inplace_reverse_idx = core.eager._get_custom_operator_inplace_map(op_name)
+
+ return (
+ params_list,
+ ins_map,
+ attrs_map,
+ outs_list,
+ in_names,
+ attr_names,
+ out_names,
+ inplace_reverse_idx,
+ )
 
 
 def _write_setup_file(

test/custom_op/test_custom_inplace.py

@@ -40,54 +40,6 @@
  verbose=True,
 )
 
-# Temporarily assemble custom python API
-from paddle.fluid import core
-from paddle.fluid.core import CustomOpKernelContext
-from paddle.fluid.framework import in_dygraph_mode
-from paddle.fluid.layer_helper import LayerHelper
-
-
-def custom_add_vec(x_vector, y):
- # prepare inputs and outputs
- attrs = {}
- outs = {}
- out_names = ["Out@VECTOR"]
-
- # The output variable's dtype use default value 'float32',
- # and the actual dtype of output variable will be inferred in runtime.
- if in_dygraph_mode():
- ctx = CustomOpKernelContext()
- for i in [x_vector, y]:
- ctx.add_inputs(i)
- for out_name in out_names:
- outs[out_name] = [core.eager.Tensor() for _ in range(len(x_vector))]
- ctx.add_outputs(outs[out_name])
- core.eager._run_custom_op(ctx, "custom_add_vec", True)
- else:
- ins = {}
- for key, value in dict({"X@VECTOR": x_vector, "Y": y}).items():
- # handle optional inputs
- if value is not None:
- ins[key] = value
- helper = LayerHelper("custom_add_vec", **locals())
- for out_name in out_names:
- outs[out_name] = [
- helper.create_variable(dtype='float32')
- for _ in range(len(x_vector))
- ]
-
- helper.append_op(
- type="custom_add_vec", inputs=ins, outputs=outs, attrs=attrs
- )
-
- res = [outs[out_name] for out_name in out_names]
-
- return res[0] if len(res) == 1 else res
-
-
-# Set custom python API manually
-custom_inplace.custom_add_vec = custom_add_vec
-
 
 def inplace_dynamic_add(phi_func, device, dtype, np_x, np_y):
  paddle.set_device(device)