[dynamic shapes] add sym_and, sym_or (pytorch#150456)

This has been pretty helpful for the size-oblivious rewrite. Wanted the variadic args version to avoid `sym_or(a, sym_or(b, sym_or(c, d)))` in favor of `sym_or(a, b, c, d)`. Happy to change this to ban the 1-arg version.

This is better than plain and/or because the whole symbolic expression gets preserved, and if we guard on it or defer as a runtime assert, we preserve all branches.

Pull Request resolved: pytorch#150456
Approved by: https://github.com/laithsakka

Author: pianpwk

docs/source/fx.experimental.rst

@@ -43,7 +43,9 @@ torch.fx.experimental.symbolic_shapes
     guard_or_true
     guard_or_false
     guard_size_oblivious
+    sym_and
     sym_eq
+    sym_or
     constrain_range
     constrain_unify
     canonicalize_bool_expr

test/export/test_export.py

@@ -6990,6 +6990,29 @@ def forward(self, a, b, mode):
             _ = exported.module()(torch.randn(4, 4), torch.randn(4), "floor")
         self.assertTrue(torch.allclose(exported.module()(*inps), foo(*inps)))
 
+    def test_sym_or_sym_and(self):
+        from torch.fx.experimental.symbolic_shapes import sym_and, sym_or
+
+        class Foo(torch.nn.Module):
+            def forward(self, xs):
+                u0, u1, u2 = xs.tolist()
+                torch._check(sym_or(u0 == 2, u0 == 4, u0 == 6))
+                torch._check(sym_and(u1 >= 4, u1 <= 8, u2 == 5))
+                return u0 + u1 + u2
+
+        ep = export(Foo(), (torch.tensor([2, 6, 5]),), strict=False)
+        ep.module()(torch.tensor([2, 6, 5]))
+        ep.module()(torch.tensor([4, 7, 5]))
+        ep.module()(torch.tensor([6, 5, 5]))
+        with self.assertRaisesRegex(
+            RuntimeError, r".* expression Eq\(u0, 2\) \| Eq\(u0, 4\) \| Eq\(u0, 6\) .*"
+        ):
+            ep.module()(torch.tensor([3, 6, 5]))
+        with self.assertRaisesRegex(
+            RuntimeError, r".* expression Eq\(u2, 5\) & \(4 <= u1\) & \(u1 <= 8\) .*"
+        ):
+            ep.module()(torch.tensor([6, 6, 6]))
+
     def test_redundant_assert_max_upper_bound(self):
         class M(torch.nn.Module):
             def forward(self, x):

torch/_export/serde/serialize.py

@@ -183,6 +183,8 @@ def _reverse_map(d: dict[Any, Enum]):
     operator.gt,
     operator.neg,
     operator.pos,
+    operator.and_,
+    operator.or_,
     math.trunc,
     torch.sym_not,
     operator.mul,

torch/fx/experimental/symbolic_shapes.py

@@ -143,7 +143,9 @@ class PendingUnbackedSymbolNotFound(RuntimeError):
     "CURRENT_NODE_KEY",
     "has_free_symbols",
     "has_free_unbacked_symbols",
+    "sym_and",
     "sym_eq",
+    "sym_or",
     "SymbolicContext",
     "StatelessSymbolicContext",
     "StatefulSymbolicContext",
@@ -1299,17 +1301,19 @@ def statically_known_true(x: Union[bool, SymBool]) -> bool:
         return result
 
 
-# When a or b is evaluated, a is evaluated eagerly first then b. This causes
-# a data dependent error for an expression “if u0==1 or True”. or over guarding for
-# “if s0==1 or True”.
-
-# On the other hand, when we use operator.or_, then dynamo will generate
-# a sympy expression Sympy.Or(u0==1, True) without evaluating the args first.
-
-# When the whole expression is passed to evaluation in that case, we do not throw a
-# data dependent error or guard because we can statically know the result is True
-# before unpacking the symbols.
-sym_or = operator.or_
+def sym_and(
+    x: Union[bool, SymBool], *others: Union[bool, SymBool]
+) -> Union[bool, SymBool]:
+    """
+    and, but for symbolic expressions, without bool casting.
+    """
+    assert isinstance(x, (bool, SymBool))
+    if len(others) == 0:
+        return x
+    for y in others:
+        assert isinstance(y, (bool, SymBool))
+        x = operator.and_(x, y)
+    return x
 
 
 def sym_eq(x: _T, y: _T) -> Union[bool, SymBool]:
@@ -1329,6 +1333,21 @@ def sym_eq(x: _T, y: _T) -> Union[bool, SymBool]:
         raise AssertionError(f"unexpected sym_eq between {type(x)} {type(y)}")
 
 
+def sym_or(
+    x: Union[bool, SymBool], *others: Union[bool, SymBool]
+) -> Union[bool, SymBool]:
+    """
+    or, but for symbolic expressions, without bool casting.
+    """
+    assert isinstance(x, (bool, SymBool))
+    if len(others) == 0:
+        return x
+    for y in others:
+        assert isinstance(y, (bool, SymBool))
+        x = operator.or_(x, y)
+    return x
+
+
 def guard_scalar(
     a: Union[SymBool, SymInt, SymFloat, int, bool, float]
 ) -> Union[bool, int, float]: