chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,236 @@
|
||||
// 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
|
||||
Reference in New Issue
Block a user