sum_op.cc 8.5 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11
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. */

Y
Yi Wang 已提交
12
#include "paddle/fluid/operators/sum_op.h"
13

14 15
#include <algorithm>
#include <string>
16
#include <vector>
17

Y
Yi Wang 已提交
18 19
#include "paddle/fluid/framework/var_type_inference.h"
#include "paddle/fluid/operators/detail/safe_ref.h"
20

21 22 23 24
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif

25 26 27 28 29 30 31 32
namespace paddle {
namespace operators {
using framework::Tensor;

class SumOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

33
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
qiaolongfei 已提交
34
    PADDLE_ENFORCE(ctx->HasInputs("X"), "Inputs(X) should not be null");
35

Q
Qiao Longfei 已提交
36 37
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of SumOp should not be null.");
38 39
    if (ctx->IsRuntime() &&
        ctx->GetOutputsVarType("Out")[0] ==
40
            framework::proto::VarType::LOD_TENSOR_ARRAY) {
41 42
      return;  // skip runtime infershape when is tensor array;
    }
43

44
    auto x_var_types = ctx->GetInputsVarType("X");
45
    auto x_dims = ctx->GetInputsDim("X");
46

Q
Qiao Longfei 已提交
47
    size_t N = x_dims.size();
48 49
    PADDLE_ENFORCE_GT(N, 0, "Input tensors count should > 0.");
    if (N == 1) {
M
minqiyang 已提交
50
      VLOG(3) << "Warning: sum have only one input, may waste memory";
51
    }
Q
qiaolongfei 已提交
52

53
    framework::DDim in_dim({0});
54
    for (size_t i = 0; i < x_dims.size(); ++i) {
55 56 57 58
      auto& x_dim = x_dims[i];
      // x_dim.size() == 1 means the real dim of selected rows is [0]
      if (x_var_types[i] == framework::proto::VarType::SELECTED_ROWS &&
          x_dim.size() == 1) {
59 60
        continue;
      }
61 62 63 64 65 66 67 68
      if (framework::product(x_dim) == 0) {
        continue;
      }
      if (framework::product(in_dim) == 0) {
        in_dim = x_dim;
      } else {
        PADDLE_ENFORCE_EQ(in_dim, x_dim, "Input tensors must have same shape");
      }
Q
qijun 已提交
69
    }
Q
Qiao Longfei 已提交
70 71
    ctx->SetOutputDim("Out", in_dim);
    ctx->ShareLoD("X", /*->*/ "Out");
72
  }
73 74

 protected:
75
  framework::OpKernelType GetExpectedKernelType(
76 77
      const framework::ExecutionContext& ctx) const override {
    auto x_vars = ctx.MultiInputVar("X");
C
chengduo 已提交
78
    auto x_vars_name = ctx.Inputs("X");
79 80 81 82 83 84 85 86 87 88 89 90

    framework::LibraryType library{framework::LibraryType::kPlain};
    framework::DataLayout layout{framework::DataLayout::kAnyLayout};

#ifdef PADDLE_WITH_MKLDNN
    if (library == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library = framework::LibraryType::kMKLDNN;
      layout = framework::DataLayout::kMKLDNN;
    }
#endif

91
    if (x_vars[0]->IsType<framework::LoDTensor>()) {
92
      int dtype = -1;
C
chengduo 已提交
93 94 95
      for (size_t idx = 0; idx < x_vars.size(); ++idx) {
        PADDLE_ENFORCE(x_vars[idx] != nullptr,
                       "Input var[%s] should not be nullptr", x_vars_name[idx]);
C
chengduo 已提交
96 97
        auto tensor =
            framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_vars[idx]);
98
        if (tensor->numel() == 0) {
99 100 101
          continue;
        }
        if (dtype == -1) {
Y
Yu Yang 已提交
102
          dtype = tensor->type();
103
        } else {
Y
Yu Yang 已提交
104
          PADDLE_ENFORCE_EQ(dtype, tensor->type());
105 106 107 108 109
        }
      }
      PADDLE_ENFORCE_NE(dtype, -1,
                        "Sum operator should have at least one tensor");

110
      return framework::OpKernelType(
111 112
          static_cast<framework::proto::VarType::Type>(dtype), ctx.GetPlace(),
          layout, library);
113
    } else if (x_vars[0]->IsType<framework::SelectedRows>()) {
114 115 116
      for (auto& var : x_vars) {
        auto& value = var->Get<framework::SelectedRows>().value();
        if (value.IsInitialized()) {
Y
Yu Yang 已提交
117 118
          return framework::OpKernelType(value.type(), ctx.device_context(),
                                         layout, library);
119 120 121 122
        }
      }
      // if input sparse vars are not initialized, use an default kernel type.
      return framework::OpKernelType(framework::proto::VarType::FP32,
123
                                     ctx.device_context(), layout, library);
124
    } else if (x_vars[0]->IsType<framework::LoDTensorArray>()) {
Y
Yang Yang(Tony) 已提交
125 126 127 128
      for (auto& x_var : x_vars) {
        auto& array = x_var->Get<framework::LoDTensorArray>();
        for (auto& each : array) {
          if (each.numel() != 0) {
Y
Yu Yang 已提交
129 130
            return framework::OpKernelType(each.type(), ctx.device_context(),
                                           layout, library);
Y
Yang Yang(Tony) 已提交
131
          }
132 133
        }
      }
Y
Yang Yang(Tony) 已提交
134
      PADDLE_THROW("Cannot find the input data type by all input data");
135 136
    }
    PADDLE_THROW("Unexpected branch. Input type is %s",
S
sneaxiy 已提交
137
                 framework::ToTypeName(x_vars[0]->Type()));
138
  }
139 140 141 142
};

class SumOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
143
  void Make() override {
144 145
    AddInput("X", "(vector<Tensor>) The input tensors of sum operator.")
        .AsDuplicable();
146
    AddOutput("Out", "(Tensor) The output tensor of sum operator.");
147 148 149
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
150
    AddComment(R"DOC(
151
Sum operator.
152

153 154
This operators sums the input tensors. All the inputs can carry the
LoD (Level of Details) information. However, the output only shares
155
the LoD information with the first input.
156
)DOC");
157 158 159
  }
};

Q
QI JUN 已提交
160 161
class SumOpVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
162 163
  void operator()(framework::InferVarTypeContext& ctx) const override {
    auto& inputs = ctx.Input("X");
164
    auto var_type = framework::proto::VarType::SELECTED_ROWS;
M
minqiyang 已提交
165 166
    for (auto& name : ctx.Input("X")) {
      VLOG(10) << name << " " << ctx.GetType(name);
Y
Yang Yang(Tony) 已提交
167 168
    }

Q
QI JUN 已提交
169
    bool any_input_is_lod_tensor = std::any_of(
M
minqiyang 已提交
170 171
        inputs.begin(), inputs.end(), [ctx](const std::string& name) {
          return ctx.GetType(name) == framework::proto::VarType::LOD_TENSOR;
Q
QI JUN 已提交
172
        });
173

M
minqiyang 已提交
174 175
    auto is_tensor_array = [ctx](const std::string& name) {
      return ctx.GetType(name) == framework::proto::VarType::LOD_TENSOR_ARRAY;
176 177 178 179 180 181 182 183
    };

    bool any_input_is_tensor_array =
        std::any_of(inputs.begin(), inputs.end(), is_tensor_array);
    bool all_inputs_are_tensor_array =
        std::all_of(inputs.begin(), inputs.end(), is_tensor_array);

    if (any_input_is_tensor_array) {
Y
Yang Yang(Tony) 已提交
184 185 186
      if (!all_inputs_are_tensor_array) {
        std::ostringstream os;
        for (auto& each : inputs) {
M
minqiyang 已提交
187
          os << "    " << each << " type is " << ctx.GetType(each) << "\n";
Y
Yang Yang(Tony) 已提交
188 189 190 191
        }
        PADDLE_ENFORCE(all_inputs_are_tensor_array,
                       "Not all inputs are tensor array:\n%s", os.str());
      }
192
      var_type = framework::proto::VarType::LOD_TENSOR_ARRAY;
193
    } else if (any_input_is_lod_tensor) {
194
      var_type = framework::proto::VarType::LOD_TENSOR;
Q
QI JUN 已提交
195 196
    }

M
minqiyang 已提交
197 198 199
    auto out_var_name = ctx.Output("Out").front();
    ctx.SetType(out_var_name, var_type);
    ctx.SetDataType(out_var_name, ctx.GetDataType(inputs.front()));
Q
QI JUN 已提交
200 201 202
  }
};

203
class SumGradMaker : public framework::GradOpDescMakerBase {
204
 public:
205
  using framework::GradOpDescMakerBase::GradOpDescMakerBase;
206

Y
Yu Yang 已提交
207
  std::vector<std::unique_ptr<framework::OpDesc>> operator()() const override {
208
    auto x_grads = InputGrad("X", false);
Y
Yu Yang 已提交
209
    std::vector<std::unique_ptr<framework::OpDesc>> grad_ops;
210 211 212 213
    grad_ops.reserve(x_grads.size());
    auto og = OutputGrad("Out");
    std::transform(x_grads.begin(), x_grads.end(), std::back_inserter(grad_ops),
                   [&og](const std::string& x_grad) {
Y
Yu Yang 已提交
214
                     auto* grad_op = new framework::OpDesc();
Y
Yu Yang 已提交
215 216 217 218
                     grad_op->SetType("scale");
                     grad_op->SetInput("X", og);
                     grad_op->SetOutput("Out", {x_grad});
                     grad_op->SetAttr("scale", 1.0f);
Y
Yu Yang 已提交
219
                     return std::unique_ptr<framework::OpDesc>(grad_op);
220 221
                   });
    return grad_ops;
222 223 224 225 226 227 228
  }
};

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
229

Q
QI JUN 已提交
230 231
REGISTER_OPERATOR(sum, ops::SumOp, ops::SumOpMaker, ops::SumGradMaker,
                  ops::SumOpVarTypeInference);
232

Q
QI JUN 已提交
233 234 235 236 237
REGISTER_OP_CPU_KERNEL(
    sum, ops::SumKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SumKernel<paddle::platform::CPUDeviceContext, double>,
    ops::SumKernel<paddle::platform::CPUDeviceContext, int>,
    ops::SumKernel<paddle::platform::CPUDeviceContext, int64_t>);