Skip to content

Commit

Permalink
[GLCC]Part-1: Add pylayer op to Support @to_static (#56108)
Browse files Browse the repository at this point in the history
  • Loading branch information
@MarioLulab
MarioLulab authored Aug 21, 2023
1 parent c6757bd commit 7577a67
Show file tree
Hide file tree
Showing 11 changed files with 1,186 additions and 125 deletions.
125 changes: 0 additions & 125 deletions paddle/fluid/operators/controlflow/conditional_block_op.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,8 @@ limitations under the License. */
#include <array>

#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"

#ifdef PADDLE_WITH_DNNL
#include "paddle/fluid/platform/mkldnn_helper.h"
Expand Down Expand Up @@ -226,129 +224,6 @@ class ConditionalBlockGradOp : public ConditionalOp {
mutable std::shared_ptr<Executor> exec_{nullptr};
mutable std::unique_ptr<ExecutorPrepareContext> ctx_{nullptr};
mutable std::shared_ptr<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));
}
// Assign zero to the grad_vars that are in outside_grads but not in
// inside_grads
AssignZeroToParentScope(
place, parent_scope, assign_zero_inputs, assign_zero_outside_grads);
}

void AssignZeroToParentScope(
const platform::Place &place,
const framework::Scope &scope,
const std::vector<std::string> &inputs,
const std::vector<std::string> &outside_grads) const {
for (size_t i = 0; i < outside_grads.size(); ++i) {
const std::string &outside_grad_name = outside_grads[i];
const std::string &input_name = inputs[i];
VLOG(4) << "[assign zero]"
<< "input_name = " << input_name
<< ", outside_grad_name = " << outside_grad_name;
framework::Variable *input_var = scope.FindVar(input_name);
if (input_var == nullptr) {
continue;
}
framework::Variable *outside_var = scope.FindVar(outside_grad_name);
if (outside_var == nullptr) {
continue;
}

if (input_var->IsType<phi::DenseTensor>()) {
PADDLE_ENFORCE_EQ(
outside_var->IsType<phi::DenseTensor>(),
true,
platform::errors::InvalidArgument(
"Type of outside_var %s is NOT phi::DenseTensor, which "
"doesn't match input_var %s.",
outside_grad_name,
input_name));
AssignZeroToOutsideTensor(place,
scope,
input_var->Get<phi::DenseTensor>(),
outside_var->GetMutable<phi::DenseTensor>());
} else if (input_var->IsType<framework::LoDTensorArray>()) {
PADDLE_ENFORCE_EQ(outside_var->IsType<framework::LoDTensorArray>(),
true,
platform::errors::InvalidArgument(
"Type of outside_var %s is NOT LoDTensorArray, "
"which doesn't match input_var %s.",
outside_grad_name,
input_name));
const auto &input_tensors = input_var->Get<framework::LoDTensorArray>();
auto *outside_tensors =
outside_var->GetMutable<framework::LoDTensorArray>();
if (outside_tensors->empty()) {
outside_tensors->resize(input_tensors.size());
}
PADDLE_ENFORCE_EQ(input_tensors.size(),
outside_tensors->size(),
platform::errors::InvalidArgument(
"LoDTensorArray outside_var %s doen't have same "
"size as input_var %s.",
outside_grad_name,
input_name));
for (size_t j = 0; j < input_tensors.size(); ++j) {
AssignZeroToOutsideTensor(
place, scope, input_tensors[j], &((*outside_tensors)[j]));
}
} else {
// TODO(huihuangzheng): add support for SelectedRows
PADDLE_THROW(platform::errors::InvalidArgument(
"Conditional block grad op doesn't support non-phi::DenseTensor "
"output "
"now."));
}
}
}

void AssignZeroToOutsideTensor(const platform::Place &place,
const framework::Scope &cur_scope,
const phi::DenseTensor &input_tensor,
phi::DenseTensor *outside_tensor) const {
if (!input_tensor.IsInitialized() || input_tensor.numel() == 0) {
return;
}
VLOG(4) << "Assigning zero to " << outside_tensor;
outside_tensor->Resize(input_tensor.dims());
outside_tensor->mutable_data(place, input_tensor.dtype());
const platform::DeviceContext *dev_ctx =
platform::DeviceContextPool::Instance().Get(place);
phi::funcs::set_constant(*dev_ctx, outside_tensor, 0.0f);
outside_tensor->set_lod(input_tensor.lod());
}
};

template <class T>
Expand Down
120 changes: 120 additions & 0 deletions paddle/fluid/operators/controlflow/control_flow_op_helper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
#pragma once

#include "paddle/fluid/framework/new_executor/standalone_executor.h"
#include "paddle/fluid/operators/assign_op.h"
#include "paddle/phi/kernels/funcs/math_function.h"

namespace paddle {
namespace operators {
Expand Down Expand Up @@ -54,5 +56,123 @@ static void BuildScopeForControlFlowOp(
}
}

static void AssignZeroToOutsideTensor(const platform::Place &place,
const framework::Scope &cur_scope,
const phi::DenseTensor &input_tensor,
phi::DenseTensor *outside_tensor) {
if (!input_tensor.IsInitialized() || input_tensor.numel() == 0) {
return;
}
VLOG(4) << "Assigning zero to " << outside_tensor;
outside_tensor->Resize(input_tensor.dims());
outside_tensor->mutable_data(place, input_tensor.dtype());
const platform::DeviceContext *dev_ctx =
platform::DeviceContextPool::Instance().Get(place);
phi::funcs::set_constant(*dev_ctx, outside_tensor, 0.0f);
outside_tensor->set_lod(input_tensor.lod());
}

static void AssignZeroToParentScope(
const platform::Place &place,
const framework::Scope &scope,
const std::vector<std::string> &inputs,
const std::vector<std::string> &outside_grads) {
for (size_t i = 0; i < outside_grads.size(); ++i) {
const std::string &outside_grad_name = outside_grads[i];
const std::string &input_name = inputs[i];
VLOG(4) << "[assign zero]"
<< "input_name = " << input_name
<< ", outside_grad_name = " << outside_grad_name;
framework::Variable *input_var = scope.FindVar(input_name);
if (input_var == nullptr) {
continue;
}
framework::Variable *outside_var = scope.FindVar(outside_grad_name);
if (outside_var == nullptr) {
continue;
}

if (input_var->IsType<phi::DenseTensor>()) {
PADDLE_ENFORCE_EQ(
outside_var->IsType<phi::DenseTensor>(),
true,
platform::errors::InvalidArgument(
"Type of outside_var %s is NOT phi::DenseTensor, which "
"doesn't match input_var %s.",
outside_grad_name,
input_name));
AssignZeroToOutsideTensor(place,
scope,
input_var->Get<phi::DenseTensor>(),
outside_var->GetMutable<phi::DenseTensor>());
} else if (input_var->IsType<framework::LoDTensorArray>()) {
PADDLE_ENFORCE_EQ(outside_var->IsType<framework::LoDTensorArray>(),
true,
platform::errors::InvalidArgument(
"Type of outside_var %s is NOT LoDTensorArray, "
"which doesn't match input_var %s.",
outside_grad_name,
input_name));
const auto &input_tensors = input_var->Get<framework::LoDTensorArray>();
auto *outside_tensors =
outside_var->GetMutable<framework::LoDTensorArray>();
if (outside_tensors->empty()) {
outside_tensors->resize(input_tensors.size());
}
PADDLE_ENFORCE_EQ(input_tensors.size(),
outside_tensors->size(),
platform::errors::InvalidArgument(
"LoDTensorArray outside_var %s doen't have same "
"size as input_var %s.",
outside_grad_name,
input_name));
for (size_t j = 0; j < input_tensors.size(); ++j) {
AssignZeroToOutsideTensor(
place, scope, input_tensors[j], &((*outside_tensors)[j]));
}
} else {
// TODO(huihuangzheng): add support for SelectedRows
PADDLE_THROW(platform::errors::InvalidArgument(
"Conditional block grad op doesn't support non-phi::DenseTensor "
"output "
"now."));
}
}
}

static 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) {
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));
}
// Assign zero to the grad_vars that are in outside_grads but not in
// inside_grads
AssignZeroToParentScope(
place, parent_scope, assign_zero_inputs, assign_zero_outside_grads);
}
} // namespace operators
} // namespace paddle
Loading

0 comments on commit 7577a67

Please sign in to comment.