[SOT][DynamicShape] Fix SymbolicVariable get_py_value reversed call and break graph on infermeta encount TypeError (#70009)

SigureMo · Copilot · gouzil · web-flow · commit de50a920be62 · 2024-12-10T00:14:13.000+08:00
---------

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
Co-authored-by: gouzil &lt;66515297+gouzil@users.noreply.github.com&gt;
diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -52,6 +52,7 @@
     map_if,
     switch_symbol_registry,
 )
+from ...utils.exceptions import BreakGraphError
 from ..instruction_utils import get_instructions
 from .guard import Guard, StringifiedExpression, make_guard
 from .mutable_data import MutationDel, MutationNew, MutationSet
@@ -661,7 +662,9 @@ def try_infer_meta_fn(args, kwargs) -> Any:
                         for arg in flatten_vars
                     ):
                         # TODO(zrr1999): maybe we can continue to fallback to all args are constant.
-                        raise e
+                        raise BreakGraphError(
+                            f"InferMeta encount {type(e)}, but all args are not symbolic."
+                        )
 
                     args, kwargs = map_if(
                         (args, kwargs),
@@ -686,7 +689,9 @@ def try_infer_meta_fn(args, kwargs) -> Any:
                         isinstance(arg, SymbolicVariable)
                         for arg in flatten_vars
                     ):
-                        raise e
+                        raise BreakGraphError(
+                            f"InferMeta encount {type(e)}, but all args are not symbolic."
+                        )
 
                     args, kwargs = map_structure(
                         replace_symbolic_var_with_constant_var, (args, kwargs)
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -106,6 +106,13 @@
 STATIC_DIM_FREQ_THRESHOLD = 5
 
 
+def method_to_reverse_method(method_name: str) -> str | None:
+    if not method_name.startswith("__") or not method_name.endswith("__"):
+        return None
+    name = method_name[2:-2]
+    return f"__r{name}__"
+
+
 class ConstantVariable(VariableBase):
     """
     ConstantVariable is a subclass of VariableBase used to wrap a Variable of the const type.
@@ -802,11 +809,30 @@ def get_py_value(self, allow_tensor: bool = False) -> bool | int | float:
             ), f"self.value is None, but tracker is not SymbolicOperationTracker. tracker: {self.tracker}"
             inputs = self.tracker.inputs
             assert len(inputs) >= 1
-            other_inputs_value = [x.get_py_value() for x in inputs[1:]]
-            self.value = getattr(
-                inputs[0].get_py_value(), self.tracker.method_name
-            )(*other_inputs_value)
-            assert isinstance(self.value, (bool, int, float))
+            input_values = [x.get_py_value() for x in inputs]
+            value = getattr(input_values[0], self.tracker.method_name)(
+                *input_values[1:]
+            )
+            # TODO(SigureMo): A Temporary solution for the case that the method is not implemented.
+            # e.g. In user code, we have `1 * 0.1`, the lhs is a SymbolicVariable, and the rhs is a float.
+            # We trace the method `__mul__` from the lhs, but actually, python use `float.__rmul__`,
+            # `int.__mul__(float)` is not implemented. So we get NotImplemented here.
+            # We need to find a better way to handle this case.
+            if isinstance(value, type(NotImplemented)):
+                reversed_method = method_to_reverse_method(
+                    self.tracker.method_name
+                )
+                if reversed_method is None:
+                    raise InnerError(
+                        f"Unsupported method {self.tracker.method_name} for SymbolicVariable"
+                    )
+                value = getattr(input_values[1], reversed_method)(
+                    input_values[0], *input_values[2:]
+                )
+            self.value = value
+            assert isinstance(
+                self.value, (bool, int, float)
+            ), f"SymbolicVariable.get_py_value() should return bool, int or float, but got {type(self.value)}"
         return self.value
 
     def get_py_type(self):
diff --git a/test/sot/test_16_paddle_api.py b/test/sot/test_16_paddle_api.py
@@ -38,6 +38,12 @@ def paddle_api_function_call_concat(
     return paddle.concat([x, y], axis=axis)
 
 
+def paddle_api_function_breakgraph_when_type_error(
+    x: paddle.Tensor, axis: paddle.Tensor
+):
+    return paddle.nn.functional.softmax(x, axis=axis)
+
+
 class TestPaddleApiCall(TestCaseBase):
     def test_paddle_api_method_call(self):
         self.assert_results(paddle_api_method_call, paddle.to_tensor(2.0))
@@ -55,6 +61,13 @@ def test_paddle_api_function_call_concat(self):
         self.assert_results(paddle_api_function_call_concat, a, b, 0)
         self.assert_results(paddle_api_function_call_concat, a, b, 1)
 
+    def test_paddle_api_function_breakgraph_when_type_error(self):
+        x = paddle.to_tensor([[1, 2], [3, 4]], dtype=paddle.float32)
+        axis = paddle.to_tensor(1)
+        self.assert_results(
+            paddle_api_function_breakgraph_when_type_error, x, axis
+        )
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/sot/test_sot_dynamic_shape.py b/test/sot/test_sot_dynamic_shape.py
@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import math
 import unittest
 
 from test_case_base import (
@@ -58,6 +59,12 @@ def dynamic_shape_in_list(x, shape):
     return x.reshape(shape)
 
 
+def dynamic_shape_int_mul_float(x):
+    y = x * 0.5
+    z = math.sin(y)  # Trigger get_py_value
+    return z
+
+
 class CustomConv(paddle.nn.Conv2D):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -193,6 +200,13 @@ def test_pad_dynamic_shape_fallback(self):
                 self.assert_results(pad_func, paddle.randn([1, 3, 224, 224]), i)
                 self.assertEqual(ctx.translate_count, i)
 
+    def test_dynamic_shape_int_mul_float(self):
+        with allow_dynamic_shape_guard(
+            True
+        ), test_instruction_translator_cache_context() as ctx:
+            for i in range(1, 6):
+                self.assert_results(dynamic_shape_int_mul_float, i)
+
 
 if __name__ == '__main__':
     unittest.main()
