Skip to content

Commit

Permalink
[PIR] add python api for pylayer op (#60359)
Browse files Browse the repository at this point in the history
* add for code reading

* add pylayer op def and api

* add api and construct program success

* add pylayer fwd execute test case

* add translation pass

* add ins partial

* add pylayer instruction

* forward success

* update testcases

* code polish

* fix
  • Loading branch information
MarioLulab authored Feb 21, 2024
1 parent 85e37db commit 6037071
Show file tree
Hide file tree
Showing 19 changed files with 683 additions and 46 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -198,27 +198,6 @@ IfInstruction::~IfInstruction() {
}
}

void IfInstruction::CopyBranchOutput(const std::vector<std::string>& var_names,
const PirInterpreter* inter) {
for (size_t i = 0; i < var_names.size(); ++i) {
auto* inner_var = inter->InnerScope()->GetVar(var_names[i]);

if (inner_var->IsType<phi::DenseTensor>()) {
output_vars_[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
inner_var->Get<phi::DenseTensor>());

} 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(
phi::errors::Unimplemented("unsupported type %d", inner_var->Type()));
}
}
}

void IfInstruction::Run() {
bool cond = true;
if (cond_var_->IsType<phi::DenseTensor>()) {
Expand Down Expand Up @@ -257,7 +236,8 @@ void IfInstruction::Run() {
paddle::platform::DontClearMKLDNNCache(true_branch_inter_->GetPlace());
#endif
true_branch_inter_->Run({}, false);
CopyBranchOutput(true_branch_outputs_, true_branch_inter_);
CopyBranchOutput(
true_branch_outputs_, output_vars_, true_branch_inter_->InnerScope());
} else {
#ifdef PADDLE_WITH_DNNL
// Executor on being destroyed clears oneDNN cache and resets
Expand All @@ -266,7 +246,8 @@ void IfInstruction::Run() {
paddle::platform::DontClearMKLDNNCache(false_branch_inter_->GetPlace());
#endif
false_branch_inter_->Run({}, false);
CopyBranchOutput(false_branch_outputs_, false_branch_inter_);
CopyBranchOutput(
false_branch_outputs_, output_vars_, false_branch_inter_->InnerScope());
}
// copy output
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -49,9 +49,6 @@ class IfInstruction : public InstructionBase {
PirInterpreter* FalseBranchInterpreter() const { return false_branch_inter_; }

private:
void CopyBranchOutput(const std::vector<std::string>& var_names,
const PirInterpreter* inter);

::pir::Operation* op_;

std::string cond_name_{"if_instruction"};
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,173 @@
// 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/framework/new_executor/instruction/control_flow/pylayer_instruction.h"

#include "paddle/fluid/framework/new_executor/interpreter/interpreter_util.h"
#include "paddle/fluid/framework/new_executor/interpreter/stream_analyzer.h"
#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
#include "paddle/fluid/framework/new_executor/pir_interpreter.h"
#include "paddle/fluid/framework/scope.h"
#include "paddle/fluid/pir/dialect/operator/interface/infermeta.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
#include "paddle/fluid/pir/dialect/operator/utils/op_yaml_info_parser.h"
#include "paddle/fluid/platform/collective_helper.h"
#include "paddle/fluid/platform/device_context.h"
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/core/meta_tensor.h"
#include "paddle/phi/core/type_defs.h"

#include "paddle/pir/include/core/builtin_attribute.h"
#include "paddle/pir/include/core/operation.h"
#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"

#ifdef PADDLE_WITH_DNNL
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif

namespace paddle {
namespace framework {

PyLayerInstruction::PyLayerInstruction(
size_t id,
const platform::Place& place,
pir::Operation* op,
ValueExecutionInfo* value_exec_info,
interpreter::ExecutionConfig execution_config)
: InstructionBase(id, place) {
PADDLE_ENFORCE(op->isa<paddle::dialect::PyLayerOp>(),
phi::errors::PreconditionNotMet(
"Cond instruction only support pylayer op"));
auto pylayer_op = op->dyn_cast<paddle::dialect::PyLayerOp>();
op_ = op;

SetKernelType(AnalyseOpFuncType(op, place));
VLOG(6) << "finish process analyse kernel type";

for (size_t i = 0; i < pylayer_op.num_results(); ++i) {
output_vars_.push_back(value_exec_info->GetScope()->GetVar(
value_exec_info->GetValue2VarName().at(pylayer_op.result(i))));
}
VLOG(6) << "finish process output_vars";

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 =
GetExternalInputs(&fwd_block, *value_exec_info, &inputs);

// NOTE(chenxi67): the variable corresponding to container value if a
// <VariableRefArray> Type. It will recursively get the ID of internal
// variables when use GetValueId() method. However, the copy_var pushed into
// the tuple does not have a corresponding ID, and will insert a -1. Here we
// remove the value of -1.
for (auto& item : inputs) {
auto& var_vec = item.second;
for (auto it = var_vec.begin(); it != var_vec.end();) {
if (*it == -1) {
it = var_vec.erase(it);
} else {
++it;
}
}
}
SetInputs(inputs);

std::unordered_map<pir::Value, std::vector<int>> outputs;
for (size_t i = 0; i < op->num_results(); i++) {
pir::Value value = op->result(i);
if (value && value.type()) {
PADDLE_ENFORCE_EQ(
value_exec_info->HasValue(value),
true,
phi::errors::PreconditionNotMet(
"output should in name map, [%d] 'th output of [%s] op",
i,
"pylayer op"));
outputs.emplace(value, GetValueIds(value, *value_exec_info));
}
}

for (auto& item : outputs) {
auto& var_vec = item.second;
for (auto it = var_vec.begin(); it != var_vec.end();) {
if (*it == -1) {
it = var_vec.erase(it);
} else {
++it;
}
}
}

SetOutputs(outputs);
VLOG(6) << "finish process inputs outputs index";

Scope* fwd_scope = &(value_exec_info->GetScope()->NewScope());
auto skip_gc_vars = execution_config.skip_gc_vars;
execution_config.skip_gc_vars.clear();
execution_config.create_local_scope = true;
fwd_inter_ = new PirInterpreter(place,
{},
&fwd_block,
fwd_scope,
value_exec_info->NewChild(fwd_scope),
execution_config);

std::set<std::string> fwd_skip_gc_names_set;
for (auto value : GetYiedOpInputs(&fwd_block)) {
fwd_outputs_.push_back(fwd_inter_->GetNameByValue(value));
fwd_skip_gc_names_.push_back(fwd_inter_->GetNameByValue(value));
fwd_skip_gc_names_set.insert(fwd_inter_->GetNameByValue(value));
}

// NOTE(zhangbo): According to the concept of control flow, child scopes
// should not control the lifecycle of parent scope variables.
for (auto value : fwd_outside_inputs) {
fwd_skip_gc_names_.push_back(fwd_inter_->GetNameByValue(value));
fwd_skip_gc_names_set.insert(fwd_inter_->GetNameByValue(value));
}
for (const auto& var_name : skip_gc_vars) {
fwd_skip_gc_names_.push_back(var_name);
fwd_skip_gc_names_set.insert(var_name);
}

fwd_inter_->SetSkipGcVars(fwd_skip_gc_names_set);
VLOG(6) << "finish process forward block interpreter";
}

PyLayerInstruction::~PyLayerInstruction() {
if (fwd_inter_ != nullptr) {
delete fwd_inter_;
}
}

void PyLayerInstruction::Run() {
VLOG(6) << "start pylayer forward block interpreter";

#ifdef PADDLE_WITH_DNNL
// Executor on being destroyed clears oneDNN cache and resets
// registered model data layout. This is unwanted for nested
// Executors (executors declared inside control ops)
paddle::platform::DontClearMKLDNNCache(fwd_inter_->GetPlace());
#endif
fwd_inter_->Run({}, false);
CopyBranchOutput(fwd_outputs_, output_vars_, fwd_inter_->InnerScope());
}

} // namespace framework
} // namespace paddle
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
// 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.

#pragma once

#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;
class Value;
class PirInterpreter;
class ValueExecutionInfo;

class PyLayerInstruction : public InstructionBase {
public:
PyLayerInstruction(size_t id,
const platform::Place& place,
::pir::Operation* op,
ValueExecutionInfo* value_exe_info,
interpreter::ExecutionConfig execution_config);

~PyLayerInstruction();

void Run() override;

const std::string& Name() const override { return name_; }

::pir::Operation* Operation() const override { return op_; }

PirInterpreter* ForwardInterpreter() const { return fwd_inter_; }

private:
::pir::Operation* op_;

std::string name_{"pylayer_instruction"};

std::vector<Variable*> output_vars_;

PirInterpreter* fwd_inter_ = nullptr;

std::vector<std::string> fwd_outputs_;

std::vector<std::string> fwd_skip_gc_names_;
};

} // namespace framework
} // namespace paddle
Original file line number Diff line number Diff line change
Expand Up @@ -418,5 +418,27 @@ bool GetCondData(const phi::DenseTensor& cond) {
return cpu_cond->data<bool>()[0];
}

void CopyBranchOutput(const std::vector<std::string>& var_names,
const std::vector<Variable*>& output_vars,
Scope* inner_scope) {
for (size_t i = 0; i < var_names.size(); ++i) {
auto* inner_var = inner_scope->GetVar(var_names[i]);

if (inner_var->IsType<phi::DenseTensor>()) {
output_vars[i]->GetMutable<phi::DenseTensor>()->ShareDataWith(
inner_var->Get<phi::DenseTensor>());

} 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(
phi::errors::Unimplemented("unsupported type %d", inner_var->Type()));
}
}
}

} // namespace framework
} // namespace paddle
Original file line number Diff line number Diff line change
Expand Up @@ -66,5 +66,10 @@ void InsertInplacedExternalInputsToOuts(
std::unordered_map<pir::Value, std::vector<int>>* outputs);

bool GetCondData(const phi::DenseTensor& cond);

void CopyBranchOutput(const std::vector<std::string>& var_names,
const std::vector<Variable*>& output_vars,
Scope* inner_scope);

} // namespace framework
} // namespace paddle
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,7 @@ const std::unordered_set<std::string> SpecialOps = {
paddle::dialect::DataOp::name(),
pir::ShadowOutputOp::name(),
paddle::dialect::IfOp::name(),
paddle::dialect::PyLayerOp::name(),
paddle::dialect::WhileOp::name(),
pir::StackCreateOp::name(),
};
Expand Down Expand Up @@ -648,6 +649,13 @@ void HandleForSpecialOp(pir::Operation* op,
auto if_op_out_value = if_op->result(i);
BuildValue(if_op_out_value, var_name_prefix, value_exe_info);
}
} else if (op->isa<paddle::dialect::PyLayerOp>()) {
auto pylayer_op = op->dyn_cast<paddle::dialect::PyLayerOp>();

for (size_t i = 0; i < pylayer_op->num_results(); ++i) {
auto pylayer_op_out_value = pylayer_op->result(i);
BuildValue(pylayer_op_out_value, var_name_prefix, value_exe_info);
}
} else if (op->isa<paddle::dialect::WhileOp>()) {
auto while_op = op->dyn_cast<paddle::dialect::WhileOp>();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,12 @@ namespace framework {

class IfInstruction;
class WhileInstruction;
class PyLayerInstruction;
class ValueExecutionInfo {
public:
friend class IfInstruction;
friend class WhileInstruction;
friend class PyLayerInstruction;

explicit ValueExecutionInfo(Scope* scope) : scope_(scope) {}

Expand Down
9 changes: 9 additions & 0 deletions paddle/fluid/framework/new_executor/pir_interpreter.cc
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
#include "paddle/fluid/framework/new_executor/instruction/control_flow/assert_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/has_elements_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/if_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/pylayer_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/select_input_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/select_output_instruction.h"
#include "paddle/fluid/framework/new_executor/instruction/control_flow/tuple_pop_instruction.h"
Expand Down Expand Up @@ -712,6 +713,14 @@ void PirInterpreter::BuildInstruction() {
{&op.dyn_cast<paddle::dialect::IfOp>().false_block(),
dynamic_cast<IfInstruction*>(vec_instruction_base_.back().get())
->FalseBranchInterpreter()});
} else if (op.isa<paddle::dialect::PyLayerOp>()) {
vec_instruction_base_.emplace_back(std::make_unique<PyLayerInstruction>(
op_idx++, place_, &op, value_exe_info_.get(), execution_config_));
sub_blocks_.insert(
{&op.dyn_cast<paddle::dialect::PyLayerOp>().forward_block(),
dynamic_cast<PyLayerInstruction*>(
vec_instruction_base_.back().get())
->ForwardInterpreter()});
} else if (op.isa<paddle::dialect::WhileOp>()) {
vec_instruction_base_.emplace_back(std::make_unique<WhileInstruction>(
op_idx++, place_, &op, value_exe_info_.get(), execution_config_));
Expand Down
Loading

0 comments on commit 6037071

Please sign in to comment.