chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,154 @@
|
||||
// Copyright (c) 2025 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/index_elementwise_get_kernel.h"
|
||||
|
||||
#include "paddle/phi/backends/xpu/xpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/kernels/funcs/index_elementwise.h"
|
||||
#include "paddle/phi/kernels/funcs/stride_utils.h"
|
||||
|
||||
namespace phi {
|
||||
template <typename T, typename Context, typename IndexT = int>
|
||||
void XPUIndexElementwiseGetGradKernel(
|
||||
const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const DenseTensor& value,
|
||||
const std::vector<const DenseTensor*>& index,
|
||||
const std::vector<int64_t>& input_dims,
|
||||
const std::vector<int64_t>& input_strides,
|
||||
const std::vector<int64_t>& index_dims,
|
||||
const std::vector<int64_t>& index_strides,
|
||||
const int64_t slice_offset,
|
||||
const bool accumulate,
|
||||
DenseTensor* output) {
|
||||
int64_t numel = 0;
|
||||
int64_t num_indices = 0;
|
||||
std::vector<int64_t> shape_tmp;
|
||||
std::vector<int64_t> stride_tmp;
|
||||
funcs::cal_shape_stride(index_dims, &num_indices, &shape_tmp, &stride_tmp);
|
||||
|
||||
auto sizes = std::array<int64_t, DDim::kMaxRank + 1>{};
|
||||
auto strides = std::array<int64_t, DDim::kMaxRank + 1>{};
|
||||
for (int64_t i = 0; i < num_indices; i++) {
|
||||
sizes[i] = index_dims[i];
|
||||
strides[i] = index_strides[i];
|
||||
}
|
||||
auto index_ptrs = funcs::GetIndexDataPtrs<IndexT>(index);
|
||||
|
||||
std::array<int64_t*, 3> strides_array;
|
||||
std::vector<int64_t> desired_shape;
|
||||
std::array<std::vector<int64_t>, 3> strides_vec;
|
||||
|
||||
funcs::IndexPutStride<3>(input_dims,
|
||||
input_strides,
|
||||
phi::SizeOf(input.dtype()),
|
||||
vectorize<int64_t>(value.dims()),
|
||||
vectorize<int64_t>(value.strides()),
|
||||
phi::SizeOf(value.dtype()),
|
||||
shape_tmp,
|
||||
stride_tmp,
|
||||
phi::SizeOf(index[0]->dtype()),
|
||||
&desired_shape,
|
||||
&strides_array,
|
||||
&numel,
|
||||
strides_vec);
|
||||
|
||||
using XPUType = typename XPUTypeTrait<T>::Type;
|
||||
using XPUTypeIndexT = typename XPUTypeTrait<IndexT>::Type;
|
||||
|
||||
const XPUType* value_ptr = reinterpret_cast<const XPUType*>(value.data<T>());
|
||||
std::vector<const XPUTypeIndexT*> index_list_vec;
|
||||
std::vector<int64_t> index_numel;
|
||||
for (int i = 0; i < num_indices; i++) {
|
||||
index_list_vec.push_back(
|
||||
reinterpret_cast<const XPUTypeIndexT*>(index[i]->data<IndexT>()));
|
||||
index_numel.push_back(index[i]->numel());
|
||||
}
|
||||
std::vector<int64_t> sizes_vec =
|
||||
std::vector<int64_t>(sizes.begin(), sizes.begin() + num_indices);
|
||||
std::vector<int64_t> orig_strides_vec =
|
||||
std::vector<int64_t>(strides.begin(), strides.begin() + num_indices);
|
||||
std::vector<std::vector<int64_t>> strides_vec_vec =
|
||||
std::vector<std::vector<int64_t>>(strides_vec.begin(), strides_vec.end());
|
||||
|
||||
XPUType* output_ptr = reinterpret_cast<XPUType*>(output->data<T>());
|
||||
|
||||
// call xpu kernel
|
||||
int r = xpu::index_elementwise_get_grad<XPUType, XPUTypeIndexT>(
|
||||
dev_ctx.x_context(),
|
||||
value_ptr,
|
||||
input_dims,
|
||||
index_list_vec,
|
||||
index_numel,
|
||||
desired_shape,
|
||||
sizes_vec,
|
||||
orig_strides_vec,
|
||||
strides_vec_vec,
|
||||
slice_offset,
|
||||
numel,
|
||||
accumulate,
|
||||
output_ptr);
|
||||
PADDLE_ENFORCE_XDNN_SUCCESS(r, "index_elementwise_get_grad");
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void IndexElementwiseGetGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const std::vector<const DenseTensor*>& index,
|
||||
const DenseTensor& out_grad,
|
||||
const std::vector<int64_t>& input_dims,
|
||||
const std::vector<int64_t>& input_strides,
|
||||
const std::vector<int64_t>& index_dims,
|
||||
const std::vector<int64_t>& index_strides,
|
||||
const int64_t slice_offset,
|
||||
const bool accumulate,
|
||||
const bool is_combined,
|
||||
DenseTensor* x_grad) {
|
||||
dev_ctx.template Alloc<T>(x_grad);
|
||||
funcs::set_constant(dev_ctx, x_grad, static_cast<float>(0));
|
||||
if (out_grad.numel() == 0) return;
|
||||
|
||||
const auto& index_type = index[0]->dtype();
|
||||
PADDLE_ENFORCE_EQ(index_type == DataType::INT64,
|
||||
true,
|
||||
common::errors::InvalidArgument(
|
||||
"Index holds the wrong type, it holds [%s], but "
|
||||
"desires to be [%s].",
|
||||
index_type,
|
||||
DataType::INT64));
|
||||
|
||||
XPUIndexElementwiseGetGradKernel<T, Context, int64_t>(dev_ctx,
|
||||
x,
|
||||
out_grad,
|
||||
index,
|
||||
input_dims,
|
||||
input_strides,
|
||||
index_dims,
|
||||
index_strides,
|
||||
slice_offset,
|
||||
accumulate,
|
||||
x_grad);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(index_elementwise_get_grad,
|
||||
XPU,
|
||||
ALL_LAYOUT,
|
||||
phi::IndexElementwiseGetGradKernel,
|
||||
float,
|
||||
int,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
Reference in New Issue
Block a user