【pir】add Refresh stopgradient pass before backward (#58942)

xiaoguoguo626807 · web-flow · commit 92b4a12c56b3 · 2023-11-14T10:31:29.000+08:00
* add refresh stopgradint

* add refresh stopgradint

* modofy
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -23,7 +23,6 @@
 
 #include "paddle/fluid/framework/ir/pass.h"
 #include "paddle/fluid/pybind/pybind_variant_caster.h"
-#include "paddle/pir/core/builtin_op.h"
 
 #include "paddle/fluid/framework/program_desc.h"
 #include "paddle/fluid/ir_adaptor/translator/program_translator.h"
@@ -46,6 +45,7 @@
 #include "paddle/pir/core/attribute.h"
 #include "paddle/pir/core/block.h"
 #include "paddle/pir/core/builtin_attribute.h"
+#include "paddle/pir/core/builtin_op.h"
 #include "paddle/pir/core/program.h"
 #include "paddle/pir/core/type.h"
 #include "paddle/pir/core/value.h"
@@ -239,6 +239,19 @@ void BindProgram(py::module *m) {
           });
 }
 
+void RefreshOpStopgradients(Operation *op) {
+  if (op->num_operands() == 0 || op->isa<pir::GetParameterOp>() ||
+      op->isa<paddle::dialect::UniformOp>()) {
+    return;
+  } else if (op->isa<pir::SliceOp>()) {
+    op->dyn_cast<pir::SliceOp>().RefreshStopGradients();
+  } else if (op->isa<pir::SplitOp>()) {
+    op->dyn_cast<pir::SplitOp>().RefreshStopGradients();
+  } else {
+    RefreshStopGradientsDefaultly(op);
+  }
+}
+
 void BindBlock(py::module *m) {
   py::class_<Block> block(*m, "Block", R"DOC(
     In IR, a Block has a list of Operation and can represent a sub computational graph.
@@ -295,24 +308,30 @@ void BindBlock(py::module *m) {
               None
 
         )DOC")
-      .def("all_parameters", [](Block &self) -> py::list {
-        py::list param_list;
+      .def("all_parameters",
+           [](Block &self) -> py::list {
+             py::list param_list;
+             for (auto iter = self.begin(); iter != self.end(); iter++) {
+               auto op = *iter;
+               if (op->HasAttribute(kAttrIsPersisable)) {
+                 auto attrs = op->attribute(kAttrIsPersisable)
+                                  .dyn_cast<pir::ArrayAttribute>()
+                                  .AsVector();
+                 for (uint32_t i = 0; i < attrs.size(); i++) {
+                   bool is_persistable =
+                       attrs[i].dyn_cast<pir::BoolAttribute>().data();
+                   if (is_persistable) {
+                     param_list.append(op->result(i));
+                   }
+                 }
+               }
+             }
+             return param_list;
+           })
+      .def("refresh_stopgradient", [](Block &self) {
         for (auto iter = self.begin(); iter != self.end(); iter++) {
-          auto op = *iter;
-          if (op->HasAttribute(kAttrIsPersisable)) {
-            auto attrs = op->attribute(kAttrIsPersisable)
-                             .dyn_cast<pir::ArrayAttribute>()
-                             .AsVector();
-            for (uint32_t i = 0; i < attrs.size(); i++) {
-              bool is_persistable =
-                  attrs[i].dyn_cast<pir::BoolAttribute>().data();
-              if (is_persistable) {
-                param_list.append(op->result(i));
-              }
-            }
-          }
+          RefreshOpStopgradients(*iter);
         }
-        return param_list;
       });
 }
 
diff --git a/paddle/pir/core/builtin_op.cc b/paddle/pir/core/builtin_op.cc
@@ -50,6 +50,32 @@ void PassStopGradientsDefaultly(OperationArgument &argument) {  // NOLINT
       pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
 }
 
+void RefreshStopGradientsDefaultly(Operation *op) {
+  bool stop_gradient = true;
+  for (auto value : op->operands_source()) {
+    auto input = value.dyn_cast<OpResult>();
+    if (!input) continue;
+    auto *defining_op = input.owner();
+    bool input_stop_gradient = true;
+    if (defining_op->HasAttribute(kStopGradientAttrName)) {
+      auto attrs = defining_op->attribute(kStopGradientAttrName)
+                       .dyn_cast<pir::ArrayAttribute>()
+                       .AsVector();
+      input_stop_gradient =
+          attrs[input.index()].dyn_cast<pir::BoolAttribute>().data();
+    }
+    if (!input_stop_gradient) {
+      stop_gradient = false;
+      break;
+    }
+  }
+  std::vector<pir::Attribute> outs_stop_gradient(
+      op->results().size(),
+      pir::BoolAttribute::get(pir::IrContext::Instance(), stop_gradient));
+  op->set_attribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
+}
 Program *ModuleOp::program() {
   const AttributeMap &attr = this->attributes();
   auto iter = attr.find("program");
@@ -260,6 +286,27 @@ void SliceOp::PassStopGradients(OperationArgument &argument, int index) {
       pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
 }
 
+void SliceOp::RefreshStopGradients() {
+  std::vector<pir::Attribute> outs_stop_gradient(
+      1, pir::BoolAttribute::get(pir::IrContext::Instance(), true));
+  auto index = attribute("index").dyn_cast<pir::Int32Attribute>().data();
+  if (auto input = (*this)->operand_source(0).dyn_cast<pir::OpResult>()) {
+    auto *defining_op = input.owner();
+    if (defining_op && defining_op->isa<CombineOp>()) {
+      IR_ENFORCE(defining_op->HasAttribute(kStopGradientAttrName),
+                 "Required CombineOp must have attribute %s",
+                 kStopGradientAttrName);
+      auto attrs = defining_op->attribute(kStopGradientAttrName)
+                       .dyn_cast<pir::ArrayAttribute>()
+                       .AsVector();
+      outs_stop_gradient[0] = attrs[static_cast<int>(index)];
+    }
+  }
+  (*this)->set_attribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
+}
+
 void SliceOp::VerifySig() const {
   // inputs.size() == 1
   auto input_size = num_operands();
@@ -364,6 +411,53 @@ void SplitOp::PassStopGradients(OperationArgument &argument) {
       pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
 }
 
+void SplitOp::RefreshStopGradients() {
+  std::vector<bool> default_stop_gradients((*this)->num_results(), true);
+  if (auto input = (*this)->operand_source(0).dyn_cast<OpResult>()) {
+    auto *defining_op = input.owner();
+    if (defining_op && defining_op->isa<CombineOp>()) {
+      IR_ENFORCE((*this)->num_results(),
+                 defining_op->num_operands(),
+                 "Required SplitOp.output.size() == CombineOp.input.size(), "
+                 "but received %d != %d",
+                 (*this)->num_results(),
+                 defining_op->num_operands());
+      for (uint32_t i = 0; i < defining_op->num_operands(); ++i) {
+        auto value = defining_op->operand_source(i);
+        if (!value) continue;
+        auto *operand_defining_op = value.dyn_cast<OpResult>().owner();
+        if (operand_defining_op->HasAttribute(kStopGradientAttrName)) {
+          auto attrs = operand_defining_op->attribute(kStopGradientAttrName)
+                           .dyn_cast<pir::ArrayAttribute>()
+                           .AsVector();
+          default_stop_gradients[i] = attrs[value.dyn_cast<OpResult>().index()]
+                                          .dyn_cast<pir::BoolAttribute>()
+                                          .data();
+        }
+      }
+    } else if (defining_op &&
+               defining_op->HasAttribute(kStopGradientAttrName)) {
+      bool stop_gradient = defining_op->attribute(kStopGradientAttrName)
+                               .dyn_cast<pir::ArrayAttribute>()
+                               .AsVector()[0]
+                               .dyn_cast<pir::BoolAttribute>()
+                               .data();
+      default_stop_gradients.assign(default_stop_gradients.size(),
+                                    stop_gradient);
+    }
+  }
+
+  std::vector<pir::Attribute> outs_stop_gradient;
+  outs_stop_gradient.reserve(num_results());
+  for (auto stop_gradient : default_stop_gradients) {
+    outs_stop_gradient.push_back(
+        pir::BoolAttribute::get(pir::IrContext::Instance(), stop_gradient));
+  }
+  (*this)->set_attribute(
+      kStopGradientAttrName,
+      pir::ArrayAttribute::get(pir::IrContext::Instance(), outs_stop_gradient));
+}
+
 void SplitOp::VerifySig() const {
   // inputs.size() == 1
   IR_ENFORCE(num_operands() == 1u, "The size of inputs must be equal to 1.");
diff --git a/paddle/pir/core/builtin_op.h b/paddle/pir/core/builtin_op.h
@@ -145,6 +145,7 @@ class IR_API SliceOp : public pir::Op<SliceOp> {
 
   void VerifySig() const;
   pir::Value input() { return operand_source(0); }
+  void RefreshStopGradients();
 
  private:
   static void PassStopGradients(OperationArgument &argument,  // NOLINT
@@ -177,6 +178,7 @@ class IR_API SplitOp : public pir::Op<SplitOp> {
     }
     return res;
   }
+  void RefreshStopGradients();
 
  private:
   static void PassStopGradients(OperationArgument &argument);  // NOLINT
@@ -210,7 +212,7 @@ class IR_API ConstantOp : public Op<ConstantOp, ConstantLikeTrait> {
 };
 
 void PassStopGradientsDefaultly(OperationArgument &argument);  // NOLINT
-
+void RefreshStopGradientsDefaultly(Operation *Op);
 }  // namespace pir
 
 IR_EXPORT_DECLARE_EXPLICIT_TYPE_ID(pir::ModuleOp)
diff --git a/python/paddle/autograd/ir_backward.py b/python/paddle/autograd/ir_backward.py
@@ -631,7 +631,9 @@ def remove_op(block, op, state):
 
 def calc_gradient_helper(outputs, inputs, grad_outputs, no_grad_set):
     block = outputs[0].get_defining_op().get_parent_block()
+    block.refresh_stopgradient()
     state = State(block.program)
+
     # check all inputs and outputs in the same block
     check_all_puts(block, inputs, outputs)
     # update no_grad_set if some value stop_gradient=True
diff --git a/test/ir/pir/test_ir_backward.py b/test/ir/pir/test_ir_backward.py
@@ -249,5 +249,29 @@ def test_basic_network(self):
             input_grad = grad(res, x)
 
 
+class TestBackward_refrash_stopgradients(unittest.TestCase):
+    def test_refreash_stopgradients(self):
+        import numpy as np
+
+        program = paddle.pir.core.default_main_program()
+        with paddle.pir_utils.IrGuard(), paddle.pir.core.program_guard(program):
+            data1 = paddle.static.data('data1', [3, 4, 5], np.float32)
+            data2 = paddle.static.data('data2', [3, 4, 5], np.float32)
+            out = paddle.add_n([data1, data2])
+            data1_arr = np.random.uniform(-1, 1, data1.shape).astype(np.float32)
+            data2_arr = np.random.uniform(-1, 1, data2.shape).astype(np.float32)
+            self.assertEqual(
+                program.global_block().ops[3].result(0).stop_gradient, True
+            )
+
+            data1.stop_gradient = False
+            data2.stop_gradient = False
+
+            dout = grad(out, [data1, data2])
+            self.assertEqual(
+                program.global_block().ops[3].result(0).stop_gradient, False
+            )
+
+
 if __name__ == "__main__":
     unittest.main()
