[SOT] Fix NumPy array guard (#71106)

SigureMo · web-flow · commit c11cac94d98b · 2025-02-13T10:30:11.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py b/python/paddle/jit/sot/opcode_translator/executor/executor_cache.py
@@ -132,7 +132,7 @@ def lookup(
                 if guard_result:
                     log(
                         2,
-                        f"[Cache]: Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                        f"[Cache] Cache hit, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
                     )
                     return custom_code
                 else:
@@ -142,14 +142,18 @@ def lookup(
                     )
                     log(
                         2,
-                        f"[Cache]: Cache miss, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                        f"[Cache] Cache miss, Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
                     )
                     log_do(
                         2,
                         self.analyse_guard_error(guard_fn, frame),
                     )
             except Exception as e:
-                log(2, f"[Cache]: Guard function error: {e}\n")
+                log(2, f"[Cache] Guard function error: {e}\n")
+                log(
+                    2,
+                    f"[Cache] Guard is \n{getattr(guard_fn, 'expr', 'None')}\n",
+                )
                 log_do(
                     2,
                     self.analyse_guard_error(guard_fn, frame),
@@ -207,7 +211,7 @@ def inner():
                     result = guard(frame)
                 except Exception as e:
                     print(
-                        f"[Cache]: Error occurred when checking guard {guard_str}: {e}"
+                        f"[Cache] Error occurred when checking guard {guard_str}: {e}"
                     )
                     return
                 if result is False:
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -1204,29 +1204,41 @@ def get_py_value(self, allow_tensor=False) -> Any:
 
     @check_guard
     def make_stringified_guard(self) -> list[StringifiedExpression]:
-        if isinstance(self.get_py_value(), np.number):
-            frame_value_tracer = self.tracker.trace_value_from_frame()
+        frame_value_tracer = self.tracker.trace_value_from_frame()
+        obj_free_var_name = f"__{self.id}"
 
-            def format_dtype(dtype: np.dtype):
-                return f"np.{dtype}"
+        def format_dtype(dtype: np.dtype):
+            return f"np.{dtype}"
 
-            def format_number(number: np.number):
-                return f"{format_dtype(number.dtype)}({number.item()})"
+        def format_number(number: np.number):
+            return f"{format_dtype(number.dtype)}({number.item()})"
 
+        dtype_guard = StringifiedExpression(
+            f"{{}}.dtype == {format_dtype(self.get_py_value().dtype)}",
+            [frame_value_tracer],
+            union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+        )
+        if isinstance(self.get_py_value(), np.number):
             return [
+                dtype_guard,
                 StringifiedExpression(
                     f"{{}} == {format_number(self.get_py_value())}",
                     [frame_value_tracer],
                     union_free_vars(frame_value_tracer.free_vars, {"np": np}),
                 ),
+            ]
+        else:
+            return [
+                dtype_guard,
                 StringifiedExpression(
-                    f"{{}}.dtype == {format_dtype(self.get_py_value().dtype)}",
+                    f"({{}} == {obj_free_var_name}).all()",
                     [frame_value_tracer],
-                    union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+                    union_free_vars(
+                        frame_value_tracer.free_vars,
+                        {obj_free_var_name: self.get_py_value()},
+                    ),
                 ),
             ]
-        else:
-            return object_equal_stringified_guard(self)
 
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
diff --git a/test/sot/test_numpy.py b/test/sot/test_numpy.py
@@ -15,7 +15,10 @@
 import unittest
 
 import numpy as np
-from test_case_base import TestCaseBase
+from test_case_base import (
+    TestCaseBase,
+    test_instruction_translator_cache_context,
+)
 
 import paddle
 from paddle.jit.sot.utils import strict_mode_guard
@@ -35,6 +38,10 @@ def large_numpy_array_to_tensor(x):
     return paddle.to_tensor(x)
 
 
+def normal_numpy_array_to_tensor(x):
+    return paddle.to_tensor(x)
+
+
 class TestNumpy(TestCaseBase):
     @strict_mode_guard(False)
     def test_numpy_add(self):
@@ -61,6 +68,15 @@ def test_large_numpy_array_to_tensor(self):
         x = np.random.rand(1024, 1024, 2).astype(np.float32)
         self.assert_results(large_numpy_array_to_tensor, x)
 
+    def test_numpy_array_guard(self):
+        x = np.array([1.0, 2.0])
+        with test_instruction_translator_cache_context() as ctx:
+            self.assertEqual(ctx.translate_count, 0)
+            self.assert_results(normal_numpy_array_to_tensor, x)
+            self.assertEqual(ctx.translate_count, 1)
+            self.assert_results(normal_numpy_array_to_tensor, x)
+            self.assertEqual(ctx.translate_count, 1)
+
 
 if __name__ == "__main__":
     unittest.main()
