【Hackathon 6th No.50】 support backward for pylayer op

luq.py

@@ -0,0 +1,44 @@
+from contextlib import contextmanager
+
+import paddle
+from paddle.autograd.py_layer import PyLayer
+
+
+class scaled_layer_1(PyLayer):
+    @staticmethod
+    def forward(ctx, x):
+        y = x * 3
+        return y
+
+    @staticmethod
+    def backward(ctx, dy):
+        dx = paddle.sin(dy)
+        return dx
+
+@contextmanager
+def enable_to_static_guard(flag: bool):
+    program_translator = paddle.jit.api.ProgramTranslator()
+    original_flag_value = program_translator.enable_to_static
+    program_translator.enable(flag)
+    try:
+        yield
+    finally:
+        program_translator.enable(original_flag_value)
+
+@paddle.jit.to_static(full_graph=True)
+def test_func(x):
+    y = scaled_layer_1.apply(x)
+    return y
+    # y.mean().backward()
+
+
+inp = paddle.randn([1, 32])
+inp.stop_gradient = False
+
+with enable_to_static_guard(True):
+    result = test_func(inp)
+    loss = result.mean()
+    loss.backward()
+
+print("ok...")
+    # print(paddle.static.default_main_program())

paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.cc

@@ -1,4 +1,4 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 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.
@@ -33,8 +33,8 @@
 #include "paddle/pir/include/core/value.h"
 
 #include "paddle/fluid/framework/new_executor/instruction/instruction_util.h"
-#include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
 
 #ifdef PADDLE_WITH_DNNL
 #include "paddle/fluid/platform/onednn_helper.h"
@@ -68,7 +68,7 @@ PyLayerInstruction::PyLayerInstruction(
   auto& fwd_block = pylayer_op.forward_block();
   std::unordered_map<pir::Value, std::vector<int>> inputs;
   GetInputIds(op, *value_exec_info, &inputs);
-  auto fwd_outside_inputs =
+  const auto fwd_outside_inputs =
       GetExternalInputs(&fwd_block, *value_exec_info, &inputs);
 
   // NOTE(chenxi67): the variable corresponding to container value if a

paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h

@@ -1,4 +1,4 @@
-// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2024 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.
@@ -17,10 +17,6 @@
 #include "paddle/fluid/framework/new_executor/instruction/instruction_base.h"
 #include "paddle/fluid/framework/new_executor/interpreter/execution_config.h"
 
-namespace ir {
-class Operation;
-}  // namespace ir
-
 namespace paddle {
 namespace framework {
 class Scope;

paddle/fluid/framework/new_executor/instruction/instruction_util.cc

@@ -433,7 +433,6 @@ void CopyBranchOutput(const std::vector<std::string>& var_names,
     } else if (inner_var->IsType<phi::TensorArray>()) {
       const auto& inner_array = inner_var->Get<phi::TensorArray>();
       auto* output_array = output_vars[i]->GetMutable<phi::TensorArray>();
-      // output_array->clear();
       *output_array = inner_array;
     } else {
       PADDLE_THROW(

paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc

@@ -30,6 +30,7 @@
 #include "paddle/fluid/pir/dialect/operator/interface/op_yaml_info.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"

paddle/fluid/framework/new_executor/pir_interpreter.cc

@@ -68,6 +68,7 @@
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
 #include "paddle/fluid/pir/dialect/operator/utils/utils.h"
 #include "paddle/pir/include/core/builtin_attribute.h"
 #include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"

paddle/fluid/pir/dialect/CMakeLists.txt

@@ -232,6 +232,7 @@ set(op_dialect_srcs
     ${op_dialect_srcs}
     ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/api_builder.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/control_flow_op.cc
+    ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/manual_pylayer_op.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/ir_meta_tensor.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/ir_selected_rows.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/operator/ir/ir_tensor.cc

paddle/fluid/pir/dialect/operator/ir/control_flow_op.cc

@@ -14,8 +14,8 @@
 #ifdef GET_OP_LIST
 #undef GET_OP_LIST
 paddle::dialect::IfOp, paddle::dialect::WhileOp, paddle::dialect::HasElementsOp,
-    paddle::dialect::PyLayerOp, paddle::dialect::AssertOp,
-    paddle::dialect::SelectInputOp, paddle::dialect::SelectOutputOp
+    paddle::dialect::AssertOp, paddle::dialect::SelectInputOp,
+    paddle::dialect::SelectOutputOp
 #else
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 
@@ -377,115 +377,6 @@ bool IfOp::InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis) {
   }
 }
 
-void PyLayerOp::Build(pir::Builder &builder,             // NOLINT
-                      pir::OperationArgument &argument,  // NOLINT
-                      pir::Value combined_inputs,
-                      std::vector<pir::Type> &&output_types) {
-  argument.AddInput(combined_inputs);
-  argument.output_types.swap(output_types);
-  argument.AddRegion().emplace_back();
-}
-
-void PyLayerOp::Build(pir::Builder &builder,             // NOLINT
-                      pir::OperationArgument &argument,  // NOLINT
-                      pir::Value combined_inputs,
-                      std::unique_ptr<pir::Block> &&fwd_block) {
-  VLOG(4) << "Start build PyLayerOp";
-  if (fwd_block && !fwd_block->empty() &&
-      fwd_block->back().isa<pir::YieldOp>()) {
-    auto &op = fwd_block->back();
-
-    std::vector<pir::Attribute> outs_stop_gradient;
-    for (size_t i = 0; i < op.num_operands(); ++i) {
-      argument.AddOutput(op.operand(i).type());
-      auto bool_attr = op.operand_source(i).attribute<pir::BoolAttribute>(
-          kStopGradientAttrName);
-      outs_stop_gradient.push_back(bool_attr ? bool_attr
-                                             : builder.bool_attr(false));
-    }
-
-    argument.AddAttribute(
-        kStopGradientAttrName,
-        pir::ArrayAttribute::get(builder.ir_context(), outs_stop_gradient));
-  }
-
-  argument.AddRegion().push_back(fwd_block.release());
-  argument.AddInput(combined_inputs);
-}
-
-pir::Block &PyLayerOp::forward_block() {
-  pir::Region &region = forward_region();
-  if (region.empty()) {
-    region.emplace_back();
-  }
-
-  return region.front();
-}
-
-void PyLayerOp::Print(pir::IrPrinter &printer) {
-  auto &os = printer.os;
-  auto op = operation();
-  printer.PrintOpResult(op);
-  os << " = pd_op.pylayer";
-  printer.PrintOpOperands(op);
-  os << " -> ";
-  printer.PrintOpReturnType(op);
-  os << "{";
-  for (auto &item : forward_block()) {
-    os << "\n  ";
-    printer.PrintOperation(&item);
-  }
-  os << "\n }";
-}
-
-void PyLayerOp::VerifySig() {
-  VLOG(4) << "Start Verifying inputs, outputs and attributes for: PyLayerOp.";
-  // NOTE(MarioLulab): do nothing.
-}
-
-void PyLayerOp::VerifyRegion() {
-  VLOG(4) << "Start Verifying sub regions for: PyLayerOp.";
-  VLOG(4) << "Start Verifying forward block.";
-  PADDLE_ENFORCE_EQ((*this)->region(0).size(),
-                    1u,
-                    phi::errors::PreconditionNotMet(
-                        "The size %d of forward_region must be 1.",
-                        (*this)->region(0).size()));
-  if ((*this)->num_results() != 0) {
-    auto &fwd_last_op = (*this)->region(0).front().back();
-    PADDLE_ENFORCE_EQ(true,
-                      fwd_last_op.isa<pir::YieldOp>(),
-                      phi::errors::PreconditionNotMet(
-                          "The last of forward block must be YieldOp"));
-    PADDLE_ENFORCE_EQ(
-        fwd_last_op.num_operands(),
-        (*this)->num_results(),
-        phi::errors::PreconditionNotMet(
-            "The size of last of forward block op's input must be "
-            "equal to PyLayerOp's outputs num."));
-  }
-}
-
-void PyLayerOp::UpdateOutput() {
-  PADDLE_ENFORCE_NOT_NULL(*this,
-                          paddle::platform::errors::InvalidArgument(
-                              "The pylayer_op in PyLayerOp used to update "
-                              "output can't be nullptr"));
-  auto block = parent();
-  PADDLE_ENFORCE_NOT_NULL(
-      block,
-      paddle::platform::errors::InvalidArgument(
-          "The parent block of pylayer_op which used to update "
-          "output can't be nullptr"));
-  pir::Block::Iterator iter = **this;
-  pir::Builder builder(ir_context(), false);
-  auto new_pylayer_op =
-      builder.Build<PyLayerOp>(combined_inputs(), forward_region().TakeBack());
-  block->Assign(iter, new_pylayer_op);
-  PyLayerOp::operator=(new_pylayer_op);
-  VerifyRegion();
-}
-
 void WhileOp::Build(pir::Builder &builder,             // NOLINT
                     pir::OperationArgument &argument,  // NOLINT
                     pir::Value cond,
@@ -1139,7 +1030,6 @@ IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::AssertOp)
-IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::PyLayerOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DEFINE_EXPLICIT_TYPE_ID(paddle::dialect::SelectOutputOp)
 

paddle/fluid/pir/dialect/operator/ir/control_flow_op.h

@@ -60,33 +60,6 @@ class IfOp : public pir::Op<IfOp, VjpInterface, InferSymbolicShapeInterface> {
   bool InferSymbolicShape(pir::ShapeConstraintIRAnalysis *shape_analysis);
 };
 
-class PyLayerOp : public pir::Op<PyLayerOp> {
- public:
-  using Op::Op;
-  static const char *name() { return "pd_op.pylayer"; }
-  static constexpr const char **attributes_name = nullptr;
-  static constexpr uint32_t attributes_num = 0;
-  static void Build(pir::Builder &builder,             // NOLINT
-                    pir::OperationArgument &argument,  // NOLINT
-                    pir::Value combined_inputs,
-                    std::vector<pir::Type> &&output_types);
-
-  static void Build(pir::Builder &builder,             // NOLINT
-                    pir::OperationArgument &argument,  // NOLINT
-                    pir::Value combined_inputs,
-                    std::unique_ptr<pir::Block> &&fwd_block);
-
-  pir::Value combined_inputs() { return operand_source(0); }
-  pir::Block &forward_block();
-  pir::Region &forward_region() { return (*this)->region(0); }
-
-  void Print(pir::IrPrinter &printer);  // NOLINT
-  void VerifySig();
-  void VerifyRegion();
-
-  void UpdateOutput();
-};
-
 ///
 /// \brief The WhileOp is an operation that iterates over a loop body based on a
 /// condition. It takes two inputs: cond_value and loop_vars. The output of the
@@ -224,6 +197,5 @@ IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::IfOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::WhileOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::HasElementsOp);
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::AssertOp);
-IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::PyLayerOp);
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectInputOp)
 IR_DECLARE_EXPLICIT_TYPE_ID(paddle::dialect::SelectOutputOp)