Files
paddlepaddle--paddle/paddle/phi/kernels/cpu/pyramid_hash_kernel.cc
T
2026-07-13 12:40:42 +08:00

270 lines
9.6 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 <xxhash.h>
#include <algorithm>
#include <cmath>
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/funcs/math/bloomfilter.h"
#include "paddle/phi/kernels/funcs/search_compute.h"
namespace phi {
#ifndef _WIN32
bool should_use_term(math::bloomfilter* _filter,
math::bloomfilter* _black_filter,
const float* word_repr,
int len) {
return (!_filter || 1 == math::bloomfilter_get(
_filter, word_repr, len * sizeof(float))) &&
(!_black_filter ||
0 == math::bloomfilter_get(
_black_filter, word_repr, len * sizeof(float)));
}
template <typename T>
void hash_embedding_ff(const float* hash_id,
int len,
T* top_pos,
const T* weights,
int _num_emb,
int _rand_len,
int _space_len) {
unsigned int pos1 = XXH32(hash_id, len * sizeof(float), 0) % _space_len;
unsigned int pos2 =
XXH32(hash_id, len * sizeof(float), _rand_len) % _space_len;
for (int j = 0; j != _num_emb; j += _rand_len) {
if (j + _rand_len < _num_emb) {
__builtin_prefetch(weights + pos2);
__builtin_prefetch(top_pos + j + _rand_len);
}
unsigned int pos3 =
XXH32(hash_id, len * sizeof(float), j + 2 * _rand_len) % _space_len;
memcpy(top_pos + j, const_cast<T*>(weights + pos1), _rand_len * sizeof(T));
pos1 = pos2;
pos2 = pos3;
}
}
template <typename T, typename Context>
void CPUPyramidHashOPKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& w,
const DenseTensor& white_list,
const DenseTensor& black_list,
int num_emb,
int space_len,
int pyramid_layer,
int rand_len,
float drop_out_percent,
int is_training,
bool use_filter,
int white_list_len,
int black_list_len,
int seed,
float lr,
const std::string& distribute_update_vars,
DenseTensor* out,
DenseTensor* drop_pos,
DenseTensor* x_temp_out) {
auto* bottom = &x;
auto* _blobs_0 = &w;
auto* _blobs_1 = &white_list;
auto* _blobs_2 = &black_list;
auto* top = out;
int _num_emb = num_emb;
int _pyramid_layer = pyramid_layer;
int _is_training = is_training;
unsigned int _seed = (unsigned int)seed;
int _rand_len = rand_len;
int _space_len = space_len;
float _drop_out_percent = drop_out_percent;
const auto& offset = bottom->lod()[0];
const auto* bottom_data_ori = bottom->data<int32_t>();
auto* buff = x_temp_out;
buff->Resize({bottom->dims()[0], bottom->dims()[1]});
float* bottom_data = dev_ctx.template Alloc<float>(buff);
for (int i = 0; i < bottom->dims()[0]; i++) {
bottom_data[i] = bottom_data_ori[i]; // NOLINT
}
const auto* weights = _blobs_0->data<T>();
std::vector<size_t> top_offset;
top_offset.resize(offset.size());
top_offset[0] = 0;
math::bloomfilter* _filter = nullptr;
math::bloomfilter* _black_filter = nullptr;
if (use_filter) {
if (white_list_len != 0) {
_filter = (math::bloomfilter*)_blobs_1->data<float>();
PADDLE_ENFORCE_EQ(
math::bloomfilter_check(_filter),
1,
common::errors::PreconditionNotMet(
"The white filter is not loaded successfully, please make sure "
"'white_list_len': %d is valid for Input(WhiteList).",
white_list_len));
}
if (black_list_len != 0) {
_black_filter = (math::bloomfilter*)_blobs_2->data<float>();
PADDLE_ENFORCE_EQ(
math::bloomfilter_check(_black_filter),
1,
common::errors::PreconditionNotMet(
"The black filter is not loaded successfully, please make sure "
"'black_list_len': %d is valid for Input(BlackList).",
black_list_len));
}
}
drop_pos->Resize(
make_ddim({bottom->dims()[0] * bottom->dims()[1] * _pyramid_layer, 1}));
std::vector<size_t> drop_pos_offset;
drop_pos_offset.resize(offset.size());
drop_pos_offset[0] = 0;
int* iter = dev_ctx.template Alloc<int>(drop_pos);
int* iter_end = iter;
for (size_t i = 0; i < top_offset.size() - 1; ++i) {
int w = static_cast<int>(offset[i + 1] - offset[i]);
int nsentense_with_pyramid = 0;
if (w < 2) {
nsentense_with_pyramid = 0;
} else {
for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
for (int l = 0; l < w - ilayer; ++l) {
if (should_use_term(_filter,
_black_filter,
(const float*)(bottom_data + offset[i] + l),
ilayer + 1)) {
if (_is_training != 0) {
unsigned int rand_val = rand_r(&_seed);
double rate = static_cast<double>(rand_val) / (RAND_MAX);
*(iter_end++) = (rate < _drop_out_percent ? 0 : 1);
} else {
*(iter_end++) = 1;
}
} else {
*(iter_end++) = 0;
}
}
}
nsentense_with_pyramid = static_cast<int>(std::count(iter, iter_end, 1));
iter = iter_end;
}
drop_pos_offset[i + 1] = drop_pos_offset[i] + nsentense_with_pyramid;
top_offset[i + 1] =
top_offset[i] +
(nsentense_with_pyramid == 0 ? 1 : nsentense_with_pyramid);
}
int top_l = static_cast<int>(top_offset[top_offset.size() - 1]);
LegacyLoD top_lod;
top_lod.push_back(top_offset);
top->set_lod(top_lod);
top->Resize({top_l, _num_emb});
auto* top_data = dev_ctx.template Alloc<T>(top);
LegacyLoD drop_pos_lod;
drop_pos_lod.push_back(drop_pos_offset);
drop_pos->set_lod(drop_pos_lod);
iter = dev_ctx.template Alloc<int>(drop_pos);
int top_counter = 0;
for (size_t i = 0; i < offset.size() - 1; ++i) {
int w_drop = static_cast<int>(drop_pos_offset[i + 1] - drop_pos_offset[i]);
int w = static_cast<int>(offset[i + 1] - offset[i]);
if (w_drop == 0) {
if (w >= 2) {
for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
for (int l = 0; l < w - ilayer; ++l) {
iter++;
}
}
}
auto* top_pos = top_data + top_counter++ * _num_emb;
memset(top_pos, 0, _num_emb * sizeof(T));
continue;
}
if (w >= 2) {
for (int ilayer = 1; ilayer < _pyramid_layer && ilayer < w; ++ilayer) {
for (int l = 0; l < w - ilayer; ++l) {
if (*(iter++) == 0) {
// do nothing
} else {
auto* top_pos = top_data + top_counter++ * _num_emb;
hash_embedding_ff<T>((const float*)(bottom_data + offset[i] + l),
ilayer + 1,
top_pos,
weights,
_num_emb,
_rand_len,
_space_len);
}
}
}
}
}
if (iter != iter_end) {
exit(1);
}
auto weight_type = _blobs_0->dtype();
if (_is_training == 0 && weight_type != DataType::INT8) {
funcs::axpy_noadd(
top_data, top_data, top->dims()[0] * top->dims()[1], _drop_out_percent);
}
}
#endif
#ifdef _WIN32
template <typename T, typename Context>
void CPUPyramidHashOPKernel(const Context& dev_ctx,
const DenseTensor& x,
const DenseTensor& w,
const DenseTensor& white_list,
const DenseTensor& black_list,
int num_emb,
int space_len,
int pyramid_layer,
int rand_len,
float drop_out_percent,
int is_training,
bool use_filter,
int white_list_len,
int black_list_len,
int seed,
float lr,
const std::string& distribute_update_vars,
DenseTensor* out,
DenseTensor* drop_pos,
DenseTensor* x_temp_out) {}
#endif
} // namespace phi
PD_REGISTER_KERNEL(
pyramid_hash, CPU, ALL_LAYOUT, phi::CPUPyramidHashOPKernel, float, int8_t) {
kernel->InputAt(0).SetDataType(phi::DataType::INT32);
kernel->OutputAt(1).SetDataType(phi::DataType::INT32);
kernel->OutputAt(2).SetDataType(phi::DataType::FLOAT32);
}