【pir】 add array_write and array_read grad case (#60360)

xiaoguoguo626807 · winter-wang · web-flow · commit 2515e7c5c1dc · 2023-12-27T11:05:01.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

* modify array_write grad info

---------

Co-authored-by: winter-wang &lt;1030748926@qq.com&gt;
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -1260,7 +1260,7 @@ OpInfoTuple ArrayReadOp::GetOpInfo() {
                   false,
                   false,
                   false,
-                  false),
+                  true),
       OpInputInfo(
           "i", "paddle::dialect::ScalarAttribute", false, false, true, false)};
 
@@ -1370,6 +1370,7 @@ void ArrayReadOp::Build(pir::Builder &builder,
       dense_out.lod());
   argument_outputs.push_back(out_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void ArrayReadOp::VerifySig() {
@@ -1428,7 +1429,7 @@ OpInfoTuple ArrayWrite_Op::GetOpInfo() {
                   false,
                   false),
       OpInputInfo(
-          "x", "paddle::dialect::DenseTensorType", false, false, false, false),
+          "x", "paddle::dialect::DenseTensorType", false, false, false, true),
       OpInputInfo(
           "i", "paddle::dialect::ScalarAttribute", false, false, true, false)};
 
@@ -1493,6 +1494,7 @@ void ArrayWrite_Op::Build(pir::Builder &builder,
       dense_out.layout());
   argument_outputs.push_back(out_type);
   argument.AddOutputs(argument_outputs.begin(), argument_outputs.end());
+  ::pir::PassStopGradientsDefaultly(argument);
 }
 
 void ArrayWrite_Op::VerifySig() {
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op_vjp.cc
@@ -256,8 +256,13 @@ std::vector<std::vector<pir::OpResult>> ArrayReadOp::Vjp(
           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]);
+
+  paddle::dialect::DenseTensorType outgrad_type =
+      out_grads[0][0].type().dyn_cast<paddle::dialect::DenseTensorType>();
+  pir::Value new_array = paddle::dialect::create_array(
+      paddle::dialect::TransToPhiDataType(outgrad_type.dtype()));
+  pir::OpResult tensor_res =
+      paddle::dialect::array_write_(new_array, out_grads[0][0], inputs_[1][0]);
 
   std::vector<std::vector<pir::OpResult>> res{{tensor_res}};
   if (stop_gradients[0][0]) {
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -273,6 +273,9 @@ std::string GetValueDataType(const pir::Value& value) {
   } else if (value.type().isa<paddle::dialect::SelectedRowsType>()) {
     return phi::DataTypeToString(dialect::TransToPhiDataType(
         value.type().dyn_cast<paddle::dialect::SelectedRowsType>().dtype()));
+  } else if (value.type().isa<DenseTensorArrayType>()) {
+    return phi::DataTypeToString(dialect::TransToPhiDataType(
+        value.type().dyn_cast<DenseTensorArrayType>().dtype()));
   } else {
     PADDLE_THROW(
         phi::errors::InvalidType("Currently, we can only get dtype for "
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -562,6 +562,9 @@ phi::DataType GetValueDtype(Value value) {
   } else if (value.type().isa<SelectedRowsType>()) {
     return paddle::dialect::TransToPhiDataType(
         value.type().dyn_cast<SelectedRowsType>().dtype());
+  } else if (value.type().isa<DenseTensorArrayType>()) {
+    return paddle::dialect::TransToPhiDataType(
+        value.type().dyn_cast<DenseTensorArrayType>().dtype());
   } else {
     PADDLE_THROW(phi::errors::InvalidArgument(
         "Currently, we can only get phi::DataType from DenseTensorType and "
diff --git a/python/paddle/tensor/array.py b/python/paddle/tensor/array.py
@@ -220,9 +220,7 @@ def array_write(x, i, array=None):
     elif in_pir_mode():
         check_variable_and_dtype(i, 'i', ['int64'], 'array_write')
         if not isinstance(x, paddle.pir.Value):
-            raise TypeError(
-                f"x should be pir.OpResult, but recevied {type(x)}."
-            )
+            raise TypeError(f"x should be pir.Value, but recevied {type(x)}.")
         if array is not None:
             if (
                 not isinstance(array, paddle.pir.Value)
@@ -304,7 +302,7 @@ def create_array(dtype, initialized_list=None):
     for val in array:
         if not isinstance(val, (Variable, paddle.pir.Value)):
             raise TypeError(
-                "All values in `initialized_list` should be Variable or pir.OpResult, but recevied {}.".format(
+                "All values in `initialized_list` should be Variable or pir.Value, but recevied {}.".format(
                     type(val)
                 )
             )
diff --git a/test/legacy_test/test_array_read_write_op.py b/test/legacy_test/test_array_read_write_op.py
@@ -22,6 +22,7 @@
 from paddle.base.backward import append_backward
 from paddle.base.executor import Executor
 from paddle.base.framework import default_main_program
+from paddle.pir_utils import test_with_pir_api
 
 
 def _test_read_write(x):
@@ -191,6 +192,56 @@ def test_array(self):
             fetched_out1, np.ones([1, 3], dtype="float32") * 6
         )
 
+    @test_with_pir_api
+    def test_array_backward(self):
+        np.random.seed(2013)
+        main_program = paddle.static.Program()
+        startup_program = paddle.static.Program()
+        with paddle.static.program_guard(main_program, startup_program):
+            d0 = paddle.static.data(name='d0', shape=[10], dtype='float32')
+            d0.stop_gradient = False
+            d0.persistable = True
+            i = paddle.zeros(shape=[1], dtype='int64')
+            mem_array = paddle.tensor.array_write(x=d0, i=i)
+            mem_array.stop_gradient = False
+            mem_array.persistable = True
+            out = paddle.tensor.array_read(array=mem_array, i=i)
+            mean = paddle.mean(out)
+            grad_list = append_backward(mean)
+
+            place = (
+                base.CUDAPlace(0)
+                if core.is_compiled_with_cuda()
+                else base.CPUPlace()
+            )
+            d = np.random.random(size=[10]).astype('float32')
+            exe = base.Executor(place)
+
+            if paddle.framework.in_pir_mode():
+                for p, g in grad_list:
+                    if p.is_same(d0):
+                        dd0 = g
+                    if p.is_same(mem_array):
+                        dmem_array = g
+                res = exe.run(
+                    main_program,
+                    feed={'d0': d},
+                    fetch_list=[mean, dd0],  # dmem_array
+                )
+                # pir not support fetch tensorarray
+            else:
+                res = exe.run(
+                    main_program,
+                    feed={'d0': d},
+                    fetch_list=[mean.name, d0.grad_name, mem_array.grad_name],
+                )
+                np.testing.assert_allclose(res[2], [[0.1] * 10], rtol=1e-05)
+
+            mean = 0.6097253
+            x_grad = [0.1] * 10
+            np.testing.assert_allclose(res[0], mean, rtol=1e-05)
+            np.testing.assert_allclose(res[1], x_grad, rtol=1e-05)
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/legacy_test/test_while_loop_op.py b/test/legacy_test/test_while_loop_op.py
@@ -418,6 +418,7 @@ def internal_body(j, x, mem_array):
             mem_array = paddle.tensor.array_write(x=init, i=i)
             data_array = paddle.tensor.array_write(x=d0, i=i)
             mem_array.stop_gradient = False
+            mem_array.persistable = True
             i = paddle.increment(i)
             paddle.tensor.array_write(d1, i, array=data_array)
             i = paddle.increment(i)
