sum_op.h 5.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
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
#include "paddle/framework/eigen.h"
14
#include "paddle/framework/lod_tensor_array.h"
15
#include "paddle/framework/op_registry.h"
16 17
#include "paddle/operators/math/math_function.h"
#include "paddle/operators/math/selected_rows_functor.h"
18 19 20 21 22

namespace paddle {
namespace operators {

using Tensor = framework::Tensor;
23 24
using SelectedRows = framework::SelectedRows;
using LoDTensor = framework::LoDTensor;
25 26 27 28
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;

Q
QI JUN 已提交
29
template <typename DeviceContext, typename T>
30
class SumKernel : public framework::OpKernel<T> {
31
 public:
32
  void Compute(const framework::ExecutionContext &context) const override {
33
    auto in_vars = context.MultiInputVar("X");
34 35 36
    int N = in_vars.size();
    auto out_var = context.OutputVar("Out");

37 38
    bool in_place = out_var == in_vars[0];

39
    if (out_var->IsType<framework::LoDTensor>()) {
Y
Update  
Yang Yu 已提交
40
      auto *out = context.Output<LoDTensor>("Out");
41
      if (!in_place) {
Y
Refine  
Yang Yu 已提交
42
        out->mutable_data<T>(context.GetPlace());
Y
Update  
Yang Yu 已提交
43 44 45
      }
      auto result = EigenVector<T>::Flatten(*out);
      if (!in_place) {
Q
QI JUN 已提交
46 47 48
        math::SetConstant<DeviceContext, T> constant_functor;
        constant_functor(context.template device_context<DeviceContext>(), out,
                         0.0);
49
      }
50

Q
QI JUN 已提交
51 52 53
      math::SelectedRowsAddToTensor<DeviceContext, T> functor;
      auto &place =
          *context.template device_context<DeviceContext>().eigen_device();
54 55
      // If in_place, just skip the first tensor
      for (int i = in_place ? 1 : 0; i < N; i++) {
56
        if (in_vars[i]->IsType<framework::LoDTensor>()) {
57
          auto &in_t = in_vars[i]->Get<framework::LoDTensor>();
58 59 60
          if (in_t.numel() == 0) {
            continue;
          }
61 62 63
          auto in = EigenVector<T>::Flatten(in_t);
          result.device(place) = result + in;
        } else if (in_vars[i]->IsType<framework::SelectedRows>()) {
64
          auto &in_t = in_vars[i]->Get<framework::SelectedRows>();
Q
QI JUN 已提交
65
          functor(context.template device_context<DeviceContext>(), in_t, out);
66 67 68 69 70
        } else {
          PADDLE_THROW("Variable type must be LoDTensor/SelectedRows.");
        }
      }
    } else if (out_var->IsType<framework::SelectedRows>()) {
71
      PADDLE_ENFORCE(!in_place, "SelectedRows not support inplace sum now");
72 73
      auto *out = context.Output<SelectedRows>("Out");
      auto *out_value = out->mutable_value();
74 75 76 77 78 79 80 81 82 83 84 85 86

      // Runtime InferShape
      size_t first_dim = 0;
      for (int i = 0; i < N; i++) {
        first_dim += in_vars[i]->Get<SelectedRows>().rows().size();
      }
      auto in_dim = in_vars[0]->Get<SelectedRows>().value().dims();
      auto in_dim_vec = framework::vectorize(in_dim);
      in_dim_vec[0] = static_cast<int64_t>(first_dim);

      out_value->Resize(framework::make_ddim(in_dim_vec));
      out_value->mutable_data<T>(context.GetPlace());

Q
QI JUN 已提交
87
      math::SelectedRowsAddTo<DeviceContext, T> functor;
88 89 90 91

      int64_t offset = 0;
      for (int i = 0; i < N; i++) {
        PADDLE_ENFORCE_EQ(out->height(),
92
                          in_vars[i]->Get<SelectedRows>().height());
Q
QI JUN 已提交
93 94
        functor(context.template device_context<DeviceContext>(),
                in_vars[i]->Get<SelectedRows>(), offset, out);
95 96
        offset += in_vars[i]->Get<SelectedRows>().value().numel();
      }
97 98 99 100 101 102 103 104 105 106 107 108 109
    } else if (out_var->IsType<framework::LoDTensorArray>()) {
      auto &out_array = *out_var->GetMutable<framework::LoDTensorArray>();
      for (size_t i = in_place ? 1 : 0; i < in_vars.size(); ++i) {
        PADDLE_ENFORCE(in_vars[i]->IsType<framework::LoDTensorArray>(),
                       "Only support all inputs are TensorArray");
        auto &in_array = in_vars[i]->Get<framework::LoDTensorArray>();

        for (size_t i = 0; i < in_array.size(); ++i) {
          if (in_array[i].numel() != 0) {
            if (i >= out_array.size()) {
              out_array.resize(i + 1);
            }
            if (out_array[i].numel() == 0) {
D
dzhwinter 已提交
110 111
              framework::CopyFrom(in_array[i], in_array[i].place(),
                                  context.device_context(), &out_array[i]);
112 113 114 115 116
              out_array[i].set_lod(in_array[i].lod());
            } else {
              PADDLE_ENFORCE(out_array[i].lod() == in_array[i].lod());
              auto in = EigenVector<T>::Flatten(in_array[i]);
              auto result = EigenVector<T>::Flatten(out_array[i]);
Q
QI JUN 已提交
117 118
              result.device(*context.template device_context<DeviceContext>()
                                 .eigen_device()) = result + in;
119 120 121 122 123 124 125
            }
          }
        }
      }
    } else {
      PADDLE_THROW("Unexpected branch, output variable type is %s",
                   out_var->Type().name());
126 127 128 129 130
    }
  }
};
}  // namespace operators
}  // namespace paddle
反馈
建议
客服 返回
顶部