mul_op.cc 12.9 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 38 39 40 41 42 43
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("X"), true,
        platform::errors::NotFound("Input(X) of MulOp should not be null."));
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("Y"), true,
        platform::errors::NotFound("Input(Y) of MulOp should not be null."));
    PADDLE_ENFORCE_EQ(
        ctx->HasOutput("Out"), true,
        platform::errors::NotFound("Output(Out) of MulOp should not be null."));
Q
Qiao Longfei 已提交
44 45 46

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

Q
Qiao Longfei 已提交
48 49
    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 已提交
50

M
minqiyang 已提交
51 52 53
    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 已提交
54

55
    PADDLE_ENFORCE_NE(framework::product(y_dims), 0,
56
                      platform::errors::PreconditionNotMet(
57
                          "The Input variable Y(%s) has not "
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75
                          "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));
76

F
fengjiayi 已提交
77 78
    auto x_mat_dims = framework::flatten_to_2d(x_dims, x_num_col_dims);
    auto y_mat_dims = framework::flatten_to_2d(y_dims, y_num_col_dims);
79

80 81
    PADDLE_ENFORCE_EQ(
        x_mat_dims[1], y_mat_dims[0],
82 83 84 85 86 87 88 89 90
        platform::errors::InvalidArgument(
            "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 "
            "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 已提交
91 92 93 94 95 96 97 98 99 100 101 102 103
    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]);
    }

    ctx->SetOutputDim("Out", framework::make_ddim(output_dims));
Q
Qiao Longfei 已提交
104
    ctx->ShareLoD("X", /*->*/ "Out");
105
  }
P
Physher 已提交
106 107 108 109 110 111 112

  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;
113
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
P
Physher 已提交
114 115 116 117 118 119
#ifdef PADDLE_WITH_MKLDNN
    if (library == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library = framework::LibraryType::kMKLDNN;
      layout = framework::DataLayout::kMKLDNN;

120 121
      if (input_data_type == framework::DataTypeTrait<int8_t>::DataType() ||
          input_data_type == framework::DataTypeTrait<uint8_t>::DataType()) {
P
Physher 已提交
122 123 124 125 126 127 128 129
        customized_type_value = kMULMKLDNNINT8;
      }
    }
#endif

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

D
dongzhihong 已提交
132
class MulOpMaker : public framework::OpProtoAndCheckerMaker {
133
 public:
Y
Yu Yang 已提交
134
  void Make() override {
C
caoying03 已提交
135 136 137
    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 已提交
138 139 140
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
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")
P
Physher 已提交
176
        .SetDefault(1.0f);
177 178 179 180 181
    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")
P
Physher 已提交
182 183 184 185 186 187 188 189 190 191
        .SetDefault({1.0f});
    AddAttr<float>("scale_out",
                   "scale_out to be used for int8 output data."
                   "Only used with MKL-DNN INT8")
        .SetDefault(1.0f);
    AddAttr<bool>(
        "force_fp32_output",
        "(bool, default false) Force quantize kernel output FP32, only "
        "used in quantized MKL-DNN.")
        .SetDefault(false);
192
    AddComment(R"DOC(
C
caoying03 已提交
193
Mul Operator.
K
kexinzhao 已提交
194

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

197 198
The equation is:

C
caoying03 已提交
199
$$Out = X * Y$$
200

C
caoying03 已提交
201 202
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 已提交
203

204 205 206 207
)DOC");
  }
};

C
chengduo 已提交
208 209 210 211 212 213 214 215
class MulOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
  std::unordered_map<std::string, std::string> GetInputOutputWithSameType()
      const override {
    return std::unordered_map<std::string, std::string>{{"X", /*->*/ "Out"}};
  }
};

216
class MulGradOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
217 218 219
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

220
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
221 222 223 224 225 226
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
    PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should not be null");
    auto x_dims = ctx->GetInputDim("X");
    auto y_dims = ctx->GetInputDim("Y");
227

Q
Qiao Longfei 已提交
228 229 230 231 232 233 234 235 236
    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 已提交
237 238 239
  }
};

H
hong 已提交
240 241
template <typename T>
class MulOpGradMaker : public framework::SingleGradOpMaker<T> {
S
sneaxiy 已提交
242
 public:
H
hong 已提交
243
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
S
sneaxiy 已提交
244 245

 protected:
H
hong 已提交
246 247
  std::unique_ptr<T> Apply() const override {
    std::unique_ptr<T> retv(new T());
S
sneaxiy 已提交
248
    retv->SetType("mul_grad");
H
hong 已提交
249 250 251 252 253 254
    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 已提交
255 256 257 258
    return retv;
  }
};

259 260 261 262 263 264 265 266 267
class MulDoubleGradOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

  void InferShape(framework::InferShapeContext* ctx) const override {
    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
    PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
    PADDLE_ENFORCE(ctx->HasInput("DOut"), "Input(DOut) should not be null");

L
lvmengsi 已提交
268 269
    if (ctx->HasOutput("DDOut") &&
        (ctx->HasInput("DDX") || (ctx->HasInput("DDY")))) {
270 271 272
      ctx->ShareDim("DOut", "DDOut");
    }
    if (ctx->HasOutput("DX") && ctx->HasInput("DDY")) {
273 274
      ctx->ShareDim("X", "DX");
    }
275
    if (ctx->HasOutput("DY") && ctx->HasInput("DDX")) {
276 277 278 279 280
      ctx->ShareDim("Y", "DY");
    }
  }
};

H
hong 已提交
281 282
template <typename T>
class MulDoubleGradMaker : public framework::SingleGradOpMaker<T> {
283
 public:
H
hong 已提交
284
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
285 286

 protected:
H
hong 已提交
287 288
  std::unique_ptr<T> Apply() const override {
    std::unique_ptr<T> retv(new T());
289 290
    retv->SetType("mul_grad_grad");

H
hong 已提交
291 292 293 294 295
    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")));
296

H
hong 已提交
297 298
    auto ddx = this->OutputGrad(framework::GradVarName("X"));
    auto ddw = this->OutputGrad(framework::GradVarName("Y"));
299

L
lvmengsi 已提交
300
    if (!ddx.empty() || !ddw.empty()) {
H
hong 已提交
301
      retv->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
L
lvmengsi 已提交
302
    }
H
hong 已提交
303 304
    retv->SetOutput("DX", ddw.empty() ? this->Empty() : this->InputGrad("X"));
    retv->SetOutput("DY", ddx.empty() ? this->Empty() : this->InputGrad("Y"));
305

H
hong 已提交
306
    retv->SetAttrMap(this->Attrs());
307 308 309 310
    return retv;
  }
};

311 312 313
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
314
namespace ops = paddle::operators;
C
chengduo 已提交
315
REGISTER_OPERATOR(mul, ops::MulOp, ops::MulOpMaker, ops::MulOpInferVarType,
H
hong 已提交
316 317
                  ops::MulOpGradMaker<paddle::framework::OpDesc>,
                  ops::MulOpGradMaker<paddle::imperative::OpBase>);
P
Physher 已提交
318

H
hong 已提交
319 320 321
REGISTER_OPERATOR(mul_grad, ops::MulGradOp,
                  ops::MulDoubleGradMaker<paddle::framework::OpDesc>,
                  ops::MulDoubleGradMaker<paddle::imperative::OpBase>);
P
Physher 已提交
322

323
REGISTER_OPERATOR(mul_grad_grad, ops::MulDoubleGradOp);
P
Physher 已提交
324

Q
QI JUN 已提交
325
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
326 327
    mul, ops::MulKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulKernel<paddle::platform::CPUDeviceContext, double>);
P
Physher 已提交
328

Q
QI JUN 已提交
329
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
330 331
    mul_grad, ops::MulGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulGradKernel<paddle::platform::CPUDeviceContext, double>);
P
Physher 已提交
332

333 334 335 336
REGISTER_OP_CPU_KERNEL(
    mul_grad_grad,
    ops::MulDoubleGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::MulDoubleGradKernel<paddle::platform::CPUDeviceContext, double>);