[PIR] refine onednn add_n (#62471)

wanghuancoder · web-flow · commit 1128c78b68d6 · 2024-03-07T10:59:21.000+08:00
* refine onednn add_n

* refine
diff --git a/paddle/fluid/ir_adaptor/translator/op_translator.cc b/paddle/fluid/ir_adaptor/translator/op_translator.cc
@@ -1367,19 +1367,10 @@ struct ShadowOutputOpTranscriber : public OpTranscriber {
 struct AddNOpTranscriber : public OpTranscriber {
   pir::OpInfo LookUpOpInfo(pir::IrContext* ctx,
                            const OpDesc& op_desc) override {
-    auto prefix = GetPrefix(ctx, op_desc);
-    std::string target_op_name;
-#ifdef PADDLE_WITH_DNNL
-    if (prefix == kOneDNNTargetDialectPrefix) {
-      target_op_name = std::string(kOneDNNTargetDialectPrefix) + "add_n_onednn";
-    } else  // NOLINT
-#endif
-    {
-      target_op_name =
-          GetPrefix(ctx, op_desc) + OpNameCompatibleMapping(op_desc.Type());
-      if (IsInplace(op_desc)) {
-        target_op_name += "_";
-      }
+    std::string target_op_name =
+        GetPrefix(ctx, op_desc) + OpNameCompatibleMapping(op_desc.Type());
+    if (IsInplace(op_desc)) {
+      target_op_name += "_";
     }
 
     const auto& op_info = ctx->GetRegisteredOpInfo(target_op_name);
diff --git a/paddle/fluid/pir/dialect/operator/ir/onednn.yaml b/paddle/fluid/pir/dialect/operator/ir/onednn.yaml
@@ -1,13 +1,3 @@
-- op : add_n_onednn
-  args : (Tensor[] inputs)
-  output : Tensor(out)
-  infer_meta:
-    func: AddNInferMeta
-    param: [inputs]
-  kernel:
-    func: add_n
-    param: [inputs]
-
 - op : dequantize
   args : (Tensor input, float scale=1.0, float shift=0.0)
   output : Tensor(output)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -28,12 +28,16 @@
     support_trans_dtype : x, y
   interfaces : paddle::dialect::InferSymbolicShapeInterface
 
-# this add_n is only for ops_api_gen.py
+# this add_n is only for ops_api_gen.py and onednn
 - op : add_n
   args : (Tensor[] inputs)
-  output : Tensor
-  invoke : add_n_impl(inputs)
-  backward : add_n_grad
+  output : Tensor(out)
+  infer_meta:
+    func: AddNInferMeta
+    param: [inputs]
+  kernel:
+    func: add_n
+    param: [inputs]
 
 - op : all
   args : (Tensor x, int64_t[] axis={}, bool keepdim=false)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -15,7 +15,7 @@
 
 - op : abs_grad
 
-- op : add_n_onednn
+- op : add_n
   extra_args : str mkldnn_data_type="float32"
 
 - op : batch_norm
