sum_op.cc 12.7 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
#include <algorithm>
M
minqiyang 已提交
15
#include <memory>
16
#include <string>
17
#include <unordered_map>
18
#include <vector>
19

Y
Yi Wang 已提交
20
#include "paddle/fluid/framework/var_type_inference.h"
21

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

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

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

34
  void InferShape(framework::InferShapeContext* ctx) const override {
35 36
    OP_INOUT_CHECK(ctx->HasInputs("X"), "Input", "X", "sum");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "sum");
37 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

47 48
    auto N = x_dims.size();
    PADDLE_ENFORCE_GT(
49 50 51 52 53
        N, 0, platform::errors::InvalidArgument(
                  "The input tensor X's dimensions of SumOp "
                  "should be larger than 0. But received X's dimensions %d, "
                  "X's shape = [%s].",
                  N, &x_dims));
54
    if (N == 1) {
55
      VLOG(3) << "Warning: SumOp have only one input, may waste memory";
56
    }
Q
qiaolongfei 已提交
57

58
    framework::DDim in_dim({0});
59
    for (size_t i = 0; i < x_dims.size(); ++i) {
60 61 62 63
      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) {
64 65
        continue;
      }
66 67 68 69 70 71
      if (framework::product(x_dim) == 0) {
        continue;
      }
      if (framework::product(in_dim) == 0) {
        in_dim = x_dim;
      } else {
Z
zhaoyuchen 已提交
72
        if (ctx->IsRuntime()) {
73 74 75 76 77 78
          PADDLE_ENFORCE_EQ(in_dim, x_dim,
                            platform::errors::InvalidArgument(
                                "The input tensor X of SumOp must"
                                " have same shape. But received X[0]'s shape = "
                                "[%s], X[%d]'s shape = [%s].",
                                in_dim, i, x_dim));
Z
zhaoyuchen 已提交
79
        } else {
80 81
          PADDLE_ENFORCE_EQ(
              in_dim.size(), x_dim.size(),
82 83 84 85 86 87
              platform::errors::InvalidArgument(
                  "The input tensor X of SumOp must have same "
                  "dimensions. But received X[0]'s dimensions = %d, X[0]'s "
                  "shape = "
                  "[%s], X[%d]'s dimensions = %d, X[%d]'s shape = [%s].",
                  in_dim.size(), in_dim, i, x_dim.size(), i, x_dim));
Z
zhaoyuchen 已提交
88
          // if in_dim or x_dim has -1, not check equal
89 90
          for (int j = 0; j < x_dim.size(); ++j) {
            if (x_dim[j] == -1 || in_dim[j] == -1) {
Z
zhaoyuchen 已提交
91 92
              continue;
            }
93 94
            PADDLE_ENFORCE_EQ(
                in_dim[j], x_dim[j],
95 96 97 98 99
                platform::errors::InvalidArgument(
                    "The input tensor X of SumOp must have same shape "
                    "if not -1."
                    "But received X[0]'s shape = [%s], X[%d]'s shape = [%s].",
                    in_dim, i, x_dim));
Z
zhaoyuchen 已提交
100 101
          }
        }
102
      }
Q
qijun 已提交
103
    }
Q
Qiao Longfei 已提交
104 105
    ctx->SetOutputDim("Out", in_dim);
    ctx->ShareLoD("X", /*->*/ "Out");
106
  }
107 108

 protected:
109
  framework::OpKernelType GetExpectedKernelType(
110 111
      const framework::ExecutionContext& ctx) const override {
    auto x_vars = ctx.MultiInputVar("X");
H
hong 已提交
112
    auto x_vars_name = ctx.InputNames("X");
113 114 115 116

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

L
Leo Chen 已提交
117 118 119 120 121 122 123
    PADDLE_ENFORCE_GT(x_vars.size(), 0, platform::errors::InvalidArgument(
                                            "Input[X] should not be empty"));

    PADDLE_ENFORCE_NOT_NULL(
        x_vars[0], platform::errors::NotFound(
                       "Input var[%s] should not be nullptr", x_vars_name[0]));

124
    if (x_vars[0]->IsType<framework::LoDTensor>()) {
125
      int dtype = -1;
C
chengduo 已提交
126
      for (size_t idx = 0; idx < x_vars.size(); ++idx) {
127 128 129 130
        PADDLE_ENFORCE_NOT_NULL(
            x_vars[idx],
            platform::errors::NotFound("Input var[%s] should not be nullptr",
                                       x_vars_name[idx]));
C
chengduo 已提交
131 132
        auto tensor =
            framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_vars[idx]);
133
        if (tensor->numel() <= 0 || (!tensor->IsInitialized())) {
134 135 136
          continue;
        }
        if (dtype == -1) {
Y
Yu Yang 已提交
137
          dtype = tensor->type();
138
        } else {
139 140 141
          PADDLE_ENFORCE_EQ(dtype, tensor->type(),
                            platform::errors::InvalidArgument(
                                "The inputs type of sum op must be same"));
142 143 144
        }
      }
      PADDLE_ENFORCE_NE(dtype, -1,
145 146
                        platform::errors::InvalidArgument(
                            "Sum operator should have at least one tensor"));
147

148 149
#ifdef PADDLE_WITH_MKLDNN
      if (library == framework::LibraryType::kPlain &&
150
          this->CanMKLDNNBeUsed(ctx) &&
J
Jacek Czaja 已提交
151 152 153 154
          (static_cast<framework::proto::VarType::Type>(dtype) ==
               framework::proto::VarType::FP32 ||
           static_cast<framework::proto::VarType::Type>(dtype) ==
               framework::proto::VarType::BF16) &&
155 156 157 158 159 160
          ctx.OutputVar("Out")->IsType<framework::LoDTensor>()) {
        if (std::all_of(x_vars.begin(), x_vars.end(),
                        [](const framework::Variable* v) {
                          return v->IsType<framework::LoDTensor>();
                        })) {
          return framework::OpKernelType(
J
Jacek Czaja 已提交
161 162 163
              static_cast<framework::proto::VarType::Type>(dtype),
              ctx.GetPlace(), framework::DataLayout::kMKLDNN,
              framework::LibraryType::kMKLDNN);
164 165 166 167
        }
      }
#endif

168
      return framework::OpKernelType(
169 170
          static_cast<framework::proto::VarType::Type>(dtype), ctx.GetPlace(),
          layout, library);
171
    } else if (x_vars[0]->IsType<framework::SelectedRows>()) {
172 173 174
      for (auto& var : x_vars) {
        auto& value = var->Get<framework::SelectedRows>().value();
        if (value.IsInitialized()) {
Y
Yu Yang 已提交
175 176
          return framework::OpKernelType(value.type(), ctx.device_context(),
                                         layout, library);
177 178 179 180
        }
      }
      // if input sparse vars are not initialized, use an default kernel type.
      return framework::OpKernelType(framework::proto::VarType::FP32,
181
                                     ctx.device_context(), layout, library);
182
    } else if (x_vars[0]->IsType<framework::LoDTensorArray>()) {
Y
Yang Yang(Tony) 已提交
183 184 185
      for (auto& x_var : x_vars) {
        auto& array = x_var->Get<framework::LoDTensorArray>();
        for (auto& each : array) {
186
          if (each.numel() != 0 && each.IsInitialized()) {
Y
Yu Yang 已提交
187 188
            return framework::OpKernelType(each.type(), ctx.device_context(),
                                           layout, library);
Y
Yang Yang(Tony) 已提交
189
          }
190 191
        }
      }
192 193 194 195 196
      PADDLE_THROW(platform::errors::InvalidArgument(
          "Expected each tensor in Input(x) in sum op has be initialized, but "
          "some tensor in Input(x) is not be initialized, please check your "
          "code.",
          framework::ToTypeName(x_vars[0]->Type())));
197
    }
198 199 200 201 202
    PADDLE_THROW(platform::errors::InvalidArgument(
        "Expected type of Input(X) must be Tensor,  SelectedRows or "
        "LodTensorArray. But got "
        "unsupport type: %s.",
        framework::ToTypeName(x_vars[0]->Type())));
203
  }
204 205 206 207
};

class SumOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
208
  void Make() override {
209 210 211 212 213
    AddInput("X",
             "A Varaible list. The shape and data type of the list elements"
             "should be consistent. Variable can be multi-dimensional Tensor"
             "or LoDTensor, and data types can be: float32, float64, int32, "
             "int64.")
214
        .AsDuplicable();
215 216 217
    AddOutput("Out",
              "the sum of input :code:`x`. its shape and data types are "
              "consistent with :code:`x`.");
218 219 220
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
J
Jacek Czaja 已提交
221 222 223 224 225
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "bfloat16"});
226 227 228
    AddComment(R"DOC(This OP is used to sum one or more Tensor or LoDTensor
                    of the input. If the input is LoDTensor, the output only
                    shares LoD information with the first input.)DOC");
229 230 231
  }
};

Q
QI JUN 已提交
232 233
class SumOpVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
234
  void operator()(framework::InferVarTypeContext* ctx) const override {
235 236 237 238 239 240 241 242
    if (!ctx->IsDygraph()) {
      auto var_type = framework::proto::VarType::SELECTED_ROWS;
      if (VLOG_IS_ON(10)) {
        for (size_t ind = 0; ind < ctx->InputSize("X"); ++ind) {
          VLOG(10) << ctx->InputVarName("X", ind) << " "
                   << ctx->GetInputType("X", ind);
        }
      }
243

244 245 246 247 248 249 250 251 252 253 254
      if (ctx->InputTypeAnyOf("X",
                              framework::proto::VarType::LOD_TENSOR_ARRAY)) {
        if (!ctx->InputTypeAllOf("X",
                                 framework::proto::VarType::LOD_TENSOR_ARRAY)) {
          std::ostringstream os;
          for (size_t ind = 0; ind < ctx->InputSize("X"); ++ind) {
            os << "    " << ctx->InputVarName("X", ind) << " type is "
               << ctx->GetInputType("X", ind) << "\n";
          }
          PADDLE_THROW(platform::errors::InvalidArgument(
              "Not all inputs are tensor array:\n%s", os.str()));
Y
Yang Yang(Tony) 已提交
255
        }
256 257 258 259
        var_type = framework::proto::VarType::LOD_TENSOR_ARRAY;
      } else if (ctx->InputTypeAnyOf("X",
                                     framework::proto::VarType::LOD_TENSOR)) {
        var_type = framework::proto::VarType::LOD_TENSOR;
Y
Yang Yang(Tony) 已提交
260
      }
Q
QI JUN 已提交
261

262 263 264
      ctx->SetOutputType("Out", var_type);
      ctx->SetOutputDataType("Out", ctx->GetInputDataType("X"));
    }
Q
QI JUN 已提交
265 266 267
  }
};

H
hong 已提交
268
class SumGradDescMaker : public framework::GradOpDescMakerBase {
269
 public:
270
  using framework::GradOpDescMakerBase::GradOpDescMakerBase;
271

Y
Yu Yang 已提交
272
  std::vector<std::unique_ptr<framework::OpDesc>> operator()() const override {
273
    auto x_grads = InputGrad("X", false);
Y
Yu Yang 已提交
274
    std::vector<std::unique_ptr<framework::OpDesc>> grad_ops;
275 276 277 278
    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 已提交
279
                     auto* grad_op = new framework::OpDesc();
Y
Yu Yang 已提交
280 281 282 283
                     grad_op->SetType("scale");
                     grad_op->SetInput("X", og);
                     grad_op->SetOutput("Out", {x_grad});
                     grad_op->SetAttr("scale", 1.0f);
Y
Yu Yang 已提交
284
                     return std::unique_ptr<framework::OpDesc>(grad_op);
285
                   });
H
hong 已提交
286 287 288 289 290 291 292 293 294

    return grad_ops;
  }
};

class SumGradOpBaseMaker : public imperative::GradOpBaseMakerBase {
 public:
  using imperative::GradOpBaseMakerBase::GradOpBaseMakerBase;

295
  std::shared_ptr<imperative::GradOpNode> operator()() const override {
H
hong 已提交
296
    auto x_grads = InputGrad("X", false);
297 298
    using InputGradsType = decltype(x_grads);

299 300 301 302 303 304 305 306 307 308 309 310 311 312 313
    if (!x_grads.empty()) {
      auto node = this->NewGradNode();
      node->reserve(x_grads.size());
      auto og = OutputGrad("Out");
      for (auto& x_grad : x_grads) {
        imperative::TracedGradOp op(node);
        op.SetType("scale");
        op.SetInput("X", og);
        op.SetOutput("Out", InputGradsType{x_grad});
        op.SetAttr("scale", 1.0f);
      }
      return node;
    } else {
      return nullptr;
    }
314 315 316
  }
};

317
DECLARE_INPLACE_OP_INFERER(SumInplaceInferer, {"X", "Out"});
318

319 320 321 322
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
323

H
hong 已提交
324 325
REGISTER_OPERATOR(sum, ops::SumOp, ops::SumOpMaker, ops::SumGradDescMaker,
                  ops::SumGradOpBaseMaker, ops::SumOpVarTypeInference,
326
                  ops::SumInplaceInferer);
327

Q
QI JUN 已提交
328 329 330 331 332
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>);