[SOT] Add a quick guard implementation for Tensor in PIR mode (#68031)

SigureMo · web-flow · commit 30eb9e582b15 · 2024-09-06T02:19:50.000+08:00
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -112,19 +112,15 @@ def dynamic_axes(self):
 
     @staticmethod
     def _handle_legacy_ir_amp_dtype(dtype):
-        expected_dtype_class = (
-            paddle.core.DataType
-            if paddle.framework.use_pir_api()
-            else paddle.core.VarDesc.VarType
-        )
-        assert isinstance(dtype, expected_dtype_class)
-
-        # TODO(@xiongkun) remove after pir become default state.
+        # TODO(cleanup-legacy-ir) remove after pir become default state.
         # We always use float32 in simulation if AMP is enabled.
+        if use_pir_api():
+            return dtype
+        assert isinstance(dtype, paddle.core.VarDesc.VarType)
+
         current_amp_state = amp_state()
         if (
-            not use_pir_api()
-            and dtype == paddle.float16
+            dtype == paddle.float16
             and current_amp_state is not None
             and current_amp_state["dtype"] == "float16"
         ):
diff --git a/python/paddle/jit/sot/opcode_translator/executor/guard.py b/python/paddle/jit/sot/opcode_translator/executor/guard.py
@@ -104,10 +104,9 @@ def make_guard(stringified_guards: list[StringifiedExpression]) -> Guard:
             guard.original_guard = guard
             return guard
 
-        free_vars = {}
-        for expr in stringified_guards:
-            free_vars = union_free_vars(free_vars, expr.free_vars)
-
+        free_vars = union_free_vars(
+            *(expr.free_vars for expr in stringified_guards)
+        )
         inlined_guard_expr = "lambda frame: " + " and ".join(
             [expr.inlined_expr for expr in stringified_guards]
         )
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -276,16 +276,30 @@ def make_stringified_guard(self) -> list[StringifiedExpression]:
                 self.tracker.obj.tracker.trace_value_from_frame()
             )
             dtype_str, dtype_free_vars = stringify_pyobject(self.value)
+            # TODO(cleanup-legacy-ir): Remove this branch after we remove legacy IR
+            if not paddle.framework.use_pir_api():
+                return [
+                    StringifiedExpression(
+                        f"MetaInfo.from_tensor({{}}).dtype == {dtype_str}",
+                        [tensor_value_tracer],
+                        union_free_vars(
+                            {"MetaInfo": MetaInfo},
+                            tensor_value_tracer.free_vars,
+                            dtype_free_vars,
+                        ),
+                    )
+                ]
             return [
                 StringifiedExpression(
-                    f"MetaInfo.from_tensor({{}}).dtype == {dtype_str}",
+                    f"{{}}.dtype == {dtype_str}",
                     [tensor_value_tracer],
                     union_free_vars(
-                        {"MetaInfo": MetaInfo},
+                        tensor_value_tracer.free_vars,
                         dtype_free_vars,
                     ),
                 )
             ]
+
         else:
             return object_equal_stringified_guard(self)
 
@@ -418,19 +432,55 @@ def _reconstruct(self, codegen: PyCodeGen):
     def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
 
-        if ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
-            str_left_expr = f"MetaInfo.from_tensor({{}}, dynamic_axes={self.meta.dynamic_axes}).guard_str()"
-        else:
-            str_left_expr = "MetaInfo.from_tensor({}).guard_str()"
+        # TODO(cleanup-legacy-ir): Remove this branch after we remove legacy IR
+        if not paddle.framework.use_pir_api():
+            if ENV_SOT_ALLOW_DYNAMIC_SHAPE.get():
+                str_left_expr = f"MetaInfo.from_tensor({{}}, dynamic_axes={self.meta.dynamic_axes}).guard_str()"
+            else:
+                str_left_expr = "MetaInfo.from_tensor({}).guard_str()"
+            return [
+                StringifiedExpression(
+                    f"{str_left_expr} == '{self.origin_meta.guard_str()}'",
+                    [frame_value_tracer],
+                    union_free_vars(
+                        {"MetaInfo": MetaInfo},
+                        frame_value_tracer.free_vars,
+                    ),
+                )
+            ]
+
+        # A quick check path for PIR, we don't need dtype conversion for AMP in PIR
+        meta = self.origin_meta
+        dtype_str, dtype_free_vars = stringify_pyobject(meta.dtype)
         return [
+            # Check rank
             StringifiedExpression(
-                f"{str_left_expr} == '{self.origin_meta.guard_str()}'",
+                f"len({{}}.shape) == {len(meta.shape)}",
                 [frame_value_tracer],
-                union_free_vars(
-                    {"MetaInfo": MetaInfo},
-                    frame_value_tracer.free_vars,
-                ),
-            )
+                union_free_vars(frame_value_tracer.free_vars),
+            ),
+            # Check each dim except dynamic dim
+            *[
+                StringifiedExpression(
+                    f"{{}}.shape[{i}] == {meta.shape[i]}",
+                    [frame_value_tracer],
+                    union_free_vars(frame_value_tracer.free_vars),
+                )
+                for i in range(len(meta.shape))
+                if not isinstance(meta.shape[i], SymbolicInt)
+            ],
+            # Check dtype
+            StringifiedExpression(
+                f"{{}}.dtype == {dtype_str}",
+                [frame_value_tracer],
+                union_free_vars(frame_value_tracer.free_vars, dtype_free_vars),
+            ),
+            # Check stop_gradient
+            StringifiedExpression(
+                f"{{}}.stop_gradient == {meta.stop_gradient!r}",
+                [frame_value_tracer],
+                union_free_vars(frame_value_tracer.free_vars),
+            ),
         ]
 
     def get_iter(self):
