trace_grad_kernel_impl.h 4.5 KB
Newer Older
H
hong 已提交
1
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
L
Li Fuchen 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15
//
// 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
H
hong 已提交
16 17 18 19 20 21

#if defined(__NVCC__) || defined(__HIPCC__)
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#endif

L
Li Fuchen 已提交
22
#include <algorithm>
H
hong 已提交
23

C
Chen Weihang 已提交
24
#include "paddle/phi/kernels/funcs/for_range.h"
25
#include "paddle/phi/kernels/funcs/math_function.h"
L
Li Fuchen 已提交
26

27
namespace phi {
L
Li Fuchen 已提交
28 29 30

template <typename T>
struct TraceGradFunctor {
H
hong 已提交
31 32 33 34 35 36 37 38 39
  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)
L
Li Fuchen 已提交
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82
      : 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_;
};

H
hong 已提交
83 84 85
template <typename T, typename Context>
void TraceGradKernel(const Context& ctx,
                     const DenseTensor& x,
P
update  
phlrain 已提交
86
                     const DenseTensor& out_grad,
H
hong 已提交
87 88 89 90 91
                     int offset,
                     int axis1,
                     int axis2,
                     DenseTensor* in_grad) {
  auto input_dims = in_grad->dims();
92
  auto input_stride = phi::stride(input_dims);
H
hong 已提交
93
  auto output_dims = out_grad.dims();
94
  auto output_stride = phi::stride(output_dims);
L
Li Fuchen 已提交
95

H
hong 已提交
96
  auto* out_data = out_grad.data<T>();
97
  T* x_data = ctx.template Alloc<T>(in_grad);
L
Li Fuchen 已提交
98

99
  phi::funcs::SetConstant<Context, T> set_zero;
L
Li Fuchen 已提交
100

H
hong 已提交
101 102 103 104 105 106 107 108 109
  set_zero(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[std::min(dim1_, dim2_)];
  auto len2 = input_dims[std::max(dim1_, dim2_)];
  auto stride1 = input_stride[std::min(dim1_, dim2_)];
  auto stride2 = input_stride[std::max(dim1_, dim2_)];
L
Li Fuchen 已提交
110

H
hong 已提交
111 112 113 114 115 116 117
  int offset_stride = 0;
  if (offset >= 0) {
    offset_stride = stride2;
    len2 -= offset;
  } else {
    offset_stride = stride1;
    len1 += offset;
L
Li Fuchen 已提交
118
  }
H
hong 已提交
119 120 121
  int64_t diag_size = len2 < len1 ? len2 : len1;
  int64_t pos = std::abs(offset) * offset_stride;
  if (diag_size > 0) {
122
#if defined(__NVCC__) || defined(__HIPCC__)
H
hong 已提交
123 124 125 126
    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());
L
Li Fuchen 已提交
127 128

#else
H
hong 已提交
129 130
    const auto* output_arr = output_stride.Get();
    const auto* input_arr = input_stride.Get();
L
Li Fuchen 已提交
131 132
#endif

C
Chen Weihang 已提交
133
    phi::funcs::ForRange<Context> for_range(ctx, in_grad->numel());
H
hong 已提交
134 135 136 137 138 139 140 141 142 143
    TraceGradFunctor<T> functor(out_data,
                                output_arr,
                                input_arr,
                                pos,
                                input_dims.size(),
                                dim1_,
                                dim2_,
                                diag_size,
                                x_data);
    for_range(functor);
L
Li Fuchen 已提交
144
  }
H
hong 已提交
145
}
L
Li Fuchen 已提交
146

147
}  // namespace phi