wehub-resource-sync
153 lines
4.6 KiB
C++
153 lines
4.6 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
#include <thrust/device_vector.h>
|
|
#include <thrust/host_vector.h>
|
|
#endif
|
|
|
|
#include <algorithm>
|
|
|
|
#include "paddle/phi/kernels/funcs/for_range.h"
|
|
#include "paddle/phi/kernels/funcs/math_function.h"
|
|
|
|
namespace phi {
|
|
|
|
template <typename T>
|
|
struct TraceGradFunctor {
|
|
TraceGradFunctor(const T* d_out,
|
|
const int64_t* out_stride,
|
|
const int64_t* x_strides,
|
|
int64_t pos,
|
|
int64_t dim_size,
|
|
int64_t dim1,
|
|
int64_t dim2,
|
|
int64_t diag_size,
|
|
T* d_x)
|
|
: d_out_(d_out),
|
|
out_stride_(out_stride),
|
|
x_strides_(x_strides),
|
|
pos_(pos),
|
|
dim_size_(dim_size),
|
|
dim1_(dim1),
|
|
dim2_(dim2),
|
|
diag_size_(diag_size),
|
|
d_x_(d_x) {}
|
|
|
|
HOSTDEVICE void operator()(size_t idx) const {
|
|
int64_t num = idx - pos_;
|
|
int64_t position = 0;
|
|
if (num >= 0) {
|
|
int64_t dim1 = 0;
|
|
int64_t dim2 = 0;
|
|
int64_t out_idx = 0;
|
|
for (int64_t i = 0; i < dim_size_; i++) {
|
|
if (i != dim1_ && i != dim2_) {
|
|
position += num / x_strides_[i] * out_stride_[out_idx++];
|
|
} else if (i == dim1_) {
|
|
dim1 = num / x_strides_[i];
|
|
} else {
|
|
dim2 = num / x_strides_[i];
|
|
}
|
|
num = num % x_strides_[i];
|
|
}
|
|
if (dim1 == dim2 && dim1 < diag_size_) {
|
|
d_x_[idx] = d_out_[position];
|
|
}
|
|
}
|
|
}
|
|
const T* d_out_;
|
|
const int64_t* out_stride_;
|
|
const int64_t* x_strides_;
|
|
int64_t pos_;
|
|
int64_t dim_size_;
|
|
int64_t dim1_;
|
|
int64_t dim2_;
|
|
int64_t diag_size_;
|
|
T* d_x_;
|
|
};
|
|
|
|
template <typename T, typename Context>
|
|
void TraceGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& x UNUSED,
|
|
const DenseTensor& out_grad,
|
|
int offset,
|
|
int axis1,
|
|
int axis2,
|
|
DenseTensor* in_grad) {
|
|
if (in_grad && in_grad->numel() == 0) {
|
|
dev_ctx.template Alloc<T>(in_grad);
|
|
return;
|
|
}
|
|
auto input_dims = in_grad->dims();
|
|
auto input_stride = common::stride(input_dims);
|
|
auto output_dims = out_grad.dims();
|
|
auto output_stride =
|
|
output_dims.size() == 0 ? DDim(output_dims) : common::stride(output_dims);
|
|
|
|
auto* out_data = out_grad.data<T>();
|
|
T* x_data = dev_ctx.template Alloc<T>(in_grad);
|
|
|
|
funcs::SetConstant<Context, T> set_zero;
|
|
|
|
set_zero(dev_ctx, in_grad, static_cast<T>(0.0));
|
|
auto dim1 = axis1;
|
|
auto dim2 = axis2;
|
|
auto dim1_ = dim1 < 0 ? input_dims.size() + dim1 : dim1;
|
|
auto dim2_ = dim2 < 0 ? input_dims.size() + dim2 : dim2;
|
|
auto len1 = input_dims[dim1_];
|
|
auto len2 = input_dims[dim2_];
|
|
auto stride1 = input_stride[dim1_];
|
|
auto stride2 = input_stride[dim2_];
|
|
|
|
int offset_stride = 0;
|
|
if (offset >= 0) {
|
|
offset_stride = stride2;
|
|
len2 -= offset;
|
|
} else {
|
|
offset_stride = stride1;
|
|
len1 += offset;
|
|
}
|
|
int64_t diag_size = len2 < len1 ? len2 : len1;
|
|
int64_t pos = std::abs(offset) * offset_stride;
|
|
if (diag_size > 0) {
|
|
#if defined(__NVCC__) || defined(__HIPCC__)
|
|
thrust::device_vector<int64_t> output_vec(vectorize(output_stride));
|
|
const int64_t* output_arr = thrust::raw_pointer_cast(output_vec.data());
|
|
thrust::device_vector<int64_t> input_vec(vectorize(input_stride));
|
|
const int64_t* input_arr = thrust::raw_pointer_cast(input_vec.data());
|
|
|
|
#else
|
|
const auto* output_arr = output_stride.Get();
|
|
const auto* input_arr = input_stride.Get();
|
|
#endif
|
|
|
|
funcs::ForRange<Context> for_range(dev_ctx, in_grad->numel());
|
|
TraceGradFunctor<T> functor(out_data,
|
|
output_arr,
|
|
input_arr,
|
|
pos,
|
|
input_dims.size(),
|
|
dim1_,
|
|
dim2_,
|
|
diag_size,
|
|
x_data);
|
|
for_range(functor);
|
|
}
|
|
}
|
|
|
|
} // namespace phi
|