Support pir for python get set (#59928)

zoooo0820 · 2742195759 · SigureMo · web-flow · commit 0d2c1eb3045f · 2023-12-20T15:20:10.000+08:00
* support __getitem__ &amp; __setitem__ in PIR

* dy2static consider inplace mapping.

* fix DtypeError in case15

* disable AST+PIR only

* add missing decorator test_legacy_and_pt_and_pir

* modify Py_Dict to ValueDict

* fix test_top_p

---------

Co-authored-by: xiongkun &lt;xiongkun03@baidu.com&gt;
Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc b/paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
@@ -456,11 +456,9 @@ void HandleForSpecialOp(pir::Operation* op,
     // change opreand name to param_name
     auto orig_name = value_exe_info->GetValue2VarName().at(value);
 
-    PADDLE_ENFORCE_NE(
-        param_name,
-        orig_name,
-        phi::errors::PreconditionNotMet(
-            "SetParamer param name should not equal with var name"));
+    if (param_name == orig_name) {
+      return;
+    }
 
     if (value_exe_info->GetScope()->root()->FindVar(param_name) == nullptr) {
       const_cast<Scope*>(value_exe_info->GetScope()->root())
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -39,6 +39,8 @@
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/dialect/operator/trait/inplace.h"
+#include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_parser.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/fluid/pir/transforms/dead_code_elimination_pass.h"
 #include "paddle/fluid/pir/transforms/fusion/attention_fuse_pass.h"
@@ -1396,8 +1398,45 @@ pir::Type CreateSelectedRowsTypeByDenseTensor(pir::Type dense_tensor_type) {
   }
 }
 
+void ResetParameterName(pir::Operation *op, const std::string &name) {
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  if (op->isa<pir::SetParameterOp>()) {
+    op->set_attribute("parameter_name", pir::StrAttribute::get(ctx, name));
+  }
+}
+
+std::map<int, int> GetOpInplaceInfo(const pir::Operation *op) {
+  std::map<int, int> inplace_info;
+  if (!op->HasTrait<paddle::dialect::InplaceTrait>()) {
+    return inplace_info;
+  }
+  pir::IrContext *ctx = pir::IrContext::Instance();
+  std::string op_name = op->name();
+  if (op->attributes().count("op_name")) {
+    op_name =
+        op->attributes().at("op_name").dyn_cast<pir::StrAttribute>().AsString();
+  }
+
+  pir::OpInfo op_info = ctx->GetRegisteredOpInfo(op_name);
+  paddle::dialect::OpYamlInfoParser yaml_parser(
+      op_info.GetInterfaceImpl<paddle::dialect::OpYamlInfoInterface>()
+          ->get_op_info_(),
+      paddle::dialect::IsLegacyOp(op_name));
+
+  for (size_t i = 0; i < op->num_results(); ++i) {
+    std::string value_name = yaml_parser.OutputNames()[i];
+    if (yaml_parser.HasInplace(value_name)) {
+      const std::string &inplace_name = yaml_parser.InplaceName(value_name);
+      inplace_info[i] = yaml_parser.InputName2Id().at(inplace_name);
+    }
+  }
+  return inplace_info;
+}
+
 void BindUtils(pybind11::module *m) {
   m->def("clone_program", CloneProgram);
+  m->def("get_op_inplace_info", GetOpInplaceInfo);
+  m->def("reset_parameter_name", ResetParameterName);
   m->def("split_program", SplitForwardBackward);
   m->def("append_set_parameters", AppendSetParameters);
   m->def("fake_op_result", FakeOpResult);
diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc
@@ -3627,6 +3627,7 @@ void SetValueInferMeta(const MetaTensor& x, MetaTensor* out) {
           "The rank of input should be less than 7, but received %d.",
           in_dims.size()));
   out->set_dims(in_dims);
+  out->set_dtype(x.dtype());
 }
 
 void ShapeInferMeta(const MetaTensor& input, MetaTensor* out) {
diff --git a/python/paddle/base/variable_index.py b/python/paddle/base/variable_index.py
@@ -89,15 +89,6 @@ def is_integer_or_scalar_tensor(ele):
     if type(ele) is int:
         return True
     elif isinstance(ele, Variable):
-        # NOTE(zoooo0820): For compatibility, if FLAGS_set_to_1d is set to True,
-        # 1-D tensor is still treated as a scalar, which means basic indexing.
-        # This will be removed in future.
-        if paddle.get_flags('FLAGS_set_to_1d')['FLAGS_set_to_1d']:
-            if len(ele.shape) == 1 and ele.shape[0] == 1:
-                warnings.warn(
-                    "1-D Tensor will be treat as advanced indexing in future version. Currently, 1-D Tensor means a scalar, not vector, and please modify it to 0-D Tensor. If advanced indexing is needed, please use `export FLAGS_set_to_1d=False` to set the flag."
-                )
-                return True
         if len(ele.shape) == 0 and ele.dtype != paddle.bool:
             return True
     elif isinstance(ele, paddle.pir.Value):
@@ -106,19 +97,6 @@ def is_integer_or_scalar_tensor(ele):
     return False
 
 
-def is_bool_tensor(ele):
-    from .framework import Variable
-
-    if isinstance(ele, Variable) and ele.dtype == paddle.bool:
-        return True
-    elif (
-        isinstance(ele, paddle.pir.Value)
-        and ele.dtype == paddle.base.libpaddle.BOOL
-    ):
-        return True
-    return False
-
-
 def deal_attrs(attrs, attr, attr_name, tensor_attr_name, inputs, infer_flags):
     from .framework import Variable
 
@@ -414,7 +392,7 @@ def _setitem_static(x, indices, values):
         values(Tensor|Number|Ndarray): values to be assigned to the x.
     """
     from . import in_dynamic_or_pir_mode
-    from .framework import Variable, default_main_program
+    from .framework import Variable, default_main_program, in_pir_mode
 
     if x.type == paddle.base.core.VarDesc.VarType.LOD_TENSOR_ARRAY:
         return _setitem_for_tensor_array(x, indices, values)
@@ -491,8 +469,19 @@ def _setitem_static(x, indices, values):
 
         # step3.1: Only basic indexing, use OP set_value to set value.
         if in_dynamic_or_pir_mode():
+            if in_pir_mode():
+                if isinstance(starts, (list, tuple)):
+                    if paddle.utils._contain_var(starts):
+                        starts = paddle.utils.get_int_tensor_list(starts)
+                if isinstance(ends, (list, tuple)):
+                    if paddle.utils._contain_var(ends):
+                        ends = paddle.utils.get_int_tensor_list(ends)
+                if isinstance(steps, (list, tuple)):
+                    if paddle.utils._contain_var(steps):
+                        ends = paddle.utils.get_int_tensor_list(steps)
+
             if value_tensor is None:
-                return paddle._C_ops.set_value_(
+                output = paddle._C_ops.set_value_(
                     x,
                     starts,
                     ends,
@@ -504,7 +493,7 @@ def _setitem_static(x, indices, values):
                     values,
                 )
             else:
-                return paddle._C_ops.set_value_with_tensor_(
+                output = paddle._C_ops.set_value_with_tensor_(
                     x,
                     value_tensor,
                     starts,
@@ -514,6 +503,14 @@ def _setitem_static(x, indices, values):
                     decrease_axes,
                     none_axes,
                 )
+            if in_pir_mode():
+                # map var to the new output, for dy2static
+                from paddle.jit.pir_dy2static.parameter_recorder import (
+                    _global_inplace_map,
+                )
+
+                _global_inplace_map.add(default_main_program(), x, output)
+            return output
         else:
             helper = paddle.base.layer_helper.LayerHelper(
                 'set_value', **locals()
@@ -585,8 +582,31 @@ def _setitem_static(x, indices, values):
         inputs["ValueTensor"] = transback_sub_tensor
 
         if in_dynamic_or_pir_mode():
-            x._bump_inplace_version()
-            output = x
+            if in_pir_mode():
+                if isinstance(starts, (list, tuple)):
+                    if paddle.utils._contain_var(starts):
+                        starts = paddle.utils.get_int_tensor_list(starts)
+                if isinstance(ends, (list, tuple)):
+                    if paddle.utils._contain_var(ends):
+                        ends = paddle.utils.get_int_tensor_list(ends)
+                if isinstance(steps, (list, tuple)):
+                    if paddle.utils._contain_var(steps):
+                        ends = paddle.utils.get_int_tensor_list(steps)
+            output = paddle._C_ops.set_value_with_tensor_(
+                x,
+                transback_sub_tensor,
+                starts,
+                ends,
+                steps,
+                axes,
+                decrease_axes,
+                none_axes,
+            )
+            from paddle.jit.pir_dy2static.parameter_recorder import (
+                _global_inplace_map,
+            )
+
+            _global_inplace_map.add(default_main_program(), x, output)
         else:
             helper = paddle.base.layer_helper.LayerHelper(
                 'set_value', **locals()
@@ -600,16 +620,15 @@ def _setitem_static(x, indices, values):
                 output = helper.create_variable_for_type_inference(
                     dtype=x.dtype
                 )
-        cur_block = default_main_program().current_block()
-        cur_block.append_op(
-            type="set_value",
-            inputs=inputs,
-            outputs={'Out': output},
-            attrs=attrs,
-            inplace_map={"Input": "Out"},
-        )
+            cur_block = default_main_program().current_block()
+            cur_block.append_op(
+                type="set_value",
+                inputs=inputs,
+                outputs={'Out': output},
+                attrs=attrs,
+                inplace_map={"Input": "Out"},
+            )
 
-        if not in_dynamic_or_pir_mode():
             # map var to the new output
             paddle.jit.api.ProgramTranslator.get_instance()._inplace_map.add(
                 cur_block.program, x.desc.id(), output
@@ -808,7 +827,10 @@ def parse_bool_and_broadcast_indices(indices):
     # deal with multiple Tensors and translating bool tensor to int tensor.
     # In static mode, bool-tensor cannot be broadcasted since its corressponding int tensor's shape cannot be infered.
     for i, indice in enumerate(indices):
-        if indice.dtype == paddle.bool:
+        if (
+            indice.dtype == paddle.bool
+            or indice.dtype == paddle.base.libpaddle.BOOL
+        ):
             indices[i] = paddle.nonzero(indice)[:, 0]
     if len(indices) > 1:
         indices = paddle.broadcast_tensors(indices)
diff --git a/python/paddle/jit/dy2static/convert_operators.py b/python/paddle/jit/dy2static/convert_operators.py
@@ -69,7 +69,7 @@ def convert_load(x):
                 _global_inplace_map,
             )
 
-            new_var = _global_inplace_map.get(cur_block.program, id(x))
+            new_var = _global_inplace_map.get(cur_block.program, x)
             if new_var is not None:
                 return new_var
 
diff --git a/python/paddle/jit/dy2static/pir_partial_program.py b/python/paddle/jit/dy2static/pir_partial_program.py
@@ -91,6 +91,29 @@ def __getitem__(self, item):
         return self._var_list[item]
 
 
+class UnionFindSet:
+    def __init__(self):
+        self.father = ValueDict()
+
+    def union(self, x, y):
+        # x -> y
+        father_x = self.find_root(x)
+        father_y = self.find_root(y)
+        if not (father_x.is_same(father_y)):
+            self.father[father_x] = father_y
+
+    def find_root(self, x):
+        if not self.father.__contains__(x):
+            self.father[x] = x
+        if self.father[x].is_same(x):
+            return x
+        self.father[x] = self.find_root(self.father[x])
+        return self.father[x]
+
+    def iter_elements(self):
+        yield from self.father.keys()
+
+
 class RunableProgram:
     """a pir program ready for run_program_op to run. constructed by 3 parts:
     - pir program (pir::Program)
@@ -120,6 +143,20 @@ def _get_value_name_map_from_program(cls, program):
                 ret[op.result(0)] = op.attrs()["parameter_name"]
         return ret
 
+    @classmethod
+    def _get_name_defining_op(cls, program, value):
+        for op in program.global_block().ops:
+            if op.name() == "pd_op.data":
+                if value.is_same(op.result(0)):
+                    return op
+            if op.name() == "builtin.set_parameter":
+                if value.is_same(op.operand(0).source()):
+                    return op
+            if op.name() == "builtin.parameter":
+                if value.is_same(op.result(0)):
+                    return op
+        return None
+
     @cached_property
     def get_name_value_map(self):
         return {v: k for k, v in self.get_value_name_map.items()}
@@ -280,8 +317,39 @@ def program_attr(self):
             else:
                 raise ValueError(f"Unknown program attr: {k}")
             value_program_attr[k] = values
+        self.deal_inplace_values(self.forward_program)
+        self.deal_inplace_values(self.backward_program)
         return value_program_attr
 
+    def deal_inplace_values(self, program):
+        # deal inplace op and modify program inplacely.
+        value2name = self._get_value_name_map_from_program(program)
+
+        def has_name(value):
+            if self._get_name_defining_op(program, value) is not None:
+                return True
+            return False
+
+        ufset = UnionFindSet()
+        for op in program.global_block().ops:
+            for out_idx, in_idx in paddle.core.pir.get_op_inplace_info(
+                op
+            ).items():
+                left = op.result(out_idx)
+                right = op.operand(in_idx).source()
+                if has_name(left):
+                    ufset.union(right, left)
+                else:
+                    ufset.union(left, right)
+
+        for value in ufset.iter_elements():
+            if has_name(ufset.find_root(value)):
+                name_defining_op = self._get_name_defining_op(program, value)
+                if name_defining_op:
+                    paddle.core.pir.reset_parameter_name(
+                        name_defining_op, value2name[ufset.find_root(value)]
+                    )
+
     @cached_property
     def program_name_attr(self):
         origin_attr = self._forward_backward_program[1]
diff --git a/python/paddle/jit/pir_dy2static/parameter_recorder.py b/python/paddle/jit/pir_dy2static/parameter_recorder.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import paddle
+from paddle.autograd.backward_utils import ValueDict
 
 from ..dy2static.program_translator import _program_hash, synchronized
 
@@ -65,28 +66,26 @@ def __init__(self):
         self.params_dict = {}
 
     @synchronized
-    def add(self, program, id, param):
-        """use the default_program as key, append param the parameter list."""
+    def add(self, program, origin_value, new_value):
         key = _program_hash(program)
         if key not in self.params_dict:
-            self.params_dict[key] = {}
+            self.params_dict[key] = ValueDict()
+        inplace_dict = self.params_dict[key]
+        inplace_dict[origin_value] = new_value
 
-        params = self.params_dict[key]
-        params[id] = param
-
-    def get(self, program, id):
-        params = self.params_dict.get(_program_hash(program))
-        if params is None:
+    def get(self, program, value):
+        inplace_dict = self.params_dict.get(_program_hash(program))
+        if inplace_dict is None:
             return None
-        if id not in params:
+        if value not in inplace_dict:
             return None
-        root_var = params[id]
+        root_var = inplace_dict[value]
         saved = []
-        while id(root_var) in params.keys():
+        while inplace_dict.__contains__(root_var):
             saved.append(root_var)
-            root_var = params[id(root_var)]
+            root_var = inplace_dict[root_var]
         for var in saved:
-            params[id(var)] = root_var
+            inplace_dict[var] = root_var
         return root_var
 
     def restore_checkpoint(self, checkpoint):
diff --git a/python/paddle/pir/math_op_patch.py b/python/paddle/pir/math_op_patch.py
@@ -644,8 +644,9 @@ def value_hash(self):
                 setattr(Value, magic_method, impl)
 
         # Handling __getitem__
-        from ..base.variable_index import _getitem_static
+        from ..base.variable_index import _getitem_static, _setitem_static
 
         Value.__getitem__ = _getitem_static
+        Value.__setitem__ = _setitem_static
 
         _already_patch_value = True
diff --git a/test/dygraph_to_static/test_jit_setitem.py b/test/dygraph_to_static/test_jit_setitem.py
diff --git a/test/indexing/test_getitem.py b/test/indexing/test_getitem.py
diff --git a/test/indexing/test_setitem.py b/test/indexing/test_setitem.py

Original file line number	Diff line number	Diff line change
`@@ -3627,6 +3627,7 @@ void SetValueInferMeta(const MetaTensor& x, MetaTensor* out) {`
`3627`	`3627`	`"The rank of input should be less than 7, but received %d.",`
`3628`	`3628`	`in_dims.size()));`
`3629`	`3629`	`out->set_dims(in_dims);`
	`3630`	`+ out->set_dtype(x.dtype());`
`3630`	`3631`	`}`
`3631`	`3632`
`3632`	`3633`	`void ShapeInferMeta(const MetaTensor& input, MetaTensor* out) {`
Original file line number	Diff line number	Diff line change
`@@ -69,7 +69,7 @@ def convert_load(x):`
`69`	`69`	`_global_inplace_map,`
`70`	`70`	`)`
`71`	`71`
`72`		`- new_var = _global_inplace_map.get(cur_block.program, id(x))`
	`72`	`+ new_var = _global_inplace_map.get(cur_block.program, x)`
`73`	`73`	`if new_var is not None:`
`74`	`74`	`return new_var`
`75`	`75`