phi autoparallel fallback (PaddlePaddle#59241)

SecretXV · Nov 28, 2023 · 8f40a4a · 8f40a4a
1 parent b4591d5
commit 8f40a4a
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 11 deletions.
diff --git a/paddle/phi/api/yaml/generator/dist_api_gen.py b/paddle/phi/api/yaml/generator/dist_api_gen.py
@@ -53,12 +53,13 @@
       // 4. Select Kernel{}
       // 5. Reshard Input{}\n
       // 6. PrepareData (DataTransform & Prepare Dense Input){}
-      // RecordOpInfoSupplement{}
-      // 7. Infer Local DenseTensor Meta{}
-      // 8. DenseTensor Kernel Call{}
+      // 7. RecordOpInfoSupplement{}
+      // 8. Infer Local DenseTensor Meta{}
+      // 9. DenseTensor Kernel Call{}
+      // 10. Fallback{}
     }}\n
-    // 9. Set Output Dist Attr For Default Impl{}\n
-    // 10. Return
+    // 11. Set Output Dist Attr For Default Impl{}\n
+    // 12. Return
     {}
   }}
 """
@@ -1462,6 +1463,26 @@ def generate_kernel_call_code(self) -> str:
                     result += MULTI_SINGLE_SET_DIST_OUT_DIMS.format(i, i)
         return result
 
+    def generate_fallback_code(self) -> str:
+        fallback_code = ""
+        fallback_code += """
+      if (kernel_result.has_fallback_cpu) {"""
+        for kernel_out in self.dense_output_args:
+            fallback_code += f"""
+        TransDataBackend({kernel_out}, kernel_backend, {kernel_out});"""
+
+        inplace_flag = False
+        if len(self.inplace_map) > 0:
+            inplace_flag = True
+
+        fallback_code += self.reset_view_after_fallback(
+            self.outputs['types'], '        ', inplace_flag
+        )
+
+        fallback_code += """
+      }"""
+        return fallback_code
+
     def generate_output_dist_attr_setting(self) -> str:
         set_out_dist_attr_code = ""
         if self.generate_general_infer_spmd is True:
@@ -1547,6 +1568,7 @@ def generate_auto_paralel_branch(self) -> str:
         )
         infer_meta_code = self.generate_infer_meta_code()
         kernel_call_code = self.generate_kernel_call_code()
+        fallback_code = self.generate_fallback_code()
         output_dist_attr_setting = self.generate_output_dist_attr_setting()
         return_code = self.generate_return_code()
 
@@ -1560,6 +1582,7 @@ def generate_auto_paralel_branch(self) -> str:
             record_op_info_supplement_code,
             infer_meta_code,
             kernel_call_code,
+            fallback_code,
             output_dist_attr_setting,
             return_code,
         )

diff --git a/paddle/phi/api/yaml/generator/dist_bw_api_gen.py b/paddle/phi/api/yaml/generator/dist_bw_api_gen.py
@@ -33,13 +33,13 @@
       // 5. Select Kernel{}
       // 6. Reshard Input{}\n
       // 7. PrepareData (DataTransform & Prepare Dense Input){}
-      // RecordOpInfoSupplement
-      {}
-      // 8. Infer Local DenseTensor Meta{}
-      // 9. DenseTensor Kernel Call{}
+      // 8. RecordOpInfoSupplement{}
+      // 9. Infer Local DenseTensor Meta{}
+      // 10. DenseTensor Kernel Call{}
+      // 11. Fallback{}
     }}
-    // 10. Reshard Kernel Output to API output{}\n
-    // 11. Return
+    // 12. Reshard Kernel Output to API output{}\n
+    // 13. Return
     {}
   }}
 """
@@ -310,6 +310,7 @@ def generate_auto_paralel_branch(self) -> str:
         )
         infer_meta_code = self.generate_infer_meta_code()
         kernel_call_code = self.generate_kernel_call_code()
+        fallback_code = self.generate_fallback_code()
         reshard_output_code = self.generate_reshard_output_code()
         return_code = self.generate_return_code()
 
@@ -324,6 +325,7 @@ def generate_auto_paralel_branch(self) -> str:
             record_op_info_supplement_code,
             infer_meta_code,
             kernel_call_code,
+            fallback_code,
             reshard_output_code,
             return_code,
         )