wehub-resource-sync
179 lines
6.7 KiB
C++
179 lines
6.7 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.
|
|
|
|
#include <algorithm>
|
|
#include <vector>
|
|
#include "glog/logging.h"
|
|
#include "paddle/phi/backends/cpu/cpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/utils/data_type.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T,
|
|
typename Context,
|
|
typename InfoT = int,
|
|
typename OutT = int>
|
|
void TDMChildInner(const Context &dev_ctx,
|
|
const DenseTensor &input,
|
|
const DenseTensor &tree_info,
|
|
int child_nums,
|
|
DenseTensor *child,
|
|
DenseTensor *mask) {
|
|
auto info_dims = tree_info.dims();
|
|
int64_t node_nums = info_dims[0];
|
|
int64_t length = info_dims[1];
|
|
|
|
int64_t input_ids_num = input.numel();
|
|
|
|
VLOG(4) << "TDM child op: input numel -> " << input_ids_num;
|
|
|
|
std::vector<OutT> child_vec{};
|
|
std::vector<OutT> item_mask_vec{};
|
|
|
|
auto *input_data = input.data<T>();
|
|
auto *tree_info_data = tree_info.data<InfoT>();
|
|
|
|
// TreeInfo: node_id : item_id; layer_id; ancestor_id; child_id
|
|
for (int64_t input_ids = 0; input_ids < input_ids_num; ++input_ids) {
|
|
PADDLE_ENFORCE_LT(
|
|
input_data[input_ids],
|
|
node_nums,
|
|
common::errors::InvalidArgument(
|
|
"input id of OP(paddle.incubate.layers.tdm_child) "
|
|
"expected >= 0 and < %ld, but got %ld. Please check input "
|
|
"value.",
|
|
node_nums,
|
|
input_data[input_ids]));
|
|
PADDLE_ENFORCE_LE(
|
|
0,
|
|
input_data[input_ids],
|
|
common::errors::InvalidArgument(
|
|
"input id of OP(paddle.incubate.layers.tdm_child) "
|
|
"expected >= 0 and < %ld, but got %ld. Please check input "
|
|
"value.",
|
|
node_nums,
|
|
input_data[input_ids]));
|
|
|
|
// TODO(large-tensor): array index not support int64
|
|
PADDLE_ENFORCE_LE_INT_MAX(input_data[input_ids], "input_data[input_ids]");
|
|
bool has_child =
|
|
(input_data[input_ids] == 0 ||
|
|
tree_info_data[static_cast<int64_t>(input_data[input_ids]) * length +
|
|
3] == 0)
|
|
? false
|
|
: true;
|
|
|
|
if (has_child) {
|
|
for (int child_ids = 0; child_ids < child_nums; ++child_ids) {
|
|
OutT child_id = static_cast<OutT>(
|
|
tree_info_data[static_cast<int64_t>(input_data[input_ids]) *
|
|
length +
|
|
3 + child_ids]);
|
|
child_vec.push_back(child_id);
|
|
// TODO(large-tensor): array index not support int64
|
|
PADDLE_ENFORCE_LE_INT_MAX(child_id, "child_id");
|
|
OutT child_is_item = static_cast<OutT>(
|
|
tree_info_data[static_cast<int64_t>(child_id) * length] == 0 ? 0
|
|
: 1);
|
|
item_mask_vec.push_back(child_is_item);
|
|
}
|
|
} else {
|
|
for (int child_ids = 0; child_ids < child_nums; ++child_ids) {
|
|
child_vec.push_back(0);
|
|
item_mask_vec.push_back(0);
|
|
}
|
|
}
|
|
}
|
|
|
|
int output_nums = child_vec.size();
|
|
auto *child_data = dev_ctx.template Alloc<OutT>(child);
|
|
auto *leaf_mask_data = dev_ctx.template Alloc<OutT>(mask);
|
|
|
|
memcpy(child_data, &child_vec[0], sizeof(OutT) * output_nums);
|
|
memcpy(leaf_mask_data, &item_mask_vec[0], sizeof(OutT) * output_nums);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void TDMChildKernel(const Context &dev_ctx,
|
|
const DenseTensor &x,
|
|
const DenseTensor &tree_info,
|
|
int child_nums,
|
|
DataType dtype,
|
|
DenseTensor *child,
|
|
DenseTensor *leaf_mask) {
|
|
const auto &input_type = x.dtype();
|
|
bool input_type_match =
|
|
input_type == DataType::INT32 || input_type == DataType::INT64;
|
|
PADDLE_ENFORCE_EQ(input_type_match,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Input(X) holds the wrong type, it holds %s, but "
|
|
"desires to be %s or %s",
|
|
DataTypeToString(input_type),
|
|
DataTypeToString(DataType::INT32),
|
|
DataTypeToString(DataType::INT64)));
|
|
|
|
const auto &info_type = tree_info.dtype();
|
|
bool info_type_match =
|
|
info_type == DataType::INT32 || info_type == DataType::INT64;
|
|
PADDLE_ENFORCE_EQ(
|
|
info_type_match,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Input(TreeInfo) holds the wrong type, it holds %s, but "
|
|
"desires to be %s or %s",
|
|
DataTypeToString(info_type),
|
|
DataTypeToString(DataType::INT32),
|
|
DataTypeToString(DataType::INT64)));
|
|
|
|
auto output_type = dtype;
|
|
bool out_type_match =
|
|
output_type == DataType::INT32 || output_type == DataType::INT64;
|
|
PADDLE_ENFORCE_EQ(out_type_match,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"Output(Child) & Output(LeafMask) holds the wrong "
|
|
"type, it holds %s, but "
|
|
"desires to be %s or %s",
|
|
DataTypeToString(output_type),
|
|
DataTypeToString(DataType::INT32),
|
|
DataTypeToString(DataType::INT64)));
|
|
|
|
if (info_type == DataType::INT32 && output_type == DataType::INT32) {
|
|
TDMChildInner<T, Context, int, int>(
|
|
dev_ctx, x, tree_info, child_nums, child, leaf_mask);
|
|
} else if (info_type == DataType::INT64 && output_type == DataType::INT32) {
|
|
TDMChildInner<T, Context, int64_t, int>(
|
|
dev_ctx, x, tree_info, child_nums, child, leaf_mask);
|
|
} else if (info_type == DataType::INT32 && output_type == DataType::INT64) {
|
|
TDMChildInner<T, Context, int, int64_t>(
|
|
dev_ctx, x, tree_info, child_nums, child, leaf_mask);
|
|
} else if (info_type == DataType::INT64 && output_type == DataType::INT64) {
|
|
TDMChildInner<T, Context, int64_t, int64_t>(
|
|
dev_ctx, x, tree_info, child_nums, child, leaf_mask);
|
|
}
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(tdm_child,
|
|
CPU,
|
|
ALL_LAYOUT,
|
|
phi::TDMChildKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t) {}
|