[Dy2St] Run PT in SOT mode only (#59658)

Author: SigureMo

python/paddle/jit/dy2static/partial_program.py

@@ -238,7 +238,9 @@ def __call__(self, inputs):
         in_vars, in_var_names = self._prepare_inputs(inputs)
         out_vars = self._prepare_outputs()
         self._cast_fp16_if_pure_fp16(in_vars)
-        attrs = self._prepare_attributes()
+        # TODO(dev): Currently AST + PT has some issues in control flow, so we only
+        # enable SOT + PT in 2.6, we will fix it later.
+        attrs = self._prepare_attributes(force_not_use_pt=True)
         attrs.extend(["x_names", in_var_names])
 
         self._sync_lr_value_with_scheduler()
@@ -777,7 +779,7 @@ def _cast_fp16_if_pure_fp16(self, in_vars):
                     in_vars[i] = var.astype('float16')
                     in_vars[i].name = name
 
-    def _prepare_attributes(self):
+    def _prepare_attributes(self, force_not_use_pt=False):
         attrs = [
             'forward_global_block',
             self.forward_program.desc.block(0),
@@ -822,6 +824,8 @@ def _prepare_attributes(self):
         is_cinn_enabled = self._build_strategy.build_cinn_pass
         if is_prim_enabled or in_cinn_backend or is_cinn_enabled:
             in_pir_pt_mode = False
+        if force_not_use_pt:
+            in_pir_pt_mode = False
         attrs.extend(['in_pir_pt_mode', in_pir_pt_mode])
 
         return attrs

python/paddle/jit/sot/symbolic/compile_cache.py

@@ -104,16 +104,13 @@ def __call__(self, *args, **kwargs):
                 ),
             )
             if self.partial_program is None:
-                with EventGuard("FallbackWrapper: call compiled_fn"):
-                    outputs = self.compiled_fn(*args, **kwargs)
+                with EventGuard("FallbackWrapper: get_concrete_program"):
                     (
                         self.concrete_program,
                         self.partial_program,
                     ) = self.compiled_fn.get_concrete_program(*args, **kwargs)
-            else:
-                # Speed up Resnet from 0.0068 --> 0.0057
-                with EventGuard("FallbackWrapper: call partial_program"):
-                    outputs = self.partial_program.sot_call(*args, **kwargs)
+            with EventGuard("FallbackWrapper: sot call partial_program"):
+                outputs = self.partial_program.sot_call(*args, **kwargs)
 
             clear_eager_tensor_name(outputs)
             log_do(