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
      auto data_type = static_cast<framework::proto::VarType::Type>(dtype);
149 150
#ifdef PADDLE_WITH_MKLDNN
      if (library == framework::LibraryType::kPlain &&
151 152 153
          this->CanMKLDNNBeUsed(ctx, data_type) &&
          (data_type == framework::proto::VarType::FP32 ||
           data_type == framework::proto::VarType::BF16) &&
154 155 156 157 158
          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>();
                        })) {
159 160 161
          return framework::OpKernelType(data_type, ctx.GetPlace(),
                                         framework::DataLayout::kMKLDNN,
                                         framework::LibraryType::kMKLDNN);
162 163 164 165
        }
      }
#endif

166 167
      return framework::OpKernelType(data_type, ctx.GetPlace(), layout,
                                     library);
168
    } else if (x_vars[0]->IsType<pten::SelectedRows>()) {
169
      for (auto& var : x_vars) {
170
        auto& value = var->Get<pten::SelectedRows>().value();
171
        if (value.IsInitialized()) {
Y
Yu Yang 已提交
172 173
          return framework::OpKernelType(value.type(), ctx.device_context(),
                                         layout, library);
174 175 176 177
        }
      }
      // if input sparse vars are not initialized, use an default kernel type.
      return framework::OpKernelType(framework::proto::VarType::FP32,
178
                                     ctx.device_context(), layout, library);
179
    } else if (x_vars[0]->IsType<framework::LoDTensorArray>()) {
Y
Yang Yang(Tony) 已提交
180 181 182
      for (auto& x_var : x_vars) {
        auto& array = x_var->Get<framework::LoDTensorArray>();
        for (auto& each : array) {
183
          if (each.numel() != 0 && each.IsInitialized()) {
Y
Yu Yang 已提交
184 185
            return framework::OpKernelType(each.type(), ctx.device_context(),
                                           layout, library);
Y
Yang Yang(Tony) 已提交
186
          }
187 188
        }
      }
189 190 191 192 193
      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())));
194
    }
195 196 197 198 199
    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())));
200
  }
201 202 203 204
};

class SumOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
205
  void Make() override {
206 207 208 209 210
    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.")
211
        .AsDuplicable();
212 213 214
    AddOutput("Out",
              "the sum of input :code:`x`. its shape and data types are "
              "consistent with :code:`x`.");
215 216 217
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
J
Jacek Czaja 已提交
218 219 220 221 222
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "bfloat16"});
223 224 225
    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");
226 227 228
  }
};

Q
QI JUN 已提交
229 230
class SumOpVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
231
  void operator()(framework::InferVarTypeContext* ctx) const override {
232 233 234 235 236 237 238 239
    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);
        }
      }
240

241 242 243 244 245 246 247 248 249 250 251
      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) 已提交
252
        }
253 254 255 256
        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) 已提交
257
      }
Q
QI JUN 已提交
258

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

H
hong 已提交
265
class SumGradDescMaker : public framework::GradOpDescMakerBase {
266
 public:
267
  using framework::GradOpDescMakerBase::GradOpDescMakerBase;
268

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

    return grad_ops;
  }
};

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

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

296 297 298 299 300 301 302 303 304 305 306 307 308 309 310
    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;
    }
311 312 313
  }
};

314
DECLARE_INPLACE_OP_INFERER(SumInplaceInferer, {"X", "Out"});
315

316 317 318 319
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
320

H
hong 已提交
321 322
REGISTER_OPERATOR(sum, ops::SumOp, ops::SumOpMaker, ops::SumGradDescMaker,
                  ops::SumGradOpBaseMaker, ops::SumOpVarTypeInference,
323
                  ops::SumInplaceInferer);
324

Q
QI JUN 已提交
325 326 327 328
REGISTER_OP_CPU_KERNEL(
    sum, ops::SumKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SumKernel<paddle::platform::CPUDeviceContext, double>,
    ops::SumKernel<paddle::platform::CPUDeviceContext, int>,
329 330
    ops::SumKernel<paddle::platform::CPUDeviceContext,
                   paddle::platform::bfloat16>,
Q
QI JUN 已提交
331
    ops::SumKernel<paddle::platform::CPUDeviceContext, int64_t>);