wehub-resource-sync
237 lines
9.0 KiB
C++
237 lines
9.0 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 "paddle/ap/include/axpr/anf_expr_util.h"
|
|
#include "paddle/ap/include/axpr/atomic.h"
|
|
#include "paddle/ap/include/axpr/data_type_util.h"
|
|
#include "paddle/ap/include/axpr/lambda_expr_builder.h"
|
|
#include "paddle/ap/include/code_gen/arg_source_ctx.h"
|
|
#include "paddle/ap/include/drr/drr_graph_descriptor.h"
|
|
#include "paddle/ap/include/drr/drr_node_descriptor.h"
|
|
#include "paddle/ap/include/drr/res_ptn_packed_ir_op_declare_data.h"
|
|
#include "paddle/ap/include/drr/result_pattern_helper.h"
|
|
#include "paddle/ap/include/drr/value.h"
|
|
#include "paddle/ap/include/graph/graph_helper.h"
|
|
#include "paddle/ap/include/ir_match/graph_match_ctx.h"
|
|
#include "paddle/ap/include/ir_match/graph_matcher.h"
|
|
#include "paddle/ap/include/ir_match/ir_match_ctx.h"
|
|
|
|
namespace ap::code_gen {
|
|
|
|
template <typename BirNode /* backend ir node*/>
|
|
struct MatchedResultPatternHelper {
|
|
using DrrNode = drr::Node;
|
|
|
|
using DrrCtx = drr::DrrCtx;
|
|
|
|
using DrrNativeIrValue = drr::NativeIrValue<DrrNode>;
|
|
using DrrPackedIrValue = drr::PackedIrValue<DrrNode>;
|
|
using DrrIrValue = drr::IrValue;
|
|
|
|
using DrrNativeIrOp = drr::NativeIrOp<DrrNode>;
|
|
using DrrNativeIrOpOperand = drr::NativeIrOpOperand<DrrNode>;
|
|
using DrrNativeIrOpResult = drr::NativeIrOpResult<DrrNode>;
|
|
using DrrPackedIrOp = drr::PackedIrOp<DrrNode>;
|
|
using DrrPackedIrOpOperand = drr::PackedIrOpOperand<DrrNode>;
|
|
using DrrPackedIrOpResult = drr::PackedIrOpResult<DrrNode>;
|
|
using DrrOptPackedIrOp = drr::OptPackedIrOp<DrrNode>;
|
|
using DrrOptPackedIrOpOperand = drr::OptPackedIrOpOperand<DrrNode>;
|
|
using DrrOptPackedIrOpResult = drr::OptPackedIrOpResult<DrrNode>;
|
|
|
|
using DrrIrOpImpl = std::variant<DrrNativeIrOp, DrrPackedIrOp>;
|
|
|
|
using IrMatchCtx = ir_match::IrMatchCtx<BirNode>;
|
|
|
|
using GraphMatchCtx = ir_match::GraphMatchCtx<BirNode>;
|
|
|
|
const GraphMatchCtx& match_ctx_;
|
|
const DrrCtx& drr_ctx_;
|
|
|
|
template <typename DoEachT>
|
|
adt::Result<adt::Ok> VisitMatchedBirInputOfRestPtnPackedIrOp(
|
|
const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
|
|
auto CollectInput =
|
|
[&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
|
|
return VisitMatchedBirValueOfResPtnIrValue(drr_ir_value, DoEach);
|
|
};
|
|
ADT_RETURN_IF_ERR(
|
|
VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, CollectInput));
|
|
return adt::Ok{};
|
|
}
|
|
|
|
template <typename DoEachT>
|
|
adt::Result<adt::Ok> VisitMatchedBirOutputOfRestPtnPackedIrOp(
|
|
const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
|
|
auto DoEachDrrIrValue =
|
|
[&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
|
|
return VisitMatchedBirValueOfResPtnIrValue(drr_ir_value, DoEach);
|
|
};
|
|
ADT_RETURN_IF_ERR(
|
|
VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachDrrIrValue));
|
|
return adt::Ok{};
|
|
}
|
|
|
|
template <typename DoEachT>
|
|
adt::Result<adt::Ok> VisitResPtnInputIrValueByResPtnIrOp(
|
|
const DrrPackedIrOp& res_ptn_ir_op, const DoEachT& DoEach) const {
|
|
drr::ResultPatternHelper helper{drr_ctx_};
|
|
return helper.VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, DoEach);
|
|
}
|
|
|
|
template <typename DoEachT>
|
|
adt::Result<adt::Ok> VisitMatchedBirValueOfResPtnIrValue(
|
|
const DrrIrValue& res_ptn_ir_value, const DoEachT& DoEach) const {
|
|
using Ok = adt::Result<adt::Ok>;
|
|
const auto& opt_ir_value = SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
|
|
ADT_CHECK(opt_ir_value.has_value());
|
|
const auto& ir_value = opt_ir_value.value();
|
|
ADT_RETURN_IF_ERR(
|
|
match_ctx_->VisitBigGraphIrValueNode(ir_value.node(), DoEach));
|
|
return adt::Ok{};
|
|
}
|
|
|
|
template <typename DoEachIndexT, typename DoEachSliceT>
|
|
adt::Result<adt::Ok> VisitApKernelInputIndexOrSlice(
|
|
const DrrPackedIrOp& res_ptn_ir_op,
|
|
const DoEachIndexT& DoEachIndex,
|
|
const DoEachSliceT& DoEachSlice) const {
|
|
std::size_t start = 0;
|
|
using Ok = adt::Result<adt::Ok>;
|
|
auto DoEachIrValue = [&](const DrrIrValue& drr_ir_value) -> Ok {
|
|
ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
|
|
ADT_RETURN_IF_ERR(drr_ir_value.Match(
|
|
[&](const DrrNativeIrValue&) -> Ok {
|
|
ADT_RETURN_IF_ERR(DoEachIndex(start));
|
|
return adt::Ok{};
|
|
},
|
|
[&](const DrrPackedIrValue&) -> Ok {
|
|
ADT_RETURN_IF_ERR(DoEachSlice(start, start + num_ir_values));
|
|
return adt::Ok{};
|
|
}));
|
|
start += num_ir_values;
|
|
return adt::Ok{};
|
|
};
|
|
ADT_RETURN_IF_ERR(
|
|
VisitResPtnInputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachIrValue));
|
|
return adt::Ok{};
|
|
}
|
|
|
|
template <typename DoEachIndexT, typename DoEachSliceT>
|
|
adt::Result<adt::Ok> VisitApKernelOutputIndexOrSlice(
|
|
const DrrPackedIrOp& res_ptn_ir_op,
|
|
const DoEachIndexT& DoEachIndex,
|
|
const DoEachSliceT& DoEachSlice) const {
|
|
std::size_t start = 0;
|
|
using Ok = adt::Result<adt::Ok>;
|
|
auto DoEachIrValue = [&](const DrrIrValue& drr_ir_value) -> Ok {
|
|
ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
|
|
ADT_RETURN_IF_ERR(drr_ir_value.Match(
|
|
[&](const DrrNativeIrValue&) -> Ok {
|
|
ADT_RETURN_IF_ERR(DoEachIndex(start));
|
|
return adt::Ok{};
|
|
},
|
|
[&](const DrrPackedIrValue&) -> Ok {
|
|
ADT_RETURN_IF_ERR(DoEachSlice(start, start + num_ir_values));
|
|
return adt::Ok{};
|
|
}));
|
|
start += num_ir_values;
|
|
return adt::Ok{};
|
|
};
|
|
ADT_RETURN_IF_ERR(
|
|
VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachIrValue));
|
|
return adt::Ok{};
|
|
}
|
|
|
|
adt::Result<std::size_t> GetApKernelNumOutputs(
|
|
const DrrPackedIrOp& res_ptn_ir_op) const {
|
|
std::size_t num_outputs = 0;
|
|
auto AccNumOutputs =
|
|
[&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
|
|
ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
|
|
num_outputs += num_ir_values;
|
|
return adt::Ok{};
|
|
};
|
|
ADT_RETURN_IF_ERR(
|
|
VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, AccNumOutputs));
|
|
return num_outputs;
|
|
}
|
|
|
|
template <typename T, typename IrOpT, typename DoEachT>
|
|
adt::Result<adt::Ok> VisitEachMatchedDrrIrValueAndOutputSlice(
|
|
const std::vector<T>& output_values,
|
|
const IrOpT& res_ptn_ir_op,
|
|
const DoEachT& DoEach) const {
|
|
std::size_t offset = 0;
|
|
auto DoEachSlice =
|
|
[&](const DrrIrValue& drr_ir_value) -> adt::Result<adt::Ok> {
|
|
ADT_LET_CONST_REF(num_ir_values, GetResPtnNumBirValues(drr_ir_value));
|
|
ADT_CHECK(offset + num_ir_values <= output_values.size());
|
|
std::vector<T> slice{output_values.begin() + offset,
|
|
output_values.begin() + offset + num_ir_values};
|
|
ADT_RETURN_IF_ERR(DoEach(drr_ir_value, slice));
|
|
offset += num_ir_values;
|
|
return adt::Ok{};
|
|
};
|
|
return VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEachSlice);
|
|
}
|
|
|
|
template <typename IrOpT, typename DoEachT>
|
|
adt::Result<adt::Ok> VisitResPtnOutputIrValueByResPtnIrOp(
|
|
const IrOpT& res_ptn_ir_op, const DoEachT& DoEach) const {
|
|
drr::ResultPatternHelper helper{drr_ctx_};
|
|
return helper.VisitResPtnOutputIrValueByResPtnIrOp(res_ptn_ir_op, DoEach);
|
|
}
|
|
|
|
adt::Result<std::size_t> GetResPtnNumBirValues(
|
|
const DrrIrValue& res_ptn_ir_value) const {
|
|
const auto& opt_src_ptn_ir_value =
|
|
SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
|
|
if (!opt_src_ptn_ir_value.has_value()) {
|
|
// internal ir value in result pattern.
|
|
return 1U;
|
|
}
|
|
return match_ctx_->GetNumBigGraphIrValueNodes(
|
|
opt_src_ptn_ir_value.value().node());
|
|
}
|
|
|
|
std::optional<DrrIrValue> SrcPtnIrValue4ResPtnIrValue(
|
|
const DrrIrValue& res_ptn_ir_value) const {
|
|
drr::ResultPatternHelper helper{drr_ctx_};
|
|
return helper.SrcPtnIrValue4ResPtnIrValue(res_ptn_ir_value);
|
|
}
|
|
|
|
using BirNativeIrValue = typename BirNode::native_value_type;
|
|
|
|
adt::Result<BirNativeIrValue> CastToBirNativeIrValue(
|
|
const BirNode& bir_node) const {
|
|
using RetT = adt::Result<BirNativeIrValue>;
|
|
return bir_node.Match(
|
|
[&](const typename BirNode::native_value_type& bir_value) -> RetT {
|
|
return bir_value;
|
|
},
|
|
[&](const typename BirNode::ref_value_type& ref_value) -> RetT {
|
|
return ref_value.GetOwnerNativeIrValue();
|
|
},
|
|
[&](const auto&) -> RetT {
|
|
return adt::errors::TypeError{
|
|
"bir_node is not an PirNode::native_value_type or "
|
|
"BirNode::ref_value_type"};
|
|
});
|
|
}
|
|
};
|
|
|
|
} // namespace ap::code_gen
|