Update on "Add support for nonzero, some improvements to reduce guards"

ezyang · ezyang · commit 8a25bc652b27 · 2023-02-23T11:51:25.000-08:00
This takes the strategy described in https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit# It is essentially #95222 but squashed and with changes that are unnecessary given that we assume nonzero returns > 1. What's in the PR: * nonzero now supports meta propagation. When `capture_dynamic_output_shape_ops`, it will return a tensor with an unbacked SymInt representing the size in question. * The unbacked SymInt is UNSOUNDLY assumed to be not equal to 0/1. We will still error if you guard otherwise. * PrimTorch pointwise operators are updated to use empty_permuted, to avoid guarding on unbacked SymInt from empty_strided (tested in `test_dynamic_pointwise_scalar`) * Convolution is updated to skip backend selection if batch is unbacked, to avoid guarding on unbacked SymInt (tested in `test_unbacked_batch_resnet`) * I kept the helper utilities like `definitely_true` for working with possibly unbacked SymInts. They're not used right now but maybe someone will find them useful. * Added `constrain_unify` to let you specify two unbacked SymInts must have the same value Signed-off-by: Edward Z. Yang <ezyangmeta.com> cc soumith voznesenskym yanboliang penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx desertfire [ghstack-poisoned]
diff --git a/torch/_decomp/decompositions.py b/torch/_decomp/decompositions.py
@@ -19,11 +19,7 @@
     _safe_copy_out,
     out_wrapper,
 )
-from torch.fx.experimental.symbolic_shapes import (
-    definitely_true,
-    guard_int,
-    tensor_has_hints,
-)
+from torch.fx.experimental.symbolic_shapes import guard_int
 from torch.utils._pytree import tree_flatten, tree_map
 
 DispatchKey = torch._C.DispatchKey  # type: ignore[attr-defined]
diff --git a/torch/_prims_common/__init__.py b/torch/_prims_common/__init__.py
@@ -77,6 +77,7 @@ def getnvFuserDtype(dtype: Union[torch.dtype, NumberTypeType]):
     torch.Tensor.device.__get__,  # type: ignore[attr-defined]
     torch.Tensor.requires_grad.__get__,  # type: ignore[attr-defined]
     torch.Tensor.layout.__get__,  # type: ignore[attr-defined]
+    torch.Tensor.is_contiguous,
     # For TorchRefsMode only
     torch.Tensor.__format__,
     torch.Tensor.__repr__,
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -426,11 +426,22 @@ def nonzero(fake_mode, func, arg):
         raise DynamicOutputShapeException(func)
     nnz = fake_mode.shape_env.create_unbacked_symint()
 
-    from torch.fx.experimental.symbolic_shapes import constrain_range
+    from torch.fx.experimental.symbolic_shapes import (
+        constrain_range,
+        definitely_true,
+        guard_int,
+    )
+
     # This is unsound, but it works well in practice
     # See https://docs.google.com/document/d/1lFRYAJo5nrfxRhwIzGnfi2pbLpU6T4ytSRSuLJ5qebI/edit#
     # TODO: Add a config knob to turn off this unsound behavior
-    constrain_range(nnz, min=min(2, arg.numel()))
+    lower = 2
+    upper = None
+    # But don't give totally unsatisfiable bounds if we know it's too small!
+    if definitely_true(arg.numel() < 2):
+        lower = 0
+        upper = guard_int(arg.numel())
+    constrain_range(nnz, min=lower, max=upper)
 
     return arg.new_empty((nnz, arg.dim()), dtype=torch.int64)
 
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
@@ -1355,12 +1355,10 @@ def create_symbol(self, val: int, source: Source, dyn=False) -> "sympy.Expr":
             if not dyn:
                 # Non explicitly marked dynamic dims register to val_to_var to get duck shaped
                 self.val_to_var[val] = sympy_expr
-                # We also infer that they must not be 0/1
-                lower = 2 if self.specialize_zero_one else 0
-                self.var_to_range[sympy_expr] = ValueRanges(lower, sympy.oo)
-            else:
-                # Avoid up front 0/1 specializing dynamic dims
-                self.var_to_range[sympy_expr] = ValueRanges(0, sympy.oo)
+
+            # We also infer that it must be not 0/1
+            lower = 2 if self.specialize_zero_one else 0
+            self.var_to_range[sympy_expr] = ValueRanges(lower, sympy.oo)
 
         if not dyn and self.duck_shape:
             # This implements duck-shaping: input sizes that match are assigned
