[PIR] support infer spmd auto gen. (#62547)

winter-wang · web-flow · commit b11f7f571997 · 2024-03-09T00:16:28.000+08:00
diff --git a/paddle/fluid/pir/dialect/CMakeLists.txt b/paddle/fluid/pir/dialect/CMakeLists.txt
@@ -95,7 +95,8 @@ execute_process(
     --op_compat_yaml_file ${op_compat_yaml_file} --namespaces ${op_namespace}
     --dialect_name ${dialect_name} --op_def_h_file ${op_header_file_tmp}
     --op_info_file ${op_info_file_tmp} --op_def_cc_file ${op_src_files_tmp}
-    --op_vjp_cc_file ${op_vjp_src_file_tmp})
+    --op_vjp_cc_file ${op_vjp_src_file_tmp} --with_distributed
+    ${WITH_DISTRIBUTE})
 
 set(generated_files_pd_op
     "${op_header_file}"
@@ -141,7 +142,7 @@ if(WITH_MKLDNN)
       --op_def_h_file ${onednn_op_header_file_tmp} --op_info_file
       ${op_onednn_info_file_tmp} --op_def_cc_file ${onednn_op_source_file_tmp}
       --onednn_yaml_file ${pir_op_onednn_yaml} --ops_onednn_extra_yaml_file
-      ${pd_ops_onednn_extra_yaml_file})
+      ${pd_ops_onednn_extra_yaml_file} --with_distributed ${WITH_DISTRIBUTE})
 
   set(generated_files_onednn_pd_op
       "${onednn_op_header_file}" "${onednn_op_source_file}"
diff --git a/paddle/fluid/pir/dialect/op_generator/op_all_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_all_func_gen.py
@@ -0,0 +1,39 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from op_infer_spmd_func_gen import gen_op_infer_spmd_func
+from op_infermeta_gen import gen_op_infermeta_func
+from op_member_access_func_gen import gen_op_member_access_func
+from op_vjp_interface_func_gen import gen_op_vjp_interface_func
+
+all_gen_op_func_list = [
+    gen_op_infer_spmd_func,
+    gen_op_infermeta_func,
+    gen_op_member_access_func,
+    gen_op_vjp_interface_func,
+]
+
+
+def gen_op_all_func(args, op_info, op_info_items):
+    interface_list = []
+    declare_list = []
+    impl_list = []
+    for func in all_gen_op_func_list:
+        interface, declare, impl = func(args, op_info, op_info_items)
+        interface_list += interface
+        if declare is not None:
+            declare_list.append(declare)
+        if impl is not None:
+            impl_list.append(impl)
+    return interface_list, declare_list, impl_list
diff --git a/paddle/fluid/pir/dialect/op_generator/op_gen.py b/paddle/fluid/pir/dialect/op_generator/op_gen.py
@@ -17,10 +17,12 @@
 import os
 import pathlib
 import sys
+from distutils.util import strtobool
 
 import yaml
 from decomp_interface_gen_op_list import decomp_interface_declare_gen_op_list
 from infer_symbolic_shape_gen import gen_infer_symbolic_shape_str
+from op_all_func_gen import gen_op_all_func
 from op_build_gen import gen_build_func_str, gen_build_func_str_by_invoke
 from op_infermeta_gen import (
     gen_infermeta_by_invoke_func_str,
@@ -32,7 +34,6 @@
     gen_op_vjp_str,
 )
 from op_kerneltype_gen import gen_kernel_type_for_var_str
-from op_member_func_gen import gen_op_get_inputs_outputs_str
 from op_verify_gen import gen_verify_func_str
 from ops_onednn_extra_parser import parse_data_format_tensors, parse_extra_args
 from parse_kernel_key_gen import gen_parse_kernel_key_str
@@ -107,6 +108,9 @@
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/operator/trait/custom_vjp.h"
 #include "paddle/phi/core/infermeta_utils.h"
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/infermeta/spmd_rules/rules.h"
+#endif
 {only_pd_op_header_files}
 
 {other_info}
@@ -147,7 +151,6 @@ class {TEST_API} {op_name} : public pir::Op<{op_name}{interfaces}{traits}> {{
 {get_kernel_type_for_var_declare}
 {parse_kernel_key_declare}
 {infer_symbolic_shape_declare}
-{get_inputs_and_outputs}
 {exclusive_interface}
 }};
 """
@@ -503,8 +506,13 @@ def __init__(self, op_yaml_item, op_compat_item):
         # parse infermeta && kernel
         self.infer_meta_map = self.parse_infer_meta_map()
         self.invoke_map = self.parse_invoke_map()
+        self.spmd_rule_func = None
         if 'infer_meta' in self.op_yaml_item:
             self.infer_meta_func = self.op_yaml_item['infer_meta']["func"]
+            if 'spmd_rule' in self.op_yaml_item['infer_meta']:
+                self.spmd_rule_func = self.op_yaml_item['infer_meta'][
+                    'spmd_rule'
+                ]
         else:
             self.infer_meta_func = None
 
@@ -1233,7 +1241,9 @@ def GenOneDnnExtraAttrsDefaultValue(onednn_extra_args):
     return attr_str
 
 
-def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
+def AutoCodeGen(
+    args, op_info_items, all_op_info_items, namespaces, dialect_name
+):
     # (3) CodeGen: Traverse op_info_items and generate
     ops_name_list = []  # all op class name store in this list
     ops_declare_list = []  # all op class declare store in this list
@@ -1291,23 +1301,17 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
         op_traits = op_info.traits_list
         op_interfaces = op_info.interfaces_list
         op_interfaces += ["paddle::dialect::OpYamlInfoInterface"]
-
-        if op_info.infer_meta_func:
-            op_interfaces += ["paddle::dialect::InferMetaInterface"]
-        elif op_invoke_map and op_invoke_map['func'] in op_info_items:
-            if op_info_items[op_invoke_map['func']].infer_meta_func:
-                op_interfaces += ["paddle::dialect::InferMetaInterface"]
-
-        if (
-            op_info.backward_name
-            and op_info.op_phi_name[0] not in vjp_interface_black_list
-            and dialect_name != "onednn_op"
-        ):
-            op_interfaces += ["paddle::dialect::VjpInterface"]
         exclusive_interface_str = gen_exclusive_interface_str(
             op_info, op_info_items
         )
 
+        interface_list, declare_list, impl_list = gen_op_all_func(
+            args, op_info, op_info_items
+        )
+        op_interfaces += interface_list
+        exclusive_interface_str += '\n' + '\n'.join(declare_list)
+        ops_defined_list += impl_list
+
         if dialect_name == "pd_op" or dialect_name == "onednn_op":
             op_interfaces += ["paddle::dialect::GetKernelTypeForVarInterface"]
 
@@ -1409,15 +1413,6 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
                         kernel_func_name
                     ]
 
-                # =================================== #
-                #  gen get input/output methods str   #
-                # =================================== #
-                op_get_inputs_outputs_str = gen_op_get_inputs_outputs_str(
-                    op_input_name_list,
-                    op_mutable_attribute_name_list,
-                    op_output_name_list,
-                )
-
                 # =================================== #
                 #         gen Build methods str       #
                 # =================================== #
@@ -1581,7 +1576,6 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
                         build_mutable_attr_is_input=build_mutable_attr_is_input,
                         build_attr_num_over_1=build_attr_num_over_1,
                         build_mutable_attr_is_input_attr_num_over_1=build_mutable_attr_is_input_attr_num_over_1,
-                        get_inputs_and_outputs=op_get_inputs_outputs_str,
                         exclusive_interface=exclusive_interface_str,
                         get_kernel_type_for_var_declare=get_kernel_type_for_var_declare_str,
                         parse_kernel_key_declare=parse_kernel_key_str,
@@ -1605,7 +1599,6 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
                         build_mutable_attr_is_input=build_mutable_attr_is_input,
                         build_attr_num_over_1=build_attr_num_over_1,
                         build_mutable_attr_is_input_attr_num_over_1=build_mutable_attr_is_input_attr_num_over_1,
-                        get_inputs_and_outputs=op_get_inputs_outputs_str,
                         exclusive_interface=exclusive_interface_str,
                         get_kernel_type_for_var_declare=get_kernel_type_for_var_declare_str,
                         parse_kernel_key_declare=parse_kernel_key_str,
@@ -2059,6 +2052,7 @@ def AutoCodeGen(op_info_items, all_op_info_items, namespaces, dialect_name):
 
 
 def OpGenerator(
+    args,
     op_yaml_files,
     op_compat_yaml_file,
     namespaces,
@@ -2206,7 +2200,9 @@ def OpGenerator(
             source_file_str,
             op_to_multi_kernels_list,
             vjp_source_file_str,
-        ) = AutoCodeGen(items, all_op_info_items, namespaces, dialect_name)
+        ) = AutoCodeGen(
+            args, items, all_op_info_items, namespaces, dialect_name
+        )
         op_list_strs.append(op_list_str)
         declare_type_id_strs.append(declare_type_id_str)
         define_type_id_strs.append(define_type_id_str)
@@ -2360,6 +2356,7 @@ def ParseArguments():
     parser.add_argument('--op_vjp_cc_file', type=str)
     parser.add_argument('--onednn_yaml_file', type=str)
     parser.add_argument('--ops_onednn_extra_yaml_file', type=str)
+    parser.add_argument('--with_distributed', type=strtobool)
     return parser.parse_args()
 
 
@@ -2384,6 +2381,7 @@ def ParseArguments():
 
     # auto code generate
     OpGenerator(
+        args,
         op_yaml_files,
         op_compat_yaml_file,
         namespaces,
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infer_spmd_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infer_spmd_func_gen.py
@@ -0,0 +1,68 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+OP_INFER_SPMD_TEMPLATE = """
+  static phi::distributed::SpmdInfo InferSpmd({infer_spmd_args}) {{
+    return phi::distributed::{func}({args});
+  }}
+"""
+
+
+def gen_op_infer_spmd_func(args, op_info, op_info_items):
+    if not args.with_distributed or not op_info.spmd_rule_func:
+        return [], None, None
+    input_types_map = {
+        'paddle::dialect::DenseTensorType': 'const phi::distributed::DistMetaTensor&',
+        'pir::VectorType<paddle::dialect::DenseTensorType>': 'const std::vector<phi::distributed::DistMetaTensor>&',
+    }
+    input_name_list = op_info.input_name_list
+    input_type_list = op_info.input_type_list
+    input_name_type_dict = {}
+    for attr_idx in range(len(input_name_list)):
+        input_name_type_dict[input_name_list[attr_idx]] = input_types_map[
+            input_type_list[attr_idx]
+        ]
+
+    attr_name_list = op_info.attribute_name_list
+    attr_type_list = op_info.attribute_gen_arg_type_list
+    attr_name_type_dict = {}
+    for attr_idx in range(len(attr_type_list)):
+        attr_name_type_dict[attr_name_list[attr_idx]] = attr_type_list[attr_idx]
+
+    spmd_params = input_name_list + attr_name_list
+    if op_info.kernel_map is not None:
+        spmd_params = op_info.kernel_map['param']
+    args_list_with_type = []
+    args_list = []
+    for param in spmd_params:
+        # is input
+        if param in op_info.input_name_list:
+            args_list_with_type.append(
+                input_name_type_dict[param] + " " + param
+            )
+            args_list.append(param)
+        # is attribute
+        else:
+            param_type = attr_name_type_dict[param]
+            if param_type == "phi::IntArray":
+                param_type = "const std::vector<int64_t>&"
+            args_list_with_type.append(param_type + " " + param)
+            args_list.append(param)
+
+    declare_str = OP_INFER_SPMD_TEMPLATE.format(
+        infer_spmd_args=', '.join(args_list_with_type),
+        func=op_info.infer_meta_map["spmd_rule"],
+        args=', '.join(args_list),
+    )
+    return [], declare_str, None
diff --git a/paddle/fluid/pir/dialect/op_generator/op_infermeta_gen.py b/paddle/fluid/pir/dialect/op_generator/op_infermeta_gen.py
@@ -674,3 +674,13 @@ def gen_infermeta_by_invoke_func_str(op_class_name, invoke_class_name):
     return OP_INFERMETA_BY_INVOKE_TEMPLATE.format(
         op_name=op_class_name, invoke_class=invoke_class_name
     )
+
+
+def gen_op_infermeta_func(args, op_info, op_info_items):
+    interface = []
+    if op_info.infer_meta_func:
+        interface = ["paddle::dialect::InferMetaInterface"]
+    elif op_info.invoke_map and op_info.invoke_map['func'] in op_info_items:
+        if op_info_items[op_info.invoke_map['func']].infer_meta_func:
+            interface = ["paddle::dialect::InferMetaInterface"]
+    return interface, None, None
diff --git a/paddle/fluid/pir/dialect/op_generator/op_member_access_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_member_access_func_gen.py
@@ -20,9 +20,13 @@
 """
 
 
-def gen_op_get_inputs_outputs_str(
-    op_input_name_list, op_mutable_attribute_name_list, op_output_name_list
-):
+# =================================== #
+#  gen get input/output methods str   #
+# =================================== #
+def gen_op_member_access_func(args, op_info, op_info_items):
+    op_input_name_list = op_info.input_name_list
+    op_mutable_attribute_name_list = op_info.mutable_attribute_name_list
+    op_output_name_list = op_info.output_name_list
     op_get_inputs_outputs_str = ""
     for idx in range(len(op_input_name_list)):
         op_get_inputs_outputs_str += OP_GET_INPUT_TEMPLATE.format(
@@ -39,4 +43,4 @@ def gen_op_get_inputs_outputs_str(
             output_name=op_output_name_list[idx],
             output_index=idx,
         )
-    return op_get_inputs_outputs_str
+    return [], op_get_inputs_outputs_str, None
diff --git a/paddle/fluid/pir/dialect/op_generator/op_vjp_interface_func_gen.py b/paddle/fluid/pir/dialect/op_generator/op_vjp_interface_func_gen.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from vjp_interface_black_list import vjp_interface_black_list
+
+
+def gen_op_vjp_interface_func(args, op_info, op_info_items):
+    if (
+        op_info.backward_name
+        and op_info.op_phi_name[0] not in vjp_interface_black_list
+        and args.dialect_name != "onednn_op"
+    ):
+        return ["paddle::dialect::VjpInterface"], None, None
+    else:
+        return [], None, None
