【prim】uniform Decomp list (#70204)

xiaoguoguo626807 · web-flow · commit bcbdc446063e · 2024-12-17T14:13:06.000+08:00
* modify DPA model

* modify ir_printer

* modify decomp list

* fix conflict
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -16,85 +16,37 @@
 # DecompInterface gen op list
 # =====================================
 
-# come into effect in generated file pd_op.h
-# manual decomp interface declare are located in manual_op.h
-decomp_interface_declare_gen_op_list = [
-    "add_n",
-    "addmm",
-    "any",
+
+##################### decomp_rule ####################
+# come into effect in generated file op_decomp_rule.cc
+# manual decomp interface implementation are located in manual_op_decomp.cc
+######################################################
+MANUAL_IMPL_DECOMP = [
     "batch_norm",
     "batch_norm_",
-    "bce_loss",
-    "bmm",
     "clip",
-    "diag",
-    "dropout",
-    "eye",
-    "elu",
-    "embedding",
-    "flatten",
-    "full_like",
-    "gelu",
-    "hardswish",
-    "hardsigmoid",
-    "heaviside",
-    "huber_loss",
-    "group_norm",
-    "index_sample",
-    "index_select",
-    "instance_norm",
-    "kldiv_loss",
-    "layer_norm",
-    "leaky_relu",
-    "lerp",
-    "log_loss",
-    "log_softmax",
-    "mean",
-    "mean_all",
-    "meshgrid",
-    "numel",
     "one_hot",
-    "p_norm",
-    "pow",
-    "reciprocal",
-    "relu",
-    "relu6",
-    "sigmoid_cross_entropy_with_logits",
-    "silu",
-    "swiglu",
-    "swish",
-    "softmax",
-    "softsign",
-    "square",
-    "squared_l2_norm",
-    "squeeze",
-    "stack",
-    "unbind",
-    "unsqueeze",
-    "unstack",
 ]
 
-# come into effect in generated file op_decomp_rule.cc
-# manual decomp interface implementation are located in manual_op_decomp.cc
-decomp_interface_implementation_gen_op_list = [
-    "any",
+GENERATE_IMPL_DECOMP = [
     "add_n",
     "addmm",
+    "any",
     "bce_loss",
     "bmm",
     "diag",
     "dropout",
-    "eye",
     "elu",
     "embedding",
+    "eye",
     "flatten",
     "full_like",
     "gelu",
-    "hardswish",
+    "group_norm",
     "hardsigmoid",
+    "hardswish",
     "heaviside",
     "huber_loss",
-    "group_norm",
     "index_sample",
     "index_select",
     "instance_norm",
@@ -115,85 +67,75 @@
     "relu6",
     "sigmoid_cross_entropy_with_logits",
     "silu",
-    "swiglu",
-    "swish",
     "softmax",
     "softsign",
     "square",
     "squared_l2_norm",
     "squeeze",
     "stack",
+    "swiglu",
+    "swish",
     "unbind",
     "unsqueeze",
     "unstack",
 ]
-
+decomp_rule_interface_declare_gen_op_list = (
+    GENERATE_IMPL_DECOMP + MANUAL_IMPL_DECOMP
+)
+decomp_rule_interface_implementation_gen_op_list = GENERATE_IMPL_DECOMP
 # xshape output will no longer used after decomp, but return none to keep output num the same as origin op
 decomp_ops_contain_unused_output = ["squeeze", "unsqueeze"]
 
-# prim op with one input and one output, with no attribute
-UNARY_PRIM_VJP_OPS = [
-    'abs_grad',
-    'erf_grad',
-    'exp_grad',
-    'floor_grad',
-    'log_grad',
-    'rsqrt_grad',
-    'sin_grad',
-    'cos_grad',
-    'tanh_grad',
-]
+##################### decomp_vjp ####################
+# come into effect in generated file op_decomp_vjp.cc
+# manual decomp interface implementation are located in manual_op_decomp_vjp.cc
+####################################################
 
-# prim op with two inputs and one output, with no attribute
-BINARY_PRIM_VJP_OPS = [
-    'matmul_grad',
+GENERATE_IMPL_VJP = [
+    'abs_grad',
     'add_grad',
+    'bce_loss_grad',
+    'cos_grad',
     'divide_grad',
-    'subtract_grad',
-    'multiply_grad',
     'elementwise_pow_grad',
-    'maximum_grad',
-    'reduce_as_grad',
-]
-
-OTHER_PRIM_VJP_OPS = [
-    'sum_grad',
-    'reshape_grad',
-    'roll_grad',
-    'transpose_grad',
-    'max_grad',
-    'squeeze_grad',
-    'unsqueeze_grad',
-]
-
-
-CUSTOM_VJP = [
-    'bce_loss_grad',
+    'erf_grad',
+    'exp_grad',
+    'floor_grad',
     'gelu_grad',
-    'hardsigmoid_grad',
+    'group_norm_grad',
     'hardswish_grad',
     'leaky_relu_grad',
+    'layer_norm_grad',
+    'log_grad',
+    'matmul_grad',
+    'max_grad',
+    'maximum_grad',
     'mean_grad',
     'minimum_grad',
+    'multiply_grad',
     'pow_grad',
+    'reduce_as_grad',
     'relu_grad',
-    'relu6_grad',
+    'reshape_grad',
+    'roll_grad',
+    'rsqrt_grad',
     'sigmoid_grad',
     'silu_grad',
+    'sin_grad',
     'softmax_grad',
     'sqrt_grad',
+    'squeeze_grad',
+    'subtract_grad',
+    'sum_grad',
     'swiglu_grad',
-    'layer_norm_grad',
-    'group_norm_grad',
-]  # custom vjp list of composite op
+    'tanh_grad',
+    'transpose_grad',
+    'unsqueeze_grad',
+]
 
 # declare belongs to codegen, but implementation not
-OTHER_VJP = ["concat_grad", "stack_grad", 'slice_grad']
-
-vjp_list = (
-    UNARY_PRIM_VJP_OPS + BINARY_PRIM_VJP_OPS + CUSTOM_VJP + OTHER_PRIM_VJP_OPS
-)
+MANUAL_IMPL_VJP = ["concat_grad", "stack_grad", 'slice_grad']
 
-decomp_vjp_interface_declare_gen_op_list = vjp_list + OTHER_VJP
+decomp_vjp_interface_declare_gen_op_list = GENERATE_IMPL_VJP + MANUAL_IMPL_VJP
 
-decomp_vjp_interface_implementation_gen_op_list = vjp_list
+decomp_vjp_interface_implementation_gen_op_list = GENERATE_IMPL_VJP
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -23,7 +23,7 @@
 
 import yaml
 from decomp_interface_gen_op_list import (
-    decomp_interface_declare_gen_op_list,
+    decomp_rule_interface_declare_gen_op_list,
     decomp_vjp_interface_declare_gen_op_list,
 )
 from gen_utils import attr_types_map, to_pascal_case
@@ -1407,8 +1407,9 @@ def AutoCodeGen(
 
             for kernel_func_name in func_list:
                 if (
-                    op_name in decomp_interface_declare_gen_op_list
-                    and kernel_func_name in decomp_interface_declare_gen_op_list
+                    op_name in decomp_rule_interface_declare_gen_op_list
+                    and kernel_func_name
+                    in decomp_rule_interface_declare_gen_op_list
                     and dialect_name != "onednn_op"
                 ):
                     if decomp_interface_str not in op_interfaces:
diff --git a/paddle/fluid/primitive/codegen/decomp_rule_gen.py b/paddle/fluid/primitive/codegen/decomp_rule_gen.py
@@ -37,8 +37,8 @@
 )
 
 from decomp_interface_gen_op_list import (
-    decomp_interface_implementation_gen_op_list,
     decomp_ops_contain_unused_output,
+    decomp_rule_interface_implementation_gen_op_list,
 )
 from op_gen import attr_types_map, to_pascal_case
 
@@ -207,7 +207,7 @@ def gen(
         templates_dir,
         destination_dir,
         apis=apis,
-        decomp_white_list=decomp_interface_implementation_gen_op_list,
+        decomp_white_list=decomp_rule_interface_implementation_gen_op_list,
         decomp_ops_list_contain_unused_output=decomp_ops_contain_unused_output,
     )
 
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py

Original file line number	Diff line number	Diff line change
`@@ -37,8 +37,8 @@`
`37`	`37`	`)`
`38`	`38`
`39`	`39`	`from decomp_interface_gen_op_list import (`
`40`		`- decomp_interface_implementation_gen_op_list,`
`41`	`40`	`decomp_ops_contain_unused_output,`
	`41`	`+ decomp_rule_interface_implementation_gen_op_list,`
`42`	`42`	`)`
`43`	`43`	`from op_gen import attr_types_map, to_pascal_case`
`44`	`44`
`@@ -207,7 +207,7 @@ def gen(`
`207`	`207`	`templates_dir,`
`208`	`208`	`destination_dir,`
`209`	`209`	`apis=apis,`
`210`		`- decomp_white_list=decomp_interface_implementation_gen_op_list,`
	`210`	`+ decomp_white_list=decomp_rule_interface_implementation_gen_op_list,`
`211`	`211`	`decomp_ops_list_contain_unused_output=decomp_ops_contain_unused_output,`
`212`	`212`	`)`
`213`	`213`