sum_op.cc 9.9 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. */

12
#include <algorithm>
M
minqiyang 已提交
13
#include <memory>
14
#include <string>
15
#include <unordered_map>
16
#include <vector>
17

Y
YuanRisheng 已提交
18 19
#include "paddle/fluid/framework/infershape_utils.h"
#include "paddle/fluid/framework/op_registry.h"
Y
Yi Wang 已提交
20
#include "paddle/fluid/framework/var_type_inference.h"
Y
YuanRisheng 已提交
21 22
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/multiary.h"
23

24 25 26
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
27
#include "paddle/fluid/framework/convert_utils.h"
28

29 30 31 32 33 34 35
namespace paddle {
namespace operators {

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

36
 protected:
37
  framework::OpKernelType GetExpectedKernelType(
38 39
      const framework::ExecutionContext& ctx) const override {
    auto x_vars = ctx.MultiInputVar("X");
H
hong 已提交
40
    auto x_vars_name = ctx.InputNames("X");
41

42
    PADDLE_ENFORCE_GT(
43 44
        x_vars.size(),
        0,
45
        platform::errors::InvalidArgument("Input[X] should not be empty"));
L
Leo Chen 已提交
46 47

    PADDLE_ENFORCE_NOT_NULL(
48 49 50
        x_vars[0],
        platform::errors::NotFound("Input var[%s] should not be nullptr",
                                   x_vars_name[0]));
L
Leo Chen 已提交
51

52
    if (x_vars[0]->IsType<framework::LoDTensor>()) {
53
      int dtype = -1;
C
chengduo 已提交
54
      for (size_t idx = 0; idx < x_vars.size(); ++idx) {
55 56 57 58
        PADDLE_ENFORCE_NOT_NULL(
            x_vars[idx],
            platform::errors::NotFound("Input var[%s] should not be nullptr",
                                       x_vars_name[idx]));
C
chengduo 已提交
59 60
        auto tensor =
            framework::GetLoDTensorOrSelectedRowsValueFromVar(*x_vars[idx]);
61
        if (tensor->numel() <= 0 || (!tensor->IsInitialized())) {
62 63 64
          continue;
        }
        if (dtype == -1) {
65
          dtype = framework::TransToProtoVarType(tensor->dtype());
66
        } else {
67 68
          PADDLE_ENFORCE_EQ(dtype,
                            framework::TransToProtoVarType(tensor->dtype()),
69 70
                            platform::errors::InvalidArgument(
                                "The inputs type of sum op must be same"));
71 72
        }
      }
73 74
      PADDLE_ENFORCE_NE(dtype,
                        -1,
75 76
                        platform::errors::InvalidArgument(
                            "Sum operator should have at least one tensor"));
77

78
      auto data_type = static_cast<framework::proto::VarType::Type>(dtype);
79
#ifdef PADDLE_WITH_MKLDNN
80
      if (this->CanMKLDNNBeUsed(ctx, data_type) &&
81 82
          (data_type == framework::proto::VarType::FP32 ||
           data_type == framework::proto::VarType::BF16) &&
83
          ctx.OutputVar("Out")->IsType<framework::LoDTensor>()) {
84 85 86 87 88 89
        if (std::all_of(
                x_vars.begin(), x_vars.end(), [](const framework::Variable* v) {
                  return v->IsType<framework::LoDTensor>();
                })) {
          return framework::OpKernelType(data_type,
                                         ctx.GetPlace(),
90 91
                                         framework::DataLayout::kMKLDNN,
                                         framework::LibraryType::kMKLDNN);
92 93 94
        }
      }
#endif
95
      return framework::OpKernelType(data_type, ctx.GetPlace());
96
    } else if (x_vars[0]->IsType<phi::SelectedRows>()) {
97
      for (auto& var : x_vars) {
98
        auto& value = var->Get<phi::SelectedRows>().value();
99
        if (value.IsInitialized()) {
100 101
          return framework::OpKernelType(
              framework::TransToProtoVarType(value.dtype()),
102
              ctx.device_context());
103 104 105 106
        }
      }
      // if input sparse vars are not initialized, use an default kernel type.
      return framework::OpKernelType(framework::proto::VarType::FP32,
107
                                     ctx.device_context());
108
    } else if (x_vars[0]->IsType<framework::LoDTensorArray>()) {
Y
Yang Yang(Tony) 已提交
109 110 111
      for (auto& x_var : x_vars) {
        auto& array = x_var->Get<framework::LoDTensorArray>();
        for (auto& each : array) {
112
          if (each.numel() != 0 && each.IsInitialized()) {
113 114
            return framework::OpKernelType(
                framework::TransToProtoVarType(each.dtype()),
115
                ctx.device_context());
Y
Yang Yang(Tony) 已提交
116
          }
117 118
        }
      }
119 120 121 122 123
      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())));
124
    }
125 126 127 128 129
    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())));
130
  }
131 132 133 134
};

class SumOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
135
  void Make() override {
136 137 138 139 140
    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.")
141
        .AsDuplicable();
142 143 144
    AddOutput("Out",
              "the sum of input :code:`x`. its shape and data types are "
              "consistent with :code:`x`.");
145 146 147
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
J
Jacek Czaja 已提交
148 149 150 151 152
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "bfloat16"});
153 154 155
    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");
156 157 158
  }
};

Q
QI JUN 已提交
159 160
class SumOpVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
161
  void operator()(framework::InferVarTypeContext* ctx) const override {
162 163 164 165 166 167 168 169
    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);
        }
      }
170

171 172 173 174 175 176 177 178 179 180 181
      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) 已提交
182
        }
183 184 185 186
        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) 已提交
187
      }
Q
QI JUN 已提交
188

189 190 191
      ctx->SetOutputType("Out", var_type);
      ctx->SetOutputDataType("Out", ctx->GetInputDataType("X"));
    }
Q
QI JUN 已提交
192 193 194
  }
};

H
hong 已提交
195
class SumGradDescMaker : public framework::GradOpDescMakerBase {
196
 public:
197
  using framework::GradOpDescMakerBase::GradOpDescMakerBase;
198

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

    return grad_ops;
  }
};

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

224
  std::shared_ptr<imperative::GradOpNode> operator()() const override {
H
hong 已提交
225
    auto x_grads = InputGrad("X", false);
226 227
    using InputGradsType = decltype(x_grads);

228 229 230 231 232 233 234 235 236 237
    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);
238
        op.SetDefaultAttrsMap(DefaultAttrsMap());
239 240 241 242 243
      }
      return node;
    } else {
      return nullptr;
    }
244 245 246
  }
};

247
DECLARE_INPLACE_OP_INFERER(SumInplaceInferer, {"X", "Out"});
248

249 250 251 252
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
253

Y
YuanRisheng 已提交
254 255 256 257 258
namespace ops = paddle::operators;
DECLARE_INFER_SHAPE_FUNCTOR(sum,
                            AddNInferShapeFunctor,
                            PD_INFER_META(phi::AddNTensorArrayInferMeta));

259 260 261 262 263 264
REGISTER_OPERATOR(sum,
                  ops::SumOp,
                  ops::SumOpMaker,
                  ops::SumGradDescMaker,
                  ops::SumGradOpBaseMaker,
                  ops::SumOpVarTypeInference,
Y
YuanRisheng 已提交
265 266
                  ops::SumInplaceInferer,
                  AddNInferShapeFunctor);