[PHI] Move oneDNN helper classes to new location (#45626)

* gaussian random

* mkldnn to onednn renaming

* fix merge conflicts

* remove fluid code

* onednn renaming

* Move classes from mkldnn_reuse.h to onednn_reuse.h

* Move more functions from mkldnn_helper.h to onednn_helpper.h

* Change MKLDNN to OneDNN in VLOG message
Co-authored-by: NSilv3S <slawomir.siwek@intel.com>

paddle/fluid/framework/new_executor/data_transfer.cc

@@ -212,10 +212,10 @@ std::shared_ptr<OperatorBase> TransferLayout(const std::string& var_name,
     out_layout = framework::DataLayout::kNCHW;
   }
 
-  if (in_layout == framework::DataLayout::MKLDNN &&
-      out_layout != framework::DataLayout::MKLDNN) {
+  if (in_layout == framework::DataLayout::ONEDNN &&
+      out_layout != framework::DataLayout::ONEDNN) {
     auto target_layout = phi::OneDNNContext::tls().get_cur_paddle_data_layout();
-    VLOG(4) << "TransDataLayoutFromMKLDNN: " << in_layout << "->"
+    VLOG(4) << "TransDataLayoutFromOneDNN: " << in_layout << "->"
             << target_layout;
 
     if (out_layout == DataLayout::kNCHW &&

paddle/fluid/framework/phi_utils_test.cc

@@ -75,7 +75,7 @@ TEST(PhiUtils, TransOpKernelTypeToPhiKernelKey) {
   auto kernel_key_mkldnn =
       paddle::framework::TransOpKernelTypeToPhiKernelKey(op_kernel_type_mkldnn);
   ASSERT_EQ(kernel_key_mkldnn.dtype(), phi::DataType::FLOAT32);
-  ASSERT_EQ(kernel_key_mkldnn.layout(), phi::DataLayout::MKLDNN);
+  ASSERT_EQ(kernel_key_mkldnn.layout(), phi::DataLayout::ONEDNN);
   ASSERT_EQ(kernel_key_mkldnn.backend(), phi::Backend::ONEDNN);
 #endif
 

paddle/phi/api/lib/kernel_dispatch.cc

@@ -56,7 +56,7 @@ BackendSet GetTensorBackendSet(const phi::TensorBase& t) {
   if (HasAllocation(t) && t.place().GetType() != AllocationType::UNDEFINED) {
     BackendSet backend_set(phi::TransToPhiBackend(t.place()));
     switch (t.layout()) {
-      case DataLayout::MKLDNN:
+      case DataLayout::ONEDNN:
         backend_set = backend_set | BackendSet(Backend::ONEDNN);
         break;
       default:

paddle/phi/backends/onednn/onednn_helper.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include <thread>
 #include "dnnl.hpp"  // NOLINT
 #include "glog/logging.h"
 
@@ -94,6 +95,106 @@ inline dnnl::memory::format_tag GetPlainOneDNNFormat(int tensor_rank) {
   }
 }
 
+template <typename Type>
+dnnl::memory::data_type oneDNNGetDataType() {
+  return dnnl::memory::data_type::undef;
+}
+
+template <>
+inline dnnl::memory::data_type oneDNNGetDataType<float>() {
+  return dnnl::memory::data_type::f32;
+}
+template <>
+inline dnnl::memory::data_type oneDNNGetDataType<int32_t>() {
+  return dnnl::memory::data_type::s32;
+}
+template <>
+inline dnnl::memory::data_type oneDNNGetDataType<int8_t>() {
+  return dnnl::memory::data_type::s8;
+}
+template <>
+inline dnnl::memory::data_type oneDNNGetDataType<uint8_t>() {
+  return dnnl::memory::data_type::u8;
+}
+
+template <>
+inline dnnl::memory::data_type oneDNNGetDataType<dtype::bfloat16>() {
+  return dnnl::memory::data_type::bf16;
+}
+
+inline std::vector<std::vector<int64_t>> ToOneDNNPadding(
+    const std::vector<int64_t>& paddings) {
+  if (paddings.size() == 6) {
+    int padding_front = paddings[0];
+    int padding_back = paddings[1];
+    int padding_top = paddings[2];
+    int padding_bottom = paddings[3];
+    int padding_left = paddings[4];
+    int padding_right = paddings[5];
+
+    return {{padding_front, padding_top, padding_left},
+            {padding_back, padding_bottom, padding_right}};
+  } else {
+    int padding_top = paddings[0];
+    int padding_bottom = paddings[1];
+    int padding_left = paddings[2];
+    int padding_right = paddings[3];
+
+    return {{padding_top, padding_left}, {padding_bottom, padding_right}};
+  }
+}
+
+template <typename T>
+inline void AppendKey(std::string* key, const T& num) {
+  key->append(std::to_string(num));
+}
+
+template <>
+inline void AppendKey(std::string* key,
+                      const dnnl::memory::format_tag& format) {
+  key->append(std::to_string(static_cast<int>(format)));
+}
+
+template <>
+inline void AppendKey(std::string* key,
+                      const dnnl::memory::data_type& data_type) {
+  key->append(std::to_string(static_cast<int>(data_type)));
+}
+
+template <>
+inline void AppendKey(std::string* key, const dnnl::algorithm& algorithm) {
+  key->append(std::to_string(static_cast<int>(algorithm)));
+}
+
+template <>
+inline void AppendKey(std::string* key,
+                      const dnnl::normalization_flags& flags) {
+  key->append(std::to_string(static_cast<int>(flags)));
+}
+
+inline void AppendKey(std::string* key, const std::string& str) {
+  key->append(str);
+}
+
+inline void AppendKey(std::string* key, const char* str) { key->append(str); }
+
+template <typename T>
+inline void AppendKey(std::string* key, const std::vector<T>& dims) {
+  for (size_t i = 0; i < dims.size(); i++) {
+    AppendKey(key, std::to_string(dims[i]));
+  }
+}
+
+template <typename... ArgTypes>
+inline std::string CreateKey(const OneDNNContext& dev_ctx, ArgTypes&&... args) {
+  std::string key;
+  key.reserve(64);
+  using expand_type = int[];
+  expand_type{0, (AppendKey(&key, std::forward<ArgTypes>(args)), 0)...};
+  key += OneDNNContext::tls().get_key_suffix();
+  return key;
+}
+
 inline void MatchShapeToLayout(DenseTensor* tensor_in,
                                DataLayout from,
                                DataLayout to) {
@@ -117,28 +218,28 @@ inline void MatchShapeToLayout(DenseTensor* tensor_in,
   // at last nhwC, so for dim==2 these layouts are the same and nothing should
   // be done. Similarly for dim==1 when you have just one possible combination.
   if (tensor_in->dims().size() < 3) {
-    VLOG(3) << "Keeping MKLDNN/NHWC/NDHWC output_shape"
+    VLOG(3) << "Keeping ONEDNN/NHWC/NDHWC output_shape"
             << print_dims(phi::vectorize<int>(tensor_in->dims()));
     return;
   }
 
   switch (from) {
-    case DataLayout::MKLDNN:
+    case DataLayout::ONEDNN:
       if ((to == DataLayout::NHWC) || (to == DataLayout::NDHWC)) {
         auto dims = phi::vectorize<int>(tensor_in->dims());
         std::rotate(dims.begin() + 1, dims.begin() + 2, dims.end());
         tensor_in->Resize(phi::make_ddim(dims));
-        VLOG(3) << "Rotating Shape from: MKLDNN to: NHWC/NDHWC output_shape"
+        VLOG(3) << "Rotating Shape from: ONEDNN to: NHWC/NDHWC output_shape"
                 << print_dims(dims);
       }
       break;
     case DataLayout::NHWC:
     case DataLayout::NDHWC:
-      if (to == DataLayout::MKLDNN) {
+      if (to == DataLayout::ONEDNN) {
         auto dims = phi::vectorize<int>(tensor_in->dims());
         std::rotate(dims.begin() + 1, dims.end() - 1, dims.end());
         tensor_in->Resize(phi::make_ddim(dims));
-        VLOG(3) << "Rotating Shape from: NHWC/NDHWC to: MKLDNN output_shape"
+        VLOG(3) << "Rotating Shape from: NHWC/NDHWC to: ONEDNN output_shape"
                 << print_dims(dims);
       }
       break;
@@ -158,5 +259,22 @@ inline dnnl::memory::desc OneDNNMemDesc(const std::vector<int64_t>& dims,
   return dnnl::memory::desc({dims}, data_type, format);
 }
 
+inline std::string ThreadIDasStr(void) {
+  return std::to_string(
+      std::hash<std::thread::id>()(std::this_thread::get_id()));
+}
+
+inline std::string ExtendKeyWithThreadInfoIfNeeded(const OneDNNContext& dev_ctx,
+                                                   const std::string& key) {
+  return (OneDNNContext::tls().is_tid_used_in_key() == true)
+             ? key + "-t:" + ThreadIDasStr()
+             : key;
+}
+
+template <typename T>
+bool constexpr is_int8() {
+  return std::is_same<T, int8_t>::value || std::is_same<T, uint8_t>::value;
+}
+
 }  // namespace funcs
 }  // namespace phi

paddle/phi/common/backend.h

@@ -32,7 +32,7 @@ namespace experimental {
  * more specific, we need to distinguish the calculation method.
  *
  * Such as the kernel for CPU device, it can be a native CPU kernel,
- * or a kernel implemented by MKLDNN library.
+ * or a kernel implemented by oneDNN library.
  *
  * Note(chenweihang): HIP is not needed now, we can added it if needed
  * in the future

paddle/phi/common/layout.h

@@ -40,7 +40,7 @@ enum class DataLayout {
   NCHW,
   NCDHW,
   NDHWC,
-  MKLDNN,
+  ONEDNN,
   SPARSE_COO,
   SPARSE_CSR,
   PSTRING_UNION,
@@ -62,7 +62,7 @@ enum class DataLayout {
   kAnyLayout = ANY,
   kNHWC = NHWC,
   kNCHW = NCHW,
-  kMKLDNN = MKLDNN,  // all layouts supported by MKLDNN internally
+  kMKLDNN = ONEDNN,  // all layouts supported by ONEDNN internally
   kNDHWC = NDHWC,
   kNCDHW = NCDHW,
 };

paddle/phi/kernels/onednn/log_softmax_kernel.cc

@@ -14,8 +14,8 @@
 
 #include "paddle/phi/kernels/log_softmax_kernel.h"
 
-#include "paddle/fluid/platform/mkldnn_reuse.h"
 #include "paddle/phi/backends/onednn/onednn_context.h"
+#include "paddle/phi/backends/onednn/onednn_reuse.h"
 #include "paddle/phi/common/bfloat16.h"
 #include "paddle/phi/common/place.h"
 #include "paddle/phi/core/kernel_registry.h"
@@ -23,16 +23,15 @@
 namespace phi {
 
 template <typename T>
-class LogSoftmaxMKLDNNHandler
-    : public paddle::platform::
-          MKLDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward> {
+class LogSoftmaxOneDNNHandler
+    : public funcs::OneDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward> {
  public:
-  LogSoftmaxMKLDNNHandler(const dnnl::engine mkldnn_engine,
+  LogSoftmaxOneDNNHandler(const dnnl::engine onednn_engine,
                           Place cpu_place,
                           const DenseTensor& x,
                           const int axis)
-      : paddle::platform::MKLDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward>(
-            mkldnn_engine, cpu_place) {
+      : funcs::OneDNNHandlerNoCachingT<T, dnnl::logsoftmax_forward>(
+            onednn_engine, cpu_place) {
     this->AcquireForwardPrimitiveDescriptor(
         dnnl::prop_kind::forward_inference, x.mem_desc(), axis);
   }
@@ -43,11 +42,11 @@ void LogSoftmaxKernel(const Context& dev_ctx,
                       const DenseTensor& x,
                       int axis,
                       DenseTensor* out) {
-  const auto& mkldnn_engine = dev_ctx.GetEngine();
+  const auto& onednn_engine = dev_ctx.GetEngine();
   axis = axis >= 0 ? axis : x.dims().size() + axis;
 
-  LogSoftmaxMKLDNNHandler<T> handler(
-      mkldnn_engine, dev_ctx.GetPlace(), x, axis);
+  LogSoftmaxOneDNNHandler<T> handler(
+      onednn_engine, dev_ctx.GetPlace(), x, axis);
 
   auto src_memory_p = handler.AcquireSrcMemory(&x);
   auto dst_memory_p = handler.AcquireDstMemory(out);

paddle/phi/kernels/transfer_layout_kernel.cc

@@ -97,7 +97,7 @@ void TransferLayoutMKLDNN(const Context& dev_ctx,
 
   // NOTE(zhiqiu): to handle the special case in ApplyDataTransform() in
   // data_transfer.cc
-  if (!x.IsInitialized() && src_layout == DataLayout::MKLDNN &&
+  if (!x.IsInitialized() && src_layout == DataLayout::ONEDNN &&
       dst_layout == DataLayout::NHWC) {
     VLOG(4) << src_layout << "->" << dst_layout << " " << x.layout();
     out->Resize(x.dims());
@@ -106,7 +106,7 @@ void TransferLayoutMKLDNN(const Context& dev_ctx,
     return;
   }
 
-  if (src_layout != DataLayout::MKLDNN && dst_layout == DataLayout::MKLDNN) {
+  if (src_layout != DataLayout::ONEDNN && dst_layout == DataLayout::ONEDNN) {
     // Case1 - transform from Non-MKLDNN OPKernel to MKLDNN OPKernel
     // Just set layout/format. No real transform occur
     auto out_format = funcs::OneDNNFormatForSize(
@@ -121,16 +121,16 @@ void TransferLayoutMKLDNN(const Context& dev_ctx,
       OneDNNContext::tls().set_cur_paddle_data_layout(src_layout);
     }
 
-    out->set_layout(DataLayout::MKLDNN);
+    out->set_layout(DataLayout::ONEDNN);
     out->set_format(out_format);
-  } else if (src_layout == DataLayout::MKLDNN &&
-             dst_layout != DataLayout::MKLDNN) {
+  } else if (src_layout == DataLayout::ONEDNN &&
+             dst_layout != DataLayout::ONEDNN) {
     // Case2 - transfrom from MKLDNN OPKernel to Non-MKLDNN OPKernel
     // Do transform via MKLDNN lib
     funcs::innerTransDataLayoutFromOneDNN(
         src_layout, dst_layout, x, out, dev_ctx.GetPlace());
-  } else if (src_layout == DataLayout::MKLDNN &&
-             dst_layout == DataLayout::MKLDNN) {
+  } else if (src_layout == DataLayout::ONEDNN &&
+             dst_layout == DataLayout::ONEDNN) {
     PADDLE_ENFORCE_NE(
         src_layout,
         dst_layout,

paddle/phi/tests/common/test_data_layout.cc

@@ -37,7 +37,7 @@ TEST(DataLayout, OStream) {
   oss << phi::DataLayout::NCHW;
   EXPECT_EQ(oss.str(), "NCHW");
   oss.str("");
-  oss << phi::DataLayout::MKLDNN;
+  oss << phi::DataLayout::ONEDNN;
   EXPECT_EQ(oss.str(), "MKLDNN");
   oss.str("");
   try {

paddle/phi/tests/kernels/test_transfer_layout_dev_api.cc

@@ -40,7 +40,7 @@ TEST(DEV_API, transfer_layout) {
   DenseTensor x;
   MetaTensor meta_x(&x);
   meta_x.set_dtype(DataType::FLOAT32);
-  meta_x.set_layout(DataLayout::MKLDNN);
+  meta_x.set_layout(DataLayout::ONEDNN);
   meta_x.set_dims(make_ddim({n, c, h, w}));
 
   DenseTensor out;