【pir】modify array_to_tensor , tensor_to_array kernel. (#60953)

xiaoguoguo626807 · web-flow · commit 0c4bb9e4f792 · 2024-01-19T10:13:09.000+08:00
* modify increment vjp and append_backward

* modify prune in appendbackward

* modfiy forward prune

* Update test/legacy_test/test_increment.py

* delete print

* tensor_to_array kernel modify

* delete stdcout
diff --git a/paddle/fluid/pir/dialect/operator/ir/manual_op.cc b/paddle/fluid/pir/dialect/operator/ir/manual_op.cc
@@ -3068,7 +3068,7 @@ OpInfoTuple TensorToArrayOp::GetOpInfo() {
       paddle::dialect::OpRunTimeInfo("TensorToArrayInferMeta",
                                      {"x", "axis", "use_stack"},
                                      "tensor_to_array",
-                                     {"x", "axis", "use_stack"},
+                                     {"x", "out_grad", "axis", "use_stack"},
                                      {"x"},
                                      {},
                                      {},
diff --git a/paddle/phi/kernels/array_grad_kernel.cc b/paddle/phi/kernels/array_grad_kernel.cc
@@ -0,0 +1,85 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/array_grad_kernel.h"
+
+#include "paddle/common/layout.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/concat_grad_kernel.h"
+#include "paddle/phi/kernels/stack_grad_kernel.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void TensorToArrayKernel(const Context& dev_ctx,
+                         const TensorArray& x,
+                         const DenseTensor& out_grad,
+                         int axis,
+                         bool use_stack,
+                         TensorArray* x_grad) {
+  std::vector<DenseTensor> tmp_inputs(x.size());
+  std::vector<const DenseTensor*> inputs;
+
+  std::vector<DenseTensor*> inputs_grad;
+  std::vector<DenseTensor> tmp_inputs_grad(x.size());
+
+  for (size_t i = 0; i < x.size(); i++) {
+    tmp_inputs[i].ShareDataWith(x[i]);
+    inputs.push_back(&tmp_inputs[i]);
+    inputs_grad.push_back(&tmp_inputs_grad[i]);
+    inputs_grad[i]->set_meta(x[i].meta());
+  }
+
+  if (use_stack) {
+    StackGradKernel<T, Context>(dev_ctx, out_grad, axis, inputs_grad);
+  } else {
+    ConcatGradKernel<T, Context>(dev_ctx, inputs, out_grad, axis, inputs_grad);
+  }
+
+  for (size_t i = 0; i < tmp_inputs_grad.size(); i++) {
+    inputs_grad[i] = nullptr;
+    x_grad->push_back(tmp_inputs_grad[i]);
+  }
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(tensor_to_array,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::TensorToArrayKernel,
+                   bool,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+PD_REGISTER_KERNEL(tensor_to_array,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::TensorToArrayKernel,
+                   bool,
+                   int,
+                   int64_t,
+                   float,
+                   double,
+                   phi::dtype::float16,
+                   phi::dtype::complex<float>,
+                   phi::dtype::complex<double>) {}
+#endif
diff --git a/paddle/phi/kernels/array_grad_kernel.h b/paddle/phi/kernels/array_grad_kernel.h
@@ -0,0 +1,31 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/tensor_array.h"
+
+namespace phi {
+
+template <typename T, typename Context>
+void TensorToArrayKernel(const Context& dev_ctx,
+                         const TensorArray& x,
+                         const DenseTensor& out_grad,
+                         int axis,
+                         bool use_stack,
+                         TensorArray* x_grad);
+}  // namespace phi
diff --git a/paddle/phi/kernels/array_kernel.cc b/paddle/phi/kernels/array_kernel.cc
@@ -17,10 +17,8 @@
 #include "paddle/common/layout.h"
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
-#include "paddle/phi/kernels/concat_grad_kernel.h"
 #include "paddle/phi/kernels/concat_kernel.h"
 #include "paddle/phi/kernels/full_kernel.h"
-#include "paddle/phi/kernels/stack_grad_kernel.h"
 #include "paddle/phi/kernels/stack_kernel.h"
 
 namespace phi {
@@ -99,18 +97,6 @@ void ArrayToTensorKernel(const Context& dev_ctx,
                                    "but the received is %d",
                                    n));
 
-  auto out_dims = x[0].dims();
-  size_t in_zero_dims_size = out_dims.size();
-  for (size_t i = 1; i < n; i++) {
-    for (size_t j = 0; j < in_zero_dims_size; j++) {
-      if (j == static_cast<size_t>(axis)) {
-        out_dims[axis] += x[i].dims()[static_cast<int>(j)];
-      }
-    }
-  }
-  auto vec = common::vectorize<int>(out_dims);
-  vec.insert(vec.begin() + axis, x.size());  // NOLINT
-  out->Resize(common::make_ddim(vec));
   std::vector<DenseTensor> tmp_inputs(x.size());
   std::vector<const DenseTensor*> inputs;
 
@@ -126,41 +112,29 @@ void ArrayToTensorKernel(const Context& dev_ctx,
   }
 
   if (use_stack) {
+    auto vec = common::vectorize<int>(x[0].dims());
+    vec.insert(vec.begin() + axis, x.size());  // NOLINT
+    out->Resize(common::make_ddim(vec));
     StackKernel<T, Context>(dev_ctx, inputs, axis, out);
   } else {
+    auto out_dims = x[0].dims();
+    size_t in_zero_dims_size = out_dims.size();
+    for (size_t i = 1; i < n; i++) {
+      for (size_t j = 0; j < in_zero_dims_size; j++) {
+        if (j == static_cast<size_t>(axis)) {
+          out_dims[axis] += x[i].dims()[static_cast<int>(j)];
+        }
+      }
+    }
+    auto vec = common::vectorize<int>(out_dims);
+    out->Resize(common::make_ddim(vec));
     ConcatKernel<T, Context>(dev_ctx, inputs, axis, out);
   }
 
   out_index->Resize(common::make_ddim({static_cast<int>(x.size())}));
   StackKernel<int, Context>(dev_ctx, indexs, 0, out_index);
 }
 
-template <typename T, typename Context>
-void TensorToArrayKernel(const Context& dev_ctx,
-                         const TensorArray& x,
-                         const DenseTensor& out_grad,
-                         int axis,
-                         bool use_stack,
-                         TensorArray* x_grad) {
-  std::vector<DenseTensor> tmp_inputs(x.size());
-  std::vector<const DenseTensor*> inputs;
-
-  std::vector<DenseTensor*> inputs_grad(x.size());
-
-  for (size_t i = 0; i < x.size(); i++) {
-    tmp_inputs[i].ShareDataWith(x[i]);
-    inputs.push_back(&tmp_inputs[i]);
-    x_grad->at(i).Resize(x[i].dims());
-    inputs_grad[i]->ShareDataWith(x_grad->at(i));
-  }
-
-  if (use_stack) {
-    StackGradKernel<T, Context>(dev_ctx, out_grad, axis, inputs_grad);
-  } else {
-    ConcatGradKernel<T, Context>(dev_ctx, inputs, out_grad, axis, inputs_grad);
-  }
-}
-
 }  // namespace phi
 PD_REGISTER_KERNEL(create_array,
                    CPU,
@@ -325,31 +299,3 @@ PD_REGISTER_KERNEL(array_to_tensor,
                    phi::dtype::complex<float>,
                    phi::dtype::complex<double>) {}
 #endif
-
-PD_REGISTER_KERNEL(tensor_to_array,
-                   CPU,
-                   ALL_LAYOUT,
-                   phi::TensorToArrayKernel,
-                   bool,
-                   int,
-                   int64_t,
-                   float,
-                   double,
-                   phi::dtype::float16,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}
-
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
-PD_REGISTER_KERNEL(tensor_to_array,
-                   GPU,
-                   ALL_LAYOUT,
-                   phi::TensorToArrayKernel,
-                   bool,
-                   int,
-                   int64_t,
-                   float,
-                   double,
-                   phi::dtype::float16,
-                   phi::dtype::complex<float>,
-                   phi::dtype::complex<double>) {}
-#endif
diff --git a/paddle/phi/kernels/array_kernel.h b/paddle/phi/kernels/array_kernel.h
@@ -57,11 +57,4 @@ void ArrayToTensorKernel(const Context& dev_ctx,
                          DenseTensor* out,
                          DenseTensor* out_index);
 
-template <typename T, typename Context>
-void TensorToArrayKernel(const Context& dev_ctx,
-                         const TensorArray& x,
-                         const DenseTensor& out_grad,
-                         int axis,
-                         bool use_stack,
-                         TensorArray* x_grad);
 }  // namespace phi
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -1033,9 +1033,22 @@ def remove_op(block, op, state):
                 )
 
 
+def all_stop_gradient_true(block):
+    for op in block.ops:
+        for value in op.results():
+            if value.stop_gradient is False:
+                return False
+    return True
+
+
 def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
     block = outputs[0].get_defining_op().get_parent_block()
     state = State(block)
+    if all_stop_gradient_true(block):
+        logging.warning(
+            "all op in block stop_grdient is True, no grad will be calculate"
+        )
+        return state.value_to_valuegrad
 
     total_ops = []
     if block.parent_block is not None:
diff --git a/python/paddle/tensor/array.py b/python/paddle/tensor/array.py
@@ -314,7 +314,7 @@ def create_array(dtype, initialized_list=None):
             dtype = paddle.base.framework.convert_np_dtype_to_dtype_(dtype)
         out = paddle._pir_ops.create_array(dtype)
         for val in array:
-            out = paddle._pir_ops.array_write_(out, val, array_length(out))
+            paddle._pir_ops.array_write_(out, val, array_length(out))
         return out
     else:
         helper = LayerHelper("array", **locals())
diff --git a/test/legacy_test/test_dynamic_rnn_stop_gradient.py b/test/legacy_test/test_dynamic_rnn_stop_gradient.py
@@ -42,10 +42,9 @@ def build_and_run_program(place, batch_size, beam_size, stop_gradient=False):
     cond = paddle.less_than(x=step_idx, y=max_len)
     while_op = paddle.static.nn.control_flow.While(cond)
     scores = paddle.tensor.array_write(x, step_idx)
-    scores.stop_gradient = False
     with while_op.block():
         bs = paddle.cast(paddle.shape(x)[0], "int64")
-        for _ in range(20):
+        for _ in range(2):
             bs = paddle.cast(bs, 'int64')
         bs.stop_gradient = stop_gradient
         batch_pos = paddle.expand(
@@ -67,29 +66,31 @@ def build_and_run_program(place, batch_size, beam_size, stop_gradient=False):
     data = np.random.random_integers(
         low=0, high=beam_size - 1, size=(batch_size, beam_size)
     ).astype("int64")
+    exe.run(paddle.static.default_startup_program())
     (loss_val,) = exe.run(feed={"indices": data}, fetch_list=[loss])
 
     return loss_val
 
 
 class TestDynRNNStopGradient(unittest.TestCase):
     def setUp(self):
-        self.batch_size = 20
-        self.beam_size = 64
+        self.batch_size = 2
+        self.beam_size = 2
 
     # @test_with_pir_api
     def run_main(self, place):
-        main_program = paddle.static.Program()
-        startup_program = paddle.static.Program()
-        with paddle.static.program_guard(main_program, startup_program):
-            with base.scope_guard(base.Scope()):
-                value1 = build_and_run_program(
-                    place, self.batch_size, self.beam_size, False
-                )
-                value2 = build_and_run_program(
-                    place, self.batch_size, self.beam_size, True
-                )
-                np.testing.assert_array_equal(value1, value2)
+        with paddle.pir_utils.IrGuard():
+            main_program = paddle.static.Program()
+            startup_program = paddle.static.Program()
+            with paddle.static.program_guard(main_program, startup_program):
+                with base.scope_guard(base.Scope()):
+                    value1 = build_and_run_program(
+                        place, self.batch_size, self.beam_size, False
+                    )
+                    value2 = build_and_run_program(
+                        place, self.batch_size, self.beam_size, True
+                    )
+                    np.testing.assert_array_equal(value1, value2)
 
     def test_check_main(self):
         places = [base.CPUPlace()]
diff --git a/test/legacy_test/test_tensor_array_to_tensor.py b/test/legacy_test/test_tensor_array_to_tensor.py
