【pir】add array_write/read vjp and api (#60338)

xiaoguoguo626807 · winter-wang · web-flow · commit 5937fb04eb5f · 2023-12-26T16:52:23.000+08:00
* optimize backward

* [PIR] add vjp interface for while op

* [PIR] fix ci error.

* modify while stopgradient

* merge

* modify while grad bug

* modify while grad op

* modify

* increment vp

* [PIR] add get_used_external_value interface for block.

* while case

* delete print

* delete print

* Update python/paddle/autograd/ir_backward.py

* [PIR] add unit_test for get_used_external_value

* modify while_loop

* code_style

* modofy ci bug

* modify while api

* modify ci

* modify array

* Update python/paddle/autograd/ir_backward.py

* Update test/legacy_test/test_cond.py

* update

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/op_generator/op_build_gen.py b/paddle/fluid/pir/dialect/op_generator/op_build_gen.py
@@ -386,7 +386,7 @@ def GenBuildOutputs(
 """
 
     CREATE_INTARRAY_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  phi::IntArray {name};
-  if ({name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullIntArrayOp>()) {{
+  if ({name}_.dyn_cast<pir::OpResult>() && {name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullIntArrayOp>()) {{
     {name} = std::move(phi::IntArray(paddle::dialect::GetInt64Vector(
                           {name}_.dyn_cast<pir::OpResult>().owner()
                           ->dyn_cast<paddle::dialect::FullIntArrayOp>()
@@ -408,7 +408,7 @@ def GenBuildOutputs(
   }}\n"""
 
     CREATE_VECTOR_INT_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  std::vector<int64_t> {name};
-  if ({name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullIntArrayOp>()) {{
+  if ({name}_.dyn_cast<pir::OpResult>() && {name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullIntArrayOp>()) {{
     {name} = paddle::dialect::GetInt64Vector(
                     {name}_.dyn_cast<pir::OpResult>().owner()
                     ->dyn_cast<paddle::dialect::FullIntArrayOp>()
@@ -428,7 +428,7 @@ def GenBuildOutputs(
   }}\n"""
 
     CREATE_SCALAR_MUTABLE_ATTRIBUE_WITH_UNKONW_DATA_TEMPLATE = """  phi::Scalar {name};
-  if ({name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullOp>()) {{
+  if ({name}_.dyn_cast<pir::OpResult>() && {name}_.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullOp>()) {{
     {name} = std::move(phi::Scalar({name}_.dyn_cast<pir::OpResult>().owner()
                                   ->dyn_cast<paddle::dialect::FullOp>()
                                   .attribute("value")
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -1340,7 +1340,8 @@ void ArrayReadOp::Build(pir::Builder &builder,
   paddle::dialect::IrMetaTensor meta_array(&dense_array);
 
   phi::Scalar i_scalar;
-  if (i.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullOp>()) {
+  if (i.dyn_cast<pir::OpResult>() &&
+      i.dyn_cast<pir::OpResult>().owner()->isa<paddle::dialect::FullOp>()) {
     i_scalar =
         std::move(phi::Scalar(i.dyn_cast<pir::OpResult>()
                                   .owner()
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.h b/paddle/fluid/pir/dialect/operator/ir/manual_op.h
@@ -215,8 +215,10 @@ class ArrayLengthOp
   static void InferMeta(phi::InferMetaContext *infer_meta);
 };
 
-class ArrayReadOp
-    : public pir::Op<ArrayReadOp, OpYamlInfoInterface, InferMetaInterface> {
+class ArrayReadOp : public pir::Op<ArrayReadOp,
+                                   OpYamlInfoInterface,
+                                   paddle::dialect::VjpInterface,
+                                   InferMetaInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.array_read"; }
@@ -236,10 +238,17 @@ class ArrayReadOp
   pir::Value i() { return operand_source(1); }
   pir::OpResult out() { return result(0); }
   static void InferMeta(phi::InferMetaContext *infer_meta);
+  static std::vector<std::vector<pir::OpResult>> Vjp(
+      pir::Operation *op,
+      const std::vector<std::vector<pir::Value>> &inputs_,
+      const std::vector<std::vector<pir::Value>> &outputs,
+      const std::vector<std::vector<pir::Value>> &out_grads,
+      const std::vector<std::vector<bool>> &stop_gradients);
 };
 
 class ArrayWrite_Op : public pir::Op<ArrayWrite_Op,
                                      OpYamlInfoInterface,
+                                     paddle::dialect::VjpInterface,
                                      InferMetaInterface,
                                      InplaceTrait> {
  public:
@@ -259,6 +268,12 @@ class ArrayWrite_Op : public pir::Op<ArrayWrite_Op,
   pir::Value i() { return operand_source(2); }
   pir::OpResult out() { return result(0); }
   static void InferMeta(phi::InferMetaContext *infer_meta);
+  static std::vector<std::vector<pir::OpResult>> Vjp(
+      pir::Operation *op,
+      const std::vector<std::vector<pir::Value>> &inputs_,
+      const std::vector<std::vector<pir::Value>> &outputs,
+      const std::vector<std::vector<pir::Value>> &out_grads,
+      const std::vector<std::vector<bool>> &stop_gradients);
 };
 
 class ArrayToTensorOp
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
@@ -170,7 +170,7 @@ std::vector<std::vector<pir::OpResult>> Increment_Op::Vjp(
       inputs_.size(),
       1,
       platform::errors::InvalidArgument(
-          "Increment_ op's inputs size should be 2, but now is %d.",
+          "Increment_ op's inputs size should be 1, but now is %d.",
           inputs_.size()));
   PADDLE_ENFORCE_EQ(
       outputs.size(),
@@ -192,5 +192,78 @@ std::vector<std::vector<pir::OpResult>> Increment_Op::Vjp(
   return res;
 }
 
+std::vector<std::vector<pir::OpResult>> ArrayWrite_Op::Vjp(
+    pir::Operation* op,
+    const std::vector<std::vector<pir::Value>>& inputs_,
+    const std::vector<std::vector<pir::Value>>& outputs,
+    const std::vector<std::vector<pir::Value>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      3,
+      platform::errors::InvalidArgument(
+          "ArrayWrite_ op's inputs size should be 3, but now is %d.",
+          inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "ArrayWrite_ op's outputs size should be 1, but now is %d.",
+          outputs.size()));
+
+  PADDLE_ENFORCE_EQ(
+      out_grads.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "ArrayWrite_ op's outputs size should be 1, but now is %d.",
+          outputs.size()));
+
+  VLOG(6) << "Vjp prepare call  ArrayWrite_'s vjp inteface";
+  pir::OpResult tensor_res =
+      paddle::dialect::array_read(out_grads[0][0], inputs_[2][0]);
+
+  std::vector<std::vector<pir::OpResult>> res{{tensor_res}};
+  if (stop_gradients[0][0]) {
+    res = {{}};
+  }
+  return res;
+}
+
+std::vector<std::vector<pir::OpResult>> ArrayReadOp::Vjp(
+    pir::Operation* op,
+    const std::vector<std::vector<pir::Value>>& inputs_,
+    const std::vector<std::vector<pir::Value>>& outputs,
+    const std::vector<std::vector<pir::Value>>& out_grads,
+    const std::vector<std::vector<bool>>& stop_gradients) {
+  PADDLE_ENFORCE_EQ(
+      inputs_.size(),
+      2,
+      platform::errors::InvalidArgument(
+          "Array_read op's inputs size should be 2, but now is %d.",
+          inputs_.size()));
+  PADDLE_ENFORCE_EQ(
+      outputs.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "Array_read op's outputs size should be 1, but now is %d.",
+          outputs.size()));
+
+  PADDLE_ENFORCE_EQ(
+      out_grads.size(),
+      1,
+      platform::errors::InvalidArgument(
+          "Array_read op's outputs size should be 1, but now is %d.",
+          outputs.size()));
+
+  VLOG(6) << "Vjp prepare call  Array_read's vjp inteface";
+  pir::OpResult tensor_res = paddle::dialect::array_write_(
+      inputs_[0][0], out_grads[0][0], inputs_[1][0]);
+
+  std::vector<std::vector<pir::OpResult>> res{{tensor_res}};
+  if (stop_gradients[0][0]) {
+    res = {{}};
+  }
+  return res;
+}
 }  // namespace dialect
 }  // namespace paddle
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -188,6 +188,45 @@ bool PyObject_CheckIROpResult(PyObject* obj) {
   return PyObject_TypeCheck(obj, g_ir_opresult_pytype);
 }
 
+bool PyObject_CheckIRValue(PyObject* obj) {
+  return PyObject_TypeCheck(obj, g_ir_value_pytype);
+}
+
+bool PyObject_CheckIRVectorOfValue(PyObject* obj) {
+  if (PyList_Check(obj)) {
+    Py_ssize_t len = PyList_Size(obj);
+    PyObject* item = nullptr;
+    // if obj is [], parse it as std::vector<scalar>
+    if (len == 0) {
+      return false;
+    }
+    for (Py_ssize_t i = 0; i < len; i++) {
+      item = PyList_GetItem(obj, i);
+      if (!PyObject_CheckIRValue(item)) {
+        return false;
+      }
+    }
+    return true;
+  } else if (PyTuple_Check(obj)) {
+    Py_ssize_t len = PyTuple_Size(obj);
+    PyObject* item = nullptr;
+    if (len == 0) {
+      return false;
+    }
+    for (Py_ssize_t i = 0; i < len; i++) {
+      item = PyTuple_GetItem(obj, i);
+      if (!PyObject_CheckIRValue(item)) {
+        return false;
+      }
+    }
+    return true;
+  } else if (PyObject_TypeCheck(obj, g_ir_value_pytype)) {
+    return true;
+  } else {
+    return false;
+  }
+}
+
 bool PyObject_CheckIRVectorOfOpResult(PyObject* obj) {
   if (PyList_Check(obj)) {
     Py_ssize_t len = PyList_Size(obj);
diff --git a/paddle/fluid/pybind/eager_utils.h b/paddle/fluid/pybind/eager_utils.h
@@ -72,7 +72,9 @@ bool PyObject_CheckLongOrConvertToLong(PyObject** obj);
 bool PyObject_CheckFloatOrConvertToFloat(PyObject** obj);
 bool PyObject_CheckStr(PyObject* obj);
 bool PyObject_CheckIROpResult(PyObject* obj);
+bool PyObject_CheckIRValue(PyObject* obj);
 bool PyObject_CheckIRVectorOfOpResult(PyObject* obj);
+bool PyObject_CheckIRVectorOfValue(PyObject* obj);
 bool CastPyArg2AttrBoolean(PyObject* obj, ssize_t arg_pos);
 int CastPyArg2AttrInt(PyObject* obj, ssize_t arg_pos);
 int64_t CastPyArg2AttrLong(PyObject* obj, ssize_t arg_pos);
diff --git a/paddle/fluid/pybind/manual_static_op_function.h b/paddle/fluid/pybind/manual_static_op_function.h
@@ -181,7 +181,7 @@ static PyObject *static_api_array_read(PyObject *self,
 
     PyObject *i_obj = PyTuple_GET_ITEM(args, 1);
     pir::Value i;
-    if (PyObject_CheckIROpResult(i_obj)) {
+    if (PyObject_CheckIRValue(i_obj)) {
       i = CastPyArg2Value(i_obj, "array_read", 1);
     } else {
       int64_t i_tmp = CastPyArg2Int(i_obj, "array_read", 1);
@@ -215,7 +215,7 @@ static PyObject *static_api_array_write_(PyObject *self,
     auto x = CastPyArg2Value(x_obj, "array_write_", 1);
     PyObject *i_obj = PyTuple_GET_ITEM(args, 2);
     pir::Value i;
-    if (PyObject_CheckIROpResult(i_obj)) {
+    if (PyObject_CheckIRValue(i_obj)) {
       i = CastPyArg2Value(i_obj, "array_write_", 2);
     } else {
       int64_t i_tmp = CastPyArg2Int(i_obj, "array_write_", 2);
diff --git a/python/paddle/tensor/array.py b/python/paddle/tensor/array.py
@@ -56,7 +56,7 @@ def array_length(array):
         return len(array)
     elif in_pir_mode():
         if (
-            not isinstance(array, paddle.pir.OpResult)
+            not isinstance(array, paddle.pir.Value)
             or not array.is_dense_tensor_array_type()
         ):
             raise TypeError(
@@ -137,7 +137,7 @@ def array_read(array, i):
         return array[i]
     elif in_pir_mode():
         if (
-            not isinstance(array, paddle.pir.OpResult)
+            not isinstance(array, paddle.pir.Value)
             or not array.is_dense_tensor_array_type()
         ):
             raise TypeError(
@@ -219,13 +219,13 @@ def array_write(x, i, array=None):
         return array
     elif in_pir_mode():
         check_variable_and_dtype(i, 'i', ['int64'], 'array_write')
-        if not isinstance(x, paddle.pir.OpResult):
+        if not isinstance(x, paddle.pir.Value):
             raise TypeError(
                 f"x should be pir.OpResult, but recevied {type(x)}."
             )
         if array is not None:
             if (
-                not isinstance(array, paddle.pir.OpResult)
+                not isinstance(array, paddle.pir.Value)
                 or not array.is_dense_tensor_array_type()
             ):
                 raise TypeError("array should be tensor array vairable")
@@ -302,7 +302,7 @@ def create_array(dtype, initialized_list=None):
 
     # NOTE: Only support plain list like [x, y,...], not support nested list in static graph mode.
     for val in array:
-        if not isinstance(val, (Variable, paddle.pir.OpResult)):
+        if not isinstance(val, (Variable, paddle.pir.Value)):
             raise TypeError(
                 "All values in `initialized_list` should be Variable or pir.OpResult, but recevied {}.".format(
                     type(val)
diff --git a/test/ir/pir/test_while_api.py b/test/ir/pir/test_while_api.py
@@ -62,16 +62,13 @@ def test_while_base(self):
     def test_get_used_external_value(self):
         main_program = paddle.static.Program()
         with paddle.pir.core.program_guard(main_program):
-            print(main_program)
             i = paddle.full(shape=[1], fill_value=0)
-            print(main_program)
             x = paddle.full(shape=[1], fill_value=10)
             y = paddle.full(shape=[1], fill_value=5)
             # i, x = paddle.static.nn.while_loop(cond, body, [i, ten])
             paddle.static.nn.while_loop(
                 lambda p, q: p < q, lambda p, q: [p + y, q + i], [i, x]
             )
-            print(main_program)
         while_op = main_program.global_block().ops[-1]
         self.assertEqual(while_op.name(), "pd_op.while")
         body_block = while_op.as_while_op().body()
@@ -175,7 +172,6 @@ def test_backward(self):
                 out,
                 [i, j],
             )
-
             self.assertEqual(
                 grad_outs[0].get_defining_op().name(), "pd_op.while"
             )
@@ -189,6 +185,16 @@ def test_backward(self):
                 "cf.has_elements",
             )
 
+            self.assertEqual(
+                main_program.global_block()
+                .ops[-1]
+                .as_while_op()
+                .body()
+                .ops[-3]
+                .name(),
+                "pd_op.add_grad",
+            )
+
     def test_backward_with_loop_var_same_to_extral_var(self):
         main_program = paddle.static.Program()
         with paddle.pir.core.program_guard(main_program):
diff --git a/test/legacy_test/test_while_loop_op.py b/test/legacy_test/test_while_loop_op.py
@@ -374,7 +374,7 @@ def body(i, x):
 
 class TestApiWhileLoop_NestedWithBackwardAndLoDTensorArray(unittest.TestCase):
     # TODO(zhangbo): Support while grad exe for pir
-
+    # @test_with_pir_api
     def test_nested_net_with_backward_and_lodtensor(self):
         def external_cond(i, j, x, mem_array):
             return paddle.less_than(i, array_len)
@@ -411,6 +411,7 @@ def internal_body(j, x, mem_array):
             d2 = paddle.static.data(name='d2', shape=[10], dtype='float32')
             x = paddle.static.data(name='x', shape=[10], dtype='float32')
             x.stop_gradient = False
+            x.persistable = True
             i = paddle.zeros(shape=[1], dtype='int64')
             i.stop_gradient = True
             init = paddle.zeros(shape=[10], dtype='float32')
@@ -436,10 +437,9 @@ def internal_body(j, x, mem_array):
                 external_cond, external_body, [i, j, x, mem_array]
             )
 
-            sum_result = paddle.tensor.array_read(array=mem_array, i=j)
+            sum_result = paddle.tensor.array_read(array=out[3], i=j)
             mean = paddle.mean(sum_result)
-            append_backward(mean)
-
+            grad_list = append_backward(mean)
             place = (
                 base.CUDAPlace(0)
                 if core.is_compiled_with_cuda()
@@ -453,11 +453,21 @@ def internal_body(j, x, mem_array):
             feed_x = np.ones(10).astype('float32')
             data_sum = d[0] + d[1] + d[2] + 3 * feed_x
             x_grad = [0.3] * 10
-            res = exe.run(
-                main_program,
-                feed={'d0': d[0], 'd1': d[1], 'd2': d[2], 'x': feed_x},
-                fetch_list=[sum_result.name, x.grad_name],
-            )
+            if paddle.framework.in_pir_mode():
+                for p, g in grad_list:
+                    if p.is_same(x):
+                        dx = g
+                res = exe.run(
+                    main_program,
+                    feed={'d0': d[0], 'd1': d[1], 'd2': d[2], 'x': feed_x},
+                    fetch_list=[sum_result, dx],
+                )
+            else:
+                res = exe.run(
+                    main_program,
+                    feed={'d0': d[0], 'd1': d[1], 'd2': d[2], 'x': feed_x},
+                    fetch_list=[sum_result.name, x.grad_name],
+                )
             np.testing.assert_allclose(res[0], data_sum, rtol=1e-05)
             np.testing.assert_allclose(res[1], x_grad, rtol=1e-05)
 
