[SOT][DynamicShape] Add basic constraint mechanism (#72250)

Author: SigureMo

python/paddle/jit/sot/infer_meta.py

@@ -38,7 +38,7 @@
 from paddle.static import InputSpec
 from paddle.utils import flatten, is_sequence
 
-from .symbolic_shape import SymbolicInt
+from .symbolic_shape.symbolic_value import SymbolicInt
 from .utils import (
     Cache,
     Singleton,

python/paddle/jit/sot/opcode_translator/executor/function_graph.py

@@ -39,7 +39,7 @@
 from ...profiler import EventGuard, event_register
 from ...symbolic.statement_ir import Reference, StatementIR, Symbol
 from ...symbolic.symbolic_context import SymbolicTraceContext
-from ...symbolic_shape import SYMBOLIC_BINARY_OPS, SYMBOLIC_UNARY_OPS
+from ...symbolic_shape.operators import SYMBOLIC_BINARY_OPS, SYMBOLIC_UNARY_OPS
 from ...utils import (
     ENV_SOT_ALLOW_DYNAMIC_SHAPE,
     ENV_SOT_ENABLE_GUARD_TREE,

python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py

@@ -31,10 +31,12 @@
 
 from ...profiler import EventGuard
 from ...psdb import NO_BREAKGRAPH_CODES
+from ...symbolic_shape.constraints import LogicalNotConstraintNode
 from ...utils import (
     ENV_MIN_GRAPH_SIZE,
     ENV_SOT_FORCE_FALLBACK_SIR_IDS,
     BreakGraphError,
+    DataDependencyDynamicShapeBreak,
     FallbackError,
     InnerError,
     SotUndefinedVar,
@@ -224,12 +226,23 @@ def inner(self: OpcodeExecutorBase, instr: Instruction):
                 )(res)
 
             assert isinstance(res, (ConstantVariable, SymbolicVariable))
-            # NOTE(SigureMo): force to constant to trigger fallback to static dim
-            # to align with old behavior. In next PR we will support guard value
-            # with constraint.
             if isinstance(res, SymbolicVariable):
-                res = res.to_constant()
-            is_jump = res.get_py_value()
+                constraint_node, symbolic_vars = res.create_constraint_tree()
+                if not all(
+                    var.value.is_backed() for var in symbolic_vars.values()
+                ):
+                    raise BreakGraphError(
+                        DataDependencyDynamicShapeBreak(
+                            f"Symbolic variable {symbolic_vars} is not backed."
+                        )
+                    )
+                is_jump = res.get_example_value()
+                if not is_jump:
+                    constraint_node = LogicalNotConstraintNode(constraint_node)
+                for var in symbolic_vars.values():
+                    var.add_constraint((constraint_node, symbolic_vars))
+            else:
+                is_jump = res.get_py_value()
             assert isinstance(is_jump, bool)
             if is_jump:
                 assert instr.jump_to is not None

python/paddle/jit/sot/opcode_translator/executor/side_effects.py

@@ -79,9 +79,9 @@ def restore_state(self, state: SideEffectsState):
         assert len(self.data_id_to_proxy.values()) == len(
             state.proxy_versions
         ), "proxy_versions length not match"
-        assert len(self.mutable_variables) == len(
-            state.mutable_attrs
-        ), "mutable_attrs length not match"
+        assert sum(
+            len(var.mutable_attrs) for var in self.mutable_variables
+        ) == len(state.mutable_attrs), "mutable_attrs length not match"
 
         for proxy, version in zip(
             self.data_id_to_proxy.values(), state.proxy_versions

python/paddle/jit/sot/opcode_translator/executor/tracker.py

@@ -137,11 +137,9 @@ def __init__(self, inputs: Sequence[VariableBase], op: UnaryOp | BinaryOp):
         self.op = op
 
     def gen_instructions(self, codegen: PyCodeGen):
-        # TODO(zrr1999): implemented in #68555
         raise InnerError("SymbolicOperationTracker has no instructions")
 
     def trace_value_from_frame(self):
-        # TODO(zrr1999): implemented in #68555
         raise InnerError(
             "SymbolicOperationTracker can't trace value from frame"
         )
@@ -150,11 +148,6 @@ def __repr__(self) -> str:
         return f"SymbolicOperationTracker(num_inputs={len(self.inputs)})"
 
     def is_traceable(self):
-        # TODO(zrr1999): implemented in #68555
-        return False
-
-    def need_guard(self) -> bool:
-        # TODO(zrr1999): implemented in #68555
         return False
 
 

python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py

@@ -24,7 +24,7 @@
 
 import paddle
 
-from ...symbolic_shape import (
+from ...symbolic_shape.operators import (
     SYMBOLIC_BINARY_OPS,
     SYMBOLIC_UNARY_OPS,
     symbolic_not,

python/paddle/jit/sot/opcode_translator/executor/variables/base.py

@@ -275,7 +275,7 @@ def _find_var(key: str = "default") -> VariableBase | None:
         return var
 
 
-def infer_debug_name_from_tracker(tracker: Tracker) -> str:
+def infer_debug_name_from_tracker(tracker: Tracker) -> str | None:
     res = None
     if isinstance(tracker, (LocalTracker, GlobalTracker, BuiltinTracker)):
         res = f"{tracker.name}"