Files
2026-07-13 12:40:42 +08:00

263 lines
11 KiB
C++

// 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 <glog/logging.h>
#include <gtest/gtest.h>
#include <memory>
#include <sstream>
#include <string>
#include <tuple>
#include <unordered_map>
#include "paddle/cinn/hlir/dialect/operator/ir/cinn_op.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_attribute.h"
#include "paddle/cinn/hlir/dialect/operator/ir/op_dialect.h"
#include "paddle/cinn/hlir/dialect/runtime/ir/jit_kernel_op.h"
#include "paddle/cinn/hlir/dialect/runtime/ir/runtime_dialect.h"
#include "paddle/cinn/hlir/framework/pir/utils.h"
#include "paddle/cinn/hlir/framework/pir_compiler.h"
#include "paddle/cinn/utils/data_util.h"
#include "paddle/fluid/framework/new_executor/interpretercore.h"
#include "paddle/fluid/pir/dialect/kernel/ir/kernel_dialect.h"
#include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
#include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
#include "paddle/fluid/pir/dialect/operator/ir/pd_op.h"
#include "paddle/fluid/pir/transforms/pd_op_to_kernel_pass.h"
#include "paddle/pir/include/core/ir_context.h"
#include "paddle/pir/include/core/program.h"
#include "paddle/pir/include/dialect/control_flow/ir/cf_op.h"
using cinn::hlir::framework::pir::CompatibleInfo;
using cinn::hlir::framework::pir::OpLoweringGroup;
using cinn::hlir::framework::pir::OpLoweringGroupPtr;
bool simple_cmp(float a, float b) { return std::abs((a - b) / a) < 1e-5; }
using ProgramInfo = std::tuple<std::shared_ptr<::pir::Program>,
std::vector<OpLoweringGroupPtr>>;
ProgramInfo BuildProgram() {
::pir::IrContext* ctx = ::pir::IrContext::Instance();
ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
auto program = std::make_shared<::pir::Program>(ctx);
::pir::Builder builder = ::pir::Builder(ctx, program->block());
const float value_one = 1.0; // relu(tan(1.)) = 1.5;
const float value_two = 2.0; // relu(tan(2.)) = 0.
auto full_op_x =
builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 128},
value_one,
phi::DataType::FLOAT32,
phi::GPUPlace());
auto full_op_y =
builder.Build<paddle::dialect::FullOp>(std::vector<int64_t>{64, 128},
value_two,
phi::DataType::FLOAT32,
phi::GPUPlace());
auto tan_op_x = builder.Build<paddle::dialect::TanOp>(full_op_x->result(0));
auto relu_op_x = builder.Build<paddle::dialect::ReluOp>(tan_op_x->result(0));
auto tan_op_y = builder.Build<paddle::dialect::TanOp>(relu_op_x->result(0));
auto relu_op_y = builder.Build<paddle::dialect::ReluOp>(tan_op_y->result(0));
builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_x.result(0)});
builder.Build<pir::YieldOp>(std::vector<pir::Value>{full_op_y.result(0)});
builder.Build<pir::YieldOp>(std::vector<pir::Value>{relu_op_y.result(0)});
std::vector<OpLoweringGroupPtr> groups;
const auto full_op_x_ops =
std::initializer_list<::pir::Operation*>({full_op_x.operation()});
groups.emplace_back(std::make_shared<OpLoweringGroup>(
full_op_x_ops,
CompatibleInfo::GroupOpsName(full_op_x_ops))); // For coverage
groups[0]->mut_output_values().push_back(groups[0]->ops().back()->result(0));
const auto full_op_y_ops =
std::initializer_list<::pir::Operation*>({full_op_x.operation()});
groups.emplace_back(std::make_shared<OpLoweringGroup>(
full_op_y_ops, CompatibleInfo::GroupOpsName(full_op_y_ops)));
groups[1]->mut_output_values().push_back(groups[1]->ops().back()->result(0));
const auto vector_ops =
std::vector<::pir::Operation*>({tan_op_x.operation(),
relu_op_x.operation(),
tan_op_y.operation(),
relu_op_y.operation()});
groups.emplace_back(std::make_shared<OpLoweringGroup>(
vector_ops, CompatibleInfo::GroupOpsName(vector_ops)));
groups[2]->mut_output_values().push_back(groups[2]->ops().back()->result(0));
return {program, groups};
}
ProgramInfo BuildSoftmax() {
::pir::IrContext* ctx = ::pir::IrContext::Instance();
ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
auto program = std::make_shared<::pir::Program>(ctx);
::pir::Builder builder = ::pir::Builder(ctx, program->block());
std::vector<int64_t> axes{-1};
auto x = builder
.Build<paddle::dialect::FullOp>(std::vector<int64_t>({16, 16}),
1.0,
phi::DataType::FLOAT32,
phi::GPUPlace(0))
.result(0);
auto max = builder.Build<cinn::dialect::ReduceMaxOp>(x, axes, true).result(0);
auto broadcast_1 =
builder
.Build<cinn::dialect::BroadcastOp>(
max, std::vector<int64_t>({0, 1}), std::vector<int64_t>({16, 16}))
.result(0);
auto sub =
builder.Build<paddle::dialect::SubtractOp>(x, broadcast_1).result(0);
auto exp = builder.Build<paddle::dialect::ExpOp>(sub).result(0);
auto sum =
builder.Build<cinn::dialect::ReduceSumOp>(exp, axes, true).result(0);
auto broadcast_2 =
builder
.Build<cinn::dialect::BroadcastOp>(
sum, std::vector<int64_t>({0, 1}), std::vector<int64_t>({16, 16}))
.result(0);
auto divide =
builder.Build<paddle::dialect::DivideOp>(exp, broadcast_2).result(0);
auto yield_op = builder.Build<pir::YieldOp>(std::vector<pir::Value>{divide});
std::vector<OpLoweringGroupPtr> groups;
const auto vector_ops =
std::initializer_list<::pir::Operation*>({max.defining_op(),
broadcast_1.defining_op(),
sub.defining_op(),
exp.defining_op(),
sum.defining_op(),
broadcast_2.defining_op(),
divide.defining_op()});
groups.emplace_back(std::make_shared<OpLoweringGroup>(
vector_ops, CompatibleInfo::GroupOpsName(vector_ops)));
groups[0]->mut_output_values().push_back(groups[0]->ops().back()->result(0));
groups[0]->set_op_pattern_kind(cinn::hlir::framework::kReduction);
return {program, groups};
}
// TEST(PirCompiler, CompileSoftmax) {
// // Step 1: Construct pir::Program
// ::pir::IrContext* ctx = ::pir::IrContext::Instance();
// ctx->GetOrRegisterDialect<paddle::dialect::OperatorDialect>();
// ctx->GetOrRegisterDialect<cinn::dialect::OperatorDialect>();
// ctx->GetOrRegisterDialect<cinn::dialect::RuntimeDialect>();
// ctx->GetOrRegisterDialect<paddle::dialect::KernelDialect>();
// auto new_program = std::make_shared<::pir::Program>(ctx);
// auto prog_info = BuildSoftmax();
// std::shared_ptr<::pir::Program> program = std::get<0>(prog_info);
// std::vector<GroupPtr> groups = std::get<1>(prog_info);
// EXPECT_EQ(program->block()->size(), 9u);
// LOG(INFO) << program->block()->size();
// std::stringstream ss;
// program->Print(ss);
// LOG(INFO) << ss.str();
// // Step 2: Compiler New pir::Program into Runtime Program
// auto target = cinn::common::DefaultNVGPUTarget();
// auto scope = cinn::hlir::framework::BuildScope(target, *program);
// LOG(INFO) << scope->var_names().size();
// ASSERT_EQ(scope->var_names().size(), 8);
// cinn::hlir::framework::PirCompiler ir_compiler(*program, target, scope);
// auto fn_ptr_res = ir_compiler.BuildCUDAJITInfo(groups);
// ::pir::Builder builder = ::pir::Builder(ctx, new_program->block());
// auto x = builder
// .Build<paddle::dialect::FullOp>(std::vector<int64_t>({16,
// 16}),
// 1.0,
// phi::DataType::FLOAT32,
// phi::GPUPlace(0))
// .result(0);
// std::unordered_map<std::string, ::pir::Attribute> op_attrs{
// {cinn::dialect::JitKernelOp::kAttrName,
// cinn::dialect::CINNKernelInfoAttribute::get(ctx, fn_ptr_res[0])},
// };
// std::vector<pir::Type> vec_types;
// vec_types.push_back(groups[0]->ops.back()->result(0).type());
// std::string jit_op_name = cinn::dialect::JitKernelOp::name();
// ::pir::OpInfo op_info = ctx->GetRegisteredOpInfo(jit_op_name);
// ::pir::Operation* cinn_op =
// ::pir::Operation::Create({x}, op_attrs, vec_types, op_info);
// new_program->block()->push_back(cinn_op);
// builder.SetInsertionPointToBlockEnd(new_program->block());
// builder.Build<paddle::dialect::FetchOp>(
// cinn_op->result(cinn_op->num_results() - 1), "out", 0);
// phi::Place place = phi::GPUPlace(0);
// auto kernel_program =
// pir::PdOpLowerToKernelPass(new_program.get(), place);
// paddle::framework::Scope exe_scope;
// paddle::framework::interpreter::ExecutionConfig exe_conf;
// exe_conf.create_local_scope = false;
// paddle::framework::InterpreterCore executor(
// place, {"out@fetch"}, kernel_program->block(), &exe_scope);
// executor.Run({}, true);
// auto out_tensor =
// executor.local_scope()->FindVar("out@fetch")->Get<phi::DenseTensor>();
// bool res0 = simple_cmp(out_tensor.data<float>()[0], 1.0 / 16);
// EXPECT_EQ(res0, true);
// }
// TEST(PirCompiler, CompileGroupOps) {
// // Step 1: Construct pir::Program
// auto prog_info = BuildProgram();
// std::shared_ptr<::pir::Program> program = std::get<0>(prog_info);
// std::vector<GroupPtr> groups = std::get<1>(prog_info);
// EXPECT_EQ(program->block()->size(), 9u);
// LOG(INFO) << program->block()->size();
// std::stringstream ss;
// program->Print(ss);
// LOG(INFO) << ss.str();
// // Step 2: Compiler New pir::Program into Runtime Program
// auto target = cinn::common::DefaultNVGPUTarget();
// auto scope = cinn::hlir::framework::BuildScope(target, *program);
// ASSERT_EQ(scope->var_names().size(), 6);
// cinn::hlir::framework::PirCompiler ir_compiler(*program, target, scope);
// auto runtime_program = ir_compiler.Build(groups);
// // Step 3: Execute Runtime Instruction and check Scope.
// ASSERT_NO_THROW(runtime_program->Execute());
// for (auto& var_name : scope->var_names()) {
// std::string name = {var_name.begin(), var_name.end()};
// std::vector<float> data =
// cinn::GetTensorData<float>(scope->GetTensor(name), target);
// for (int i = 0; i < 1; ++i) {
// LOG_FIRST_N(INFO, 10) << "data: " << data[i];
// }
// }
// }