wehub-resource-sync
93 lines
3.7 KiB
C++
93 lines
3.7 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_grad_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 CEmbeddingGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& w,
|
|
const DenseTensor& ids,
|
|
const DenseTensor& out_grad,
|
|
int64_t start_index,
|
|
DenseTensor* w_grad) {
|
|
w_grad->Resize(w.dims());
|
|
dev_ctx.Alloc(w_grad, w.dtype());
|
|
const auto& index_type = ids.dtype();
|
|
if (index_type == phi::DataType::INT32 ||
|
|
index_type == phi::DataType::INT64) {
|
|
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->set_meta(w.meta());
|
|
dev_ctx.Alloc(w_tmp.get(), w_tmp->dtype());
|
|
auto out_grad_tmp = std::make_shared<DenseTensor>();
|
|
out_grad_tmp->ShareDataWith(out_grad).Resize({K, D});
|
|
paddle::Tensor x_tensor(x_tmp), w_tensor(w_tmp),
|
|
out_grad_tensor(out_grad_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 real_ids_tensor = (x_tensor - start_index_tensor)
|
|
.multiply(paddle::experimental::cast(
|
|
ids_mask_tensor, x_tensor.dtype()));
|
|
auto out_grad_tensor_mul_mask =
|
|
paddle::experimental::reshape(out_grad_tensor, {K, D})
|
|
.multiply(paddle::experimental::reshape(
|
|
paddle::experimental::cast(ids_mask_tensor, w.dtype()),
|
|
{K, 1}));
|
|
paddle::Tensor w_grad_tensor;
|
|
paddle::experimental::embedding_grad(real_ids_tensor,
|
|
w_tensor,
|
|
out_grad_tensor_mul_mask,
|
|
-1,
|
|
false,
|
|
&w_grad_tensor);
|
|
w_grad->ShareDataWith(
|
|
*reinterpret_cast<DenseTensor*>(w_grad_tensor.impl().get()));
|
|
|
|
} else {
|
|
PADDLE_THROW(common::errors::Unavailable(
|
|
"Custom Device c_embedding_grad ids only support int32 or int64."));
|
|
}
|
|
}
|
|
#endif
|
|
} // namespace phi
|
|
|
|
#ifdef PADDLE_WITH_CUSTOM_DEVICE
|
|
PD_REGISTER_KERNEL(c_embedding_grad,
|
|
Custom,
|
|
ALL_LAYOUT,
|
|
phi::CEmbeddingGradKernel,
|
|
float,
|
|
phi::float16,
|
|
phi::bfloat16) {}
|
|
#endif
|