[Auto Parallel] Try to refine dist_api_gen (#73114)

* refine code

* fix

Author: ooooo-create

paddle/phi/api/generator/dist_api_gen.py

@@ -581,6 +581,29 @@
       }}
 """
 
+CALCULATE_LOCAL_SHAPE_KERNEL_TEMPLATE = """
+
+      auto out_grad_shape = out_grad.dims();
+      std::vector<{dtype}> local_kernel_shape;
+      const auto& out_grad_dist_attr = {out_grad_dist_attr};
+      for (int i = 0; i < out_grad_shape.size(); i++) {{
+        if (out_grad_dist_attr.dims_mapping()[i] >= 0) {{
+          {dtype} shape_i = out_grad_shape[i];
+          int64_t dim = out_grad_dist_attr.dims_mapping()[i];
+          int64_t mesh_dim = out_grad_dist_attr.process_mesh().shape()[dim];
+          // TODO: Support aliquant condition.
+          PADDLE_ENFORCE(shape_i % mesh_dim == 0,
+                common::errors::InvalidArgument(
+                    "{op_name} only support local shape dim is divisible "
+                    "by the mesh dim, however local_kernel_shape[%lld] is %lld "
+                    "and shard mesh dims is %lld.", i, shape_i, mesh_dim));
+          local_kernel_shape.push_back(shape_i / mesh_dim);
+        }} else {{
+          local_kernel_shape.push_back(out_grad_shape[i]);
+        }}
+      }}
+"""
+
 # BaseAPI members:
 # inputs:
 #     names : [], list of input names
@@ -1755,7 +1778,7 @@ def generate_infer_meta_code(self) -> str:
 
         return output_decl_code + infer_meta_code
 
-    def generate_kernel_call_code(self) -> str:
+    def generate_kernel_call_code(self, is_forward=True) -> str:
         dense_input_trans_map = {
             'const Tensor&': 'const phi::DenseTensor&',
             'const std::vector<Tensor>&': 'const std::vector<const phi::DenseTensor*>&',
@@ -1773,6 +1796,7 @@ def generate_kernel_call_code(self) -> str:
         kernel_args_type_list = ['const phi::DeviceContext&']
 
         attr_names = self.attrs['names']
+        pure_kernel_args = self.kernel['param']
         kernel_args = self.kernel['param']
         if kernel_args is None:
             kernel_args = input_names + attr_names
@@ -1803,7 +1827,14 @@ def generate_kernel_call_code(self) -> str:
                     kernel_args_type_list.append('const phi::IntArray&')
                     # TODO(GhostScreaming): kernel like reshape need calculate local_shape
                     if self.infer_meta['local_shape'] is not None:
-                        arg = 'phi::IntArray(local_shape)'
+                        if is_forward or (
+                            pure_kernel_args is not None
+                            and self.infer_meta['local_shape']
+                            not in pure_kernel_args
+                        ):
+                            arg = 'phi::IntArray(local_shape)'
+                        else:
+                            arg = 'phi::IntArray(local_kernel_shape)'
                     else:
                         arg = 'phi::IntArray(' + arg + ')'
                 elif 'vector<phi::Scalar>' in self.attrs['attr_info'][arg][0]:
@@ -1818,9 +1849,15 @@ def generate_kernel_call_code(self) -> str:
                         self.attrs['attr_info'][arg][0]
                     )
                     # calculate local_shape for expand_as
-                    # TODO(ooooo): bwd reuse this function to kernel, but actually the local_shape isn't same meaning.
                     if self.infer_meta['local_shape'] is not None:
-                        arg = "local_shape"
+                        if is_forward or (
+                            pure_kernel_args is not None
+                            and self.infer_meta['local_shape']
+                            not in pure_kernel_args
+                        ):
+                            arg = 'local_shape'
+                        else:
+                            arg = 'local_kernel_shape'
                 input_args.append(arg)
             elif isinstance(arg, bool):
                 input_args.append(str(arg).lower())

paddle/phi/api/generator/dist_bw_api_gen.py

@@ -288,7 +288,7 @@ def generate_output_creation_code(self) -> str:
 
         return output_creation_code
 
-    def generate_bw_infer_local_shape_code(self) -> str:
+    def generate_bw_infer_local_shape_code(self, need_kernel=False):
         arg_name = self.infer_meta['local_shape']
         assert arg_name in self.outputs['names'], (
             f"Auto Parallel will calculate local_shape for {arg_name} "
@@ -304,7 +304,7 @@ def generate_bw_infer_local_shape_code(self) -> str:
             self.outputs['names'].index(arg_name)
         ]
         shape_type = self.get_shape_type(fw_attrs['attr_info'])
-        return dist_api_gen.CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
+        return_code = dist_api_gen.CALCULATE_LOCAL_SHAPE_TEMPLATE.format(
             out_name=dist_out_name,
             out_dist_attr=(
                 "PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.second[0]);"
@@ -314,6 +314,20 @@ def generate_bw_infer_local_shape_code(self) -> str:
             dtype=shape_type,
             op_name=self.kernel['func'][0],
         )
+        if need_kernel:
+            return (
+                dist_api_gen.CALCULATE_LOCAL_SHAPE_KERNEL_TEMPLATE.format(
+                    out_grad_dist_attr=(
+                        "PADDLE_GET_CONST(phi::distributed::TensorDistAttr, spmd_info.first[1]);"
+                        if self.infer_meta['spmd_rule']
+                        else "phi::distributed::TensorDistAttr(common::vectorize(out_grad.dims()))"
+                    ),
+                    dtype=shape_type,
+                    op_name=self.kernel['func'][0],
+                )
+                + return_code
+            )
+        return return_code
 
     def generate_infer_meta_code(self) -> str:
         (
@@ -347,7 +361,15 @@ def generate_infer_meta_code(self) -> str:
         )
         # TODO(GhostScreaming): kernel like reshape need calculate local_shape
         if self.infer_meta['local_shape'] is not None:
-            infer_meta_code += self.generate_bw_infer_local_shape_code()
+            if (
+                self.kernel['param'] is not None
+                and self.infer_meta['local_shape'] not in self.kernel['param']
+            ):
+                infer_meta_code += self.generate_bw_infer_local_shape_code()
+            else:
+                infer_meta_code += self.generate_bw_infer_local_shape_code(
+                    need_kernel=True
+                )
             infer_meta_code += SET_LOCAL_SHAPE_TEMPLATE.format(
                 meta_tensor="meta_" + self.dense_output_args[0]
             )
@@ -462,7 +484,7 @@ def generate_auto_parallel_branch(self) -> str:
             )
         )
         infer_meta_code = self.generate_infer_meta_code()
-        kernel_call_code = self.generate_kernel_call_code()
+        kernel_call_code = self.generate_kernel_call_code(is_forward=False)
         fallback_code = self.generate_fallback_code()
         reshard_output_code = self.generate_reshard_output_code()
         return_code = self.generate_return_code()