chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,203 @@
|
||||
// Copyright (c) 2022 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/set_value_grad_kernel.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_context.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/infermeta/unary.h"
|
||||
#include "paddle/phi/kernels/funcs/common_shape.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/impl/share_data_kernel_impl.h"
|
||||
#include "paddle/phi/kernels/reduce_sum_kernel.h"
|
||||
#include "paddle/phi/kernels/set_value_kernel.h"
|
||||
#include "paddle/phi/kernels/shape_kernel.h"
|
||||
#include "paddle/phi/kernels/strided_slice_kernel.h"
|
||||
|
||||
namespace phi {
|
||||
template <typename T, typename Context>
|
||||
void SetValueGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& out_grad,
|
||||
const IntArray& starts,
|
||||
const IntArray& ends,
|
||||
const IntArray& steps,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& decrease_axes,
|
||||
const std::vector<int64_t>& none_axes,
|
||||
DenseTensor* x_grad,
|
||||
DenseTensor* value_grad) {
|
||||
const int rank = out_grad.dims().size();
|
||||
std::vector<int64_t> starts_local = starts.GetData();
|
||||
std::vector<int64_t> ends_local = ends.GetData();
|
||||
std::vector<int64_t> steps_local = steps.GetData();
|
||||
|
||||
bool ellipsis_flag = true;
|
||||
for (size_t i = 0; i < axes.size(); i++) {
|
||||
auto idx = axes[i];
|
||||
if (!(starts_local[i] == 0 && ends_local[i] == out_grad.dims()[idx] &&
|
||||
steps_local[i] == 1)) {
|
||||
ellipsis_flag = false;
|
||||
}
|
||||
}
|
||||
|
||||
if (ellipsis_flag) {
|
||||
if (x_grad) {
|
||||
dev_ctx.template Alloc<T>(x_grad);
|
||||
funcs::set_constant(dev_ctx, x_grad, static_cast<float>(0.0));
|
||||
}
|
||||
if (value_grad) {
|
||||
if (value_grad->numel() == out_grad.numel()) {
|
||||
if (value_grad->dims() != out_grad.dims()) {
|
||||
DenseTensor out_grad_temp;
|
||||
ShareDataKernel<T, Context>(dev_ctx, out_grad, &out_grad_temp);
|
||||
out_grad_temp.Resize(value_grad->dims());
|
||||
Copy(dev_ctx, out_grad_temp, dev_ctx.GetPlace(), false, value_grad);
|
||||
} else {
|
||||
Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, value_grad);
|
||||
}
|
||||
} else {
|
||||
auto reduce_dim = funcs::GetReduceDims(out_grad, *value_grad);
|
||||
SumKernel<T, Context>(
|
||||
dev_ctx, out_grad, reduce_dim, out_grad.dtype(), false, value_grad);
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (x_grad) {
|
||||
Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
|
||||
SetValueKernel<T, Context>(dev_ctx,
|
||||
*x_grad,
|
||||
starts,
|
||||
ends,
|
||||
steps,
|
||||
axes,
|
||||
decrease_axes,
|
||||
none_axes,
|
||||
{1},
|
||||
std::vector<Scalar>({Scalar(0)}),
|
||||
x_grad);
|
||||
}
|
||||
|
||||
if (value_grad) {
|
||||
DenseTensor value_grad_orig;
|
||||
MetaTensor meta_out(&value_grad_orig);
|
||||
MetaTensor meta_in(out_grad);
|
||||
std::vector<int> infer_flags(axes.size(), 1);
|
||||
std::vector<int> axes_int32(axes.begin(), axes.end());
|
||||
std::vector<int> decrease_axes_int32(decrease_axes.begin(),
|
||||
decrease_axes.end());
|
||||
StridedSliceRawInferMeta(meta_in,
|
||||
axes_int32,
|
||||
starts,
|
||||
ends,
|
||||
steps,
|
||||
infer_flags,
|
||||
decrease_axes_int32,
|
||||
&meta_out,
|
||||
MetaConfig(true, false));
|
||||
if (value_grad_orig.dims() != value_grad->dims()) {
|
||||
StridedSliceRawKernel<T, Context>(dev_ctx,
|
||||
out_grad,
|
||||
axes_int32,
|
||||
starts,
|
||||
ends,
|
||||
steps,
|
||||
infer_flags,
|
||||
decrease_axes_int32,
|
||||
&value_grad_orig);
|
||||
|
||||
if (value_grad->numel() == value_grad_orig.numel()) {
|
||||
value_grad_orig.Resize(value_grad->dims());
|
||||
Copy(dev_ctx, value_grad_orig, dev_ctx.GetPlace(), false, value_grad);
|
||||
} else {
|
||||
auto reduce_dim = funcs::GetReduceDims(value_grad_orig, *value_grad);
|
||||
SumKernel<T, Context>(dev_ctx,
|
||||
value_grad_orig,
|
||||
reduce_dim,
|
||||
value_grad->dtype(),
|
||||
false,
|
||||
value_grad);
|
||||
}
|
||||
} else {
|
||||
StridedSliceRawKernel<T, Context>(dev_ctx,
|
||||
out_grad,
|
||||
axes_int32,
|
||||
starts,
|
||||
ends,
|
||||
steps,
|
||||
infer_flags,
|
||||
decrease_axes_int32,
|
||||
value_grad);
|
||||
// 0-dim will change to 1 dim so we need to set meta
|
||||
value_grad->set_meta(value_grad_orig.meta());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void SetValueWithScalarGradKernel(const Context& dev_ctx,
|
||||
const DenseTensor& out_grad,
|
||||
const IntArray& starts,
|
||||
const IntArray& ends,
|
||||
const IntArray& steps,
|
||||
const std::vector<int64_t>& axes,
|
||||
const std::vector<int64_t>& decrease_axes,
|
||||
const std::vector<int64_t>& none_axes,
|
||||
DenseTensor* x_grad) {
|
||||
SetValueGradKernel<T, Context>(dev_ctx,
|
||||
out_grad,
|
||||
starts,
|
||||
ends,
|
||||
steps,
|
||||
axes,
|
||||
decrease_axes,
|
||||
none_axes,
|
||||
x_grad,
|
||||
nullptr);
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
PD_REGISTER_KERNEL(set_value_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SetValueGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
bool,
|
||||
int16_t,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
|
||||
PD_REGISTER_KERNEL(set_value_with_scalar_grad,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::SetValueWithScalarGradKernel,
|
||||
float,
|
||||
double,
|
||||
int,
|
||||
int64_t,
|
||||
bool,
|
||||
int16_t,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
phi::float16,
|
||||
phi::bfloat16,
|
||||
phi::complex64,
|
||||
phi::complex128) {}
|
||||
Reference in New Issue
Block a user