fix pickling for BitwiseFn (pytorch#163861)

* fix pickling for BitwiseFn (pytorch#163571)

Summary:
ran into AttributeError: Can't get local object 'make_opaque_bitwise_fn.<locals>.BitwiseFn'

looks like it was fixed for UnaryFn but not BitwiseFn in pytorch#138395

Fixes pytorch#147841

Pull Request resolved: pytorch#163571
Approved by: https://github.com/jamesjwu

(cherry picked from commit cde5c9a)

* Fix lintrunner with -a

---------

Co-authored-by: dolpm <[email protected]>
Co-authored-by: Lucas Kabela <[email protected]>

Author: 3 people

test/inductor/test_compile_subprocess.py

@@ -206,7 +206,8 @@ def model_add(x, y):
 
             start = time.time()
             last_report = start
-            while _AsyncFxCompile._stat_compiled_runs < 4:
+            while True:
+                start_stat_compiled_runs = _AsyncFxCompile._stat_compiled_runs
                 # Sleep a bit so we don't drive the CPU unnecessarily.
                 time.sleep(0.25)
 
@@ -219,6 +220,9 @@ def model_add(x, y):
                 # Backward pass
                 output.sum().backward()
 
+                if _AsyncFxCompile._stat_compiled_runs - start_stat_compiled_runs == 2:
+                    break
+
                 # DEBUGGING: Print a periodic message so we know we're still
                 # running...
                 now = time.time()
@@ -231,12 +235,12 @@ def model_add(x, y):
                         "Test timed out before producing a compiled artifact."
                     )
 
-            self.assertEqual(_AsyncFxCompile._stat_compiled_runs, 4)
+            self.assertGreater(_AsyncFxCompile._stat_compiled_runs, 1)
             # Make sure we ran eager at least once. Normally this will be
             # something like 80.
             self.assertGreater(_AsyncFxCompile._stat_eager_runs, 0)
-            self.assertEqual(_AsyncFxCompile._stat_bg_started, 1)
-            self.assertEqual(_AsyncFxCompile._stat_bg_finished, 1)
+            self.assertEqual(_AsyncFxCompile._stat_bg_started, 2)
+            self.assertEqual(_AsyncFxCompile._stat_bg_finished, 2)
 
 
 if RUN_CPU:

test/test_sympy_utils.py

@@ -5,7 +5,7 @@
 import math
 import pickle
 import sys
-from typing import Callable
+from collections.abc import Callable
 
 import sympy
 
@@ -24,6 +24,7 @@
     FloorDiv,
     Identity,
     OpaqueUnaryFn_cos,
+    BitwiseFn_bitwise_and,
     simple_floordiv_gcd,
 )
 from torch.utils._sympy.interp import sympy_interp
@@ -426,7 +427,7 @@ def test_interp(self, fn):
                 # Yes, I know this is a long-winded way of saying xreplace; the
                 # point is to test sympy_interp
                 r = sympy_interp(
-                    ReferenceAnalysis, dict(zip(symbols, sargs)), sympy_expr
+                    ReferenceAnalysis, dict(zip(symbols, sargs, strict=False)), sympy_expr
                 )
                 self.assertEqual(ref_r, r)
 
@@ -501,7 +502,7 @@ def trace_f(px, py):
 
                 self.assertEqual(
                     sympy_interp(
-                        PythonReferenceAnalysis, dict(zip(symbols, args)), sympy_expr
+                        PythonReferenceAnalysis, dict(zip(symbols, args, strict=False)), sympy_expr
                     ),
                     gm(*args),
                 )
@@ -555,7 +556,7 @@ def test_tensor_interp(self, fn):
                     direct_result = tensor_fn(*tensor_args)
                     interp_result = sympy_interp(
                         TensorReferenceAnalysis,
-                        dict(zip(symbols, tensor_args)),
+                        dict(zip(symbols, tensor_args, strict=False)),
                         sympy_expr,
                     )
 
@@ -873,6 +874,10 @@ def test_pickle(self):
         r = pickle.loads(pickle.dumps(x))
         self.assertEqual(x, r)
 
+        x = BitwiseFn_bitwise_and(sympy.Symbol("a"), sympy.Symbol("b"))
+        r = pickle.loads(pickle.dumps(x))
+        self.assertEqual(x, r)
+
 
 class TestSingletonInt(TestCase):
     def test_basic(self):

torch/utils/_sympy/functions.py

@@ -3,7 +3,8 @@
 import math
 import operator
 import sys
-from typing import Callable, Optional, SupportsFloat, TYPE_CHECKING, TypeVar, Union
+from collections.abc import Callable
+from typing import Optional, SupportsFloat, TYPE_CHECKING, TypeVar, Union
 from typing_extensions import TypeVarTuple, Unpack
 
 import sympy
@@ -1192,7 +1193,8 @@ def eval(cls, *args):
             # When all strides are integral, we can sort, and the size for the
             # largest stride doesn't matter and can be arbitrarily symbolic
             s_sizes, s_strides = zip(
-                *sorted(zip(sizes, strides), key=operator.itemgetter(1))
+                *sorted(zip(sizes, strides, strict=False), key=operator.itemgetter(1)),
+                strict=False,
             )
             # Put something arbitrary in the max size spot, it'll be ignored
             if all(isinstance(a, sympy.Integer) for a in s_sizes[:-1]):
@@ -1411,7 +1413,10 @@ def eval(cls, a, b):
                 return sympy.Integer(getattr(operator, real_op_name)(int(a), int(b)))
             return None
 
-    BitwiseFn.__name__ = "BitwiseFn_" + name
+    nm = "BitwiseFn_" + name
+    BitwiseFn.__name__ = nm
+    BitwiseFn.__qualname__ = nm
+
     return BitwiseFn