sum_op.h 5.0 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"
Q
QI JUN 已提交
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;
Q
QI JUN 已提交
23 24
using SelectedRows = framework::SelectedRows;
using LoDTensor = framework::LoDTensor;
25 26 27 28 29
template <typename T, int MajorType = Eigen::RowMajor,
          typename IndexType = Eigen::DenseIndex>
using EigenVector = framework::EigenVector<T, MajorType, IndexType>;

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

Y
Yu Yang 已提交
37 38
    bool in_place = out_var == in_vars[0];

Q
QI JUN 已提交
39
    if (out_var->IsType<framework::LoDTensor>()) {
40
      auto *out = context.Output<Tensor>("Out");
Q
QI JUN 已提交
41 42 43 44
      out->mutable_data<T>(context.GetPlace());

      auto result = EigenVector<T>::Flatten(*out);

Y
Yu Yang 已提交
45 46 47 48
      if (!in_place) {
        math::SetConstant<Place, T> constant_functor;
        constant_functor(context.device_context(), out, 0.0);
      }
Q
QI JUN 已提交
49 50 51

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

      // 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());

      math::SelectedRowsAddTo<Place, T> functor;

      int64_t offset = 0;
      for (int i = 0; i < N; i++) {
        PADDLE_ENFORCE_EQ(out->height(),
90
                          in_vars[i]->Get<SelectedRows>().height());
Q
QI JUN 已提交
91 92 93 94
        functor(context.device_context(), in_vars[i]->Get<SelectedRows>(),
                offset, out);
        offset += in_vars[i]->Get<SelectedRows>().value().numel();
      }
95 96 97 98 99 100 101 102 103 104 105 106 107
    } 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 已提交
108 109
              framework::CopyFrom(in_array[i], in_array[i].place(),
                                  context.device_context(), &out_array[i]);
110 111 112 113 114 115 116 117 118 119 120 121 122
              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]);
              result.device(context.GetEigenDevice<Place>()) = result + in;
            }
          }
        }
      }
    } else {
      PADDLE_THROW("Unexpected branch, output variable type is %s",
                   out_var->Type().name());
123 124 125 126 127
    }
  }
};
}  // namespace operators
}  // namespace paddle