wehub-resource-sync
83 lines
3.1 KiB
C++
83 lines
3.1 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 "paddle/phi/kernels/c_embedding_kernel.h"
|
|
#include "glog/logging.h"
|
|
#include "paddle/phi/api/backward/backward_api_base.h"
|
|
#include "paddle/phi/api/include/api.h"
|
|
#include "paddle/phi/backends/all_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
template <typename T, typename Context>
|
|
void CEmbeddingKernel(const Context& dev_ctx,
|
|
const DenseTensor& w,
|
|
const DenseTensor& ids,
|
|
int64_t start_index,
|
|
int64_t vocab_size,
|
|
DenseTensor* out) {
|
|
const auto& index_type = ids.dtype();
|
|
if (index_type == phi::DataType::INT32 ||
|
|
index_type == phi::DataType::INT64) {
|
|
auto out_dims = out->dims();
|
|
auto K = ids.numel();
|
|
auto N = w.dims()[0];
|
|
auto D = w.dims()[1];
|
|
|
|
auto x_tmp = std::make_shared<DenseTensor>();
|
|
x_tmp->ShareDataWith(ids).Resize({K});
|
|
auto w_tmp = std::make_shared<DenseTensor>();
|
|
w_tmp->ShareDataWith(w).Resize({N, D});
|
|
paddle::Tensor x_tensor(x_tmp), w_tensor(w_tmp);
|
|
|
|
auto start_index_tensor = paddle::experimental::full_like(
|
|
x_tensor, start_index, x_tensor.dtype(), x_tensor.place());
|
|
auto end_index_tensor = paddle::experimental::full_like(
|
|
x_tensor, start_index + N, x_tensor.dtype(), x_tensor.place());
|
|
auto ids_mask_tensor = paddle::experimental::logical_and(
|
|
x_tensor.greater_equal(start_index_tensor),
|
|
x_tensor.less_than(end_index_tensor));
|
|
auto ids_tensor = (x_tensor - start_index_tensor)
|
|
.multiply(paddle::experimental::cast(
|
|
ids_mask_tensor, x_tensor.dtype()));
|
|
auto out_tensor =
|
|
paddle::experimental::reshape(
|
|
paddle::experimental::cast(ids_mask_tensor, w_tensor.dtype()),
|
|
{K, 1})
|
|
.multiply(paddle::experimental::reshape(
|
|
paddle::experimental::embedding(
|
|
ids_tensor, w_tensor, -1, false),
|
|
{K, D}));
|
|
out->ShareDataWith(*reinterpret_cast<DenseTensor*>(out_tensor.impl().get()))
|
|
.Resize(out_dims);
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unavailable(
|
|
"Custom Device c_embedding ids only support int32 or int64."));
|
|
}
|
|
}
|
|
#endif
|
|
} // namespace phi
|
|
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
PD_REGISTER_KERNEL(c_embedding,
|
|
Custom,
|
|
ALL_LAYOUT,
|
|
phi::CEmbeddingKernel,
|
|
float,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
#endif
|