mul_op.cc 14.2 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2

3 4 5
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
6

7
    http://www.apache.org/licenses/LICENSE-2.0
8

9 10 11 12 13
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. */
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/mul_op.h"
16
#include <memory>
17
#include <string>
18
#include <unordered_map>
19
#include <vector>
P
Physher 已提交
20 21 22
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
23 24 25 26

namespace paddle {
namespace operators {

27
using framework::OpKernelType;
D
dongzhihong 已提交
28 29
using framework::Tensor;

30
class MulOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
31
 public:
32 33 34
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
35 36 37
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "Mul");
    OP_INOUT_CHECK(ctx->HasInput("Y"), "Input", "Y", "Mul");
    OP_INOUT_CHECK(ctx->HasOutput("Out"), "Output", "Out", "Mul");
Q
Qiao Longfei 已提交
38 39 40

    auto x_dims = ctx->GetInputDim("X");
    auto y_dims = ctx->GetInputDim("Y");
Y
Yu Yang 已提交
41

Q
Qiao Longfei 已提交
42 43
    int x_num_col_dims = ctx->Attrs().Get<int>("x_num_col_dims");
    int y_num_col_dims = ctx->Attrs().Get<int>("y_num_col_dims");
F
WIP  
fengjiayi 已提交
44

M
minqiyang 已提交
45 46 47
    VLOG(3) << "mul operator x.shape=" << x_dims << " y.shape=" << y_dims
            << " x_num_col_dims=" << x_num_col_dims
            << " y_num_col_dims=" << y_num_col_dims;
Y
Yu Yang 已提交
48

49
    PADDLE_ENFORCE_NE(phi::product(y_dims), 0,
50
                      platform::errors::PreconditionNotMet(
51
                          "The Input variable Y(%s) has not "
52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69
                          "been initialized. You may need to confirm "
                          "if you put exe.run(startup_program) "
                          "after optimizer.minimize function.",
                          ctx->Inputs("Y").front()));
    PADDLE_ENFORCE_GT(
        x_dims.size(), x_num_col_dims,
        platform::errors::InvalidArgument(
            "The input tensor X's dimensions of MulOp "
            "should be larger than x_num_col_dims. But received X's "
            "dimensions = %d, X's shape = [%s], x_num_col_dims = %d.",
            x_dims.size(), x_dims, x_num_col_dims));
    PADDLE_ENFORCE_GT(
        y_dims.size(), y_num_col_dims,
        platform::errors::InvalidArgument(
            "The input tensor Y's dimensions of MulOp "
            "should be larger than y_num_col_dims. But received Y's "
            "dimensions = %d, Y's shape = [%s], y_num_col_dims = %d.",
            y_dims.size(), y_dims, y_num_col_dims));
70

71 72
    auto x_mat_dims = phi::flatten_to_2d(x_dims, x_num_col_dims);
    auto y_mat_dims = phi::flatten_to_2d(y_dims, y_num_col_dims);
73

74 75
    PADDLE_ENFORCE_EQ(
        x_mat_dims[1], y_mat_dims[0],
76
        platform::errors::InvalidArgument(
77 78 79
            "After flatten the input tensor X and Y to 2-D dimensions matrix "
            "X1 and Y1, the matrix X1's width must be equal with matrix Y1's "
            "height. But received X's shape = [%s], X1's shape = [%s], X1's "
80 81 82 83
            "width = %s; Y's shape = [%s], Y1's shape = [%s], Y1's height = "
            "%s.",
            x_dims, x_mat_dims, x_mat_dims[1], y_dims, y_mat_dims,
            y_mat_dims[0]));
Y
Yu Yang 已提交
84 85 86 87 88 89 90 91 92 93 94 95
    std::vector<int64_t> output_dims;
    output_dims.reserve(
        static_cast<size_t>(x_num_col_dims + y_dims.size() - y_num_col_dims));

    for (int i = 0; i < x_num_col_dims; ++i) {
      output_dims.push_back(x_dims[i]);
    }

    for (int i = y_num_col_dims; i < y_dims.size(); ++i) {
      output_dims.push_back(y_dims[i]);
    }

96
    ctx->SetOutputDim("Out", phi::make_ddim(output_dims));
Q
Qiao Longfei 已提交
97
    ctx->ShareLoD("X", /*->*/ "Out");
98
  }
P
Physher 已提交
99 100 101 102 103 104 105

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const {
    framework::LibraryType library = framework::LibraryType::kPlain;
    framework::DataLayout layout = framework::DataLayout::kAnyLayout;
    int customized_type_value =
        framework::OpKernelType::kDefaultCustomizedTypeValue;
106
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
P
Physher 已提交
107 108
#ifdef PADDLE_WITH_MKLDNN
    if (library == framework::LibraryType::kPlain &&
109
        this->CanMKLDNNBeUsed(ctx, input_data_type)) {
P
Physher 已提交
110 111 112
      library = framework::LibraryType::kMKLDNN;
      layout = framework::DataLayout::kMKLDNN;

113 114
      if (input_data_type == framework::DataTypeTrait<int8_t>::DataType() ||
          input_data_type == framework::DataTypeTrait<uint8_t>::DataType()) {
P
Physher 已提交
115
        customized_type_value = kMULMKLDNNINT8;
116 117 118 119 120 121
      } else if (input_data_type ==
                     framework::DataTypeTrait<
                         paddle::platform::bfloat16>::DataType() ||
                 input_data_type ==
                     framework::DataTypeTrait<float>::DataType()) {
        customized_type_value = kMULMKLDNNFP32;
P
Physher 已提交
122 123 124 125 126 127 128
      }
    }
#endif

    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout,
                                   library, customized_type_value);
  }
129 130
};

D
dongzhihong 已提交
131
class MulOpMaker : public framework::OpProtoAndCheckerMaker {
132
 public:
Y
Yu Yang 已提交
133
  void Make() override {
C
caoying03 已提交
134 135 136
    AddInput("X", "(Tensor), The first input tensor of mul op.");
    AddInput("Y", "(Tensor), The second input tensor of mul op.");
    AddOutput("Out", "(Tensor), The output tensor of mul op.");
P
Physher 已提交
137 138
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
139 140
        .SetDefault(false)
        .AsExtra();
F
WIP  
fengjiayi 已提交
141
    AddAttr<int>(
F
fengjiayi 已提交
142
        "x_num_col_dims",
C
caoying03 已提交
143 144 145 146 147 148 149 150 151 152 153 154 155 156 157
        R"DOC((int, default 1), The mul_op can take tensors with more than two
              dimensions as its inputs. If the input $X$ is a tensor with more
              than two dimensions, $X$ will be flattened into a two-dimensional
              matrix first. The flattening rule is: the first `num_col_dims`
              will be flattened to form the first dimension of the final matrix
              (the height of the matrix), and the rest `rank(X) - num_col_dims`
              dimensions are flattened to form the second dimension of the final
              matrix (the width of the matrix). As a result, height of the
              flattened matrix is equal to the product of $X$'s first
              `x_num_col_dims` dimensions' sizes, and width of the flattened
              matrix is equal to the product of $X$'s last `rank(x) - num_col_dims`
              dimensions' size. For example, suppose $X$ is a 6-dimensional
              tensor with the shape [2, 3, 4, 5, 6], and `x_num_col_dims` = 3.
              Thus, the flattened matrix will have a shape [2 x 3 x 4, 5 x 6] =
              [24, 30].
F
fengjiayi 已提交
158
        )DOC")
F
WIP  
fengjiayi 已提交
159
        .SetDefault(1)
F
fengjiayi 已提交
160
        .EqualGreaterThan(1);
F
WIP  
fengjiayi 已提交
161
    AddAttr<int>(
F
fengjiayi 已提交
162
        "y_num_col_dims",
C
caoying03 已提交
163 164 165 166
        R"DOC((int, default 1), The mul_op can take tensors with more than two,
              dimensions as its inputs. If the input $Y$ is a tensor with more
              than two dimensions, $Y$ will be flattened into a two-dimensional
              matrix first. The attribute `y_num_col_dims` determines how $Y$ is
C
caoying03 已提交
167
              flattened. See comments of `x_num_col_dims` for more details.
F
fengjiayi 已提交
168
        )DOC")
F
WIP  
fengjiayi 已提交
169
        .SetDefault(1)
F
fengjiayi 已提交
170
        .EqualGreaterThan(1);
171 172 173 174 175
    AddAttr<float>(
        "scale_x",
        "scale_x to be used for int8 mul input data x. scale_x has the"
        "same purpose as scale_in in OPs that support quantization."
        "Only to be used with MKL-DNN INT8")
176 177
        .SetDefault(1.0f)
        .AsExtra();
178 179 180 181 182
    AddAttr<std::vector<float>>(
        "scale_y",
        "scale_y to be used for int8 mul input data y. scale_y has the"
        "same purpose as scale_weights in OPs that support quantization."
        "Only to be used with MKL-DNN INT8")
183 184
        .SetDefault({1.0f})
        .AsExtra();
P
Physher 已提交
185 186 187
    AddAttr<float>("scale_out",
                   "scale_out to be used for int8 output data."
                   "Only used with MKL-DNN INT8")
188 189
        .SetDefault(1.0f)
        .AsExtra();
P
Physher 已提交
190 191 192 193
    AddAttr<bool>(
        "force_fp32_output",
        "(bool, default false) Force quantize kernel output FP32, only "
        "used in quantized MKL-DNN.")
194 195
        .SetDefault(false)
        .AsExtra();
196
    AddComment(R"DOC(
C
caoying03 已提交
197
Mul Operator.
K
kexinzhao 已提交
198

C
caoying03 已提交
199
This operator is used to perform matrix multiplication for input $X$ and $Y$.
200

201 202
The equation is:

C
caoying03 已提交
203
$$Out = X * Y$$
204

C
caoying03 已提交
205 206
Both the input $X$ and $Y$ can carry the LoD (Level of Details) information,
or not. But the output only shares the LoD information with input $X$.
K
kexinzhao 已提交
207

208 209 210 211
)DOC");
  }
};

C
chengduo 已提交
212 213
class MulOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
214
  std::unordered_map<std::string, std::string>& GetInputOutputWithSameType()
C
chengduo 已提交
215
      const override {
216 217
    static std::unordered_map<std::string, std::string> m{{"X", /*->*/ "Out"}};
    return m;
C
chengduo 已提交
218 219 220
  }
};

221
class MulGradOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
222 223 224
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

225
  void InferShape(framework::InferShapeContext* ctx) const override {
226 227 228 229
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "mul");
    OP_INOUT_CHECK(ctx->HasInput("Y"), "Input", "Y", "mul");
    OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Out")), "Input",
                   "Out@GRAD", "mul");
Q
Qiao Longfei 已提交
230 231
    auto x_dims = ctx->GetInputDim("X");
    auto y_dims = ctx->GetInputDim("Y");
232

Q
Qiao Longfei 已提交
233 234 235 236 237 238 239 240 241
    auto x_grad_name = framework::GradVarName("X");
    auto y_grad_name = framework::GradVarName("Y");

    if (ctx->HasOutput(x_grad_name)) {
      ctx->SetOutputDim(x_grad_name, x_dims);
    }
    if (ctx->HasOutput(y_grad_name)) {
      ctx->SetOutputDim(y_grad_name, y_dims);
    }
D
dongzhihong 已提交
242
  }
243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const {
    framework::LibraryType library = framework::LibraryType::kPlain;
    framework::DataLayout layout = framework::DataLayout::kAnyLayout;
    int customized_type_value =
        framework::OpKernelType::kDefaultCustomizedTypeValue;
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
#ifdef PADDLE_WITH_MKLDNN
    if (library == framework::LibraryType::kPlain &&
        this->CanMKLDNNBeUsed(ctx, input_data_type)) {
      library = framework::LibraryType::kMKLDNN;
      layout = framework::DataLayout::kMKLDNN;

      if (input_data_type == framework::DataTypeTrait<int8_t>::DataType() ||
          input_data_type == framework::DataTypeTrait<uint8_t>::DataType()) {
        customized_type_value = kMULMKLDNNINT8;
      } else if (input_data_type ==
                     framework::DataTypeTrait<
                         paddle::platform::bfloat16>::DataType() ||
                 input_data_type ==
                     framework::DataTypeTrait<float>::DataType()) {
        customized_type_value = kMULMKLDNNFP32;
      }
    }
#endif

    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout,
                                   library, customized_type_value);
  }
D
dongzhihong 已提交
273 274
};

H
hong 已提交
275 276
template <typename T>
class MulOpGradMaker : public framework::SingleGradOpMaker<T> {
S
sneaxiy 已提交
277
 public:
H
hong 已提交
278
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
S
sneaxiy 已提交
279 280

 protected:
281
  void Apply(GradOpPtr<T> retv) const override {
S
sneaxiy 已提交
282
    retv->SetType("mul_grad");
H
hong 已提交
283 284 285 286 287 288
    retv->SetInput("X", this->Input("X"));
    retv->SetInput("Y", this->Input("Y"));
    retv->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
    retv->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
    retv->SetOutput(framework::GradVarName("Y"), this->InputGrad("Y"));
    retv->SetAttrMap(this->Attrs());
S
sneaxiy 已提交
289 290 291
  }
};

292 293 294 295 296
class MulDoubleGradOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
297 298 299
    OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "mul");
    OP_INOUT_CHECK(ctx->HasInput("Y"), "Input", "Y", "mul");
    OP_INOUT_CHECK(ctx->HasInput("DOut"), "Input", "DOut", "mul");
300

L
lvmengsi 已提交
301 302
    if (ctx->HasOutput("DDOut") &&
        (ctx->HasInput("DDX") || (ctx->HasInput("DDY")))) {
303 304 305
      ctx->ShareDim("DOut", "DDOut");
    }
    if (ctx->HasOutput("DX") && ctx->HasInput("DDY")) {
306 307
      ctx->ShareDim("X", "DX");
    }
308
    if (ctx->HasOutput("DY") && ctx->HasInput("DDX")) {
309 310 311 312 313
      ctx->ShareDim("Y", "DY");
    }
  }
};

H
hong 已提交
314 315
template <typename T>
class MulDoubleGradMaker : public framework::SingleGradOpMaker<T> {
316
 public:
H
hong 已提交
317
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
318 319

 protected:
320
  void Apply(GradOpPtr<T> retv) const override {
321 322
    retv->SetType("mul_grad_grad");

H
hong 已提交
323 324 325 326 327
    retv->SetInput("X", this->Input("X"));
    retv->SetInput("Y", this->Input("Y"));
    retv->SetInput("DOut", this->Input(framework::GradVarName("Out")));
    retv->SetInput("DDX", this->OutputGrad(framework::GradVarName("X")));
    retv->SetInput("DDY", this->OutputGrad(framework::GradVarName("Y")));
328

H
hong 已提交
329 330
    auto ddx = this->OutputGrad(framework::GradVarName("X"));
    auto ddw = this->OutputGrad(framework::GradVarName("Y"));
331

L
lvmengsi 已提交
332
    if (!ddx.empty() || !ddw.empty()) {
H
hong 已提交
333
      retv->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
L
lvmengsi 已提交
334
    }
335 336 337 338
    retv->SetOutput(
        "DX", ddw.empty() ? this->EmptyInputGrad() : this->InputGrad("X"));
    retv->SetOutput(
        "DY", ddx.empty() ? this->EmptyInputGrad() : this->InputGrad("Y"));
339

H
hong 已提交
340
    retv->SetAttrMap(this->Attrs());
341 342 343
  }
};

344 345 346
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
347
namespace ops = paddle::operators;
C
chengduo 已提交
348
REGISTER_OPERATOR(mul, ops::MulOp, ops::MulOpMaker, ops::MulOpInferVarType,
H
hong 已提交
349 350
                  ops::MulOpGradMaker<paddle::framework::OpDesc>,
                  ops::MulOpGradMaker<paddle::imperative::OpBase>);
P
Physher 已提交
351

H
hong 已提交
352 353 354
REGISTER_OPERATOR(mul_grad, ops::MulGradOp,
                  ops::MulDoubleGradMaker<paddle::framework::OpDesc>,
                  ops::MulDoubleGradMaker<paddle::imperative::OpBase>);
P
Physher 已提交
355

356
REGISTER_OPERATOR(mul_grad_grad, ops::MulDoubleGradOp);
P
Physher 已提交
357

Q
QI JUN 已提交
358
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
359 360
    mul, ops::MulKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulKernel<paddle::platform::CPUDeviceContext, double>);
P
Physher 已提交
361

Q
QI JUN 已提交
362
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
363 364
    mul_grad, ops::MulGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulGradKernel<paddle::platform::CPUDeviceContext, double>);
P
Physher 已提交
365

366 367 368 369
REGISTER_OP_CPU_KERNEL(
    mul_grad_grad,
    ops::MulDoubleGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulDoubleGradKernel<paddle::platform::CPUDeviceContext, double>);