sum_op.cc 10.4 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

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"
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 36
namespace paddle {
namespace operators {
using framework::Tensor;

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

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

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

46
    PADDLE_ENFORCE_GT(
47 48
        x_vars.size(),
        0,
49
        platform::errors::InvalidArgument("Input[X] should not be empty"));
L
Leo Chen 已提交
50 51

    PADDLE_ENFORCE_NOT_NULL(
52 53 54
        x_vars[0],
        platform::errors::NotFound("Input var[%s] should not be nullptr",
                                   x_vars_name[0]));
L
Leo Chen 已提交
55

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

82
      auto data_type = static_cast<framework::proto::VarType::Type>(dtype);
83 84
#ifdef PADDLE_WITH_MKLDNN
      if (library == framework::LibraryType::kPlain &&
85 86 87
          this->CanMKLDNNBeUsed(ctx, data_type) &&
          (data_type == framework::proto::VarType::FP32 ||
           data_type == framework::proto::VarType::BF16) &&
88
          ctx.OutputVar("Out")->IsType<framework::LoDTensor>()) {
89 90 91 92 93 94
        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(),
95 96
                                         framework::DataLayout::kMKLDNN,
                                         framework::LibraryType::kMKLDNN);
97 98 99 100
        }
      }
#endif

101 102
      return framework::OpKernelType(
          data_type, ctx.GetPlace(), layout, library);
103
    } else if (x_vars[0]->IsType<phi::SelectedRows>()) {
104
      for (auto& var : x_vars) {
105
        auto& value = var->Get<phi::SelectedRows>().value();
106
        if (value.IsInitialized()) {
107 108
          return framework::OpKernelType(
              framework::TransToProtoVarType(value.dtype()),
109 110 111
              ctx.device_context(),
              layout,
              library);
112 113 114 115
        }
      }
      // if input sparse vars are not initialized, use an default kernel type.
      return framework::OpKernelType(framework::proto::VarType::FP32,
116 117 118
                                     ctx.device_context(),
                                     layout,
                                     library);
119
    } else if (x_vars[0]->IsType<framework::LoDTensorArray>()) {
Y
Yang Yang(Tony) 已提交
120 121 122
      for (auto& x_var : x_vars) {
        auto& array = x_var->Get<framework::LoDTensorArray>();
        for (auto& each : array) {
123
          if (each.numel() != 0 && each.IsInitialized()) {
124 125
            return framework::OpKernelType(
                framework::TransToProtoVarType(each.dtype()),
126 127 128
                ctx.device_context(),
                layout,
                library);
Y
Yang Yang(Tony) 已提交
129
          }
130 131
        }
      }
132 133 134 135 136
      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())));
137
    }
138 139 140 141 142
    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())));
143
  }
144 145 146 147
};

class SumOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
148
  void Make() override {
149 150 151 152 153
    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.")
154
        .AsDuplicable();
155 156 157
    AddOutput("Out",
              "the sum of input :code:`x`. its shape and data types are "
              "consistent with :code:`x`.");
158 159 160
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
J
Jacek Czaja 已提交
161 162 163 164 165
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "bfloat16"});
166 167 168
    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");
169 170 171
  }
};

Q
QI JUN 已提交
172 173
class SumOpVarTypeInference : public framework::VarTypeInference {
 public:
M
minqiyang 已提交
174
  void operator()(framework::InferVarTypeContext* ctx) const override {
175 176 177 178 179 180 181 182
    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);
        }
      }
183

184 185 186 187 188 189 190 191 192 193 194
      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) 已提交
195
        }
196 197 198 199
        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) 已提交
200
      }
Q
QI JUN 已提交
201

202 203 204
      ctx->SetOutputType("Out", var_type);
      ctx->SetOutputDataType("Out", ctx->GetInputDataType("X"));
    }
Q
QI JUN 已提交
205 206 207
  }
};

H
hong 已提交
208
class SumGradDescMaker : public framework::GradOpDescMakerBase {
209
 public:
210
  using framework::GradOpDescMakerBase::GradOpDescMakerBase;
211

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

    return grad_ops;
  }
};

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

237
  std::shared_ptr<imperative::GradOpNode> operator()() const override {
H
hong 已提交
238
    auto x_grads = InputGrad("X", false);
239 240
    using InputGradsType = decltype(x_grads);

241 242 243 244 245 246 247 248 249 250
    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);
251
        op.SetDefaultAttrsMap(DefaultAttrsMap());
252 253 254 255 256
      }
      return node;
    } else {
      return nullptr;
    }
257 258 259
  }
};

260
DECLARE_INPLACE_OP_INFERER(SumInplaceInferer, {"X", "Out"});
261

262 263 264 265
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
266

267 268 269 270 271
namespace ops = paddle::operators;
DECLARE_INFER_SHAPE_FUNCTOR(sum,
                            AddNInferShapeFunctor,
                            PD_INFER_META(phi::AddNTensorArrayInferMeta));

272 273 274 275 276 277
REGISTER_OPERATOR(sum,
                  ops::SumOp,
                  ops::SumOpMaker,
                  ops::SumGradDescMaker,
                  ops::SumGradOpBaseMaker,
                  ops::SumOpVarTypeInference,
278 279
                  ops::SumInplaceInferer,
                  AddNInferShapeFunctor);