Fix precedence of bitwise and/or printing (pytorch#143197)

Signed-off-by: Edward Z. Yang <[email protected]>

Pull Request resolved: pytorch#143197
Approved by: https://github.com/albanD, https://github.com/williamwen42

Author: ezyang

test/dynamo/test_repros.py

@@ -6315,6 +6315,16 @@ def make_dist_and_execute(t, d):
         for _ in range(2):
             make_dist_and_execute(torch.randn(10), SubCateg)
 
+    def test_bitwise_print_precedence(self):
+        import math
+
+        @torch.compile(fullgraph=True, dynamic=True)
+        def f(x):
+            torch._check(math.floor((x.size(0) | 3) * 4) == 12)
+            return x.sin()
+
+        f(torch.randn(2))
+
     def test_tensor_split_within_device_cm(self):
         @torch.compile(fullgraph=True)
         def split(x):

torch/utils/_sympy/functions.py

@@ -25,6 +25,7 @@
 from sympy.core.operations import LatticeOp, ShortCircuit
 from sympy.core.sorting import ordered
 from sympy.core.traversal import walk
+from sympy.printing.precedence import PRECEDENCE
 from sympy.utilities.iterables import sift
 
 from .numbers import int_oo
@@ -1350,8 +1351,16 @@ def eval(cls, a):
 
 
 def make_opaque_bitwise_fn(name, real_op_name):
+    if name == "bitwise_and":
+        prec = PRECEDENCE["BitwiseAnd"]
+    elif name == "bitwise_or":
+        prec = PRECEDENCE["BitwiseOr"]
+    else:
+        raise AssertionError(f"unrecognized {name}")
+
     class BitwiseFn(sympy.Function):
         _torch_handler_name = name
+        precedence: int = prec
 
         @classmethod
         def eval(cls, a, b):

torch/utils/_sympy/printers.py

@@ -25,10 +25,10 @@ def _print_Relational(self, expr: sympy.Expr) -> str:
         return self.stringify(expr.args, f" {expr.rel_op} ", precedence(expr))
 
     def _print_BitwiseFn_bitwise_and(self, expr: sympy.Expr) -> str:
-        return self.stringify(expr.args, " & ", PRECEDENCE["Atom"] - 0.5)
+        return self.stringify(expr.args, " & ", PRECEDENCE["BitwiseAnd"])
 
     def _print_BitwiseFn_bitwise_or(self, expr: sympy.Expr) -> str:
-        return self.stringify(expr.args, " | ", PRECEDENCE["Atom"] - 0.5)
+        return self.stringify(expr.args, " | ", PRECEDENCE["BitwiseOr"])
 
     # NB: this is OK to put here, because Mod is only defined for positive
     # numbers, and so across C/Python its behavior is consistent