Bug fixes for disabling 0/1 specialization on plain int (#129961)

ezyang · pytorchmergebot · commit d7680a564b65 · 2024-07-02T23:19:48.000Z
These bug fixes will be exercised in #128327 but I separate them from the actual policy change (which is more risky) Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: #129961 Approved by: https://github.com/lezcano
diff --git a/test/test_dynamic_shapes.py b/test/test_dynamic_shapes.py
@@ -1256,18 +1256,12 @@ def get_constant_bool(self, val):
     def test_symnode_hashing(self):
         shape_env = ShapeEnv()
 
-        # SymInt, SymBool, SymFloat are unhashable
-        unhashable = (
-            create_symint(shape_env, 3),
-            create_symbool(shape_env, True),
-            # We should be passing in float here, but create_symbol currently
-            # only supports int
-            create_symfloat(shape_env, 3.0),
-        )
-
-        for x in unhashable:
-            with self.assertRaisesRegex(TypeError, "unhashable"):
-                hash(x)
+        # These all trigger specialization when hashed
+        hash(create_symint(shape_env, 3))
+        hash(create_symbool(shape_env, True))
+        # We should be passing in float here, but create_symbol currently
+        # only supports int
+        hash(create_symfloat(shape_env, 3.0))
 
         # NestedInt (SymInt), constant SymBool, SymNode are hashable
         j1 = torch._C._get_nested_int(1, 1)
diff --git a/torch/__init__.py b/torch/__init__.py
@@ -511,8 +511,8 @@ def __hash__(self) -> builtins.int:
         if self.node.is_nested_int():
             return hash(self.node.nested_int())
         else:
-            # We could support constant SymInts as well, but not doing it for now
-            raise TypeError("unhashable type: non-nested SymInt")
+            # Force specialization
+            return hash(builtins.int(self))
 
 
 class SymFloat:
@@ -550,6 +550,9 @@ def __rfloordiv__(self, other):
     def __bool__(self):
         return self.node.bool_()
 
+    def __float__(self):
+        return self.node.guard_float("", 0)
+
     # Symbolic power does NOT work with negative base, this is to avoid
     # potential complex outputs
     def __pow__(self, other):
@@ -612,6 +615,13 @@ def is_integer(self):
     def __repr__(self):
         return self.node.str()
 
+    def __hash__(self):
+        if self.node.is_constant():
+            return hash(self.node.float_())
+        else:
+            # Force specialization
+            return hash(builtins.float(self))
+
 
 class SymBool:
     """
@@ -674,7 +684,8 @@ def __hash__(self):
         if self.node.is_constant():
             return hash(self.node.bool_())
         else:
-            raise TypeError("unhashable type: SymBool")
+            # Force specialization
+            return hash(builtins.bool(self))
 
 
 def sym_not(a):
diff --git a/torch/_dynamo/variables/lists.py b/torch/_dynamo/variables/lists.py
@@ -324,6 +324,8 @@ def call_method(
         args: List["VariableTracker"],
         kwargs: Dict[str, "VariableTracker"],
     ) -> "VariableTracker":
+        from .tensor import SymNodeVariable
+
         if name == "append" and self.mutable_local:
             assert not kwargs
             (arg,) = args
@@ -345,7 +347,10 @@ def call_method(
         elif name == "insert" and self.mutable_local:
             assert not kwargs
             idx, value = args
-            const_idx = idx.as_python_constant()
+            if isinstance(idx, SymNodeVariable):
+                const_idx = idx.evaluate_expr()
+            else:
+                const_idx = idx.as_python_constant()
             tx.output.side_effects.mutation(self)
             self.items.insert(const_idx, value)
             return ConstantVariable.create(None)
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -764,7 +764,12 @@ def squeeze(x, dim=None):
     if dim is None:
         return TensorBox(SqueezeView.create(x.data))
 
-    dim = canonicalize_dims(len(x.get_size()), dim)
+    dim = (
+        V.graph.sizevars.evaluate_static_shape(dim)
+        if isinstance(dim, (int, sympy.Expr))
+        else tuple(V.graph.sizevars.evaluate_static_shape(d) for d in dim)
+    )
+    dim = canonicalize_dims(len(x.get_size()), dim)  # type: ignore[call-overload]
     dims = set((dim,) if not isinstance(dim, tuple) else dim)
 
     new_shape = []
@@ -1589,7 +1594,7 @@ def unsqueeze_(x, dim):
 
 
 def _validate_dim(x, dim, offset=0):
-    assert isinstance(dim, int)
+    dim = V.graph.sizevars.shape_env.evaluate_expr(sympy.sympify(dim))
     ndim = len(x.get_size())
     if dim < 0:
         dim += ndim + offset
diff --git a/torch/_inductor/sizevars.py b/torch/_inductor/sizevars.py
@@ -448,7 +448,9 @@ def evaluate_max(self, left: Expr, right: Expr) -> Expr:
         min_val = self.evaluate_min(left, right)
         return right if min_val is left else left
 
-    def evaluate_static_shape(self, left: Expr) -> int:
+    def evaluate_static_shape(self, left: Union[Expr, int]) -> int:
+        if isinstance(left, int):
+            return left
         right = self.size_hint(left)
         self.guard_equals(left, sympy.Integer(right))
         return int(right)
@@ -461,7 +463,9 @@ def remove_precomputed_replacements(self, expr: Expr) -> Expr:
             return sympy_subs(expr, self.inv_precomputed_replacements)  # type: ignore[arg-type]
         return expr
 
-    def symbolic_hint(self, expr: Expr) -> Union[Expr, int]:
+    def symbolic_hint(self, expr: Union[Expr, int]) -> Union[Expr, int]:
+        if isinstance(expr, int):
+            return expr
         # Substitute all hints into expr, but leave unbacked symints alone
         expr = self.simplify(expr)
         if not isinstance(expr, Expr):
@@ -476,7 +480,9 @@ def symbolic_hint(self, expr: Expr) -> Union[Expr, int]:
         expr = self.remove_precomputed_replacements(expr)
         return sympy_subs(expr, self.var_to_val)
 
-    def size_hint(self, expr: Expr, *, fallback: Optional[int] = None) -> int:
+    def size_hint(
+        self, expr: Union[Expr, int], *, fallback: Optional[int] = None
+    ) -> int:
         out = self.symbolic_hint(expr)
         if not isinstance(out, (int, sympy.Integer)) and fallback is not None:
             # Use the provided heuristic fallback hint
