PaddlePaddle
diff --git a/‎.gitignore
Copy file name to clipboard
+1-3Lines changed: 1 addition & 3 deletions b/‎.gitignore
Copy file name to clipboard
+1-3Lines changed: 1 addition & 3 deletions
diff --git a/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+14-13Lines changed: 14 additions & 13 deletions b/‎paddle/fluid/framework/new_executor/pir_adaptor/pir_adaptor_util.cc
Copy file name to clipboard
+14-13Lines changed: 14 additions & 13 deletions
diff --git a/‎paddle/fluid/ir_adaptor/translator/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/ir_adaptor/translator/CMakeLists.txt
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/fluid/ir_adaptor/translator/op_translator.cc
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/fluid/ir_adaptor/translator/utils.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/fluid/ir_adaptor/translator/utils.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/fluid/ir_adaptor/translator/utils.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/fluid/ir_adaptor/translator/utils.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/fluid/operators/mkldnn/quantize_mkldnn_op.cc
Copy file name to clipboard
+3-5Lines changed: 3 additions & 5 deletions b/‎paddle/fluid/operators/mkldnn/quantize_mkldnn_op.cc
Copy file name to clipboard
+3-5Lines changed: 3 additions & 5 deletions
diff --git a/‎paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc
Copy file name to clipboard
+8-7Lines changed: 8 additions & 7 deletions b/‎paddle/fluid/operators/mkldnn/requantize_mkldnn_op.cc
Copy file name to clipboard
+8-7Lines changed: 8 additions & 7 deletions
diff --git a/‎paddle/fluid/pir/dialect/CMakeLists.txt
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions b/‎paddle/fluid/pir/dialect/CMakeLists.txt
Copy file name to clipboard
+3-2Lines changed: 3 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+55-18Lines changed: 55 additions & 18 deletions b/‎paddle/fluid/pir/dialect/op_generator/op_gen.py
Copy file name to clipboard
+55-18Lines changed: 55 additions & 18 deletions
@@ -109,9 +109,7 @@ paddle/fluid/pir/dialect/operator/ir/op_decomp.cc
 paddle/fluid/pir/dialect/operator/ir/pd_op_vjp.cc
 paddle/fluid/pir/dialect/operator/ir/pd_op.*
 paddle/fluid/pir/dialect/operator/ir/onednn_op.*
-paddle/fluid/pir/dialect/operator/ir/pd_onednn_op.*
-paddle/fluid/pir/dialect/operator/ir/pd_onednn_op_info.*
-paddle/fluid/pir/dialect/operator/ir/pd_onednn_op.*
+paddle/fluid/pir/dialect/operator/ir/onednn_op_info.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_bwd.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_fused.*
 paddle/fluid/pir/dialect/operator/ir/pd_op_fused_bwd.*
 
@@ -879,19 +879,20 @@ std::shared_ptr<OperatorBase> BuildOperatorBase(
   auto attr_name_list = op_yaml_info.AttrParams(true);
   for (auto& name : attr_name_list) {
     auto& val = op_attr_map.at(name);
+    auto legacy_arg_name = op_normalizer.GetLegacyAttrName(fluid_op_name, name);
 
     if (val.isa<pir::StrAttribute>()) {
-      attr_map[name] = val.dyn_cast<pir::StrAttribute>().AsString();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::StrAttribute>().AsString();
     } else if (val.isa<pir::Int32Attribute>()) {
-      attr_map[name] = val.dyn_cast<pir::Int32Attribute>().data();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::Int32Attribute>().data();
     } else if (val.isa<pir::BoolAttribute>()) {
-      attr_map[name] = val.dyn_cast<pir::BoolAttribute>().data();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::BoolAttribute>().data();
     } else if (val.isa<pir::FloatAttribute>()) {
-      attr_map[name] = val.dyn_cast<pir::FloatAttribute>().data();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::FloatAttribute>().data();
     } else if (val.isa<pir::DoubleAttribute>()) {
-      attr_map[name] = val.dyn_cast<pir::DoubleAttribute>().data();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::DoubleAttribute>().data();
     } else if (val.isa<pir::Int64Attribute>()) {
-      attr_map[name] = val.dyn_cast<pir::Int64Attribute>().data();
+      attr_map[legacy_arg_name] = val.dyn_cast<pir::Int64Attribute>().data();
     } else if (val.isa<pir::ArrayAttribute>()) {
       auto array_list = val.dyn_cast<pir::ArrayAttribute>().AsVector();
       PADDLE_ENFORCE(
@@ -902,49 +903,49 @@ std::shared_ptr<OperatorBase> BuildOperatorBase(
         for (auto attribute : array_list) {
           vec_int.push_back(attribute.dyn_cast<pir::Int32Attribute>().data());
         }
-        attr_map[name] = vec_int;
+        attr_map[legacy_arg_name] = vec_int;
       } else if (array_list[0].isa<pir::Int64Attribute>()) {
         std::vector<int> vec_int64;
         for (auto attribute : array_list) {
           vec_int64.push_back(
               attribute.dyn_cast<pir::Int64Attribute>().data());  // NOLINT
         }
-        attr_map[name] = vec_int64;
+        attr_map[legacy_arg_name] = vec_int64;
       } else if (array_list[0].isa<pir::BoolAttribute>()) {
         std::vector<int> vec_bool;
         for (auto attribute : array_list) {
           vec_bool.push_back(attribute.dyn_cast<pir::BoolAttribute>().data());
         }
-        attr_map[name] = vec_bool;
+        attr_map[legacy_arg_name] = vec_bool;
       } else if (array_list[0].isa<pir::FloatAttribute>()) {
         std::vector<int> vec_float;
         for (auto attribute : array_list) {
           vec_float.push_back(
               attribute.dyn_cast<pir::FloatAttribute>().data());  // NOLINT
         }
-        attr_map[name] = vec_float;
+        attr_map[legacy_arg_name] = vec_float;
       } else if (array_list[0].isa<pir::DoubleAttribute>()) {
         std::vector<int> vec_double;
         for (auto attribute : array_list) {
           vec_double.push_back(
               attribute.dyn_cast<pir::DoubleAttribute>().data());  // NOLINT
         }
-        attr_map[name] = vec_double;
+        attr_map[legacy_arg_name] = vec_double;
       } else if (array_list[0].isa<pir::StrAttribute>()) {
         std::vector<std::string> vec_string;
         for (auto attribute : array_list) {
           vec_string.push_back(
               attribute.dyn_cast<pir::StrAttribute>().AsString());  // NOLINT
         }
-        attr_map[name] = vec_string;
+        attr_map[legacy_arg_name] = vec_string;
       } else {
         std::stringstream ss;
         val.Print(ss);
         VLOG(1) << "type not support " << ss.str() << std::endl;
         PADDLE_THROW("Type[%s] in attribute map not support yet", ss.str());
       }
     } else if (val.isa<paddle::dialect::DataTypeAttribute>()) {
-      attr_map[name] = paddle::framework::TransToProtoVarType(
+      attr_map[legacy_arg_name] = paddle::framework::TransToProtoVarType(
           val.dyn_cast<paddle::dialect::DataTypeAttribute>().data());
     } else {
       std::stringstream ss;
 
@@ -20,4 +20,4 @@ file(GLOB PD_PROGRAM_TRANSLATOR_SRCS "*.cc")
 cc_library(
   program_translator
   SRCS ${PD_PROGRAM_TRANSLATOR_SRCS} ${op_compat_source_file}
-  DEPS proto_desc op_dialect_vjp pir framework_proto)
+  DEPS proto_desc op_dialect op_dialect_vjp pir framework_proto)
@@ -254,7 +254,8 @@ inline std::string GetPrefix(pir::IrContext* ctx, const OpDesc& op_desc) {
     return kCustomOpDialectPrefix;
   }
 #ifdef PADDLE_WITH_DNNL
-  if (op_desc.GetAttrIfExists<bool>("use_mkldnn")) {
+  if (op_desc.GetAttrIfExists<bool>("use_mkldnn") ||
+      paddle::dialect::IsOneDNNOnlyOp(op_desc.Type())) {
     if (!HasOpInfo(ctx, op_desc, kOneDNNTargetDialectPrefix)) {
       VLOG(3) << op_desc.Type()
               << "'s use_mkldnn == True, but PIR not support OneDNN for this "
 
@@ -24,6 +24,7 @@
 #include "paddle/pir/core/builtin_type.h"
 #include "paddle/pir/core/utils.h"
 #ifdef PADDLE_WITH_DNNL
+#include "paddle/fluid/pir/dialect/operator/ir/onednn_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_onednn_dialect.h"
 #endif
 
@@ -36,6 +37,13 @@ bool HaveOpToMultiKernelsMap(std::string op_name) {
 const std::vector<PdOpSig>& LegacyOpToPdOpsMapping(std::string op_name) {
   return op_to_multi_kernels_map[op_name];
 }
+
+#ifdef PADDLE_WITH_DNNL
+bool IsOneDNNOnlyOp(std::string op_name) {
+  return paddle::onednn::dialect::onednn_only_op_set.count(op_name);
+}
+#endif
+
 }  // namespace dialect
 }  // namespace paddle
 
 
@@ -43,6 +43,10 @@ bool HaveOpToMultiKernelsMap(std::string op_name);
 
 const std::vector<PdOpSig>& LegacyOpToPdOpsMapping(std::string op_name);
 
+#ifdef PADDLE_WITH_DNNL
+bool IsOneDNNOnlyOp(std::string op_name);
+#endif
+
 }  // namespace dialect
 }  // namespace paddle
 
 
@@ -24,7 +24,7 @@ namespace operators {
 
 using dnnl::memory;
 
-template <typename T>
+template <typename T, typename DeviceContext>
 class QuantOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
@@ -125,7 +125,5 @@ class QuantOpKernel : public framework::OpKernel<T> {
 }  // namespace paddle
 namespace ops = paddle::operators;
 
-REGISTER_OP_KERNEL(quantize,
-                   MKLDNN,
-                   ::phi::CPUPlace,
-                   ops::QuantOpKernel<float>);
+PD_REGISTER_STRUCT_KERNEL(quantize, OneDNN, ONEDNN, ops::QuantOpKernel, float) {
+}
@@ -32,7 +32,7 @@ inline uint8_t clip_to_uint8(float x) {
 
 }  // namespace
 
-template <typename T>
+template <typename T, typename DeviceContext>
 class ReQuantOpKernel : public framework::OpKernel<T> {
  public:
   void Compute(const framework::ExecutionContext& ctx) const override {
@@ -134,9 +134,10 @@ class ReQuantOpKernel : public framework::OpKernel<T> {
 
 namespace ops = paddle::operators;
 
-REGISTER_OP_KERNEL(requantize,
-                   MKLDNN,
-                   ::phi::CPUPlace,
-                   ops::ReQuantOpKernel<int8_t>,
-                   ops::ReQuantOpKernel<uint8_t>,
-                   ops::ReQuantOpKernel<paddle::platform::bfloat16>);
+PD_REGISTER_STRUCT_KERNEL(requantize,
+                          OneDNN,
+                          ONEDNN,
+                          ops::ReQuantOpKernel,
+                          int8_t,
+                          uint8_t,
+                          paddle::platform::bfloat16) {}
@@ -119,7 +119,7 @@ if(WITH_MKLDNN)
       ${PADDLE_SOURCE_DIR}/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
   )
 
-  set(op_onednn_info_file ${PD_DIALECT_SOURCE_DIR}/pd_onednn_op_info.cc)
+  set(op_onednn_info_file ${PD_DIALECT_SOURCE_DIR}/onednn_op_info.cc)
   set(op_onednn_info_file_tmp ${op_onednn_info_file}.tmp)
 
   set(onednn_op_namespace paddle,onednn,dialect)
@@ -248,7 +248,8 @@ set(op_dialect_srcs
     ${api_source_file})
 
 if(WITH_MKLDNN)
-  set(op_dialect_srcs ${op_dialect_srcs} ${onednn_op_source_file})
+  set(op_dialect_srcs ${op_dialect_srcs} ${onednn_op_source_file}
+                      ${op_onednn_info_file})
 endif()
 
 set(op_dialect_deps phi common pir type_info string_helper)
 
@@ -80,7 +80,7 @@
 #include "paddle/fluid/ir_adaptor/translator/utils.h"
 {only_pd_op_header_files}
 
-{op_to_multi_kernels_map}
+{other_info}
 
 {input}
 
@@ -91,6 +91,10 @@
 extern std::unordered_map<std::string, std::vector<PdOpSig>> op_to_multi_kernels_map;
 """
 
+ONEDNN_ONLY_OP_SET_H = """
+extern std::set<std::string> onednn_only_op_set;
+"""
+
 GET_OP_LIST_TEMPALTE = """{}
 """
 
@@ -196,7 +200,7 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 // This file is generated by "paddle/fluid/pir/dialect/op_generator/op_gen.py"
 #include "{h_file}"
 
-{op_to_multi_kernels_map}
+{other_info}
 #endif
 """
 
@@ -260,6 +264,12 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 }};
 """
 
+ONEDNN_ONLY_OP_SET = """
+std::set<std::string> onednn_only_op_set = {{
+{maps}
+}};
+"""
+
 scalar_type_maps = {
     'int': 'pir::Int32Attribute',
     'int64_t': 'pir::Int64Attribute',
@@ -497,15 +507,24 @@ def __init__(self, op_yaml_item, op_compat_item):
         # OneDNN info
         if "extra_args" in self.op_yaml_item:
             self.onednn_extra_args = self.op_yaml_item["extra_args"]
+        else:
+            self.onednn_extra_args = []
+
+        if "data_format_tensors" in self.op_yaml_item:
             self.onednn_data_format_tensors = self.op_yaml_item[
                 "data_format_tensors"
             ]
-            self.is_onednn_only = self.op_yaml_item["is_onednn_only"]
-            self.dynamic_fallback = self.op_yaml_item["dynamic_fallback"]
         else:
-            self.onednn_extra_args = []
             self.onednn_data_format_tensors = None
+
+        if "is_onednn_only" in self.op_yaml_item:
+            self.is_onednn_only = self.op_yaml_item["is_onednn_only"]
+        else:
             self.is_onednn_only = False
+
+        if "dynamic_fallback" in self.op_yaml_item:
+            self.dynamic_fallback = self.op_yaml_item["dynamic_fallback"]
+        else:
             self.dynamic_fallback = False
 
     def parse_op_traits(self):
@@ -1931,6 +1950,7 @@ def OpGenerator(
     op_infos = []
     all_op_info_items = {}
     first_file = True
+    onednn_only_op_list = []
     for yaml_file in op_yaml_files:
         op_yaml_items = []
         with open(yaml_file, "r") as f:
@@ -1968,8 +1988,8 @@ def OpGenerator(
                 op_compat_item['int_array'] = int_array_item
             if dialect_name == "onednn_op":
                 if first_file:
-                    first_file = False
                     op["is_onednn_only"] = True
+                    onednn_only_op_list.append("\"" + op['name'] + "\"")
                 elif op['name'] in ops_onednn_extra_map:
                     onednn_item = ops_onednn_extra_map[op['name']]
                     op["is_onednn_only"] = onednn_item["is_onednn_only"]
@@ -1984,8 +2004,11 @@ def OpGenerator(
             item = OpInfoParser(op, op_compat_item)
             op_info_items[op['name']] = item
             all_op_info_items[op['name']] = item
-
         op_infos.append(op_info_items)
+
+        if first_file:
+            first_file = False
+
     if dialect_name == "onednn_op":
         op_infos = [all_op_info_items]
 
@@ -2020,17 +2043,23 @@ def OpGenerator(
     # (4) write to files for pd_op.h.tmp (all yaml gen to pd_op.h)
     only_pd_op_header_files_str = ""
     if dialect_name == "pd_op":
-        op_to_multi_kernels_map = OP_TO_MULTI_KERNELS_MAP_H
+        other_info = OP_TO_MULTI_KERNELS_MAP_H
         for name in reversed(namespaces):
-            op_to_multi_kernels_map = NAMESPACE_GARD_TEMPLATE.format(
-                namespace=name, input=op_to_multi_kernels_map
+            other_info = NAMESPACE_GARD_TEMPLATE.format(
+                namespace=name, input=other_info
             )  # Add namespaces
         only_pd_op_header_files_str = """
 #include \"paddle/phi/common/data_type.h\"
 #include \"paddle/fluid/pir/dialect/operator/interface/get_kernel_type_for_var.h\"
             """
+    elif dialect_name == "onednn_op":
+        other_info = ONEDNN_ONLY_OP_SET_H
+        for name in reversed(namespaces):
+            other_info = NAMESPACE_GARD_TEMPLATE.format(
+                namespace=name, input=other_info
+            )  # Add namespaces
     else:
-        op_to_multi_kernels_map = ""
+        other_info = ""
 
     head_file_str = "\n".join(head_file_strs)
     declare_type_id_str = "\n".join(declare_type_id_strs)
@@ -2039,7 +2068,7 @@ def OpGenerator(
             namespace=name, input=head_file_str
         )  # Add namespaces
     head_file_str = H_FILE_TEMPLATE.format(
-        op_to_multi_kernels_map=op_to_multi_kernels_map,
+        other_info=other_info,
         input=head_file_str,
         declare_type_id=declare_type_id_str,
         only_pd_op_header_files=only_pd_op_header_files_str,
@@ -2049,15 +2078,23 @@ def OpGenerator(
 
     # (5) write to files for pd_op_info.tmp
     if dialect_name == "pd_op":
-        op_to_multi_kernels_map_str = OP_TO_MULTI_KERNELS_MAPS.format(
+        other_info_str = OP_TO_MULTI_KERNELS_MAPS.format(
             maps=", \r".join(op_to_multi_kernels_lists)
         )
         for name in reversed(namespaces):
-            op_to_multi_kernels_map_str = NAMESPACE_GARD_TEMPLATE.format(
-                namespace=name, input=op_to_multi_kernels_map_str
+            other_info_str = NAMESPACE_GARD_TEMPLATE.format(
+                namespace=name, input=other_info_str
+            )  # Add namespaces
+    elif dialect_name == "onednn_op":
+        other_info_str = ONEDNN_ONLY_OP_SET.format(
+            maps=", \r".join(onednn_only_op_list)
+        )
+        for name in reversed(namespaces):
+            other_info_str = NAMESPACE_GARD_TEMPLATE.format(
+                namespace=name, input=other_info_str
             )  # Add namespaces
     else:
-        op_to_multi_kernels_map_str = ""
+        other_info_str = ""
 
     if op_info_file is not None:
         if sys.platform == "win32":
@@ -2075,7 +2112,7 @@ def OpGenerator(
                 op_declare_second_part=",".join(second_part_op_info).replace(
                     "\n", ""
                 ),
-                op_to_multi_kernels_map=op_to_multi_kernels_map_str,
+                other_info=other_info_str,
                 h_file=op_def_h_file[:-4],
             )
         else:
@@ -2084,7 +2121,7 @@ def OpGenerator(
             )
             op_info_str = CC_OP_INFO_FILE_TEMPLATE.format(
                 op_declare=",".join(op_list_strs).replace("\n", ""),
-                op_to_multi_kernels_map=op_to_multi_kernels_map_str,
+                other_info=other_info_str,
                 h_file=op_def_h_file[:-4],
             )