sum_op.h 4.9 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>();
Q
QI JUN 已提交
56 57 58
          auto in = EigenVector<T>::Flatten(in_t);
          result.device(place) = result + in;
        } else if (in_vars[i]->IsType<framework::SelectedRows>()) {
59
          auto &in_t = in_vars[i]->Get<framework::SelectedRows>();
Q
QI JUN 已提交
60 61 62 63 64 65
          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 已提交
66
      PADDLE_ENFORCE(!in_place, "SelectedRows not support inplace sum now");
67 68
      auto *out = context.Output<SelectedRows>("Out");
      auto *out_value = out->mutable_value();
Q
QI JUN 已提交
69 70 71 72 73 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());

      math::SelectedRowsAddTo<Place, T> functor;

      int64_t offset = 0;
      for (int i = 0; i < N; i++) {
        PADDLE_ENFORCE_EQ(out->height(),
87
                          in_vars[i]->Get<SelectedRows>().height());
Q
QI JUN 已提交
88 89 90 91
        functor(context.device_context(), in_vars[i]->Get<SelectedRows>(),
                offset, out);
        offset += in_vars[i]->Get<SelectedRows>().value().numel();
      }
92 93 94 95 96 97 98 99 100 101 102 103 104
    } 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 已提交
105 106
              framework::CopyFrom(in_array[i], in_array[i].place(),
                                  context.device_context(), &out_array[i]);
107 108 109 110 111 112 113 114 115 116 117 118 119
              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());
120 121 122 123 124
    }
  }
};
}  // namespace operators
}  // namespace paddle