[PIR] add vjp interface for while op. (#59935)

winter-wang · web-flow · commit c2e20ccc99a4 · 2023-12-14T20:01:52.000+08:00
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc
@@ -20,6 +20,8 @@ paddle::dialect::IfOp, paddle::dialect::WhileOp, paddle::dialect::HasElementsOp
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
+#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
+#include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/phi/core/enforce.h"
 #include "paddle/pir/core/builder.h"
 #include "paddle/pir/core/builtin_attribute.h"
@@ -289,6 +291,84 @@ void WhileOp::Print(pir::IrPrinter &printer) {
   os << "\n }";
 }
 
+std::vector<std::vector<pir::OpResult>> WhileOp::Vjp(
+    pir::Operation *op,
+    const std::vector<std::vector<pir::Value>> &inputs,
+    const std::vector<std::vector<pir::OpResult>> &outputs,
+    const std::vector<std::vector<pir::Value>> &out_grads,
+    const std::vector<std::vector<bool>> &stop_gradients) {
+  auto fwd_op = WhileOp::dyn_cast(op);
+  PADDLE_ENFORCE_NE(
+      fwd_op,
+      nullptr,
+      phi::errors::InvalidArgument("The input op used to called WhileOp::vjp "
+                                   "must be non-nullptr while_op"));
+  TuplePushOp push_op;
+  for (auto iter = fwd_op.body().rbegin(); iter != fwd_op.body().rend();
+       ++iter) {
+    if (iter->isa<TuplePushOp>()) {
+      push_op = iter->dyn_cast<TuplePushOp>();
+      PADDLE_ENFORCE_EQ(push_op.container().use_empty(),
+                        true,
+                        phi::errors::InvalidArgument(
+                            "The last container in foward while op must used "
+                            "empty while construct while_grad op"));
+      break;
+    }
+  }
+  PADDLE_ENFORCE_NE(push_op,
+                    nullptr,
+                    phi::errors::InvalidArgument(
+                        "The forward WhileOp must include TuplePushOp, denying "
+                        "that we can't construct a reverse loop condition."));
+
+  PADDLE_ENFORCE_GT(inputs.size(),
+                    outputs.size(),
+                    phi::errors::InvalidArgument(
+                        "while op's inputs' size should greater than "
+                        "outputs' size, Now the inputs's size is %d ."
+                        "the outputs size is %d.",
+                        inputs.size(),
+                        outputs.size()));
+  PADDLE_ENFORCE_EQ(stop_gradients[0][0],
+                    true,
+                    phi::errors::InvalidArgument(
+                        "The stop_gradient of condition input must be true."));
+
+  auto &builder = *ApiBuilder::Instance().GetBuilder();
+  auto cond_val = builder.Build<HasElementsOp>(push_op.container()).out();
+
+  std::vector<pir::Type> output_types;
+  std::vector<pir::Value> loop_vars;
+  size_t index = 0;
+
+  for (; index < outputs.size(); ++index) {
+    if (!stop_gradients[index + 1][0]) {
+      loop_vars.push_back(out_grads[index][0]);
+    }
+  }
+  for (++index; index < inputs.size(); ++index) {
+    if (!stop_gradients[index][0]) {
+      auto fwd_type = inputs[index][0].type().dyn_cast<DenseTensorType>();
+      PADDLE_ENFORCE_NE(
+          fwd_type,
+          pir::Type(),
+          phi::errors::InvalidArgument(
+              "The forward value type must be dense tensor type."));
+      auto shape = vectorize(fwd_type.dims());
+      auto dtype = TransToPhiDataType(fwd_type.dtype());
+      auto full_op = builder.Build<FullOp>(shape, 0.0, dtype, phi::CPUPlace());
+      loop_vars.push_back(full_op.out());
+    }
+  }
+  auto while_grad = builder.Build<WhileOp>(cond_val, loop_vars);
+
+  std::vector<std::vector<pir::OpResult>> res(inputs.size());
+  for (size_t i = 0, j = 0; i < inputs.size(); ++i) {
+    res[i].push_back(stop_gradients[i][0] ? nullptr : while_grad.result(j++));
+  }
+  return res;
+}
 std::vector<std::vector<pir::OpResult>> TuplePushOpVjpInterfaceModel::Vjp(
     pir::Operation *op,
     const std::vector<std::vector<pir::Value>> &inputs,
@@ -318,17 +398,17 @@ std::vector<std::vector<pir::OpResult>> TuplePushOpVjpInterfaceModel::Vjp(
 
 void HasElementsOp::Build(pir::Builder &builder,             // NOLINT
                           pir::OperationArgument &argument,  // NOLINT
-                          pir::Value stack) {
-  argument.AddInput(stack);
+                          pir::Value container) {
+  argument.AddInput(container);
   argument.AddOutput(
       DenseTensorType::get(builder.ir_context(), builder.bool_type(), {1}));
 }
 void HasElementsOp::VerifySig() {
   VLOG(4) << "Verifying inputs, outputs ,attributes for: HasElementsOp.";
   // Verify inputs:
   IR_ENFORCE(num_operands() == 1u, "The size of inputs must equal to 1.");
-  IR_ENFORCE(operand_source(0).type().isa<pir::StackType>(),
-             "The first input of cf.has_elements must be stack_type.");
+  IR_ENFORCE(operand_type(0).isa<pir::ContainerType>(),
+             "The first input of cf.has_elements must be container type.");
 
   // No attributes should be verify.
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h b/paddle/fluid/pir/dialect/operator/ir/control_flow_op.h
@@ -65,7 +65,7 @@ class IfOp : public pir::Op<IfOp, VjpInterface> {
 ///      cond, outputs = body(outputs)
 ///   }
 ///
-class WhileOp : public pir::Op<WhileOp> {
+class WhileOp : public pir::Op<WhileOp, VjpInterface> {
  public:
   using Op::Op;
   static const char *name() { return "pd_op.while"; }
@@ -81,6 +81,12 @@ class WhileOp : public pir::Op<WhileOp> {
   void Print(pir::IrPrinter &printer);  // NOLINT
   void VerifySig() {}
   void VerifyRegion() {}
+  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::OpResult>> &outputs,
+      const std::vector<std::vector<pir::Value>> &out_grads,
+      const std::vector<std::vector<bool>> &stop_gradients);
 };
 
 struct TuplePushOpVjpInterfaceModel : public VjpInterface::Concept {
@@ -114,7 +120,7 @@ class HasElementsOp : public pir::Op<HasElementsOp> {
 
   static void Build(pir::Builder &builder,             // NOLINT
                     pir::OperationArgument &argument,  // NOLINT
-                    pir::Value stack);
+                    pir::Value container);
   void VerifySig();
   pir::Value input() { return operand_source(0); }
   pir::Value out() { return result(0); }
diff --git a/paddle/fluid/pybind/pir.cc b/paddle/fluid/pybind/pir.cc
@@ -79,6 +79,7 @@ using paddle::dialect::DenseTensorArrayType;
 using paddle::dialect::DenseTensorType;
 using paddle::dialect::IfOp;
 using paddle::dialect::SelectedRowsType;
+using paddle::dialect::WhileOp;
 
 using pir::Attribute;
 using pir::Block;
@@ -499,7 +500,15 @@ void BindOperation(py::module *m) {
              self.ReplaceAllUsesWith(op_results);
            })
       .def("as_if_op",
-           [](Operation &self) { return PyIfOp(self.dyn_cast<IfOp>()); });
+           [](Operation &self) { return PyIfOp(self.dyn_cast<IfOp>()); })
+      .def("as_while_op", [](Operation &self) -> WhileOp {
+        auto while_op = self.dyn_cast<WhileOp>();
+        if (!while_op) {
+          PADDLE_THROW(phi::errors::InvalidArgument(
+              "Can't cast non-while type Operation to WhileOp."));
+        }
+        return while_op;
+      });
   py::class_<Operation::BlockContainer> block_container(
       *m, "Operation_BlockContainer", R"DOC(
     The Operation_BlockContainer only use to walk all blocks in the operation.
diff --git a/paddle/pir/core/op_base.h b/paddle/pir/core/op_base.h
@@ -62,6 +62,9 @@ class IR_API OpBase {
   Value operand_source(uint32_t index) const {
     return operation()->operand_source(index);
   }
+  Type operand_type(uint32_t index) const {
+    return operation()->operand_type(index);
+  }
 
   OpResult result(uint32_t index) const { return operation()->result(index); }
 
diff --git a/paddle/pir/core/operation.h b/paddle/pir/core/operation.h
@@ -100,6 +100,7 @@ class IR_API alignas(8) Operation final
   std::vector<OpOperand> operands();
   Value operand_source(uint32_t index) const;
   std::vector<Value> operands_source() const;
+  Type operand_type(uint32_t index) const { return operand(index).type(); }
 
   ///
   /// \brief op successor related public interfaces
diff --git a/paddle/pir/dialect/control_flow/ir/cf_type.h b/paddle/pir/dialect/control_flow/ir/cf_type.h
@@ -21,6 +21,7 @@
 namespace pir {
 
 class IR_API ContainerType : public Type {
+ public:
   using Type::Type;
   static bool classof(Type);
 };
diff --git a/test/ir/pir/test_while_api.py b/test/ir/pir/test_while_api.py
@@ -15,6 +15,11 @@
 import unittest
 
 import paddle
+from paddle.base.core import call_vjp, has_vjp
+from paddle.base.libpaddle.pir import (
+    build_pipe_for_block,
+    get_used_external_value,
+)
 
 paddle.enable_static()
 
@@ -51,6 +56,37 @@ def test_while_base(self):
         self.assertEqual(last_op.name(), "pd_op.while")
         self.assertEqual(len(out), 2)
 
+    def test_while_op_vjp_interface(self):
+        main_program = self.construct_program_with_while()
+        while_op = main_program.global_block().ops[-1]
+        self.assertEqual(while_op.name(), "pd_op.while")
+        build_pipe_for_block(while_op.as_while_op().body())
+        with paddle.pir.core.program_guard(main_program):
+            out_grad = paddle.full(shape=[6, 1], dtype='float32', fill_value=3)
+            # check vjp interface for while_op
+            while_input = [
+                [input] for input in get_used_external_value(while_op)
+            ]
+            self.assertEqual(len(while_input), 4)
+            while_input_stop_graditents = [[True], [False], [True], [True]]
+            while_output = [while_op.results()]
+            while_output_grad = [[out_grad, out_grad]]
+            self.assertEqual(has_vjp(while_op), True)
+            grad_outs = call_vjp(
+                while_op,
+                while_input,
+                while_output,
+                while_output_grad,
+                while_input_stop_graditents,
+            )
+
+            self.assertEqual(grad_outs[0][0], None)
+
+            while_grad_op = grad_outs[1][0].get_defining_op()
+            self.assertEqual(while_grad_op.name(), "pd_op.while")
+            while_grad_output = while_grad_op.results()
+            self.assertEqual(len(while_grad_output), 1)
+
 
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -62,6 +62,9 @@ class IR_API OpBase {`
`62`	`62`	`Value operand_source(uint32_t index) const {`
`63`	`63`	`return operation()->operand_source(index);`
`64`	`64`	`}`
	`65`	`+ Type operand_type(uint32_t index) const {`
	`66`	`+ return operation()->operand_type(index);`
	`67`	`+ }`
`65`	`68`
`66`	`69`	`OpResult result(uint32_t index) const { return operation()->result(index); }`
`67`	`70`