matmul_op.cc 33.2 KB
Newer Older
1
/* Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
M
Markus Kliegl 已提交
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>
Y
Yu Yang 已提交
13
#include <utility>
14
#include <vector>
15

Y
Yu Yang 已提交
16
#include "paddle/fluid/framework/op_registry.h"
17
#include "paddle/fluid/framework/op_version_registry.h"
18
#include "paddle/phi/kernels/funcs/blas/blas.h"
M
Markus Kliegl 已提交
19 20 21

namespace paddle {
namespace operators {
22 23 24 25

/**
 * Printing shape information into a string is easy to use.
 */
26
inline static std::string DumpMatrixShape(
27
    const phi::funcs::MatDescriptor &desc) {
28 29 30 31 32 33
  std::stringstream buffer;
  buffer << "[" << desc.batch_size_ << ", " << desc.height_ << ", "
         << desc.width_ << "]";
  return buffer.str();
}

Y
Yu Yang 已提交
34 35 36 37
/**
 * Get row matrix shape from a vector shape. If the rank of x_dim > 1, the
 * original x_dim is returned.
 */
Y
yuyang18 已提交
38
static framework::DDim RowMatrixFromVector(const framework::DDim &x_dim) {
Y
Yu Yang 已提交
39 40 41
  if (x_dim.size() > 1) {
    return x_dim;
  }
42
  return phi::make_ddim({1, x_dim[0]});
Y
Yu Yang 已提交
43 44 45 46 47 48
}

/**
 * Get column matrix shape from a vector shape. If the ran of y_dim > 1, the
 * original y_dim is returned.
 */
Y
yuyang18 已提交
49
static framework::DDim ColumnMatrixFromVector(const framework::DDim &y_dim) {
Y
Yu Yang 已提交
50 51 52
  if (y_dim.size() > 1) {
    return y_dim;
  }
53
  return phi::make_ddim({y_dim[0], 1});
Y
Yu Yang 已提交
54 55 56 57 58
}

template <typename DeviceContext, typename T>
class MatMulKernel : public framework::OpKernel<T> {
 public:
Y
yuyang18 已提交
59
  void Compute(const framework::ExecutionContext &context) const override {
60
    auto &x = GET_DATA_SAFELY(
61
        context.Input<phi::DenseTensor>("X"), "Input", "X", "MatMul");
62
    auto &y = GET_DATA_SAFELY(
63 64
        context.Input<phi::DenseTensor>("Y"), "Input", "Y", "MatMul");
    auto *out = context.Output<phi::DenseTensor>("Out");
W
Wilber 已提交
65 66 67

    auto &dev_ctx = context.template device_context<DeviceContext>();
    dev_ctx.template Alloc<T>(out, out->numel() * sizeof(T));
Y
Yu Yang 已提交
68

69
    auto blas = phi::funcs::GetBlas<DeviceContext, T>(dev_ctx);
70
    auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(
Y
Yu Yang 已提交
71
        RowMatrixFromVector(x.dims()), 0, context.Attr<bool>("transpose_X"));
72
    auto mat_dim_b = phi::funcs::CreateMatrixDescriptor(
Y
Yu Yang 已提交
73
        ColumnMatrixFromVector(y.dims()), 0, context.Attr<bool>("transpose_Y"));
S
sneaxiy 已提交
74
    auto scale = static_cast<T>(context.Attr<float>("alpha"));
75

76
    int head_number = 1;
77 78
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
79 80 81 82 83 84 85 86 87 88 89 90
    head_number = context.Attr<int>("head_number");
#endif

    const auto &x_dims = x.dims();
    const auto &y_dims = y.dims();
    if (head_number <= 1 && x_dims.size() == 3 && y_dims.size() <= 2) {
      // the transpose_X must be false, if is true, the transpose cost much time
      if (!context.Attr<bool>("transpose_X")) {
        mat_dim_a.height_ *= mat_dim_a.batch_size_;
        mat_dim_a.batch_size_ = 0;
      }
    }
91 92
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
93 94 95
    bool split_vertical_y = (mat_dim_a.width_ != mat_dim_b.height_);

    if (head_number > 1) {
96 97 98 99 100 101 102 103 104
      blas.MatMulWithHead(x,
                          mat_dim_a,
                          y,
                          mat_dim_b,
                          scale,
                          head_number,
                          out,
                          T(0),
                          split_vertical_y);
105 106
    } else {
      blas.MatMul(x, mat_dim_a, y, mat_dim_b, scale, out, T(0));
107 108
    }
#else
S
sneaxiy 已提交
109
    blas.MatMul(x, mat_dim_a, y, mat_dim_b, scale, out, T(0));
110
#endif
Y
Yu Yang 已提交
111 112 113 114 115
  }
};

// Reshape a rank-3 tensor from P x M x N to (P * M) x N.
// Identity op if the tensor is not of rank 3.
116
static phi::DenseTensor FoldInitDims(const phi::DenseTensor &input) {
Y
Yu Yang 已提交
117 118 119 120 121 122 123 124 125 126 127 128
  auto output = input;
  auto in_dims = input.dims();
  if (in_dims.size() == 3) {
    output.Resize({in_dims[0] * in_dims[1], in_dims[2]});
  }
  return output;
}

// Reshape a rank-3 tensor from P x M x N to M x (P * N).
// (Warning: This requires transposing data and writes into new memory.)
// Identity op if the tensor is not of rank 3.
template <typename DeviceContext, typename T>
129 130
static phi::DenseTensor FoldHeadAndLastDims(const DeviceContext &context,
                                            const phi::DenseTensor &input) {
Y
Yu Yang 已提交
131 132 133 134
  auto in_dims = input.dims();
  if (in_dims.size() != 3) {
    return input;
  }
135
  phi::DenseTensor output;
Y
Yu Yang 已提交
136 137 138
  output.Resize({in_dims[1], in_dims[0], in_dims[2]});
  output.mutable_data<T>(context.GetPlace());
  std::vector<int> axis = {1, 0, 2};
139
  phi::funcs::Transpose<DeviceContext, T, 3> trans;
Y
Yu Yang 已提交
140 141
  trans(context, input, &output, axis);
  output.Resize({in_dims[1], in_dims[0] * in_dims[2]});
M
Markus Kliegl 已提交
142

Y
Yu Yang 已提交
143 144 145 146 147 148 149 150 151 152
  return output;
}

/**
 * Reshape a tensor to 3-D or 2-D tensor by matrix descriptor.
 *
 * The shape would be [BatchSize, H, W] or [H, W].
 * If transposed, `H,W` will be swapped.
 */
static void ReshapeTensorIntoMatrixSequence(
153
    phi::DenseTensor *x, const phi::funcs::MatDescriptor &descriptor) {
Y
Yu Yang 已提交
154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180
  int64_t h, w;
  h = descriptor.height_;
  w = descriptor.width_;
  if (descriptor.trans_) {
    std::swap(w, h);
  }
  if (descriptor.batch_size_) {
    x->Resize({descriptor.batch_size_, h, w});
  } else {
    x->Resize({h, w});
  }
}

/**
 * Reshape the x,y,out tensor to 3-D or 2-D tensor by matrix descriptor
 * Out = matmul(x, y)
 *
 * This method will first calculate X,Y matrix sequence, and then calculate
 * the out shape.
 *
 * Assume X = [BatchSize, H1, W1], Y = [BatchSize, H2, W2]
 * The out = [BatchSize, H1, W2]
 *
 * If there is no batch size in `X` and `Y`, the out will be [H1, W2]
 * If any of `X` and `Y` has batch size BatchSize, the out will have the
 * BatchSize.
 */
181 182 183
static void ReshapeXYOutIntoMatrixSequence(phi::DenseTensor *x,
                                           phi::DenseTensor *y,
                                           phi::DenseTensor *out,
184
                                           bool trans_x,
Y
Yu Yang 已提交
185 186 187
                                           bool trans_y) {
  auto x_dim = RowMatrixFromVector(x->dims());
  auto y_dim = ColumnMatrixFromVector(y->dims());
188 189
  auto mat_dim_x = phi::funcs::CreateMatrixDescriptor(x_dim, 0, trans_x);
  auto mat_dim_y = phi::funcs::CreateMatrixDescriptor(y_dim, 0, trans_y);
Y
Yu Yang 已提交
190 191 192 193
  if (mat_dim_x.batch_size_ == 0 && mat_dim_y.batch_size_ == 0) {
    out->Resize({mat_dim_x.height_, mat_dim_y.width_});
  } else {
    out->Resize({std::max(mat_dim_x.batch_size_, mat_dim_y.batch_size_),
194 195
                 mat_dim_x.height_,
                 mat_dim_y.width_});
Y
Yu Yang 已提交
196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229
  }

  ReshapeTensorIntoMatrixSequence(x, mat_dim_x);
  ReshapeTensorIntoMatrixSequence(y, mat_dim_y);
}

// Using dimensional constraints on matrix multiplication, it is
// straight-forward to check the following table for when X and Y
// are both matrices.
//
// transpose_X | False    | True     | False    | True
// transpose_Y | False    | False    | True     | True
// -----------+----------+----------+----------+-----------
//        dX = | dOut Y^T | Y dOut^T | dOut Y   | Y^T dOut^T
//        dY = | X^T dOut | X dOut   | dOut^T X | dOut^T X^T
//
// When X is a vector of size K, we treat it instead as a matrix of shape
// (1, K). Similarly, when Y is a vector of size K, we treat it instead as
// a matrix of shape (K, 1).
//
// When X and Y are both 3-dimensional tensors, then the first dimension
// the batch dimension can be ignored and the exact same formulas apply
// as for two matrices.
//
// Finally, when, e.g., X is a 3-dimensional tensor but Y is a matrix, we end
// up with formulas like
//
//   dY_{ij} = \sum_{p, m} X_{pmi} dOut_{pmj}
//
// To handle this sort of scenario, we reshape X : P x M x K, dOut: P x M x N
// to X: (P * M) x K, dOut: (P * M) x N.
template <typename DeviceContext, typename T>
class MatMulGradKernel : public framework::OpKernel<T> {
 public:
Y
yuyang18 已提交
230
  void MatMul(const framework::ExecutionContext &context,
231
              const phi::DenseTensor &a,
232
              bool trans_a,
233
              const phi::DenseTensor &b,
234
              bool trans_b,
235
              phi::DenseTensor *out) const {
Y
Yu Yang 已提交
236
    out->mutable_data<T>(context.GetPlace());
237 238
    auto &dev_ctx = context.template device_context<DeviceContext>();
    auto blas = phi::funcs::GetBlas<DeviceContext, T>(dev_ctx);
239 240
    auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(a.dims(), 0, trans_a);
    auto mat_dim_b = phi::funcs::CreateMatrixDescriptor(b.dims(), 0, trans_b);
241 242

    int head_number = 1;
243 244
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
245 246 247
    if (context.HasAttr("head_number")) {
      head_number = context.Attr<int>("head_number");
    }
248 249 250 251 252 253 254 255 256
#endif

    if (head_number <= 1 && a.dims().size() == 3 && b.dims().size() <= 2) {
      // the transpose_X must be false, if is true, the transpose cost much time
      if (!trans_a) {
        mat_dim_a.height_ *= mat_dim_a.batch_size_;
        mat_dim_a.batch_size_ = 0;
      }
    }
257 258 259 260 261 262 263
    blas.MatMul(a,
                mat_dim_a,
                b,
                mat_dim_b,
                static_cast<T>(context.Attr<float>("alpha")),
                out,
                T(0));
Y
Yu Yang 已提交
264 265
  }

Y
yuyang18 已提交
266
  void CalcInputGrad(const framework::ExecutionContext &context,
267
                     const phi::DenseTensor &a,
268 269
                     bool trans_a,
                     bool is_fold_init_dims_a,
270
                     const phi::DenseTensor &b,
271 272
                     bool trans_b,
                     bool is_fold_init_dims_b,
273
                     phi::DenseTensor *out) const {
Y
Yu Yang 已提交
274 275 276 277 278 279
    if (out == nullptr) return;
    bool need_combine = (a.dims().size() == 3 || b.dims().size() == 3) &&
                        out->dims().size() == 2;
    if (!need_combine) {
      MatMul(context, a, trans_a, b, trans_b, out);
    } else {
Y
yuyang18 已提交
280
      auto &ctx = context.template device_context<DeviceContext>();
281 282 283 284 285 286 287
      MatMul(
          context,
          is_fold_init_dims_a ? FoldInitDims(a)
                              : FoldHeadAndLastDims<DeviceContext, T>(ctx, a),
          trans_a,
          is_fold_init_dims_b ? FoldInitDims(b)
                              : FoldHeadAndLastDims<DeviceContext, T>(ctx, b),
288 289
          trans_b,
          out);
Y
Yu Yang 已提交
290 291 292
    }
  }

Y
yuyang18 已提交
293
  void Compute(const framework::ExecutionContext &context) const override {
294 295 296 297 298
    auto x = *context.Input<phi::DenseTensor>("X");
    auto y = *context.Input<phi::DenseTensor>("Y");
    auto dout = *context.Input<phi::DenseTensor>(framework::GradVarName("Out"));
    auto *dx = context.Output<phi::DenseTensor>(framework::GradVarName("X"));
    auto *dy = context.Output<phi::DenseTensor>(framework::GradVarName("Y"));
Y
Yu Yang 已提交
299 300 301 302
    bool transpose_x = context.Attr<bool>("transpose_X");
    bool transpose_y = context.Attr<bool>("transpose_Y");

    ReshapeXYOutIntoMatrixSequence(&x, &y, &dout, transpose_x, transpose_y);
303
    phi::DDim dx_dims;
Y
Yu Yang 已提交
304 305 306 307 308 309 310
    if (dx) {
      dx_dims = dx->dims();
      if (dx_dims != x.dims()) {
        dx->Resize(x.dims());
      }
    }

311
    phi::DDim dy_dims;
Y
Yu Yang 已提交
312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344
    if (dy) {
      dy_dims = dy->dims();
      if (dy_dims != y.dims()) {
        dy->Resize(y.dims());
      }
    }

    if (transpose_x && transpose_y) {
      CalcInputGrad(context, y, true, true, dout, true, false, dx);
      CalcInputGrad(context, dout, true, true, x, true, false, dy);
    } else if (transpose_x) {
      CalcInputGrad(context, y, false, false, dout, true, false, dx);
      CalcInputGrad(context, x, false, false, dout, false, true, dy);
    } else if (transpose_y) {
      CalcInputGrad(context, dout, false, false, y, false, true, dx);
      CalcInputGrad(context, dout, true, true, x, false, true, dy);
    } else {
      CalcInputGrad(context, dout, false, false, y, true, false, dx);
      CalcInputGrad(context, x, true, true, dout, false, true, dy);
    }

    if (dx) {
      if (dx_dims != x.dims()) {
        dx->Resize(dx_dims);
      }
    }
    if (dy) {
      if (dy_dims != y.dims()) {
        dy->Resize(dy_dims);
      }
    }
  }
};
M
Markus Kliegl 已提交
345

346 347
phi::DDim GetDimForInput(const framework::InferShapeContext &ctx,
                         std::string input_name) {
348 349 350
  auto dim = ctx.GetInputDim(input_name);
  PADDLE_ENFORCE_GT(dim.size(),
                    0,
351
                    phi::errors::InvalidArgument(
352 353 354 355 356 357
                        "The Input(%s) has not been initialized properly. The "
                        "shape of Input(%s) = [%s].",
                        dim));
  return dim;
}

358 359 360 361
template <typename DeviceContext, typename T>
class MatMulDoubleGradKernel : public framework::OpKernel<T> {
 public:
  void MatMul(const framework::ExecutionContext &context,
362
              const phi::DenseTensor &a,
363
              bool trans_a,
364
              const phi::DenseTensor &b,
365 366
              bool trans_b,
              bool flag,
367
              phi::DenseTensor *out) const {
368
    out->mutable_data<T>(context.GetPlace());
369 370
    auto &dev_ctx = context.template device_context<DeviceContext>();
    auto blas = phi::funcs::GetBlas<DeviceContext, T>(dev_ctx);
371 372
    auto mat_dim_a = phi::funcs::CreateMatrixDescriptor(a.dims(), 0, trans_a);
    auto mat_dim_b = phi::funcs::CreateMatrixDescriptor(b.dims(), 0, trans_b);
373 374

    int head_number = 1;
375 376
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
377 378 379 380 381 382 383 384 385 386
    head_number = context.Attr<int>("head_number");
#endif

    if (head_number <= 1 && a.dims().size() == 3 && b.dims().size() <= 2) {
      // the transpose_X must be false, if is true, the transpose cost much time
      if (!trans_a) {
        mat_dim_a.height_ *= mat_dim_a.batch_size_;
        mat_dim_a.batch_size_ = 0;
      }
    }
387 388 389 390 391 392
    blas.MatMul(a,
                mat_dim_a,
                b,
                mat_dim_b,
                static_cast<T>(context.Attr<float>("alpha")),
                out,
393 394 395 396
                static_cast<T>(flag));
  }

  void CalcInputGrad(const framework::ExecutionContext &context,
397
                     const phi::DenseTensor &a,
398 399
                     bool trans_a,
                     bool is_fold_init_dims_a,
400
                     const phi::DenseTensor &b,
401 402 403
                     bool trans_b,
                     bool is_fold_init_dims_b,
                     bool flag,
404
                     phi::DenseTensor *out) const {
405 406 407 408 409 410 411
    if (out == nullptr) return;
    bool need_combine = (a.dims().size() == 3 || b.dims().size() == 3) &&
                        out->dims().size() == 2;
    if (!need_combine) {
      MatMul(context, a, trans_a, b, trans_b, flag, out);
    } else {
      auto &ctx = context.template device_context<DeviceContext>();
412 413 414 415 416 417 418
      MatMul(
          context,
          is_fold_init_dims_a ? FoldInitDims(a)
                              : FoldHeadAndLastDims<DeviceContext, T>(ctx, a),
          trans_a,
          is_fold_init_dims_b ? FoldInitDims(b)
                              : FoldHeadAndLastDims<DeviceContext, T>(ctx, b),
419 420 421
          trans_b,
          flag,
          out);
422 423 424 425
    }
  }

  void Compute(const framework::ExecutionContext &context) const override {
426 427
    auto x = *context.Input<phi::DenseTensor>("X");
    auto y = *context.Input<phi::DenseTensor>("Y");
428 429 430
    auto dout = *context.Input<phi::DenseTensor>("DOut");
    auto *ddx = context.Input<phi::DenseTensor>("DDX");
    auto *ddy = context.Input<phi::DenseTensor>("DDY");
431

432 433 434
    auto *dx = context.Output<phi::DenseTensor>("DX");
    auto *dy = context.Output<phi::DenseTensor>("DY");
    auto *ddout = context.Output<phi::DenseTensor>("DDOut");
435 436 437 438 439 440

    bool transpose_x = context.Attr<bool>("transpose_X");
    bool transpose_y = context.Attr<bool>("transpose_Y");

    ReshapeXYOutIntoMatrixSequence(&x, &y, &dout, transpose_x, transpose_y);

441
    phi::DDim dx_dims;
442 443 444 445 446 447 448
    if (dx) {
      dx_dims = dx->dims();
      if (dx_dims != x.dims()) {
        dx->Resize(x.dims());
      }
    }

449
    phi::DDim dy_dims;
450 451 452 453 454 455 456
    if (dy) {
      dy_dims = dy->dims();
      if (dy_dims != y.dims()) {
        dy->Resize(y.dims());
      }
    }

457
    phi::DDim ddout_dims;
458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473
    if (ddout) {
      ddout_dims = ddout->dims();
      if (ddout_dims != dout.dims()) {
        ddout->Resize(dout.dims());
      }
    }

    bool ddout_flag = false;
    if (ddx) {
      auto ddx_mat = *ddx;
      if (ddx_mat.dims() != x.dims()) {
        ddx_mat.Resize(x.dims());
      }
      if (dy) {
        if (transpose_x && transpose_y) {
          // dy = dout' * ddx'
474 475
          CalcInputGrad(
              context, dout, true, true, ddx_mat, true, false, false, dy);
476 477
        } else if (transpose_x) {
          // dy = ddx * dout
478 479
          CalcInputGrad(
              context, ddx_mat, false, false, dout, false, true, false, dy);
480 481
        } else if (transpose_y) {
          // dy = dout' * ddx
482 483
          CalcInputGrad(
              context, dout, true, true, ddx_mat, false, true, false, dy);
484 485
        } else {
          // dy = ddx' * dout
486 487
          CalcInputGrad(
              context, ddx_mat, true, true, dout, false, true, false, dy);
488 489 490 491
        }
      }

      if (ddout) {
492 493 494 495 496 497 498 499 500
        CalcInputGrad(context,
                      ddx_mat,
                      transpose_x,
                      true,
                      y,
                      transpose_y,
                      false,
                      ddout_flag,
                      ddout);
501 502 503 504 505 506 507 508 509 510 511 512
        ddout_flag = true;
      }
    }

    if (ddy) {
      auto ddy_mat = *ddy;
      if (ddy_mat.dims() != y.dims()) {
        ddy_mat.Resize(y.dims());
      }
      if (dx) {
        if (transpose_x && transpose_y) {
          // dx = ddy' * dout'
513 514
          CalcInputGrad(
              context, ddy_mat, true, true, dout, true, false, false, dx);
515 516
        } else if (transpose_x) {
          // dx = ddy * dout'
517 518
          CalcInputGrad(
              context, ddy_mat, false, false, dout, true, false, false, dx);
519 520
        } else if (transpose_y) {
          // dx = dout * ddy
521 522
          CalcInputGrad(
              context, dout, false, false, ddy_mat, false, true, false, dx);
523 524
        } else {
          // dx = dout * ddy'
525 526
          CalcInputGrad(
              context, dout, false, false, ddy_mat, true, false, false, dx);
527 528 529 530
        }
      }

      if (ddout) {
531 532 533 534 535 536 537 538 539
        CalcInputGrad(context,
                      x,
                      transpose_x,
                      true,
                      ddy_mat,
                      transpose_y,
                      false,
                      ddout_flag,
                      ddout);
540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562
      }
    }

    if (dx) {
      if (dx_dims != x.dims()) {
        dx->Resize(dx_dims);
      }
    }

    if (dy) {
      if (dy_dims != y.dims()) {
        dy->Resize(dy_dims);
      }
    }

    if (ddout) {
      if (ddout_dims != dout.dims()) {
        ddout->Resize(ddout_dims);
      }
    }
  }
};

M
Markus Kliegl 已提交
563 564 565 566 567
class MatMulOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
Y
yuyang18 已提交
568
  void InferShape(framework::InferShapeContext *context) const override {
569 570 571
    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "matmul");
    OP_INOUT_CHECK(context->HasInput("Y"), "Input", "Y", "matmul");
    OP_INOUT_CHECK(context->HasOutput("Out"), "Output", "Out", "matmul");
M
Markus Kliegl 已提交
572

573 574
    auto dim_x = GetDimForInput(*context, "X");
    auto dim_y = GetDimForInput(*context, "Y");
575 576

#ifdef PADDLE_WITH_MKLDNN
577 578
    // For NHWC execution output shape needs to be
    // computed like instead x*y we are to do y*x
579 580
    bool channelwise_onednn =
        context->IsRunMKLDNNKernel() &&
581
        (phi::OneDNNContext::tls().get_cur_paddle_data_layout() ==
582
         phi::DataLayout::kNHWC);
583 584 585 586 587
    if (channelwise_onednn) {
      std::swap(dim_x, dim_y);
    }
#endif

588
    auto mat_dim_x = phi::funcs::CreateMatrixDescriptor(
589 590
        RowMatrixFromVector(dim_x),
        0,
591
        context->Attrs().Get<bool>("transpose_X"));
592
    auto mat_dim_y = phi::funcs::CreateMatrixDescriptor(
593 594
        ColumnMatrixFromVector(dim_y),
        0,
595
        context->Attrs().Get<bool>("transpose_Y"));
C
chengduoZH 已提交
596

597 598 599 600 601 602 603
    if (mat_dim_x.width_ == -1) {
      mat_dim_x.width_ = mat_dim_y.height_;
    }
    if (mat_dim_y.height_ == -1) {
      mat_dim_y.height_ = mat_dim_x.width_;
    }

P
phlrain 已提交
604
    if (context->IsRuntime()) {
605
      PADDLE_ENFORCE_EQ(
606 607
          mat_dim_x.batch_size_ == mat_dim_y.batch_size_ ||
              mat_dim_x.batch_size_ == 0 || mat_dim_y.batch_size_ == 0,
608
          true,
609
          phi::errors::InvalidArgument(
610 611 612 613 614
              "The batch size of the two matrices should be equal, or "
              "at least one is zero.\n"
              "But received X's shape: %s, Y's shape: %s.",
              DumpMatrixShape(mat_dim_x).c_str(),
              DumpMatrixShape(mat_dim_y).c_str()));
P
phlrain 已提交
615
    }
616
    int64_t dim_out_y = mat_dim_y.width_;
617 618
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
619
    int head_number = context->Attrs().Get<int>("head_number");
620
    bool split_vertical_y = (mat_dim_x.width_ != mat_dim_y.height_);
621 622
    if (context->IsRuntime()) {
      PADDLE_ENFORCE_LE(
623 624
          head_number,
          mat_dim_x.width_,
625
          phi::errors::InvalidArgument(
626 627 628
              "Unsatisfied mkl acceleration library requirements: "
              "The number of heads "
              "(%d) must be equal to X's width. But received X's shape: %s.",
629 630
              head_number,
              DumpMatrixShape(mat_dim_x).c_str()));
631 632 633 634

      if (!split_vertical_y && head_number > 0) {
        dim_out_y = head_number * mat_dim_y.width_;
      }
635
    }
636
#else
637 638
    PADDLE_ENFORCE_EQ(mat_dim_x.width_,
                      mat_dim_y.height_,
639
                      phi::errors::InvalidArgument(
640
                          "Input X's width should be equal to the Y's height, "
641
                          "but received X's shape: [%s], "
642
                          "Y's shape: [%s].",
643 644
                          dim_x,
                          dim_y));
645 646
#endif

647
    std::vector<int64_t> dim_out;
Y
Yu Yang 已提交
648
    if (mat_dim_x.batch_size_ != 0) {
649
      dim_out = phi::vectorize(dim_x);
Y
Yu Yang 已提交
650
      dim_out[dim_out.size() - 2] = mat_dim_x.height_;
651
      dim_out[dim_out.size() - 1] = dim_out_y;
Y
Yu Yang 已提交
652
    } else if (mat_dim_y.batch_size_ != 0) {
653
      dim_out = phi::vectorize(dim_y);
Y
Yu Yang 已提交
654
      dim_out[dim_out.size() - 2] = mat_dim_x.height_;
655
      dim_out[dim_out.size() - 1] = dim_out_y;
Y
Yu Yang 已提交
656
    } else {
657
      dim_out = {mat_dim_x.height_, dim_out_y};
M
Markus Kliegl 已提交
658 659
    }

Y
Yu Yang 已提交
660 661 662
    if (dim_x.size() == 1 && dim_out[dim_out.size() - 2] == 1) {
      std::swap(dim_out[dim_out.size() - 2], dim_out[dim_out.size() - 1]);
      dim_out.resize(dim_out.size() - 1);
M
Markus Kliegl 已提交
663 664
    }

Y
Yu Yang 已提交
665 666
    if (dim_y.size() == 1 && dim_out[dim_out.size() - 1] == 1) {
      dim_out.resize(dim_out.size() - 1);
M
Markus Kliegl 已提交
667 668
    }

669
    phi::DDim ddim_out = phi::make_ddim(dim_out);
670

671 672
    context->SetOutputDim("Out", ddim_out);
    context->ShareLoD("X", "Out");
M
Markus Kliegl 已提交
673
  }
674

675
  phi::KernelKey GetExpectedKernelType(
676
      const framework::ExecutionContext &ctx) const override {
677 678
    auto input_data_type =
        OperatorWithKernel::IndicateOrPromoteVarDataTypes(ctx, "X", "Y");
679
    return phi::KernelKey(input_data_type, ctx.GetPlace());
680
  }
681

682
  phi::KernelKey GetKernelTypeForVar(
683
      const std::string &var_name,
684
      const phi::DenseTensor &tensor,
685 686
      const phi::KernelKey &expected_kernel_type) const override {
    if (framework::IsComplexType(expected_kernel_type.dtype())) {
687
      // only promote inputs’s types when contains complex input
688
      return phi::KernelKey(tensor.place(), tensor.layout(), tensor.dtype());
689
    } else {
690 691 692 693
#ifdef PADDLE_WITH_MKLDNN
      // When matmul is first oneDNN op in a chain (there was some non oneDNN op
      // previously)
      // then we also need to rotate shape NHWC -> NCWH
694
      if ((expected_kernel_type.layout() == phi::DataLayout::ONEDNN) &&
695
          (tensor.layout() != phi::DataLayout::ONEDNN) &&
696 697
          phi::OneDNNContext::tls().get_cur_paddle_data_layout() ==
              phi::DataLayout::kNHWC) {
698 699 700
        return phi::KernelKey(tensor.place(),
                              phi::DataLayout::kNHWC,
                              expected_kernel_type.dtype());
701 702
      }
#endif
703 704
      return phi::KernelKey(
          tensor.place(), tensor.layout(), expected_kernel_type.dtype());
705 706
    }
  }
M
Markus Kliegl 已提交
707 708 709 710
};

class MatMulOpMaker : public framework::OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
711
  void Make() override {
M
Markus Kliegl 已提交
712 713 714 715 716 717 718 719 720 721 722
    AddInput("X", "The first input of MatMul op");
    AddInput("Y", "The second input of MatMul op");
    AddOutput("Out", "The output of MatMul op");
    AddAttr<bool>("transpose_X",
                  R"DOC(If true, use the transpose of `X`.
        )DOC")
        .SetDefault(false);
    AddAttr<bool>("transpose_Y",
                  R"DOC(If true, use the transpose of `Y`.
        )DOC")
        .SetDefault(false);
S
sneaxiy 已提交
723
    AddAttr<float>("alpha", "The scale of Out").SetDefault(1.0f);
724 725 726
    AddAttr<bool>(
        "use_mkldnn",
        "(bool, default false) Indicates if MKL-DNN kernel will be used")
727 728
        .SetDefault(false)
        .AsExtra();
729 730 731 732
    AddAttr<bool>(
        "use_quantizer",
        "(bool, default false) "
        "This parameter is no longer used. Use 'mkldnn_data_type' instead.")
733 734
        .SetDefault(false)
        .AsExtra();
735 736 737 738
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
739 740
        .InEnum({"float32", "int8", "bfloat16"})
        .AsExtra();
741
    /* int8 parameters */
742 743
    AddAttr<float>("Scale_x",
                   "(float, default 1.0f), The quantize scale of X tensor")
744 745
        .SetDefault(1.0f)
        .AsExtra();
746 747
    AddAttr<float>("Scale_y",
                   "(float, default 1.0f), The quantize scale of Y tensor")
748 749
        .SetDefault(1.0f)
        .AsExtra();
750 751
    AddAttr<float>("Scale_out",
                   "(float, default 1.0f), The quantize scale of output data")
752 753
        .SetDefault(1.0f)
        .AsExtra();
754 755 756
    AddAttr<bool>("force_fp32_output",
                  "(bool, default false) Force INT8 kernel output FP32, only "
                  "used in MKL-DNN INT8")
757 758
        .SetDefault(false)
        .AsExtra();
759

760 761
#if defined(PADDLE_WITH_MKLML) && !defined(PADDLE_WITH_CUDA) && \
    !defined(PADDLE_WITH_HIP)
762 763 764
    AddAttr<int>("head_number", "The number of heads of the matrix")
        .SetDefault(1);
#endif
M
Markus Kliegl 已提交
765
    AddComment(R"DOC(
K
kexinzhao 已提交
766 767
MatMul Operator.
This operator is used to perform (batched) matrix multiplication
M
Markus Kliegl 已提交
768 769 770 771 772 773 774 775 776 777 778 779
over the last two dimensions of the input tensors `X` and `Y`.
If a transpose flag is specified, the last two dimensions of the
tensor are transposed. If the tensor is rank-1 of shape [D], then
for `X` it is treated as [1, D] in nontransposed form and as [D, 1]
in transposed form, whereas for `Y` it is the opposite: It is treated
as [D, 1] in nontransposed form and as [1, D] in transposed form.
Examples without transpose:
- X: [K], Y: [K] => Out: [1]
- X: [K], Y: [K, N] => Out: [N]
- X: [B, M, K], Y: [K] => Out: [B, M]
- X: [M, K], Y: [B, K, N] => Out: [B, M, N]
- X: [B, M, K], Y: [B, K, N] => Out: [B, M, N]
C
chengduoZH 已提交
780
- X: [B, ..., M, K], Y: [B, ..., K, N] => Out: [B, ..., M, N]
781 782
Example of matrix multiplication with head_number of H
- X: [B, M, K], Y: [B, K, N] => Out: [B, M, H * N]
M
Markus Kliegl 已提交
783 784
The behavior is designed to be similar to the `numpy.matmul` function.
The differences are:
C
chengduoZH 已提交
785 786
- When the rank of the input data is less than or equal to 3, it
  is similar to the `numpy.matmul` function.
C
chengduoZH 已提交
787
- When the rank of the input is greater than 3, the rank of X and
C
chengduoZH 已提交
788
  Y must be equal, and the first `rank - 2` dimensions must be equal.
M
Markus Kliegl 已提交
789
- We add `transpose_X` and `transpose_Y` flags.
790 791 792
- We add `head_number` attribute, which is used to multiple two matrixes head
  by head, and eventually concatenates the output of several (head_number)
  small matrixes multiplication.
M
Markus Kliegl 已提交
793
Both the input `X` and `Y` can carry the LoD (Level of Details) information,
K
kexinzhao 已提交
794
or not. But the output only shares the LoD information with input `X`.
M
Markus Kliegl 已提交
795 796 797 798 799 800 801 802 803
)DOC");
  }
};

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

 protected:
Y
yuyang18 已提交
804
  void InferShape(framework::InferShapeContext *context) const override {
805 806
    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "matmul");
    OP_INOUT_CHECK(context->HasInput("Y"), "Input", "Y", "matmul");
807 808 809 810
    OP_INOUT_CHECK(context->HasInput(framework::GradVarName("Out")),
                   "Input",
                   "Out@GRAD",
                   "matmul");
M
Markus Kliegl 已提交
811 812 813 814 815 816 817 818 819 820 821 822 823
    auto x_dims = context->GetInputDim("X");
    auto y_dims = context->GetInputDim("Y");

    auto x_grad_name = framework::GradVarName("X");
    auto y_grad_name = framework::GradVarName("Y");

    if (context->HasOutput(x_grad_name)) {
      context->SetOutputDim(x_grad_name, x_dims);
    }
    if (context->HasOutput(y_grad_name)) {
      context->SetOutputDim(y_grad_name, y_dims);
    }
  }
824

825
  phi::KernelKey GetExpectedKernelType(
826 827 828
      const framework::ExecutionContext &ctx) const override {
    auto input_data_type =
        OperatorWithKernel::IndicateOrPromoteVarDataTypes(ctx, "X", "Y");
829
    return phi::KernelKey(input_data_type, ctx.GetPlace());
830
  }
M
Markus Kliegl 已提交
831 832
};

H
hong 已提交
833 834
template <typename T>
class MatMulOpGradMaker : public framework::SingleGradOpMaker<T> {
Y
Yu Yang 已提交
835
 public:
H
hong 已提交
836
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
Y
Yu Yang 已提交
837 838

 protected:
839
  void Apply(GradOpPtr<T> retv) const override {
Y
Yu Yang 已提交
840
    retv->SetType("matmul_grad");
H
hong 已提交
841 842 843 844 845 846
    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());
Y
Yu Yang 已提交
847 848
  }
};
849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885 886 887 888 889 890 891 892 893 894 895 896 897 898 899 900 901 902 903

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

 protected:
  void InferShape(framework::InferShapeContext *context) const override {
    OP_INOUT_CHECK(context->HasInput("X"), "Input", "X", "matmul");
    OP_INOUT_CHECK(context->HasInput("Y"), "Input", "Y", "matmul");
    OP_INOUT_CHECK(context->HasInput("DOut"), "Input", "DOut", "matmul");

    if (context->HasOutput("DX") && context->HasInput("DDY")) {
      context->ShareDim("X", "DX");
    }

    if (context->HasOutput("DY") && context->HasInput("DDX")) {
      context->ShareDim("Y", "DY");
    }

    if (context->HasOutput("DDOut") &&
        (context->HasInput("DDY") || context->HasInput("DDX"))) {
      context->ShareDim("DOut", "DDOut");
    }
  }
};

template <typename T>
class MatMulOpDoubleGradMaker : public framework::SingleGradOpMaker<T> {
 public:
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;

 protected:
  void Apply(GradOpPtr<T> retv) const override {
    retv->SetType("matmul_grad_grad");
    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")));

    auto ddx = this->OutputGrad(framework::GradVarName("X"));
    auto ddy = this->OutputGrad(framework::GradVarName("Y"));

    if (!ddx.empty() || !ddy.empty()) {
      retv->SetOutput("DDOut", this->InputGrad(framework::GradVarName("Out")));
    }
    retv->SetOutput(
        "DX", ddy.empty() ? this->EmptyInputGrad() : this->InputGrad("X"));
    retv->SetOutput(
        "DY", ddx.empty() ? this->EmptyInputGrad() : this->InputGrad("Y"));

    retv->SetAttrMap(this->Attrs());
  }
};

M
Markus Kliegl 已提交
904 905 906 907
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
908 909 910
REGISTER_OPERATOR(matmul,
                  ops::MatMulOp,
                  ops::MatMulOpMaker,
H
hong 已提交
911 912
                  ops::MatMulOpGradMaker<paddle::framework::OpDesc>,
                  ops::MatMulOpGradMaker<paddle::imperative::OpBase>);
913 914
REGISTER_OPERATOR(matmul_grad,
                  ops::MatMulOpGrad,
915 916 917
                  ops::MatMulOpDoubleGradMaker<paddle::framework::OpDesc>,
                  ops::MatMulOpDoubleGradMaker<paddle::imperative::OpBase>);
REGISTER_OPERATOR(matmul_grad_grad, ops::MatMulOpDoubleGrad);
L
Leo Chen 已提交
918 919 920 921 922 923 924 925 926 927
REGISTER_OP_CPU_KERNEL(matmul,
                       ops::MatMulKernel<phi::CPUContext, float>,
                       ops::MatMulKernel<phi::CPUContext, double>);
REGISTER_OP_CPU_KERNEL(matmul_grad,
                       ops::MatMulGradKernel<phi::CPUContext, float>,
                       ops::MatMulGradKernel<phi::CPUContext, double>);

REGISTER_OP_CPU_KERNEL(matmul_grad_grad,
                       ops::MatMulDoubleGradKernel<phi::CPUContext, float>,
                       ops::MatMulDoubleGradKernel<phi::CPUContext, double>);
928

929
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
Y
Yu Yang 已提交
930
REGISTER_OP_CUDA_KERNEL(
931
    matmul,
L
Leo Chen 已提交
932 933 934
    ops::MatMulKernel<phi::GPUContext, float>,
    ops::MatMulKernel<phi::GPUContext, double>,
    ops::MatMulKernel<phi::GPUContext, paddle::platform::float16>);
Y
Yu Yang 已提交
935 936
REGISTER_OP_CUDA_KERNEL(
    matmul_grad,
L
Leo Chen 已提交
937 938 939 940 941 942
    ops::MatMulGradKernel<phi::GPUContext, float>,
    ops::MatMulGradKernel<phi::GPUContext, double>,
    ops::MatMulGradKernel<phi::GPUContext, paddle::platform::float16>);
REGISTER_OP_CUDA_KERNEL(matmul_grad_grad,
                        ops::MatMulDoubleGradKernel<phi::GPUContext, float>,
                        ops::MatMulDoubleGradKernel<phi::GPUContext, double>);
Y
Yu Yang 已提交
943
#endif
944

945 946
REGISTER_OP_VERSION(matmul).AddCheckpoint(
    R"ROC(Register matmul for adding the attribute of
947
       fused_reshape_Y)ROC",
948 949 950 951 952 953
    paddle::framework::compatible::OpVersionDesc().NewAttr(
        "fused_reshape_Y",
        "In order to support the function of fused the input Y "
        " and input X into the input X when "
        "using the operator of matmul, and get raw shape of input Y.",
        std::vector<int>{}));