[PIR] pir onednn support fused elementwise (#62126)

wanghuancoder · web-flow · commit aa86d86ab9ed · 2024-02-28T09:47:10.000+08:00
* pir onednn support fused eselemtwise
diff --git a/paddle/fluid/pir/dialect/operator/ir/onednn.yaml b/paddle/fluid/pir/dialect/operator/ir/onednn.yaml
@@ -28,6 +28,42 @@
     data_type : input
   optional : bias, residual_param
 
+- op : fused_elementwise_add
+  args : (Tensor x, Tensor y, int axis=-1, str fuse_activation="", float fuse_alpha=0.0, float fuse_beta=0.0, float fused_output_scale=1.0, int[] fused_unsqueeze2_axes={}, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+    param : [x, y]
+  kernel :
+    func : fused_elementwise_add
+
+- op : fused_elementwise_div
+  args : (Tensor x, Tensor y, int axis=-1, str fuse_activation="", float fuse_alpha=0.0, float fuse_beta=0.0, float fused_output_scale=1.0, int[] fused_unsqueeze2_axes={}, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+    param : [x, y]
+  kernel :
+    func : fused_elementwise_div
+
+- op : fused_elementwise_mul
+  args : (Tensor x, Tensor y, int axis=-1, str fuse_activation="", float fuse_alpha=0.0, float fuse_beta=0.0, float fused_output_scale=1.0, int[] fused_unsqueeze2_axes={}, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+    param : [x, y]
+  kernel :
+    func : fused_elementwise_mul
+
+- op : fused_elementwise_sub
+  args : (Tensor x, Tensor y, int axis=-1, str fuse_activation="", float fuse_alpha=0.0, float fuse_beta=0.0, float fused_output_scale=1.0, int[] fused_unsqueeze2_axes={}, float scale_x=1.0, float scale_y=1.0, float scale_out=1.0)
+  output : Tensor(out)
+  infer_meta :
+    func : ElementwiseInferMeta
+    param : [x, y]
+  kernel :
+    func : fused_elementwise_sub
+
 - op: multi_gru
   args: (Tensor x, Tensor[] weight_x, Tensor[] weight_h, Tensor[] bias, Tensor[] scale_weights, str activation="tanh", str gate_activation="sigmoid", int layers=1, bool origin_mode=false, str mkldnn_data_type="float32", float scale_data=1.0, float shift_data=1.0, bool force_fp32_output=false)
   output: Tensor(hidden)
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml b/paddle/fluid/pir/dialect/operator/ir/ops_onednn_extra.yaml
@@ -104,13 +104,13 @@
   extra_args : float fuse_alpha = 0.0, float fuse_beta = 0.0, float scale_in=1.0, float scale_out=1.0, float scale_in_eltwise=1.0, float[] scale_weights={1.0f}
   data_format_tensors : input
 
-# - op : fused_elementwise_add
+- op : fused_elementwise_add
 
-# - op : fused_elementwise_div
+- op : fused_elementwise_div
 
-# - op : fused_elementwise_mul
+- op : fused_elementwise_mul
 
-# - op : fused_elementwise_sub
+- op : fused_elementwise_sub
 
 # - op : fused_matmul
 
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -1354,6 +1354,30 @@
     attrs : [bool use_cudnn = false, float fuse_alpha = 0.0f, float fuse_beta = 0.0f, float Scale_in = 1.0f,
              float Scale_out = 1.0f, float Scale_in_eltwise = 1.0f, 'float[] Scale_weights = {1.0f}', bool use_mkldnn = true, str mkldnn_data_type = "float32"]
 
+- op : fused_elementwise_add
+  inputs :
+    {x : X, y : Y}
+  outputs :
+    {out : Out}
+
+- op : fused_elementwise_div
+  inputs :
+    {x : X, y : Y}
+  outputs :
+    {out : Out}
+
+- op : fused_elementwise_mul
+  inputs :
+    {x : X, y : Y}
+  outputs :
+    {out : Out}
+
+- op : fused_elementwise_sub
+  inputs :
+    {x : X, y : Y}
+  outputs :
+    {out : Out}
+
 - op : fused_embedding_eltwise_layernorm
   inputs :
     ids : Ids
