mul_mkldnn_op.cc 20.7 KB
Newer Older
P
Physher 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15
/* Copyright (c) 2019 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. */

#include <string>
W
wanghuancoder 已提交
16

P
Physher 已提交
17
#include "paddle/fluid/operators/mul_op.h"
18
#include "paddle/fluid/platform/mkldnn_reuse.h"
W
wanghuancoder 已提交
19

20 21 22 23
namespace pten {
class DenseTensor;
}  // namespace pten

W
wanghuancoder 已提交
24
namespace paddle {
25
namespace framework {}  // namespace framework
W
wanghuancoder 已提交
26 27 28 29
namespace platform {
class MKLDNNDeviceContext;
}  // namespace platform
}  // namespace paddle
P
Physher 已提交
30 31 32 33 34 35 36

namespace paddle {
namespace operators {

using framework::DataLayout;
using framework::DDim;
using framework::ExecutionContext;
37
using framework::LoDTensor;
P
Physher 已提交
38
using framework::Tensor;
39 40 41 42 43

using platform::MatMulV2MKLDNNHandler;
using platform::MKLDNNDeviceContext;
using platform::to_void_cast;

44 45 46 47
using dnnl::inner_product_forward;
using dnnl::memory;
using dnnl::prop_kind;
using dnnl::stream;
P
Physher 已提交
48 49 50 51

template <typename XT, typename YT, typename OT>
class MulPrimitiveFactory {
 public:
52
  explicit MulPrimitiveFactory(const dnnl::engine &engine) : engine_(engine) {}
P
Physher 已提交
53

54 55 56 57 58
  inner_product_forward CreateMulPrimitive(const Tensor *x_input,
                                           const Tensor *y_input,
                                           Tensor *output,
                                           const ExecutionContext &ctx) {
    /* check data format and reorder if need */
P
Physher 已提交
59 60 61
    int x_num_col_dims = ctx.Attr<int>("x_num_col_dims");
    int y_num_col_dims = ctx.Attr<int>("y_num_col_dims");

62 63 64 65 66 67 68 69 70 71
    // TODO(intel-minghui) : Remove the restriction that only supports Input(Y)
    // as weights
    PADDLE_ENFORCE_EQ(
        (std::is_same<YT, float>::value), true,
        platform::errors::InvalidArgument(
            "Input(Y) must be fp32 data type since only fp32 data type is "
            "supported in the current design of MKLDNN INT8."));

    auto x_matrix = UpdateDataFormat<XT>(x_input, x_num_col_dims, ctx);
    auto y_matrix = UpdateDataFormat<YT>(y_input, y_num_col_dims, ctx);
P
Physher 已提交
72 73 74 75 76 77 78 79

    auto output_dim = output->dims();
    if (output_dim.size() != 2) {
      output->Resize({x_matrix.dims()[0], y_matrix.dims()[1]});
    }

    if (mul_) {
      UpdateDataPointers(ctx, output, &x_matrix);
A
Adam 已提交
80
      Execute();
81
      return *(mul_);
P
Physher 已提交
82 83
    }

84
    auto src_desc = CreateMemDescriptor<XT>(&x_matrix, MKLDNNMemoryFormat::nc);
P
Physher 已提交
85
    x_input_ = CreateMemory<XT>(src_desc, &x_matrix);
86 87 88 89 90 91 92 93 94

    if (is_int8_) {
      const auto trans_y = TransposeInputY(&y_matrix);
      auto scale_y = ctx.Attr<std::vector<float>>("scale_y");
      y_input_ = QuantInputY(trans_y, scale_y);
    } else {
      y_input_ = TransposeInputY(&y_matrix);
    }

95
    auto dst_desc = CreateMemDescriptor<OT>(output, MKLDNNMemoryFormat::any);
P
Physher 已提交
96 97

    mul_ = CreateMulPrimitive(*x_input_, *y_input_, dst_desc, output, ctx);
A
Adam 已提交
98
    Execute();
99 100 101 102 103 104 105 106
    return *(mul_);
  }

 private:
  memory ReorderWithScale(const memory::desc &src_desc,
                          const memory::desc &dst_desc, void *src_data,
                          const std::vector<float> &scale) {
    auto mask = scale.size() > 1 ? 1 : 0;
107
    dnnl::primitive_attr attr;
108 109 110 111 112
    attr.set_output_scales(mask, scale);

    auto src_mem = memory(src_desc, engine_, src_data);
    auto dst_mem = memory(dst_desc, engine_);

113
    auto reorder_pd = dnnl::reorder::primitive_desc(src_mem, dst_mem, attr);
114

115
    auto reorder = dnnl::reorder(reorder_pd);
116

117
    auto &astream = platform::MKLDNNDeviceContext::tls().get_stream();
118
    {
C
chenjian 已提交
119 120 121
      platform::RecordEvent record_reorder(
          "int_reorder", platform::TracerEventType::UserDefined, 2,
          platform::EventRole::kUniqueOp);
122 123 124
      reorder.execute(astream, src_mem, dst_mem);
      astream.wait();
    }
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140

    return dst_mem;
  }

  memory QuantInputY(memory input_y, const std::vector<float> &scale_y) {
    const auto &dims = input_y.get_desc().data.dims;
    auto ndims = input_y.get_desc().data.ndims;
    auto y_dims = std::vector<int64_t>(dims, dims + ndims);

    auto user_y_desc = CreateMemDescriptor<YT>(y_dims, MKLDNNMemoryFormat::oi);
    auto y_desc = CreateMemDescriptor<int8_t>(y_dims, MKLDNNMemoryFormat::oi);

    return ReorderWithScale(user_y_desc, y_desc, input_y.get_data_handle(),
                            scale_y);
  }

141 142 143
  dnnl::primitive_attr CreateMulAttr(const ExecutionContext &ctx,
                                     bool force_fp32_output) {
    dnnl::primitive_attr mul_attr;
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 185 186 187 188 189

    auto scale_y_data = ctx.Attr<std::vector<float>>("scale_y");
    auto scale_x_data = ctx.Attr<float>("scale_x");
    auto scale_out_data =
        force_fp32_output ? 1.0f : ctx.Attr<float>("scale_out");

    bool is_multi_channel = scale_y_data.size() > 1;
    int count = is_multi_channel ? scale_y_data.size() : 1;
    std::vector<float> output_shift_scale(count);
    for (int i = 0; i < count; i++) {
      if (scale_y_data[i] == 0.0)
        output_shift_scale[i] = scale_out_data;
      else
        output_shift_scale[i] =
            scale_out_data / (scale_x_data * scale_y_data[i]);
    }
    int mul_mask = is_multi_channel ? 1 : 0;
    mul_attr.set_output_scales(mul_mask, output_shift_scale);

    return mul_attr;
  }

  inner_product_forward CreateMulPrimitive(const memory &x_memory,
                                           const memory &y_memory,
                                           const memory::desc &dst_desc,
                                           Tensor *output,
                                           const ExecutionContext &ctx) {
    const auto x_desc = x_memory.get_desc();
    const auto y_desc = y_memory.get_desc();
    inner_product_forward::primitive_desc mul_prim_desc;

    const auto &mul_desc = inner_product_forward::desc(
        prop_kind::forward, x_desc, y_desc, dst_desc);

    if (is_int8_) {
      bool force_fp32_output = ctx.Attr<bool>("force_fp32_output");
      auto mul_attr = CreateMulAttr(ctx, force_fp32_output);
      mul_prim_desc =
          inner_product_forward::primitive_desc(mul_desc, mul_attr, engine_);
    } else {
      mul_prim_desc = inner_product_forward::primitive_desc(mul_desc, engine_);
    }

    output_ = CreateDstMemory(mul_prim_desc, ctx, output);

    return inner_product_forward(mul_prim_desc);
P
Physher 已提交
190 191
  }

A
Adam 已提交
192
  void Execute() {
193
    auto &astream = platform::MKLDNNDeviceContext::tls().get_stream();
194 195 196
    (*mul_).execute(astream, {{DNNL_ARG_SRC, *x_input_},
                              {DNNL_ARG_WEIGHTS, *y_input_},
                              {DNNL_ARG_DST, *output_}});
A
Adam 已提交
197 198 199
    astream.wait();
  }

P
Physher 已提交
200 201 202 203 204
  template <typename T>
  Tensor UpdateDataFormat(const Tensor *data, int num_col_dims,
                          const ExecutionContext &ctx) {
    Tensor x_tmp;
    Tensor data_matrix;
205 206
    MKLDNNMemoryFormat src_fmt = data->format();
    MKLDNNMemoryFormat dst_fmt;
P
Physher 已提交
207 208 209
    auto src_mdesc = CreateMemDescriptor<T>(data, src_fmt);

    if ((data->dims().size() == 4 &&
210
         src_fmt != (dst_fmt = MKLDNNMemoryFormat::nchw)) ||
P
Physher 已提交
211
        (data->dims().size() == 5 &&
212
         src_fmt != (dst_fmt = MKLDNNMemoryFormat::ncdhw))) {
P
Physher 已提交
213 214 215 216 217 218 219
      auto dst_mdesc = CreateMemDescriptor<T>(data, dst_fmt);
      x_tmp.mutable_data<T>(ctx.GetPlace(), data->memory_size());

      Reorder(src_mdesc, dst_mdesc, to_void_cast<T>(data->data<T>()),
              to_void_cast<T>(x_tmp.data<T>()));

      x_tmp.Resize(data->dims());
A
Adam 已提交
220
      x_tmp.set_format(platform::GetMKLDNNFormat(dst_mdesc));
P
Physher 已提交
221 222 223 224 225 226 227 228 229 230 231 232 233
      data_matrix = framework::ReshapeToMatrix(x_tmp, num_col_dims);
    } else {
      data_matrix = framework::ReshapeToMatrix(*data, num_col_dims);
    }

    return data_matrix;
  }

  void UpdateDataPointers(const ExecutionContext &ctx, Tensor *out,
                          const Tensor *in) {
    x_input_->set_data_handle(to_void_cast<XT>(in->data<XT>()));
    output_->set_data_handle(out->mutable_data<OT>(ctx.GetPlace()));

A
Adam 已提交
234
    if (out->format() == MKLDNNMemoryFormat::undef) {
A
Adam 已提交
235
      auto output_format = platform::GetMKLDNNFormat(*output_);
236
      out->set_format((MKLDNNMemoryFormat)output_format);
P
Physher 已提交
237 238 239 240 241
    }
  }

  template <typename T>
  memory::desc CreateMemDescriptor(
242
      const Tensor *tensor, MKLDNNMemoryFormat format,
P
Physher 已提交
243
      memory::data_type type = platform::MKLDNNGetDataType<T>()) {
A
Adam 已提交
244
    auto dims = framework::vectorize<int64_t>(tensor->dims());
P
Physher 已提交
245 246 247 248 249
    return platform::MKLDNNMemDesc(dims, type, format);
  }

  template <typename T>
  memory::desc CreateMemDescriptor(
A
Adam 已提交
250
      const std::vector<int64_t> &dims, MKLDNNMemoryFormat format,
P
Physher 已提交
251 252 253 254 255 256
      memory::data_type type = platform::MKLDNNGetDataType<T>()) {
    return platform::MKLDNNMemDesc(dims, type, format);
  }

  template <typename T>
  memory CreateMemory(const memory::desc &desc, const Tensor *tensor) {
A
Adam 已提交
257
    return memory(desc, engine_, to_void_cast<T>(tensor->data<T>()));
P
Physher 已提交
258 259 260 261 262
  }

  memory CreateDstMemory(
      const inner_product_forward::primitive_desc &mul_prim_desc,
      const ExecutionContext &ctx, Tensor *output) {
A
Adam 已提交
263 264
    auto dst_desc = mul_prim_desc.dst_desc();
    auto buffer_size = dst_desc.get_size();
P
Physher 已提交
265 266

    OT *output_data = output->mutable_data<OT>(ctx.GetPlace(), buffer_size);
A
Adam 已提交
267 268
    output->set_format(paddle::platform::GetMKLDNNFormat(dst_desc));
    return memory(dst_desc, engine_, to_void_cast<OT>(output_data));
P
Physher 已提交
269 270 271 272
  }

  memory Reorder(const memory::desc &src_desc, const memory::desc &dst_desc,
                 void *src_data, void *dst_data = NULL) {
A
Adam 已提交
273 274 275
    auto src_mem = memory(src_desc, engine_, src_data);
    auto dst_mem = dst_data ? memory(dst_desc, engine_, dst_data)
                            : memory(dst_desc, engine_);
P
Physher 已提交
276

277
    auto reorder = dnnl::reorder(src_mem, dst_mem);
A
Adam 已提交
278

279
    auto &astream = platform::MKLDNNDeviceContext::tls().get_stream();
280
    {
C
chenjian 已提交
281 282 283
      platform::RecordEvent record_reorder(
          "int_reorder", platform::TracerEventType::UserDefined, 2,
          platform::EventRole::kUniqueOp);
284 285 286
      reorder.execute(astream, src_mem, dst_mem);
      astream.wait();
    }
P
Physher 已提交
287 288 289 290 291

    return dst_mem;
  }

  memory TransposeInputY(const Tensor *input_y) {
A
Adam 已提交
292
    auto dims = framework::vectorize<int64_t>(input_y->dims());
P
Physher 已提交
293
    std::swap(dims[0], dims[1]);  // Correct output dimensions
294 295
    auto src_desc = CreateMemDescriptor<YT>(dims, MKLDNNMemoryFormat::io);
    auto dst_desc = CreateMemDescriptor<YT>(dims, MKLDNNMemoryFormat::oi);
P
Physher 已提交
296 297 298
    return Reorder(src_desc, dst_desc, to_void_cast<YT>(input_y->data<YT>()));
  }

299
  const dnnl::engine &engine_;
300 301 302 303
  paddle::optional<memory> x_input_;
  paddle::optional<memory> y_input_;
  paddle::optional<memory> output_;
  paddle::optional<inner_product_forward> mul_;
304 305
  static constexpr bool is_int8_ =
      std::is_same<XT, int8_t>::value || std::is_same<XT, uint8_t>::value;
P
Physher 已提交
306 307 308 309 310 311 312
};

/* OT: output data type */
template <typename XT, typename YT, typename OT>
std::shared_ptr<MulPrimitiveFactory<XT, YT, OT>> GetPrimitiveFactory(
    const MKLDNNDeviceContext &dev_ctx, const ExecutionContext &ctx,
    const Tensor *input_x, const Tensor *input_y,
313
    const dnnl::engine &mkldnn_engine) {
314
  std::string key = platform::CreateKey(
315 316 317 318
      dev_ctx, framework::TransToProtoVarType(input_x->dtype()),
      framework::vectorize(input_x->dims()),
      framework::TransToProtoVarType(input_y->dtype()),
      framework::vectorize(input_y->dims()), ctx.OutputName("Out"));
319
  key = platform::ExtendKeyWithThreadInfoIfNeeded(dev_ctx, key);
P
Physher 已提交
320 321 322 323 324 325

  auto prim_creator = std::static_pointer_cast<MulPrimitiveFactory<XT, YT, OT>>(
      dev_ctx.GetBlob(key));

  if (prim_creator == nullptr) {
    prim_creator =
326
        std::make_shared<MulPrimitiveFactory<XT, YT, OT>>(mkldnn_engine);
P
Physher 已提交
327 328 329 330 331 332 333 334 335 336 337
    dev_ctx.SetBlob(key, prim_creator);
  }

  return prim_creator;
}

template <typename XT, typename YT>
inner_product_forward GetMulPrimitive(const MKLDNNDeviceContext &dev_ctx,
                                      const ExecutionContext &ctx,
                                      const Tensor *input_x,
                                      const Tensor *input_y, Tensor *output,
338
                                      const dnnl::engine &mkldnn_engine) {
339
  constexpr bool is_int8 =
P
Physher 已提交
340 341 342
      std::is_same<XT, int8_t>::value || std::is_same<XT, uint8_t>::value;
  bool force_fp32_output = ctx.Attr<bool>("force_fp32_output");

343
  if (is_int8 && !force_fp32_output) {
P
Physher 已提交
344
    return GetPrimitiveFactory<XT, YT, int8_t>(dev_ctx, ctx, input_x, input_y,
345
                                               mkldnn_engine)
P
Physher 已提交
346 347 348 349
        ->CreateMulPrimitive(input_x, input_y, output, ctx);

  } else {
    return GetPrimitiveFactory<XT, YT, float>(dev_ctx, ctx, input_x, input_y,
350
                                              mkldnn_engine)
P
Physher 已提交
351 352 353 354 355 356
        ->CreateMulPrimitive(input_x, input_y, output, ctx);
  }
}

/* XT: input x data type, YT: input y data type */
template <typename XT, typename YT>
357
class MulMKLDNNINT8Kernel : public framework::OpKernel<XT> {
P
Physher 已提交
358 359
 public:
  void Compute(const ExecutionContext &ctx) const override {
360 361 362
    PADDLE_ENFORCE_EQ(platform::is_cpu_place(ctx.GetPlace()), true,
                      paddle::platform::errors::PreconditionNotMet(
                          "Operator DNNL Mul must use CPUPlace"));
363
    platform::MKLDNNDeviceContext::tls().log_lib_version();
P
Physher 已提交
364
    auto &dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
365
    auto &mkldnn_engine = dev_ctx.GetEngine();
P
Physher 已提交
366 367 368 369 370 371 372 373 374 375 376 377

    const Tensor *x = ctx.Input<Tensor>("X");
    const Tensor *y = ctx.Input<Tensor>("Y");
    Tensor *out = ctx.Output<Tensor>("Out");
    auto out_dims = out->dims();

    auto mul = GetMulPrimitive<XT, YT>(dev_ctx, ctx, x, y, out, mkldnn_engine);

    if (out_dims.size() != 2) {
      out->Resize(out_dims);
    }
    out->set_layout(DataLayout::kMKLDNN);
A
Adam 已提交
378 379
    out->set_format(platform::MKLDNNFormatForSize(out_dims.size(),
                                                  MKLDNNMemoryFormat::nchw));
P
Physher 已提交
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 476 477 478 479 480 481 482 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
template <typename XT, typename YT>
class MulMKLDNNKernel : public framework::OpKernel<XT> {
 public:
  void Compute(const ExecutionContext &ctx) const override { RunKernel(ctx); }

 protected:
  void ExecuteMatMul(const ExecutionContext &ctx,
                     const MKLDNNDeviceContext &dev_ctx,
                     const dnnl::engine &onednn_engine,
                     const platform::Place &cpu_place, const Tensor *x,
                     const std::vector<int64_t> &x_dims, bool trans_x,
                     const Tensor *y, const std::vector<int64_t> &y_dims,
                     bool trans_y, Tensor *out) const {
    static const std::vector<int64_t> vec_placeholder;
    MatMulV2MKLDNNHandler<XT> handler(onednn_engine, ctx.GetPlace(), x_dims,
                                      trans_x, y_dims, trans_y, false,
                                      vec_placeholder, vec_placeholder);

    const auto src_memory_p = handler.AcquireSrcMemory(x);
    const auto weights_memory_p = handler.AcquireWeightsMemory(y);
    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 = MKLDNNDeviceContext::tls().get_stream();
    matmul_p->execute(astream, matmul_args);
    astream.wait();

    out->set_layout(framework::DataLayout::kMKLDNN);
    // plain output formats are enforced inside handler
    out->set_format(platform::MKLDNNFormatForSize(
        out->dims().size(), dnnl::memory::format_tag::nchw));
  }

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

    const auto *x = ctx.Input<Tensor>("X");
    const auto *y = ctx.Input<Tensor>("Y");
    auto *out = ctx.Output<Tensor>("Out");

    int x_num_col_dims = ctx.Attr<int>("x_num_col_dims");
    int y_num_col_dims = ctx.Attr<int>("y_num_col_dims");

    const Tensor x_matrix = x->dims().size() > 2
                                ? framework::ReshapeToMatrix(*x, x_num_col_dims)
                                : *x;
    const Tensor y_matrix = y->dims().size() > 2
                                ? framework::ReshapeToMatrix(*y, y_num_col_dims)
                                : *y;

    // adding mb dim because MatMulV2 handler needs it
    std::vector<int64_t> y_dims(3, 1);
    std::vector<int64_t> x_dims(3, 1);

    y_dims[1] = y_matrix.dims()[0];
    y_dims[2] = y_matrix.dims()[1];

    x_dims[1] = x_matrix.dims()[0];
    x_dims[2] = x_matrix.dims()[1];

    ExecuteMatMul(ctx, dev_ctx, onednn_engine, ctx.GetPlace(), &x_matrix,
                  x_dims, false, &y_matrix, y_dims, false, out);
  }
};

template <typename XT, typename YT>
class MulGradMKLDNNKernel : public MulMKLDNNKernel<XT, YT> {
 public:
  void Compute(const ExecutionContext &ctx) const override { RunKernel(ctx); }

 private:
  template <typename OT = XT>
  void RunKernel(const ExecutionContext &ctx) const {
    const auto &dev_ctx = ctx.template device_context<MKLDNNDeviceContext>();
    const auto &onednn_engine = dev_ctx.GetEngine();

    const auto *x = ctx.Input<LoDTensor>("X");
    const auto *y = ctx.Input<LoDTensor>("Y");
    const auto *dout = ctx.Input<Tensor>(framework::GradVarName("Out"));

    auto *dx = ctx.Output<LoDTensor>(framework::GradVarName("X"));
    auto *dy = ctx.Output<LoDTensor>(framework::GradVarName("Y"));

    int x_num_col_dims = ctx.Attr<int>("x_num_col_dims");
    int y_num_col_dims = ctx.Attr<int>("y_num_col_dims");

    const Tensor x_matrix = x->dims().size() > 2
                                ? framework::ReshapeToMatrix(*x, x_num_col_dims)
                                : static_cast<const Tensor &>(*x);
    const Tensor y_matrix = y->dims().size() > 2
                                ? framework::ReshapeToMatrix(*y, y_num_col_dims)
                                : static_cast<const Tensor &>(*y);

    Tensor dout_matrix = *dout;
    dout_matrix.Resize(
        {framework::flatten_to_2d(x->dims(), x_num_col_dims)[0],
         framework::flatten_to_2d(y->dims(), y_num_col_dims)[1]});

    // adding mb dim because MatMulV2 handler needs it
    std::vector<int64_t> x_dims(3, 1);
    std::vector<int64_t> y_dims(3, 1);
    std::vector<int64_t> dout_dims(3, 1);

    x_dims[1] = x_matrix.dims()[0];
    x_dims[2] = x_matrix.dims()[1];

    y_dims[1] = y_matrix.dims()[0];
    y_dims[2] = y_matrix.dims()[1];

    dout_dims[1] = dout_matrix.dims()[0];
    dout_dims[2] = dout_matrix.dims()[1];

    if (dx != nullptr) {
      dx->set_lod(x->lod());
      this->ExecuteMatMul(ctx, dev_ctx, onednn_engine, ctx.GetPlace(),
                          &dout_matrix, dout_dims, false, &y_matrix, y_dims,
                          true, static_cast<Tensor *>(dx));
    }
    if (dy != nullptr) {
      dy->set_lod(y->lod());
      this->ExecuteMatMul(ctx, dev_ctx, onednn_engine, ctx.GetPlace(),
                          &x_matrix, x_dims, true, &dout_matrix, dout_dims,
                          false, static_cast<Tensor *>(dy));
    }
  }
};

P
Physher 已提交
518 519 520 521 522 523
}  // namespace operators
}  // namespace paddle

namespace ops = paddle::operators;
REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(mul, MKLDNN, ::paddle::platform::CPUPlace,
                                    U8, ops::kMULMKLDNNINT8,
524
                                    ops::MulMKLDNNINT8Kernel<uint8_t, float>);
P
Physher 已提交
525 526 527

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(mul, MKLDNN, ::paddle::platform::CPUPlace,
                                    S8, ops::kMULMKLDNNINT8,
528 529 530 531 532 533 534 535 536 537
                                    ops::MulMKLDNNINT8Kernel<int8_t, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(mul, MKLDNN, ::paddle::platform::CPUPlace,
                                    FP32, ops::kMULMKLDNNFP32,
                                    ops::MulMKLDNNKernel<float, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(
    mul, MKLDNN, ::paddle::platform::CPUPlace, BF16, ops::kMULMKLDNNFP32,
    ops::MulMKLDNNKernel<paddle::platform::bfloat16,
                         paddle::platform::bfloat16>);
P
Physher 已提交
538 539

REGISTER_OP_KERNEL(mul, MKLDNN, ::paddle::platform::CPUPlace,
540 541 542 543 544 545 546 547 548 549 550 551 552 553 554
                   ops::MulMKLDNNINT8Kernel<uint8_t, float>,
                   ops::MulMKLDNNKernel<paddle::platform::bfloat16,
                                        paddle::platform::bfloat16>,
                   ops::MulMKLDNNKernel<float, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(mul_grad, MKLDNN,
                                    ::paddle::platform::CPUPlace, FP32,
                                    ops::kMULMKLDNNFP32,
                                    ops::MulGradMKLDNNKernel<float, float>);

REGISTER_OP_KERNEL_WITH_CUSTOM_TYPE(
    mul_grad, MKLDNN, ::paddle::platform::CPUPlace, BF16, ops::kMULMKLDNNFP32,
    ops::MulGradMKLDNNKernel<paddle::platform::bfloat16,
                             paddle::platform::bfloat16>,
    ops::MulGradMKLDNNKernel<float, float>);