matmul_v2_mkldnn_op.cc 45.4 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.

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. */
14
#include "paddle/fluid/operators/mkldnn/matmul_mkldnn_op.h"
15

16
namespace {
17
using dnnl::memory;
18 19 20
using paddle::framework::DataLayout;
using paddle::framework::ExecutionContext;
using paddle::platform::GetMKLDNNFormat;
21
using paddle::platform::MatMulV2MKLDNNHandler;
22
using paddle::platform::MKLDNNDeviceContext;
23
using paddle::platform::MKLDNNFormatForSize;
24 25
using paddle::platform::MKLDNNGetDataType;
using paddle::platform::to_void_cast;
26
using phi::vectorize;
27
using Tensor = phi::DenseTensor;
28
using paddle::framework::GradVarName;
29
using phi::make_ddim;
30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90

// 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.
static Tensor FoldOuterDims(const Tensor &input) {
  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 T>
static Tensor FoldFirstAndLastDims(const MKLDNNDeviceContext &dev_ctx,
                                   const Tensor *input) {
  auto input_dims = vectorize(input->dims());
  if (input_dims.size() != 3) {
    return *input;
  }

  Tensor output;
  output.Resize({input_dims[1], input_dims[0], input_dims[2]});

  auto output_dims = vectorize(output.dims());

  memory::data_type input_type = paddle::framework::ToMKLDNNDataType(
      paddle::framework::TransToProtoVarType(input->dtype()));
  paddle::platform::ReorderMKLDNNHandler reorder_handler(
      output_dims,
      paddle::framework::TransToProtoVarType(input->dtype()),
      input_type,
      dev_ctx.GetEngine());

  auto reorder_src_memory_p = reorder_handler.AcquireSrcMemory(
      memory::format_tag::abc,
      paddle::platform::to_void_cast(input->data<T>()));
  auto reorder_dst_memory_p = reorder_handler.AcquireDstMemory(
      &output, memory::format_tag::bac, dev_ctx.GetPlace());
  auto reorder_p = reorder_handler.AcquireReorder(reorder_src_memory_p,
                                                  reorder_dst_memory_p);

  auto &astream = MKLDNNDeviceContext::tls().get_stream();
  reorder_p->execute(astream, *reorder_src_memory_p, *reorder_dst_memory_p);
  astream.wait();

  output.Resize({input_dims[1], input_dims[0] * input_dims[2]});
  return output;
}

template <typename T>
constexpr bool IsInt8() {
  return std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value;
}

template <typename T>
constexpr bool IsBfloat16() {
  return std::is_same<T, paddle::platform::bfloat16>::value;
}
91

92 93
// Get row matrix shape from a vector shape. If the rank of x_dim > 1, the
// original x_dim is returned.
94 95
static paddle::framework::DDim RowMatrixDimsFromVector(
    const paddle::framework::DDim &x_dim) {
96
  return x_dim.size() > 1 ? x_dim : phi::make_ddim({1, x_dim[0]});
97 98 99 100
}

// Get column matrix shape from a vector shape. If the ran of y_dim > 1, the
// original y_dim is returned.
101 102
static paddle::framework::DDim ColumnMatrixDimsFromVector(
    const paddle::framework::DDim &y_dim) {
103
  return y_dim.size() > 1 ? y_dim : phi::make_ddim({y_dim[0], 1});
104 105
}

J
Jacek Czaja 已提交
106 107 108
phi::DDim GetDimForInput(const ExecutionContext &ctx, std::string input_name) {
  auto shape = ctx.Attr<std::vector<int>>("fused_reshape_" + input_name);
  auto axis = ctx.Attr<std::vector<int>>("fused_transpose_" + input_name);
109
  auto input_dims = ctx.Input<phi::DenseTensor>(input_name)->dims();
J
Jacek Czaja 已提交
110 111 112 113 114 115
  if (!shape.empty() && !axis.empty()) {
    return input_dims.reshape(shape).transpose(axis);
  }
  return input_dims;
}

116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 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 181 182 183 184
template <typename XT, typename YT, typename OT>
class MatMulMKLDNNHandler
    : public paddle::platform::MKLDNNHandlerNoCachingT<XT, dnnl::matmul> {
 public:
  MatMulMKLDNNHandler(const dnnl::engine engine,
                      paddle::platform::Place cpu_place,
                      Tensor *x,
                      bool trans_x,
                      Tensor *y,
                      bool trans_y,
                      Tensor *out,
                      float scale)
      : paddle::platform::MKLDNNHandlerNoCachingT<XT, dnnl::matmul>(engine,
                                                                    cpu_place) {
    auto mat_dim_x = phi::funcs::CreateMatrixDescriptor(x->dims(), 0, trans_x);
    auto mat_dim_y = phi::funcs::CreateMatrixDescriptor(y->dims(), 0, trans_y);

    memory::dim x_bs = mat_dim_x.batch_size_;
    memory::dim y_bs = mat_dim_y.batch_size_;

    memory::dim out_bs = x_bs || y_bs ? std::max(x_bs, y_bs) : 1;
    const memory::dim M = mat_dim_x.height_;
    const memory::dim N = mat_dim_y.width_;
    const memory::dim K = mat_dim_x.width_;

    memory::dims x_dims = {x_bs > 0 ? x_bs : 1, M, K};
    memory::dims y_dims = {y_bs > 0 ? y_bs : 1, K, N};
    memory::dims out_dims = {out_bs, M, N};

    memory::dims x_strides =
        !trans_x ? memory::dims{M * K, K, 1} : memory::dims{M * K, 1, M};

    memory::dims y_strides =
        !trans_y ? memory::dims{N * K, N, 1} : memory::dims{N * K, 1, K};
    memory::dims out_strides = memory::dims{M * N, N, 1};

    auto x_md = memory::desc(x_dims, MKLDNNGetDataType<XT>(), x_strides);
    auto y_md = memory::desc(y_dims, MKLDNNGetDataType<YT>(), y_strides);
    auto out_md = memory::desc(out_dims, MKLDNNGetDataType<OT>(), out_strides);

    dnnl::primitive_attr attrs;
    if (scale != 1.0f) attrs.set_output_scales(0, {scale});

    this->AcquireForwardPrimitiveDescriptor(attrs, x_md, y_md, out_md);
  }
  // Constructor for FWD MatMul
  MatMulMKLDNNHandler(const dnnl::engine engine, const ExecutionContext &ctx)
      : paddle::platform::MKLDNNHandlerNoCachingT<XT, dnnl::matmul>(
            engine, ctx.GetPlace()) {
    const dnnl::primitive_attr matmul_attrs = CreateMatmulAttrs(ctx);

    auto matmul_dims_ = GetMatmulDims(ctx);
    auto x_md = memory::desc(
        matmul_dims_.x_dims, MKLDNNGetDataType<XT>(), matmul_dims_.x_strides);
    auto y_md = memory::desc(
        matmul_dims_.y_dims, MKLDNNGetDataType<YT>(), matmul_dims_.y_strides);
    auto out_md = memory::desc(matmul_dims_.out_dims,
                               MKLDNNGetDataType<OT>(),
                               matmul_dims_.out_strides);
    this->AcquireForwardPrimitiveDescriptor(matmul_attrs, x_md, y_md, out_md);
  }

  std::shared_ptr<memory> AcquireWeightsMemory(const Tensor *input) {
    const YT *input_data = input->data<YT>();
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->weights_desc(),
                                            to_void_cast<YT>(input_data));
  }

 public:
185 186 187
  void Execute(const phi::DenseTensor *x,
               const phi::DenseTensor *y,
               phi::DenseTensor *out) {
188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209
    const auto src_memory_p = this->AcquireSrcMemory(x);
    const auto weights_memory_p = this->AcquireWeightsMemory(y);
    const auto dst_memory_p = this->AcquireDstMemory(out);

    auto matmul_p = this->AcquireForwardPrimitive();

    std::unordered_map<int, dnnl::memory> matmul_args = {
        {DNNL_ARG_SRC, *src_memory_p},
        {DNNL_ARG_WEIGHTS, *weights_memory_p},
        {DNNL_ARG_DST, *dst_memory_p}};

    auto &astream = paddle::platform::MKLDNNDeviceContext::tls().get_stream();

    // Simulate batch matmul by processing in loop
    void *x_ptr = src_memory_p->get_data_handle();
    void *y_ptr = weights_memory_p->get_data_handle();
    void *out_ptr = dst_memory_p->get_data_handle();
    auto offsets = this->GetOffsets();
    for (uint16_t i = 0; i < this->GetBatchSize(); ++i) {
      src_memory_p->set_data_handle(x_ptr);
      weights_memory_p->set_data_handle(y_ptr);
      dst_memory_p->set_data_handle(out_ptr);
210
      matmul_p->execute(astream, matmul_args);
211 212 213 214 215 216
      x_ptr = static_cast<char *>(x_ptr) + std::get<0>(offsets);
      y_ptr = static_cast<char *>(y_ptr) + std::get<1>(offsets);
      out_ptr = static_cast<char *>(out_ptr) + std::get<2>(offsets);
    }
    astream.wait();

217
    out->set_mem_desc(dst_memory_p->get_desc().reshape(out->dims()));
218 219
  }

220
  std::shared_ptr<dnnl::memory> AcquireDstMemory(phi::DenseTensor *output) {
221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242
    // We cannot use base AcquireDstMemory as it makes an allocation request
    // base on DST memory primitive size. This is fine in general, but in MatMul
    // we have primitive that covers only one batch of Data and then shift
    // pointer for every new batch. Hence Tensor size is bigger that dst memory
    // primitive size. So would we request less memory that is there and it
    // triggers an
    // assertion.  So as there is no 'any' format here we can leave default size
    // of Tensor as computed in ComputeInferShape
    OT *ptr = output->mutable_data<OT>(this->place_);
    return this->AcquireMemoryFromPrimitive(this->fwd_pd_->dst_desc(), ptr);
  }

 private:
  struct MatMulDims {
    const memory::dims x_dims, y_dims, out_dims, x_strides, y_strides,
        out_strides;
  };

  std::pair<phi::funcs::MatDescriptor, memory::dims> GetInputDimsAndStrides(
      const ExecutionContext &ctx, std::string input_name) {
    auto shape = ctx.Attr<std::vector<int>>("fused_reshape_" + input_name);
    auto axis = ctx.Attr<std::vector<int>>("fused_transpose_" + input_name);
243
    auto input_dims = ctx.Input<phi::DenseTensor>(input_name)->dims();
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 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 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 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475
    auto new_dims = input_dims;
    if (!shape.empty() && !axis.empty()) {
      new_dims = input_dims.reshape(shape).transpose(axis);
    }

    auto &MatrixDimsFromVector = input_name == "X" ? RowMatrixDimsFromVector
                                                   : ColumnMatrixDimsFromVector;
    phi::funcs::MatDescriptor mat_dim = phi::funcs::CreateMatrixDescriptor(
        MatrixDimsFromVector(new_dims),
        0,
        ctx.Attr<bool>("transpose_" + input_name));

    memory::dims strides;
    if (!shape.empty()) {
      auto shape2 = input_dims.reshape(shape);
      strides.push_back(1);
      for (auto i = shape2.size() - 1; i > 0; --i) {
        strides.insert(strides.begin(), strides.front() * shape2[i]);
      }
      strides = Transpose(strides, axis);
      if (shape.size() == 4)
        strides.erase(strides.begin());
      else if (shape.size() == 2)
        strides.insert(strides.begin(), shape[0] * shape[1]);
      mat_dim.stride_ = strides[0];
      if (mat_dim.trans_) std::swap(*strides.rbegin(), *(++strides.rbegin()));
    }
    return std::make_pair(mat_dim, strides);
  }

  float ComputeOutputScale(const ExecutionContext &ctx) {
    float scale_x = ctx.Attr<float>("Scale_x");
    float scale_y = ctx.Attr<float>("Scale_y");
    bool force_fp32_out = ctx.Attr<bool>("force_fp32_output");
    float scale_out = force_fp32_out ? 1.f : ctx.Attr<float>("Scale_out");
    float alpha = ctx.Attr<float>("alpha");
    return alpha * scale_out / (scale_x * scale_y);
  }

  bool IsInputFused(const ExecutionContext &ctx) const {
    return !(ctx.Attr<std::vector<int>>("fused_reshape_X").empty() &&
             ctx.Attr<std::vector<int>>("fused_reshape_Y").empty());
  }

  bool IsOutputFused(const ExecutionContext &ctx) const {
    auto &fused_reshape_Out = ctx.Attr<std::vector<int>>("fused_reshape_Out");
    auto &fused_transpose_Out =
        ctx.Attr<std::vector<int>>("fused_transpose_Out");
    return !fused_reshape_Out.empty() && !fused_transpose_Out.empty();
  }

  MatMulDims GetMatmulDims(const ExecutionContext &ctx) {
    phi::funcs::MatDescriptor mat_dim_x;
    memory::dims strides_x;
    std::tie(mat_dim_x, strides_x) = GetInputDimsAndStrides(ctx, "X");
    phi::funcs::MatDescriptor mat_dim_y;
    memory::dims strides_y;
    std::tie(mat_dim_y, strides_y) = GetInputDimsAndStrides(ctx, "Y");

    auto x_bs = mat_dim_x.batch_size_;
    auto y_bs = mat_dim_y.batch_size_;
    PADDLE_ENFORCE_EQ(x_bs > 0 && y_bs > 0 && x_bs != y_bs,
                      false,
                      paddle::platform::errors::InvalidArgument(
                          "If batch sizes of X and Y are positive,"
                          "they have to be equal."));

    memory::dim out_bs = x_bs || y_bs ? std::max(x_bs, y_bs) : 1;
    const memory::dim M = mat_dim_x.height_;
    const memory::dim N = mat_dim_y.width_;
    const memory::dim K = mat_dim_x.width_;

    batch_size_ = 1;
    if (out_bs > 1 && (IsOutputFused(ctx) || IsInputFused(ctx))) {
      auto x_dims = GetDimForInput(ctx, "X");
      auto y_dims = GetDimForInput(ctx, "Y");
      batch_size_ = x_bs > y_bs ? x_dims[0] : y_dims[0];
      x_bs /= batch_size_;
      y_bs /= batch_size_;
      out_bs /= batch_size_;
    }
    memory::dims x_dims = {x_bs > 0 ? x_bs : 1, M, K};
    memory::dims y_dims = {y_bs > 0 ? y_bs : 1, K, N};
    memory::dims out_dims = {out_bs, M, N};

    x_offset_ = x_bs * M * K * sizeof(XT);
    y_offset_ = y_bs * K * N * sizeof(YT);
    out_offset_ = out_bs * M * N * sizeof(OT);

    // Translate transA and transB
    if (strides_x.empty())
      strides_x = !ctx.Attr<bool>("transpose_X") ? memory::dims{M * K, K, 1}
                                                 : memory::dims{M * K, 1, M};
    if (strides_y.empty())
      strides_y = !ctx.Attr<bool>("transpose_Y") ? memory::dims{N * K, N, 1}
                                                 : memory::dims{N * K, 1, K};
    memory::dims out_strides = memory::dims{M * N, N, 1};

    CorrectStridesWhenFloatOutputFused(ctx, N, out_bs, &out_strides);

    return {x_dims, y_dims, out_dims, strides_x, strides_y, out_strides};
  }

  std::vector<int64_t> Transpose(const std::vector<int64_t> &x,
                                 const std::vector<int> &axis) {
    size_t in_rank = x.size();
    size_t axis_size = axis.size();

    auto axis_set = std::set<int>(axis.begin(), axis.end());
    PADDLE_ENFORCE_EQ(axis_set.size(),
                      axis_size,
                      paddle::platform::errors::InvalidArgument(
                          "In an axis array, elements must be unique."));

    PADDLE_ENFORCE_EQ(in_rank,
                      axis_size,
                      paddle::platform::errors::InvalidArgument(
                          "The input dimension's size "
                          "should be equal to the axis's size. "
                          "But received dimension is %d, "
                          "axis's size is %d",
                          in_rank,
                          axis_size));

    PADDLE_ENFORCE_LT(*std::max_element(axis.begin(), axis.end()),
                      axis_size,
                      paddle::platform::errors::InvalidArgument(
                          "Axis values must be ranging from 0 to (dims - 1)."));

    std::vector<int64_t> new_x(x.size());
    for (size_t i = 0; i < x.size(); i++) {
      new_x[i] = x[axis[i]];
    }
    return new_x;
  }

  void CorrectStridesWhenFloatOutputFused(const ExecutionContext &ctx,
                                          const memory::dim N,
                                          memory::dim b,
                                          memory::dims *out_strides) const {
    if (!IsInt8<OT>() && !IsBfloat16<OT>() && IsOutputFused(ctx)) {
      *out_strides = {N, b * N, 1};
    }
  }

  uint16_t GetBatchSize(void) const { return batch_size_; }

  std::tuple<uint32_t, uint32_t, uint32_t> GetOffsets() const {
    return std::make_tuple(x_offset_, y_offset_, out_offset_);
  }

  dnnl::primitive_attr CreateMatmulAttrs(const ExecutionContext &ctx) {
    dnnl::primitive_attr matmul_attrs;
    dnnl::post_ops post_operations;

    float scale_out = ComputeOutputScale(ctx);
    if (scale_out != 1.0f) {
      matmul_attrs.set_output_scales(0, {scale_out});
    }
    paddle::platform::AppendActivation(ctx, post_operations);

    matmul_attrs.set_post_ops(post_operations);
    return matmul_attrs;
  }

 private:
  uint32_t x_offset_;
  uint32_t y_offset_;
  uint32_t out_offset_;
  uint16_t batch_size_;
};

/**
 * 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 ReshapeTensorToMatrixSequence(
    Tensor *x, const phi::funcs::MatDescriptor &descriptor) {
  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.
 */
static void ReshapeXYOutToMatrixSequence(
    Tensor *x, Tensor *y, Tensor *out, bool trans_x, bool trans_y) {
  auto x_dim = RowMatrixDimsFromVector(x->dims());
  auto y_dim = ColumnMatrixDimsFromVector(y->dims());
  auto mat_dim_x = phi::funcs::CreateMatrixDescriptor(x_dim, 0, trans_x);
  auto mat_dim_y = phi::funcs::CreateMatrixDescriptor(y_dim, 0, trans_y);
  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_),
                 mat_dim_x.height_,
                 mat_dim_y.width_});
  }

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

// Choose appropriate Handler instances based on inferred
// output type (uint8, int8 or float).
template <typename XT, typename YT>
static void ExecuteMatMul(const ExecutionContext &ctx) {
  constexpr bool is_int8 = IsInt8<XT>();
  constexpr bool is_bfloat16 = IsBfloat16<XT>();
  const bool force_fp32_output = ctx.Attr<bool>("force_fp32_output");
476 477 478 479
  const bool fuse_relu =
      ctx.HasAttr("fuse_activation")
          ? ctx.Attr<std::string>("fuse_activation") == "relu"
          : false;
480 481 482
  auto *x = ctx.Input<phi::DenseTensor>("X");
  auto *y = ctx.Input<phi::DenseTensor>("Y");
  auto *out = ctx.Output<phi::DenseTensor>("Out");
483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517
  const auto &dev_ctx =
      ctx.template device_context<paddle::platform::MKLDNNDeviceContext>();
  const auto &onednn_engine = dev_ctx.GetEngine();

  if (force_fp32_output || ((!is_int8) && (!is_bfloat16))) {
    MatMulMKLDNNHandler<XT, YT, float>(onednn_engine, ctx).Execute(x, y, out);
  } else if (is_bfloat16) {
    MatMulMKLDNNHandler<XT, YT, paddle::platform::bfloat16>(onednn_engine, ctx)
        .Execute(x, y, out);
  } else if (fuse_relu) {
    MatMulMKLDNNHandler<XT, YT, uint8_t>(onednn_engine, ctx).Execute(x, y, out);
  } else {
    MatMulMKLDNNHandler<XT, YT, int8_t>(onednn_engine, ctx).Execute(x, y, out);
  }
}

template <typename T>
class MatMulMKLDNNKernel : public paddle::framework::OpKernel<T> {
 public:
  void Compute(const ExecutionContext &ctx) const override {
    if (ctx.HasAttr("head_number")) {
      PADDLE_ENFORCE_EQ(
          ctx.Attr<int>("head_number"),
          1,
          paddle::platform::errors::Unimplemented(
              "oneDNN matmul doesn't support multiple heads. Expected "
              "head_number=1. But received `head_number` is %d",
              ctx.Attr<int>("head_number")));
    }
    ExecuteMatMul<T, T>(ctx);
  }
};

static std::vector<int64_t> Transpose(const std::vector<int64_t> &x,
                                      const std::vector<int> &axis) {
518 519 520 521
  size_t in_rank = x.size();
  size_t axis_size = axis.size();

  auto axis_set = std::set<int>(axis.begin(), axis.end());
522 523
  PADDLE_ENFORCE_EQ(axis_set.size(),
                    axis_size,
524 525 526
                    paddle::platform::errors::InvalidArgument(
                        "In an axis array, elements must be unique."));

527 528
  PADDLE_ENFORCE_EQ(in_rank,
                    axis_size,
529 530 531 532 533
                    paddle::platform::errors::InvalidArgument(
                        "The input dimension's size "
                        "should be equal to the axis's size. "
                        "But received dimension is %d, "
                        "axis's size is %d",
534 535
                        in_rank,
                        axis_size));
536

537 538
  PADDLE_ENFORCE_LT(*std::max_element(axis.begin(), axis.end()),
                    axis_size,
539 540 541 542 543 544 545 546 547 548
                    paddle::platform::errors::InvalidArgument(
                        "Axis values must be ranging from 0 to (dims - 1)."));

  std::vector<int64_t> new_x(x.size());
  for (size_t i = 0; i < x.size(); i++) {
    new_x[i] = x[axis[i]];
  }
  return new_x;
}

549
std::vector<int64_t> GetInputStrides(const ExecutionContext &ctx,
550 551 552
                                     const std::string input_name) {
  auto shape = ctx.Attr<std::vector<int>>("fused_reshape_" + input_name);
  auto axis = ctx.Attr<std::vector<int>>("fused_transpose_" + input_name);
553
  auto input_dims = ctx.Input<phi::DenseTensor>(input_name)->dims();
554 555 556 557 558
  auto new_dims = input_dims;
  if (!shape.empty() && !axis.empty()) {
    new_dims = input_dims.reshape(shape).transpose(axis);
  }

559
  auto &MatrixDimsFromVector =
560
      input_name == "X" ? RowMatrixDimsFromVector : ColumnMatrixDimsFromVector;
561
  phi::funcs::MatDescriptor mat_dim = phi::funcs::CreateMatrixDescriptor(
562 563
      MatrixDimsFromVector(new_dims),
      0,
564 565 566 567
      ctx.HasAttr("trans_x")
          ? ctx.Attr<bool>(std::string("trans_") +
                           static_cast<char>(std::tolower(input_name[0])))
          : ctx.Attr<bool>(std::string("transpose_") + input_name[0]));
568 569 570 571 572 573 574 575 576 577 578 579 580 581 582 583 584 585 586

  std::vector<int64_t> strides;
  if (!shape.empty()) {
    auto shape2 = input_dims.reshape(shape);
    strides.push_back(1);
    for (auto i = shape2.size() - 1; i > 0; --i) {
      strides.insert(strides.begin(),
                     strides.front() * static_cast<int64_t>(shape2[i]));
    }
    strides = Transpose(strides, axis);
    if (shape.size() == 2)
      strides.insert(strides.begin(),
                     static_cast<int64_t>(shape[0] * shape[1]));
    mat_dim.stride_ = strides[0];
    if (mat_dim.trans_) std::swap(*strides.rbegin(), *(++strides.rbegin()));
  }
  return strides;
}

587 588 589
bool IsOutputFused(const ExecutionContext &ctx) {
  auto &fused_reshape_Out = ctx.Attr<std::vector<int>>("fused_reshape_Out");
  auto &fused_transpose_Out = ctx.Attr<std::vector<int>>("fused_transpose_Out");
590 591 592
  return !fused_reshape_Out.empty() && !fused_transpose_Out.empty();
}

593
float ComputeOutputScale(const ExecutionContext &ctx) {
594 595 596 597
  float scale_x = ctx.Attr<float>("Scale_x");
  float scale_y = ctx.Attr<float>("Scale_y");
  bool force_fp32_out = ctx.Attr<bool>("force_fp32_output");
  float scale_out = force_fp32_out ? 1.f : ctx.Attr<float>("Scale_out");
598 599
  float alpha = ctx.HasAttr("alpha") ? ctx.Attr<float>("alpha") : 1.0f;
  return alpha * scale_out / (scale_x * scale_y);
600
}
601

602
template <typename T, typename T_out>
603 604
void ExecuteMatMulV2(const ExecutionContext &ctx,
                     const MKLDNNDeviceContext &dev_ctx,
605
                     const dnnl::engine onednn_engine,
606
                     paddle::platform::Place cpu_place,
607 608
                     const Tensor *x,
                     const std::vector<int64_t> &x_dims,
609
                     bool trans_x,
610 611
                     const Tensor *y,
                     const std::vector<int64_t> &y_dims,
612
                     bool trans_y,
613 614
                     Tensor *out,
                     const std::vector<int64_t> &out_dims,
615
                     int execution_number = 0) {
616 617
  std::vector<int64_t> x_strides_override = GetInputStrides(ctx, "X");
  std::vector<int64_t> y_strides_override = GetInputStrides(ctx, "Y");
618 619 620 621 622 623 624 625 626 627
  MatMulV2MKLDNNHandler<T, T, T_out> handler(ctx,
                                             onednn_engine,
                                             ctx.GetPlace(),
                                             x_dims,
                                             trans_x,
                                             y_dims,
                                             trans_y,
                                             IsOutputFused(ctx),
                                             x_strides_override,
                                             y_strides_override);
628

629 630 631
  const auto src_memory_p = handler.AcquireSrcMemory(x);
  const auto weights_memory_p = handler.AcquireWeightsMemory(y);
  const auto dst_memory_p = handler.AcquireDstMemory(out);
632

633
  auto matmul_p = handler.AcquireForwardPrimitive();
634

635 636 637 638
  std::unordered_map<int, memory> matmul_args = {
      {DNNL_ARG_SRC, *src_memory_p},
      {DNNL_ARG_WEIGHTS, *weights_memory_p},
      {DNNL_ARG_DST, *dst_memory_p}};
639

640
  if (ctx.HasInput("ResidualData")) {
641
    auto *residual_data = ctx.Input<phi::DenseTensor>("ResidualData");
642 643 644 645 646
    const auto residual_data_memory_p = handler.AcquireSrcMemory(residual_data);
    matmul_args.insert({DNNL_ARG_ATTR_MULTIPLE_POST_OP(0) | DNNL_ARG_SRC_1,
                        *residual_data_memory_p});
  }

647
  auto &astream = MKLDNNDeviceContext::tls().get_stream();
648 649
  matmul_p->execute(astream, matmul_args);
  astream.wait();
650 651 652 653 654 655 656 657 658 659 660 661

  // TODO(jczaja): Explain why int8 format of dst is ABCD and do not need
  // permute
  if (IsOutputFused(ctx) && !IsInt8<T_out>()) {
    auto axis = ctx.Attr<std::vector<int>>("fused_transpose_Out");
    auto permuted_md = dst_memory_p->get_desc().permute_axes(axis);
    out->set_mem_desc(
        permuted_md.reshape(phi::vectorize<int64_t>(out->dims())));
  } else {
    out->set_mem_desc(
        dst_memory_p->get_desc().reshape(phi::vectorize<int64_t>(out->dims())));
  }
662 663 664 665 666
}

template <typename T>
class MatMulV2MKLDNNKernel : public paddle::framework::OpKernel<T> {
 public:
667 668 669 670 671 672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692
  void Compute(const ExecutionContext &ctx) const override {
    if (ctx.HasAttr("head_number")) {
      PADDLE_ENFORCE_EQ(
          ctx.Attr<int>("head_number"),
          1,
          paddle::platform::errors::Unimplemented(
              "oneDNN matmul doesn't support multiple heads. Expected "
              "head_number=1. But received `head_number` is %d",
              ctx.Attr<int>("head_number")));
    }
    constexpr bool is_int8 = IsInt8<T>();
    constexpr bool is_bfloat16 = IsBfloat16<T>();
    const bool force_fp32_output = ctx.HasAttr("force_fp32_output")
                                       ? ctx.Attr<bool>("force_fp32_output")
                                       : false;
    constexpr bool fuse_relu = false;  // TODO(intel): Enable eltwise fuses
    if (force_fp32_output || ((!is_int8) && (!is_bfloat16))) {
      RunKernel<float>(ctx);
    } else if (is_bfloat16) {
      RunKernel<paddle::platform::bfloat16>(ctx);
    } else if (fuse_relu) {
      RunKernel<uint8_t>(ctx);
    } else {
      RunKernel<int8_t>(ctx);
    }
  }
693

694
 private:
695 696 697 698 699 700 701
  void CalculateMatrixDims(const ExecutionContext &ctx,
                           const std::vector<int64_t> &x_dims,
                           const std::vector<int64_t> &y_dims,
                           std::vector<int64_t> *x_bd_dims,
                           std::vector<int64_t> *y_bd_dims,
                           std::vector<int64_t> *out_dims,
                           Tensor *out) const {
702
    if (x_dims.size() == 1) {
703
      (*x_bd_dims)[(*x_bd_dims).size() - 1] = x_dims[0];
704
    } else if (x_dims.size() == 2) {
705 706
      (*x_bd_dims)[(*x_bd_dims).size() - 1] = x_dims[1];
      (*x_bd_dims)[(*x_bd_dims).size() - 2] = x_dims[0];
707 708
    } else {
      for (size_t i = 0; i < x_dims.size(); ++i) {
709
        (*x_bd_dims)[(*x_bd_dims).size() - x_dims.size() + i] = x_dims[i];
710 711 712
      }
    }
    if (y_dims.size() == 1) {
713
      (*y_bd_dims)[(*x_bd_dims).size() - 2] = y_dims[0];
714
    } else if (y_dims.size() == 2) {
715 716
      (*y_bd_dims)[(*y_bd_dims).size() - 1] = y_dims[1];
      (*y_bd_dims)[(*y_bd_dims).size() - 2] = y_dims[0];
717 718
    } else {
      for (size_t i = 0; i < y_dims.size(); ++i) {
719
        (*y_bd_dims)[(*y_bd_dims).size() - y_dims.size() + i] = y_dims[i];
720 721 722
      }
    }

723
    if (!IsOutputFused(ctx) && x_dims.size() > 2 && y_dims.size() > 2) {
724
      for (size_t i = 0; i < (*x_bd_dims).size() - 2; ++i) {
725
        PADDLE_ENFORCE_EQ(
726 727
            (*x_bd_dims)[i] == (*y_bd_dims)[i] || (*x_bd_dims)[i] == 1 ||
                (*y_bd_dims)[i] == 1,
728 729 730 731 732
            true,
            paddle::platform::errors::InvalidArgument(
                "Tensor dimensions are incorrect for broadcasting."
                "Dimensions in X and Y must be same or equal to 1, but "
                "received x_dim[%d]=%d and y_dims[%d]= %d",
733 734 735 736 737
                i,
                (*x_bd_dims)[i],
                i,
                (*y_bd_dims)[i]));
        (*out_dims)[i] = std::max((*x_bd_dims)[i], (*y_bd_dims)[i]);
738
      }
739
      out->Resize(phi::make_ddim((*out_dims)));
740 741 742
    }
  }

743
  template <typename T_out>
744 745 746
  void RunKernel(const ExecutionContext &ctx) const {
    const auto &dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
    const auto &onednn_engine = dev_ctx.GetEngine();
747

748 749 750
    auto *x = ctx.Input<phi::DenseTensor>("X");
    auto *y = ctx.Input<phi::DenseTensor>("Y");
    auto *out = ctx.Output<phi::DenseTensor>("Out");
751 752 753 754
    bool trans_x = ctx.HasAttr("trans_x") ? ctx.Attr<bool>("trans_x")
                                          : ctx.Attr<bool>("transpose_X");
    bool trans_y = ctx.HasAttr("trans_y") ? ctx.Attr<bool>("trans_y")
                                          : ctx.Attr<bool>("transpose_Y");
755

756 757
    auto x_dims = vectorize(GetDimForInput(ctx, "X"));
    auto y_dims = vectorize(GetDimForInput(ctx, "Y"));
758
    auto out_dims = vectorize(out->dims());
759

760
    int ndims = std::max(x_dims.size(), y_dims.size());
761 762 763 764 765
    ndims = std::max(ndims, 3);

    std::vector<int64_t> x_bd_dims(ndims, 1);
    std::vector<int64_t> y_bd_dims(ndims, 1);

766 767 768
    CalculateMatrixDims(
        ctx, x_dims, y_dims, &x_bd_dims, &y_bd_dims, &out_dims, out);

769 770 771 772 773 774 775 776 777 778 779 780
    ExecuteMatMulV2<T, T_out>(ctx,
                              dev_ctx,
                              onednn_engine,
                              ctx.GetPlace(),
                              x,
                              x_bd_dims,
                              trans_x,
                              y,
                              y_bd_dims,
                              trans_y,
                              out,
                              out_dims);
781 782
  }
};
783

784
template <typename T>
785
class MatMulV2GradMKLDNNKernel : public paddle::framework::OpKernel<T> {
786
 public:
787
  void Compute(const ExecutionContext &ctx) const override { RunKernel(ctx); }
788

789
 private:
790 791 792 793 794 795 796
  void CalculateGradMatrixDims(const ExecutionContext &ctx,
                               Tensor *dx_tmp,
                               Tensor *dy_tmp,
                               const std::vector<int64_t> &dx_dims,
                               const std::vector<int64_t> &dy_dims,
                               std::vector<int64_t> *dx_bd_dims,
                               std::vector<int64_t> *dy_bd_dims) const {
797 798 799
    for (size_t i = 0; i < dx_dims.size() - 2; ++i) {
      if (dx_dims[i] != dy_dims[i]) {
        if (dx_dims[i] == 1) {
800
          (*dx_bd_dims)[i] = dy_dims[i];
801
        } else {
802
          (*dy_bd_dims)[i] = dx_dims[i];
803 804 805
        }
      }
    }
806

807
    dx_tmp->Resize(phi::make_ddim((*dx_bd_dims)));
808
    dx_tmp->mutable_data<T>(ctx.GetPlace());
809
    dy_tmp->Resize(phi::make_ddim((*dy_bd_dims)));
810 811 812
    dy_tmp->mutable_data<T>(ctx.GetPlace());
  }

813
  void ReduceSumForMatmulGradOutput(
814 815
      const ExecutionContext &ctx,
      const MKLDNNDeviceContext &dev_ctx,
816
      const dnnl::engine onednn_engine,
817 818 819 820
      const Tensor *dx_tmp,
      Tensor *dx,
      const std::vector<int64_t> &dx_dims,
      const std::vector<int64_t> &squeezed_dims) const {
821
    paddle::platform::ReductionMKLDNNHandler<T> handler(
822 823 824 825 826 827 828 829
        dnnl::algorithm::reduction_sum,
        0.0f,
        0.0f,
        onednn_engine,
        ctx.GetPlace(),
        dx_tmp,
        dx,
        dx_dims);
830 831 832 833 834 835

    auto src_memory_p = handler.AcquireSrcMemory(dx_tmp);
    auto dst_memory_p = handler.AcquireDstMemory(dx);

    std::unordered_map<int, dnnl::memory> reduction_args = {
        {DNNL_ARG_SRC, *src_memory_p}, {DNNL_ARG_DST, *dst_memory_p}};
836

837
    auto &astream = MKLDNNDeviceContext::tls().get_stream();
838 839 840
    auto reduction_p = handler.AcquireForwardPrimitive();

    reduction_p->execute(astream, reduction_args);
841
    astream.wait();
842

843
    dx->set_mem_desc(dst_memory_p->get_desc().reshape(squeezed_dims));
844 845
  }

846
  std::vector<int64_t> ExtendDimsWithOnes(const std::vector<int64_t> &dims,
847 848 849 850 851 852 853
                                          int new_size) const {
    std::vector<int64_t> new_dims(new_size, 1);
    for (size_t i = 0; i < dims.size(); ++i) {
      new_dims[new_size - dims.size() + i] = dims[i];
    }

    return new_dims;
854
  }
855

856 857 858
  void RunKernel(const ExecutionContext &ctx) const {
    const auto &dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
    const auto &onednn_engine = dev_ctx.GetEngine();
859

860 861
    auto *x = ctx.Input<phi::DenseTensor>("X");
    auto *y = ctx.Input<phi::DenseTensor>("Y");
862 863 864 865 866 867 868 869 870 871

    auto x_dims = vectorize(x->dims());
    auto y_dims = vectorize(y->dims());

    bool is_broadcast = true;
    if (x_dims.size() <= 2 || y_dims.size() <= 2) {
      is_broadcast = false;
    } else if (x_dims.size() != y_dims.size()) {
      is_broadcast = true;
    } else {
872 873 874
      is_broadcast = !std::equal(x_dims.cbegin(),
                                 x_dims.cbegin() + x_dims.size() - 2,
                                 y_dims.cbegin());
875 876 877 878 879
    }

    // if no broadcasting is needed, we can simply use matmul's grad and avoid
    // using reduce_sum
    if (!is_broadcast) {
880
      matmul_v1_grad_mkldnn_kernel.Compute(ctx);
881 882 883
      return;
    }

884 885 886
    auto *dout = ctx.Input<phi::DenseTensor>(GradVarName("Out"));
    auto *dx = ctx.Output<phi::DenseTensor>(GradVarName("X"));
    auto *dy = ctx.Output<phi::DenseTensor>(GradVarName("Y"));
887

888 889 890 891
    bool trans_x = ctx.HasAttr("trans_x") ? ctx.Attr<bool>("trans_x")
                                          : ctx.Attr<bool>("transpose_X");
    bool trans_y = ctx.HasAttr("trans_y") ? ctx.Attr<bool>("trans_y")
                                          : ctx.Attr<bool>("transpose_Y");
892 893
    auto dout_dims = vectorize(dout->dims());

894 895 896 897 898 899 900 901
    size_t ndims = std::max(x->dims().size(), y->dims().size());
    ndims = std::max<size_t>(ndims, 3);

    if (x_dims.size() != ndims) {
      x_dims = ExtendDimsWithOnes(x_dims, ndims);
    } else if (y_dims.size() != ndims) {
      y_dims = ExtendDimsWithOnes(y_dims, ndims);
    }
902 903 904 905 906 907 908 909

    // in broadcasting scenario new memory is required because
    // reduce sum must be calculated upon broadcasted dims
    Tensor dx_tmp, dy_tmp;

    std::vector<int64_t> dx_bd_dims(x_dims);
    std::vector<int64_t> dy_bd_dims(y_dims);

910 911
    CalculateGradMatrixDims(
        ctx, &dx_tmp, &dy_tmp, x_dims, y_dims, &dx_bd_dims, &dy_bd_dims);
912 913

    if (trans_x && trans_y) {
914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938 939
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            y,
                            y_dims,
                            true,
                            dout,
                            dout_dims,
                            true,
                            &dx_tmp,
                            dx_bd_dims,
                            1);
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            dout,
                            dout_dims,
                            true,
                            x,
                            x_dims,
                            true,
                            &dy_tmp,
                            dy_bd_dims,
                            2);
940
    } else if (trans_x) {
941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965 966
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            y,
                            y_dims,
                            false,
                            dout,
                            dout_dims,
                            true,
                            &dx_tmp,
                            dx_bd_dims,
                            1);
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            x,
                            x_dims,
                            false,
                            dout,
                            dout_dims,
                            false,
                            &dy_tmp,
                            dy_bd_dims,
                            2);
967
    } else if (trans_y) {
968 969 970 971 972 973 974 975 976 977 978 979 980 981 982 983 984 985 986 987 988 989 990 991 992 993
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            dout,
                            dout_dims,
                            false,
                            y,
                            y_dims,
                            false,
                            &dx_tmp,
                            dx_bd_dims,
                            1);
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            dout,
                            dout_dims,
                            true,
                            x,
                            x_dims,
                            false,
                            &dy_tmp,
                            dy_bd_dims,
                            2);
994
    } else {
995 996 997 998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015 1016 1017 1018 1019 1020
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            dout,
                            dout_dims,
                            false,
                            y,
                            y_dims,
                            true,
                            &dx_tmp,
                            dx_bd_dims,
                            1);
      ExecuteMatMulV2<T, T>(ctx,
                            dev_ctx,
                            onednn_engine,
                            ctx.GetPlace(),
                            x,
                            x_dims,
                            true,
                            dout,
                            dout_dims,
                            false,
                            &dy_tmp,
                            dy_bd_dims,
                            2);
1021 1022 1023
    }

    if (x_dims != dx_bd_dims) {
1024 1025 1026 1027 1028 1029 1030
      ReduceSumForMatmulGradOutput(ctx,
                                   dev_ctx,
                                   onednn_engine,
                                   &dx_tmp,
                                   dx,
                                   x_dims,
                                   phi::vectorize(x->dims()));
1031 1032 1033 1034
    } else {
      *dx = std::move(dx_tmp);
    }
    if (y_dims != dy_bd_dims) {
1035 1036 1037 1038 1039 1040 1041
      ReduceSumForMatmulGradOutput(ctx,
                                   dev_ctx,
                                   onednn_engine,
                                   &dy_tmp,
                                   dy,
                                   y_dims,
                                   phi::vectorize(y->dims()));
1042 1043 1044 1045
    } else {
      *dy = std::move(dy_tmp);
    }

1046 1047
    dx->Resize(x->dims());
    dy->Resize(y->dims());
1048
  }
1049 1050 1051

 private:
  paddle::operators::MatMulGradMKLDNNKernel<T> matmul_v1_grad_mkldnn_kernel;
1052
};
1053
}  // anonymous namespace
1054

1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080 1081 1082 1083 1084 1085 1086 1087 1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120 1121 1122 1123 1124
namespace paddle {
namespace operators {

template <typename T>
void MatMulGradMKLDNNKernel<T>::Compute(const ExecutionContext &ctx) const {
  if (ctx.HasAttr("head_number")) {
    PADDLE_ENFORCE_EQ(
        ctx.Attr<int>("head_number"),
        1,
        platform::errors::Unimplemented(
            "oneDNN matmul doesn't support multiple heads. Expected "
            "head_number=1. But received `head_number` is %d",
            ctx.Attr<int>("head_number")));
  }
  RunKernel(ctx);
}

template <typename T>
void MatMulGradMKLDNNKernel<T>::ExecuteMatMulGrad(
    const ExecutionContext &ctx,
    const MKLDNNDeviceContext &dev_ctx,
    const dnnl::engine &engine,
    Tensor *x,
    bool trans_x,
    bool is_fold_init_dims_x,
    Tensor *y,
    bool trans_y,
    bool is_fold_init_dims_y,
    Tensor *out) const {
  // gradient is calculated in a different way when broadcasting is used
  bool need_combine = (x->dims().size() == 3 || y->dims().size() == 3) &&
                      out->dims().size() == 2;

  Tensor x_combined, y_combined;
  if (!need_combine) {
    x_combined = *x;
    y_combined = *y;
  } else {
    x_combined = is_fold_init_dims_x ? FoldOuterDims(*x)
                                     : FoldFirstAndLastDims<T>(dev_ctx, x);
    y_combined = is_fold_init_dims_y ? FoldOuterDims(*y)
                                     : FoldFirstAndLastDims<T>(dev_ctx, y);
  }

  float alpha = ctx.HasAttr("alpha") ? ctx.Attr<float>("alpha") : 1.0f;

  MatMulMKLDNNHandler<T, T, T> handler(engine,
                                       ctx.GetPlace(),
                                       &x_combined,
                                       trans_x,
                                       &y_combined,
                                       trans_y,
                                       out,
                                       alpha);

  const auto src_memory_p = handler.AcquireSrcMemory(&x_combined);
  const auto weights_memory_p = handler.AcquireWeightsMemory(&y_combined);
  const auto dst_memory_p = handler.AcquireDstMemory(out);

  auto matmul_p = handler.AcquireForwardPrimitive();

  std::unordered_map<int, dnnl::memory> matmul_args = {
      {DNNL_ARG_SRC, *src_memory_p},
      {DNNL_ARG_WEIGHTS, *weights_memory_p},
      {DNNL_ARG_DST, *dst_memory_p}};

  auto &astream = platform::MKLDNNDeviceContext::tls().get_stream();
  matmul_p->execute(astream, matmul_args);
  astream.wait();

1125 1126
  out->set_mem_desc(
      dst_memory_p->get_desc().reshape(vectorize<int64_t>(out->dims())));
1127 1128 1129 1130 1131 1132 1133 1134
}

template <typename T>
void MatMulGradMKLDNNKernel<T>::RunKernel(const ExecutionContext &ctx) const {
  const auto &dev_ctx =
      ctx.template device_context<platform::MKLDNNDeviceContext>();
  const auto &onednn_engine = dev_ctx.GetEngine();

1135 1136 1137 1138 1139
  auto x = *ctx.Input<phi::DenseTensor>("X");
  auto y = *ctx.Input<phi::DenseTensor>("Y");
  auto dout = *ctx.Input<phi::DenseTensor>(framework::GradVarName("Out"));
  auto *dx = ctx.Output<phi::DenseTensor>(framework::GradVarName("X"));
  auto *dy = ctx.Output<phi::DenseTensor>(framework::GradVarName("Y"));
1140 1141 1142 1143 1144 1145 1146 1147 1148 1149 1150 1151 1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184 1185 1186 1187 1188

  bool transpose_x = ctx.HasAttr("transpose_X") ? ctx.Attr<bool>("transpose_X")
                                                : ctx.Attr<bool>("trans_x");
  bool transpose_y = ctx.HasAttr("transpose_Y") ? ctx.Attr<bool>("transpose_Y")
                                                : ctx.Attr<bool>("trans_y");

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

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

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

  if (transpose_x && transpose_y) {
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &y, true, true, &dout, true, false, dx);
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &dout, true, true, &x, true, false, dy);
  } else if (transpose_x) {
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &y, false, false, &dout, true, false, dx);
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &x, false, false, &dout, false, true, dy);
  } else if (transpose_y) {
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &dout, false, false, &y, false, true, dx);
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &dout, true, true, &x, false, true, dy);
  } else {
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &dout, false, false, &y, true, false, dx);
    this->ExecuteMatMulGrad(
        ctx, dev_ctx, onednn_engine, &x, true, true, &dout, false, true, dy);
  }

  if (dx) {
    if (dx_dims != x.dims()) {
      dx->Resize(dx_dims);
1189
      dx->set_mem_desc(x.mem_desc());
1190 1191 1192 1193 1194
    }
  }
  if (dy) {
    if (dy_dims != y.dims()) {
      dy->Resize(dy_dims);
1195
      dy->set_mem_desc(y.mem_desc());
1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207
    }
  }
}

template class MatMulGradMKLDNNKernel<float>;
template class MatMulGradMKLDNNKernel<paddle::platform::bfloat16>;

}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;

1208 1209 1210 1211 1212 1213 1214
REGISTER_OP_KERNEL(matmul,
                   MKLDNN,
                   ::paddle::platform::CPUPlace,
                   MatMulV2MKLDNNKernel<float>,
                   MatMulV2MKLDNNKernel<paddle::platform::bfloat16>,
                   MatMulV2MKLDNNKernel<int8_t>,
                   MatMulV2MKLDNNKernel<uint8_t>);
1215 1216 1217 1218 1219 1220 1221

REGISTER_OP_KERNEL(matmul_grad,
                   MKLDNN,
                   ::paddle::platform::CPUPlace,
                   ops::MatMulGradMKLDNNKernel<float>,
                   ops::MatMulGradMKLDNNKernel<paddle::platform::bfloat16>);

1222 1223 1224
REGISTER_OP_KERNEL(matmul_v2,
                   MKLDNN,
                   ::paddle::platform::CPUPlace,
1225
                   MatMulV2MKLDNNKernel<float>,
1226 1227 1228
                   MatMulV2MKLDNNKernel<paddle::platform::bfloat16>,
                   MatMulV2MKLDNNKernel<int8_t>,
                   MatMulV2MKLDNNKernel<uint8_t>);
1229

1230 1231 1232
REGISTER_OP_KERNEL(matmul_v2_grad,
                   MKLDNN,
                   ::paddle::platform::CPUPlace,
1233 1234
                   MatMulV2GradMKLDNNKernel<float>,
                   MatMulV2GradMKLDNNKernel<paddle::platform::bfloat16>);