wehub-resource-sync
264 lines
10 KiB
C++
264 lines
10 KiB
C++
// 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.
|
|
// 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 <unordered_map>
|
|
|
|
#include "paddle/fluid/custom_engine/custom_engine_ext.h"
|
|
#include "paddle/fluid/framework/new_executor/instruction/custom_engine_instruction.h"
|
|
#include "paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.h"
|
|
#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
|
|
#include "test/cpp/pir/custom_engine/custom_engine_op.h"
|
|
#include "test/cpp/pir/custom_engine/fake_cpu_engine_base.h"
|
|
|
|
C_Status RegisterCustomEngineOp() {
|
|
pir::IrContext* ctx = pir::IrContext::Instance();
|
|
pir::Dialect* custom_engine_dialect =
|
|
ctx->GetOrRegisterDialect<paddle::dialect::CustomEngineDialect>();
|
|
EXPECT_EQ(custom_engine_dialect != nullptr, true);
|
|
ctx->RegisterOpInfo(custom_engine_dialect,
|
|
pir::TypeId::get<paddle::dialect::FakeEngineOp>(),
|
|
paddle::dialect::FakeEngineOp::name(),
|
|
paddle::dialect::FakeEngineOp::interface_set(),
|
|
paddle::dialect::FakeEngineOp::GetTraitSet(),
|
|
paddle::dialect::FakeEngineOp::attributes_num,
|
|
paddle::dialect::FakeEngineOp::attributes_name,
|
|
paddle::dialect::FakeEngineOp::VerifySigInvariants,
|
|
paddle::dialect::FakeEngineOp::VerifyRegionInvariants);
|
|
return C_SUCCESS;
|
|
}
|
|
|
|
C_Status CustomEngineOpLower(C_CustomEngineLowerParams* lower_param) {
|
|
// get lower params
|
|
pir::IrContext* ctx =
|
|
reinterpret_cast<pir::IrContext*>(lower_param->ir_context);
|
|
pir::Operation* op_item =
|
|
reinterpret_cast<pir::Operation*>(lower_param->operation);
|
|
phi::KernelKey* kernel_key =
|
|
reinterpret_cast<phi::KernelKey*>(lower_param->kernel_key);
|
|
phi::Place* place = reinterpret_cast<phi::Place*>(lower_param->place);
|
|
std::unordered_map<pir::Operation*, pir::Operation*>* map_op_pair =
|
|
reinterpret_cast<std::unordered_map<pir::Operation*, pir::Operation*>*>(
|
|
lower_param->map_op_pair);
|
|
std::unordered_map<pir::Value, pir::Value>* map_value_pair =
|
|
reinterpret_cast<std::unordered_map<pir::Value, pir::Value>*>(
|
|
lower_param->map_value_pair);
|
|
pir::Block* block = reinterpret_cast<pir::Block*>(lower_param->block);
|
|
|
|
// Prepare output types
|
|
std::vector<pir::Type> op_output_types;
|
|
for (size_t i = 0; i < op_item->num_results(); ++i) {
|
|
PushBackOutputTypes(ctx,
|
|
op_item,
|
|
op_item->result(i).type(),
|
|
*place,
|
|
*kernel_key,
|
|
&op_output_types);
|
|
}
|
|
// Prepare input
|
|
std::vector<pir::Value> vec_inputs;
|
|
|
|
for (size_t i = 0; i < op_item->num_operands(); ++i) {
|
|
auto cur_in = op_item->operand_source(i);
|
|
PADDLE_ENFORCE_EQ(
|
|
map_value_pair->count(cur_in),
|
|
true,
|
|
common::errors::PreconditionNotMet(
|
|
"[%d]'s input of [%s] op MUST in map pair", i, op_item->name()));
|
|
|
|
auto new_in = map_value_pair->at(cur_in);
|
|
|
|
vec_inputs.push_back(new_in);
|
|
}
|
|
// Prepare attr
|
|
std::unordered_map<std::string, pir::Attribute> op_attribute;
|
|
auto op_attr_map = op_item->attributes();
|
|
for (auto& map_item : op_attr_map) {
|
|
op_attribute.emplace(map_item.first, map_item.second);
|
|
}
|
|
|
|
pir::OpInfo custom_engine_op_info =
|
|
ctx->GetRegisteredOpInfo(paddle::dialect::FakeEngineOp::name());
|
|
pir::Operation* op = pir::Operation::Create(vec_inputs,
|
|
op_attribute,
|
|
op_output_types,
|
|
custom_engine_op_info,
|
|
1,
|
|
{},
|
|
true);
|
|
op->set_attribute("origin_id", pir::Int64Attribute::get(ctx, op->id()));
|
|
op->set_attribute("op_name", pir::StrAttribute::get(ctx, op->name()));
|
|
|
|
VLOG(3) << "CustomEngineOpLower get op_item subgraph block.";
|
|
pir::Region& op_item_region = op_item->region(0);
|
|
PADDLE_ENFORCE_EQ(op_item_region.empty(),
|
|
false,
|
|
::common::errors::Unavailable(
|
|
"Required CustomEngineOp's region must not be empty."));
|
|
pir::Block* sub_graph_block = &(op_item_region.front());
|
|
|
|
VLOG(3) << "CustomEngineOpLower set new op subgraph block.";
|
|
pir::Region& region = op->region(0);
|
|
if (region.empty()) {
|
|
region.emplace_back();
|
|
}
|
|
pir::Block* op_block = &(region.front());
|
|
|
|
// process subgraph block
|
|
pir::ProcessBlock(
|
|
*place, sub_graph_block, op_block, ctx, map_op_pair, map_value_pair);
|
|
|
|
if (VLOG_IS_ON(3)) {
|
|
std::stringstream ss;
|
|
ss << "CustomEngineOpLower new op:";
|
|
op->Print(ss);
|
|
VLOG(3) << ss.str();
|
|
}
|
|
|
|
(*map_op_pair)[op_item] = op;
|
|
// only deal with single output
|
|
if (op_item->num_results() > 0) {
|
|
for (size_t i = 0; i < op_item->num_results(); ++i) {
|
|
(*map_value_pair)[op_item->result(i)] = op->result(i);
|
|
}
|
|
}
|
|
block->push_back(op);
|
|
return C_SUCCESS;
|
|
}
|
|
|
|
class CustomEngine {
|
|
public:
|
|
CustomEngine(std::vector<phi::DenseTensor*> tensor_args,
|
|
std::vector<phi::DenseTensor*> return_tensor)
|
|
: tensor_args_(tensor_args), return_tensor_(return_tensor) {}
|
|
~CustomEngine() {}
|
|
|
|
void Run(const phi::DeviceContext& device_ctx, const phi::Place& place) {
|
|
PADDLE_ENFORCE_EQ(
|
|
tensor_args_.size(),
|
|
2u,
|
|
common::errors::PreconditionNotMet("tensor_args.size != 2"));
|
|
PADDLE_ENFORCE_EQ(
|
|
return_tensor_.size(),
|
|
1u,
|
|
common::errors::PreconditionNotMet("return_tensor.size != 1"));
|
|
// phi::AddKernel<float, phi::DeviceContext>(device_ctx, *(tensor_args_[0]),
|
|
// *(tensor_args_[1]),return_tensor_[0]);
|
|
phi::Copy(device_ctx, *(tensor_args_[0]), place, true, return_tensor_[0]);
|
|
return;
|
|
}
|
|
|
|
private:
|
|
std::vector<phi::DenseTensor*> tensor_args_;
|
|
std::vector<phi::DenseTensor*> return_tensor_;
|
|
std::vector<phi::DenseTensor*> template_tensor_;
|
|
};
|
|
|
|
C_Status GraphEngineExecute(C_CustomEngineInstruction instruction) {
|
|
paddle::framework::CustomEngineInstruction* instruction_ =
|
|
reinterpret_cast<paddle::framework::CustomEngineInstruction*>(
|
|
instruction);
|
|
CustomEngine* customengine =
|
|
reinterpret_cast<CustomEngine*>(instruction_->CustomEngine());
|
|
|
|
customengine->Run(instruction_->DeviceContext(),
|
|
instruction_->DeviceContext().GetPlace());
|
|
return C_SUCCESS;
|
|
}
|
|
|
|
C_Status GraphEngineBuild(C_CustomEngineInstruction instruction) {
|
|
paddle::framework::CustomEngineInstruction* instruction_ =
|
|
reinterpret_cast<paddle::framework::CustomEngineInstruction*>(
|
|
instruction);
|
|
pir::Operation* op = instruction_->Operation();
|
|
const paddle::framework::ValueExecutionInfo* value_exec_info =
|
|
instruction_->GetValueExecutionInfo();
|
|
// prepare input tensors
|
|
std::vector<phi::DenseTensor*> tensor_args;
|
|
PADDLE_ENFORCE_EQ(op->num_operands(),
|
|
1u,
|
|
common::errors::PreconditionNotMet(
|
|
"custom engine op should has 1 operand"));
|
|
auto vec_in = op->operand_source(0).defining_op()->operands_source();
|
|
for (auto in : vec_in) {
|
|
auto var_name = value_exec_info->GetVarName(in);
|
|
auto tensor = value_exec_info->GetScope()
|
|
->FindVar(var_name)
|
|
->GetMutable<phi::DenseTensor>();
|
|
tensor_args.push_back(tensor);
|
|
}
|
|
|
|
// prepare output tensors
|
|
std::vector<phi::DenseTensor*> return_tensor;
|
|
PADDLE_ENFORCE_EQ(op->num_results(),
|
|
1u,
|
|
common::errors::PreconditionNotMet(
|
|
"custom engine op should has 1 result"));
|
|
pir::Value vec_result = op->result(0);
|
|
PADDLE_ENFORCE_EQ(vec_result.type().isa<pir::VectorType>(),
|
|
true,
|
|
common::errors::PreconditionNotMet(
|
|
"custom engine op result should be vectortype"));
|
|
auto vec_out = op->result(0).first_use().owner()->results();
|
|
|
|
for (auto out : vec_out) {
|
|
bool check =
|
|
out && out.type() && out.type().isa<paddle::dialect::DenseTensorType>();
|
|
PADDLE_ENFORCE_EQ(
|
|
check,
|
|
true,
|
|
common::errors::PreconditionNotMet(
|
|
"customEngine instruction only support DenseTensorType"));
|
|
auto var_name = value_exec_info->GetVarName(out);
|
|
|
|
auto tensor = value_exec_info->GetScope()
|
|
->Var(var_name)
|
|
->GetMutable<phi::DenseTensor>();
|
|
|
|
return_tensor.push_back(tensor);
|
|
auto alloc_tensor_type =
|
|
out.type().dyn_cast<paddle::dialect::AllocatedDenseTensorType>();
|
|
tensor->set_type(
|
|
paddle::dialect::TransToPhiDataType(alloc_tensor_type.dtype()));
|
|
tensor->Resize(alloc_tensor_type.dims());
|
|
}
|
|
|
|
CustomEngine* fake_engine = new CustomEngine(tensor_args, return_tensor);
|
|
auto customEngineDeleter = [](void* ptr) {
|
|
CustomEngine* customEngine = static_cast<CustomEngine*>(ptr);
|
|
|
|
if (customEngine != nullptr) {
|
|
delete customEngine;
|
|
} else {
|
|
PADDLE_THROW(
|
|
common::errors::PreconditionNotMet("customEngine is nullptr"));
|
|
}
|
|
};
|
|
instruction_->SetCustomEngine(reinterpret_cast<void*>(fake_engine));
|
|
instruction_->SetCustomEngineDeleter(customEngineDeleter);
|
|
|
|
return C_SUCCESS;
|
|
}
|
|
|
|
void InitPluginCustomEngine(CustomEngineParams* params) {
|
|
memset(reinterpret_cast<void*>(params->interface),
|
|
0,
|
|
sizeof(C_CustomEngineInterface));
|
|
|
|
params->interface->register_custom_engine_op = RegisterCustomEngineOp;
|
|
params->interface->graph_engine_build = GraphEngineBuild;
|
|
params->interface->graph_engine_execute = GraphEngineExecute;
|
|
params->interface->custom_engine_op_lower = CustomEngineOpLower;
|
|
}
|