[GLCC]Part-1: Add pylayer op to Support @to_static

paddle/fluid/framework/op_compatible_info.cc

@@ -117,6 +117,7 @@ void OpCompatibleMap::InitOpCompatibleMap() {
                                              OpCompatibleType::possible};
   op_compatible_map_["conditional_block_infer"] = {"1.6.0",
                                                    OpCompatibleType::possible};
+  op_compatible_map_["static_pylayer"] = {"1.6.0", OpCompatibleType::possible};
   op_compatible_map_["conv2d"] = {"1.6.0", OpCompatibleType::possible};
   op_compatible_map_["conv2d_transpose"] = {"1.6.0",
                                             OpCompatibleType::possible};

paddle/fluid/operators/controlflow/CMakeLists.txt

@@ -3,21 +3,25 @@ if(WITH_UNITY_BUILD)
   # Load Unity Build rules for operators in paddle/fluid/operators/controlflow.
   include(unity_build_rule.cmake)
 endif()
-register_operators(EXCLUDES conditional_block_op DEPS naive_executor
-                   standalone_executor)
+register_operators(EXCLUDES conditional_block_op static_pylayer_op DEPS
+                   naive_executor standalone_executor)
 
 cc_library(
   conditional_block_op
   SRCS conditional_block_op.cc
   DEPS executor standalone_executor)
+cc_library(
+  static_pylayer_op
+  SRCS static_pylayer_op.cc
+  DEPS executor standalone_executor)
 cc_library(
   op_variant
   SRCS op_variant.cc
   DEPS operator proto_desc)
 cc_library(
   conditional_block_op_helper
   SRCS conditional_block_op_helper.cc
-  DEPS op_variant operator conditional_block_op)
+  DEPS op_variant operator conditional_block_op static_pylayer_op)
 cc_library(
   recurrent_op_helper
   SRCS recurrent_op_helper.cc
@@ -28,9 +32,11 @@ cc_library(
   DEPS op_variant operator)
 
 if(WITH_UNITY_BUILD)
-  target_link_libraries(paddle_operators_controlflow_unity conditional_block_op)
+  target_link_libraries(paddle_operators_controlflow_unity conditional_block_op
+                        static_pylayer_op)
 else()
-  target_link_libraries(conditional_block_infer_op conditional_block_op)
+  target_link_libraries(conditional_block_infer_op conditional_block_op
+                        static_pylayer_op)
 endif()
 
 file(APPEND ${pybind_file}

paddle/fluid/operators/controlflow/conditional_block_op_helper.cc

@@ -27,6 +27,12 @@ namespace operators {
 
 static bool IsMatchedConditionalBlockOpAndConditionalBlockGradOp(
     const OpVariant &fwd_op, const OpVariant &bwd_op) {
+  // NOTE(MarioLulab): To avoide strip `static_pylayer_op.cc` when linker static
+  // link libstatic_pylayer_op.a, we should reference Variable or Function
+  // defined in libstatic_pylayer_op.a. Remind of that this is a `Temporary`
+  // method, we will use `op_library` to avoide linker strip.
+  std::string temp = std::string(StaticPyLayerOp::kInputs);
+
   return fwd_op.Outputs().at(ConditionalOp::kScope) ==
          bwd_op.Inputs().at(ConditionalOp::kScope);
 }

paddle/fluid/operators/controlflow/conditional_block_op_helper.h

@@ -20,6 +20,7 @@
 #include "paddle/fluid/framework/operator.h"
 #include "paddle/fluid/operators/controlflow/conditional_block_op.h"
 #include "paddle/fluid/operators/controlflow/op_variant.h"
+#include "paddle/fluid/operators/controlflow/static_pylayer_op.h"
 #include "paddle/fluid/string/string_helper.h"
 
 namespace paddle {

paddle/fluid/operators/controlflow/static_pylayer_op.cc

@@ -0,0 +1,279 @@
+// 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
+//
+//     http://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/fluid/operators/controlflow/static_pylayer_op.h"
+
+#include "paddle/fluid/framework/new_executor/standalone_executor.h"
+#include "paddle/fluid/operators/assign_op.h"
+#include "paddle/fluid/operators/controlflow/control_flow_op_helper.h"
+#include "paddle/phi/core/flags.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace paddle {
+namespace operators {
+
+const char StaticPyLayerOp::kInputs[] = "Input";
+const char StaticPyLayerOp::kOutputs[] = "Out";
+const char StaticPyLayerOp::kScope[] = "Scope";
+const char StaticPyLayerOp::kSkipEagerDeletionVars[] =
+    "skip_eager_deletion_vars";
+
+class StaticPyLayerForwardOp : public StaticPyLayerOp {
+ public:
+  StaticPyLayerForwardOp(const std::string &type,
+                         const framework::VariableNameMap &inputs,
+                         const framework::VariableNameMap &outputs,
+                         const framework::AttributeMap &attrs)
+      : StaticPyLayerOp(type, inputs, outputs, attrs) {}
+
+ private:
+  void RunImpl(const framework::Scope &scope,
+               const platform::Place &dev_place) const {
+    auto *scope_var = scope.FindVar(Output(kScope));
+    PADDLE_ENFORCE_NOT_NULL(
+        scope_var,
+        platform::errors::PreconditionNotMet(
+            "Expect Scope variable to be set in static_pylayer_op, but "
+            "got a null Scope variable. Please set the Scope variable."));
+
+    auto *scopes = scope_var->GetMutable<std::vector<framework::Scope *>>();
+    scopes->resize(1);
+    scopes->front() = &scope.NewScope();
+
+    auto &cur_scope = *scopes->front();
+    auto *block = Attr<framework::BlockDesc *>("forward_block");
+    VLOG(3) << "StaticPyLayer forward_block block.idx = " << block->ID()
+            << ", scope = " << &cur_scope;
+
+    auto &skip_vars = Attr<std::vector<std::string>>(kSkipEagerDeletionVars);
+
+    LOG_FIRST_N(INFO, 1)
+        << "[ControlFlow][StaticPyLayer] New Executor is Running.";
+
+    if (!core_ || !platform::is_same_place(core_->GetPlace(), dev_place)) {
+      VLOG(10) << "[interpreterCore cache]" << core_.get();
+      VLOG_IF(10, core_) << platform::is_same_place(core_->GetPlace(),
+                                                    dev_place);
+
+      framework::interpreter::ExecutionConfig execution_config;
+      execution_config.create_local_scope = false;
+      execution_config.used_for_control_flow_op = true;
+      execution_config.skip_gc_vars =
+          std::set<std::string>(skip_vars.begin(), skip_vars.end());
+
+      core_.reset(new framework::InterpreterCore(
+          dev_place, *block, &cur_scope, execution_config));
+      VLOG(10) << "[interpreterCore] created:" << core_;
+    } else {
+      // NOTE: Borrowed from
+      // `paddle/fluid/operators/controlflow/control_flow_op_helper.h`
+      // TODO(MarioLulab): Add StaticPyLayer Helper ?
+      BuildScopeForControlFlowOp(*core_, *block, &cur_scope);
+      core_->reset_scope(&cur_scope);
+    }
+
+    core_->Run({}, false);
+  }
+
+ private:
+  mutable std::shared_ptr<framework::InterpreterCore> core_{nullptr};
+};
+
+class StaticPyLayerForwardInferShape : public framework::InferShapeBase {
+ public:
+  void operator()(framework::InferShapeContext *context) const override {
+    // TODO(MarioLulab): do nothing.
+  }
+};
+
+class StaticPyLayerBackwardOp : public StaticPyLayerOp {
+ public:
+  StaticPyLayerBackwardOp(const std::string &type,
+                          const framework::VariableNameMap &inputs,
+                          const framework::VariableNameMap &outputs,
+                          const framework::AttributeMap &attrs)
+      : StaticPyLayerOp(type, inputs, outputs, attrs) {}
+
+ private:
+  void RunImpl(const framework::Scope &scope,
+               const platform::Place &dev_place) const override {
+    const auto &inputs = Inputs(StaticPyLayerOp::kInputs);
+    const auto &outside_grads =
+        Outputs(framework::GradVarName(StaticPyLayerOp::kInputs));
+    std::vector<std::string> inside_grads;
+    inside_grads.reserve(inputs.size());
+    for (auto &in : inputs) {
+      inside_grads.emplace_back(framework::GradVarName(in));
+    }
+
+    auto *scope_var = scope.FindVar(Input(StaticPyLayerOp::kScope));
+    PADDLE_ENFORCE_NOT_NULL(
+        scope_var,
+        platform::errors::PreconditionNotMet(
+            "Expect Scope variable to be set in static_pylayer_op, but "
+            "got a null Scope variable. Please set the Scope variable."));
+    auto &scopes = scope_var->Get<std::vector<framework::Scope *>>();
+    PADDLE_ENFORCE_GT(
+        scopes.size(),
+        0,
+        platform::errors::InvalidArgument(
+            "Expect Scope variable contains at least 1 scope, but got: %d",
+            scopes.size()));
+    framework::Scope &cur_scope = *(scopes[0]);
+
+    // auto *block = Attr<framework::BlockDesc *>("sub_block");
+    auto *block = Attr<framework::BlockDesc *>("forward_block");
+    VLOG(3) << "Static PyLayer backward block.idx = " << block->ID()
+            << ", scope = " << &cur_scope;
+
+    LOG_FIRST_N(INFO, 1)
+        << "[ControlFlow][StaticPyLayerBackwardOp] New Executor is Running.";
+    if (!core_ || !platform::is_same_place(core_->GetPlace(), dev_place)) {
+      VLOG(10) << "[interpreterCore cache]" << core_.get();
+      VLOG_IF(10, core_) << platform::is_same_place(core_->GetPlace(),
+                                                    dev_place);
+
+      framework::interpreter::ExecutionConfig execution_config;
+      execution_config.create_local_scope = false;
+      execution_config.used_for_control_flow_op = true;
+      execution_config.skip_gc_vars =
+          std::set<std::string>(inside_grads.begin(), inside_grads.end());
+
+      core_.reset(new framework::InterpreterCore(
+          dev_place, *block, &cur_scope, execution_config));
+      VLOG(10) << "[interpreterCore] created:" << core_;
+    } else {
+      BuildScopeForControlFlowOp(*core_, *block, &cur_scope);
+      core_->reset_scope(&cur_scope);
+    }
+    core_->Run({}, false);
+
+    // NOTE: It's neccessary. The reason of associating `inside_grads` and
+    // `outside_grads` at runtime `RunImpl` instead of `assgin` op at block is
+    // that the Var name of grad_op's outputs may be changed in the
+    // `append_backward` function (e.g. `_addup_repetitive_outputs_`).
+    AssignLocalGradientToParentScope(
+        dev_place, cur_scope, scope, inside_grads, outside_grads, inputs);
+
+    // Release the cur_scope, otherwise memory leakage occurs.
+    scope.DeleteScope(&cur_scope);
+    return;
+  }
+
+ private:
+  mutable std::shared_ptr<framework::InterpreterCore> core_{nullptr};
+
+ private:
+  void AssignLocalGradientToParentScope(
+      const platform::Place &place,
+      const framework::Scope &cur_scope,
+      const framework::Scope &parent_scope,
+      const std::vector<std::string> &inside_grads,
+      const std::vector<std::string> &outside_grads,
+      const std::vector<std::string> &inputs) const {
+    std::vector<std::string> assign_zero_outside_grads;
+    std::vector<std::string> assign_zero_inputs;
+    for (size_t i = 0; i < outside_grads.size(); ++i) {
+      const std::string &outside_grad_name = outside_grads[i];
+      const std::string &inside_grad_name = inside_grads[i];
+      VLOG(4) << "[assign local]"
+              << "inside_grad_name = " << inside_grad_name
+              << ", outside_grad_name = " << outside_grad_name;
+      framework::Variable *outside_var =
+          parent_scope.FindVar(outside_grad_name);
+      if (outside_var == nullptr) {
+        continue;
+      }
+      framework::Variable *inside_var =
+          cur_scope.FindLocalVar(inside_grad_name);
+      if (inside_var == nullptr) {
+        assign_zero_outside_grads.emplace_back(outside_grad_name);
+        assign_zero_inputs.emplace_back(inputs[i]);
+        continue;
+      }
+      platform::DeviceContext *dev_ctx =
+          platform::DeviceContextPool::Instance().Get(place);
+      framework::VisitVarType(*inside_var,
+                              AssignFunctor(outside_var, *dev_ctx));
+    }
+  }
+};
+
+class StaticPyLayerBackwardInferShape : public framework::InferShapeBase {
+ public:
+  void operator()(framework::InferShapeContext *context) const override {
+    if (context->HasInputs(StaticPyLayerOp::kInputs) &&
+        context->HasOutputs(framework::GradVarName(StaticPyLayerOp::kInputs))) {
+      context->SetOutputsDim(framework::GradVarName(StaticPyLayerOp::kInputs),
+                             context->GetInputsDim(StaticPyLayerOp::kInputs));
+    }
+  }
+};
+
+class StaticPyLayerBackwardInferVarType : public framework::VarTypeInference {
+ public:
+  void operator()(framework::InferVarTypeContext *ctx) const override {
+    auto forward_input_size = ctx->InputSize(StaticPyLayerOp::kInputs);
+    auto backward_output_size =
+        ctx->OutputSize(framework::GradVarName(StaticPyLayerOp::kInputs));
+    PADDLE_ENFORCE_EQ(forward_input_size,
+                      backward_output_size,
+                      platform::errors::InvalidArgument(
+                          "input_size and output_size should be equal for "
+                          "static_pylayer_grad_op."));
+    for (size_t i = 0; i < backward_output_size; ++i) {
+      ctx->SyncTypeAndDataType(StaticPyLayerOp::kInputs,
+                               framework::GradVarName(StaticPyLayerOp::kInputs),
+                               i);
+    }
+  }
+};
+
+template <typename T>
+class StaticPyLayerBackwardMaker : public framework::SingleGradOpMaker<T> {
+ public:
+  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
+
+ protected:
+  void Apply(GradOpPtr<T> grad_op) const override {
+    grad_op->SetType("static_pylayer_grad");
+    grad_op->SetInput(StaticPyLayerOp::kInputs,
+                      this->Input(StaticPyLayerOp::kInputs));
+    grad_op->SetInput(framework::GradVarName(StaticPyLayerOp::kOutputs),
+                      this->OutputGrad(StaticPyLayerOp::kOutputs));
+    grad_op->SetInput(StaticPyLayerOp::kScope,
+                      this->Output(StaticPyLayerOp::kScope));
+
+    auto fwd_inputs = this->InputGrad(StaticPyLayerOp::kInputs, false);
+    grad_op->SetOutput(framework::GradVarName(StaticPyLayerOp::kInputs),
+                       fwd_inputs);
+    grad_op->SetBlockAttr("forward_block",
+                          PADDLE_GET_CONST(framework::BlockDesc *,
+                                           this->GetAttr("backward_block")));
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OPERATOR(static_pylayer,
+                  ops::StaticPyLayerForwardOp,
+                  ops::StaticPyLayerForwardInferShape,
+                  ops::StaticPyLayerForwardOpProtoMaker,
+                  ops::StaticPyLayerBackwardMaker<paddle::framework::OpDesc>);
+REGISTER_OPERATOR(static_pylayer_grad,
+                  ops::StaticPyLayerBackwardOp,
+                  ops::StaticPyLayerBackwardInferShape,
+                  ops::StaticPyLayerBackwardInferVarType);