Update on "Grab Current Tracing Fake Mode in a couple spots"

eellison · eellison · commit 2e98e29dce55 · 2023-04-18T01:40:37.000Z
Fix for #99286. There were a couple locations we were instantiating new fake modes instead of grabbing the correct one from the current tracing context/inputs. cc soumith voznesenskym penguinwu anijain2305 EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx peterbell10 desertfire [ghstack-poisoned]
diff --git a/.github/workflows/_mac-build.yml b/.github/workflows/_mac-build.yml
@@ -174,9 +174,6 @@ jobs:
             export PATH="$CONDA_ENV/bin":$PATH
           fi
 
-          # NB: Same trick as Linux, there is no need to initialize sccache with the risk of getting
-          # it hangs or timeout at initialization. The cache will be started automatically
-          export SKIP_SCCACHE_INITIALIZATION=1
           ${CONDA_RUN} .ci/pytorch/macos-build.sh
 
       - name: Archive artifacts into zip
diff --git a/aten/src/ATen/native/TensorShape.cpp b/aten/src/ATen/native/TensorShape.cpp
@@ -2807,7 +2807,7 @@ static std::vector<Tensor> reshape_input_for_column_stack(TensorList tensors) {
   auto transform_lambda = [](const Tensor& input) -> Tensor {
     // reshape 0D or 1D tensor t into (t.numel(), 1)
     if (input.dim() <= 1) {
-      return input.reshape_symint({input.sym_numel(), 1});
+      return input.reshape({input.numel(), 1});
     }
     return input;
   };
diff --git a/c10/core/Scalar.h b/c10/core/Scalar.h
@@ -171,7 +171,7 @@ class C10_API Scalar {
   }
 
   bool isSymbolic() const {
-    return Tag::HAS_si == tag || Tag::HAS_sd == tag;
+    return Tag::HAS_si == tag || Tag::HAS_sd == tag || Tag::HAS_sb == tag;
   }
 
   C10_ALWAYS_INLINE Scalar& operator=(Scalar&& other) noexcept {
diff --git a/test/cpp_extensions/open_registration_extension.cpp b/test/cpp_extensions/open_registration_extension.cpp
@@ -2,16 +2,27 @@
 #include <c10/core/Allocator.h>
 
 #include <torch/csrc/Device.h>
+#include <c10/core/impl/DeviceGuardImplInterface.h>
+#include <c10/macros/Macros.h>
 #include <torch/extension.h>
 
 #include <ATen/native/cpu/Loops.h>
 #include <ATen/native/DispatchStub.h>
+#include <ATen/native/Resize.h>
 #include <ATen/EmptyTensor.h>
 #include <ATen/core/GeneratorForPrivateuseone.h>
 
 static uint64_t add_counter = 0;
 static uint64_t last_saved_value = 0;
 
+// register guard
+namespace at {
+namespace detail {
+
+C10_REGISTER_GUARD_IMPL(PrivateUse1, c10::impl::NoOpDeviceGuardImpl<DeviceType::PrivateUse1>);
+
+}} // namespace at::detail
+
 // basic dummy add function
 at::Tensor custom_add_Tensor(const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
   add_counter += 1;
@@ -79,6 +90,16 @@ at::Tensor custom_empty_strided(c10::IntArrayRef size, c10::IntArrayRef stride,
   return  at::detail::empty_strided_generic(size, stride, &global_custom_alloc, private_use_ks, dtype);
 }
 
+// Some set operations for the basic use case
+at::Tensor& custom_set_source_Storage(at::Tensor& result, c10::Storage src) {
+  int64_t new_size = static_cast<int64_t>(src.nbytes() / result.dtype().itemsize());
+  c10::IntArrayRef stride = {};
+  result.unsafeGetTensorImpl()->set_storage_offset(0);
+  at::OptionalIntArrayRef stride_opt = stride.data() != nullptr ? at::OptionalIntArrayRef(stride) : c10::nullopt;
+  at::native::resize_impl_cpu_(result.unsafeGetTensorImpl(), new_size, stride_opt, /*resize_storage=*/!result.is_meta());
+  return result;
+}
+
 // This macro does the heavy lifting.
 // With TORCH_LIBRARY_IMPL, you can register custom kernels for your backend.
 // For open registration, we're registering all of our kernels to the PrivateUse1 dispatch key.
@@ -94,6 +115,7 @@ TORCH_LIBRARY_IMPL(aten, PrivateUse1, m) {
   m.impl("fill_.Scalar", &custom_fill__scalar);
   m.impl("_copy_from", &custom__copy_from);
   m.impl("empty_strided", &custom_empty_strided);
+  m.impl("set_.source_Storage", &custom_set_source_Storage);
 }
 
 // This basic implementation doesn't bother dealing with different device indices
diff --git a/test/dynamo/test_export.py b/test/dynamo/test_export.py
@@ -2496,6 +2496,25 @@ def forward(self, input):
             tracing_mode="real",
         )
 
+    @config.patch(
+        dynamic_shapes=True,
+        capture_dynamic_output_shape_ops=True,
+        capture_scalar_outputs=True,
+        assume_static_by_default=False,
+    )
+    def test_sym_contains(self):
+        def f(x, y):
+            return x.size(0) in y
+
+        gm, _ = torch._dynamo.export(
+            f, torch.ones(2), torch.ones(3), aten_graph=True, tracing_mode="symbolic"
+        )
+
+        true_inp = (torch.Tensor([6, 4, 5]), torch.ones(6, 4).add_(5))
+        false_inp = (torch.Tensor([6, 4, 5]), torch.ones(6, 4).add_(2))
+        self.assertEqual(gm(*true_inp), f(*true_inp))
+        self.assertEqual(gm(*false_inp), f(*false_inp))
+
 
 common_utils.instantiate_parametrized_tests(ExportTests)
 
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
@@ -2509,7 +2509,9 @@ def forward(self, x):
     xfail('cdist', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('cholesky_inverse', ''),  # could not find kernel
     xfail('cholesky_solve', ''),  # could not find kernel
+    xfail('column_stack', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('combinations', ''),  # aten.masked_select.default
+    xfail('cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('cumprod', ''),  # aten.cumprod.default - couldn't find symbolic meta function/decomposition
     xfail('cumulative_trapezoid', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('diff', ''),  # aten.zeros_like.default - couldn't find symbolic meta function/decomposition
@@ -2525,6 +2527,7 @@ def forward(self, x):
     xfail('kthvalue', ''),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail('linalg.cholesky_ex', ''),  # could not find kernel for aten.linalg_solve_triangular.default
     xfail('linalg.cond', ''),  # Cannot call numel() on tensor with symbolic sizes/strides
+    xfail('linalg.cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.det', ''),  # aten._linalg_det.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.det', 'singular'),  # aten._linalg_det.default - couldn't find symbolic meta function/deco...
     xfail('linalg.eigh', ''),  # aten._linalg_eigh.default - couldn't find symbolic meta function/decomposition
diff --git a/test/test_cpp_extensions_open_device_registration.py b/test/test_cpp_extensions_open_device_registration.py
@@ -162,6 +162,28 @@ def test_open_device_registration(self):
         self.assertTrue(x.is_foo)
         self.assertTrue(hasattr(torch.nn.Module, 'foo'))
 
+        # check whether the attributes and methods for storage of the corresponding custom backend are generated correctly
+        torch.utils.generate_methods_for_privateuse1_backend(for_tensor=False, for_module=False, for_storage=True)
+
+        x = torch.empty(4, 4)
+        # check TypedStorage
+        z1 = x.storage()
+        self.assertFalse(z1.is_foo)
+        z1 = z1.foo()
+        self.assertFalse(self.module.custom_add_called())
+        self.assertTrue(z1.is_foo)
+        with self.assertRaisesRegex(RuntimeError, "Invalid device"):
+            z1.foo(torch.device("cpu"))
+        z1 = z1.cpu()
+
+        y = torch.empty(4, 4)
+        # check UntypedStorage
+        z2 = y.untyped_storage()
+        self.assertFalse(z2.is_foo)
+        z2 = z2.foo()
+        self.assertFalse(self.module.custom_add_called())
+        self.assertTrue(z2.is_foo)
+
     def test_open_device_random(self):
         torch.utils.rename_privateuse1_backend('foo')
         with self.assertRaisesRegex(RuntimeError, "Expected one of cpu"):
diff --git a/test/test_proxy_tensor.py b/test/test_proxy_tensor.py
@@ -1236,17 +1236,6 @@ def f(x, y):
         inp = (torch.randn(8)[3:], torch.randn(5))
         self.assertEqual(fx_g(*inp), f(*inp))
 
-    def test_sym_contains(self):
-        def f(x, y):
-            return x.size(0) in y
-
-        # This shouldn't raise but we need SymBool from
-        # https://github.com/pytorch/pytorch/pull/98453
-        # then modify this.  It should NOT raise a RuntimeError
-        # though!
-        with self.assertRaisesRegex(NotImplementedError, "item NYI for torch.bool"):
-            make_fx(f, tracing_mode="symbolic")(torch.randn(2), torch.randn(3))
-
     def _assert_no_guards(self, fx_g, free_symbols):
         assert _get_free_symbols(fx_g.shape_env) == free_symbols, fx_g.shape_env.var_to_val
         assert len(fx_g.shape_env.get_nontrivial_guards()) == 0, fx_g.shape_env.format_guards()
@@ -1369,7 +1358,9 @@ def f(a, b, c, d, e):
     xfail('linalg.eig'),
     xfail('linalg.eigvals'),
     xfail('cholesky_solve', ''),  # Could not run 'aten::_cholesky_solve_helper' with arguments from the 'Meta' back...
+    xfail('column_stack', ''),  # Tensors of type TensorImpl do not have numel
     xfail('combinations', ''),
+    xfail('cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('cumulative_trapezoid', ''),  # aten.slice.Tensor - couldn't find symbolic meta function/decomposition
     xfail('diff', ''),  # aten.empty_like.default - couldn't find symbolic meta function/decomposition
     xfail('dsplit', ''),  # aten.slice.Tensor - couldn't find symbolic meta function/decomposition
@@ -1386,6 +1377,7 @@ def f(a, b, c, d, e):
     xfail('kron', ''),  # aten.size.default - couldn't find symbolic meta function/decomposition
     xfail('kthvalue', ''),  # aten.kthvalue.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.cond', ''),  # Tensors of type TensorImpl do not have numel
+    xfail('linalg.cross', ''),  # aten.linalg_cross.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.eigh', ''),  # aten._linalg_eigh.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.eigvalsh', ''),  # aten._linalg_eigh.default - couldn't find symbolic meta function/decomposition
     xfail('linalg.householder_product', ''),  # aten.linalg_householder_product.default - couldn't find symbolic meta funct...
diff --git a/third_party/fbgemm b/third_party/fbgemm
@@ -1 +1 @@
-Subproject commit 03b2046676707da64504e898490ab46104d4682a
+Subproject commit e07dda2d50562f534abc682b0dc79a8cd5ce819d
diff --git a/torch/_dynamo/skipfiles.py b/torch/_dynamo/skipfiles.py
@@ -145,10 +145,7 @@ def add(import_name: str):
     if isinstance(import_name, types.ModuleType):
         return add(import_name.__name__)
     assert isinstance(import_name, str)
-    try:
-        module_spec = importlib.util.find_spec(import_name)
-    except Exception:
-        return
+    module_spec = importlib.util.find_spec(import_name)
     if not module_spec:
         return
     origin = module_spec.origin
@@ -191,7 +188,6 @@ def check(filename, allow_torch=False):
     "tensorflow",
     "tensorrt",
     "torch2trt",
-    "torchrec.distributed",
     "tqdm",
     "tree",
     "tvm",
diff --git a/torch/_dynamo/variables/builder.py b/torch/_dynamo/variables/builder.py
@@ -1088,7 +1088,7 @@ def _clone_input(value):
     ):
         proxy.node.meta["example_value"] = example_value
         return ConstantVariable(example_value, **options)
-    elif isinstance(example_value, (torch.SymInt, torch.SymFloat)):
+    elif isinstance(example_value, (torch.SymInt, torch.SymFloat, torch.SymBool)):
         proxy.node.meta["example_value"] = example_value
         return SymNodeVariable(proxy, example_value, **options)
     elif proxy.node.target in [torch.cuda.streams.Stream, torch.cuda.current_stream]:
diff --git a/torch/_dynamo/variables/tensor.py b/torch/_dynamo/variables/tensor.py
@@ -423,6 +423,16 @@ def call_method(
                 tx, [result, kwargs["value"]], {}
             )
             return self.call_method(tx, "add_", [result], {})
+        elif name == "__contains__":
+            # Rewrite __contains__ here so that downstream passes can trace through
+            # without dealing with unbacked symbool. Roughly the code we translate is:
+            # def __contains__(self, x):
+            #     return (x == self).any().item()
+            result = TorchVariable(torch.eq, **options).call_function(
+                tx, [self, args[0]], {}
+            )
+            result = TorchVariable(torch.any, **options).call_function(tx, [result], {})
+            return result.call_method(tx, "item", [], {})
         else:
             # Convert x.new(torch.Size) into x.new_empty(torch.Size),
             # as Tensor.new acts differently with a Size input versus a tuple input.
diff --git a/torch/_meta_registrations.py b/torch/_meta_registrations.py
@@ -88,26 +88,6 @@ def meta_take(self, index, *, out=None):
     return result
 
 
-@register_meta([aten.linalg_cross.default, aten.linalg_cross.out])
-@out_wrapper()
-def linalg_cross(self, other, *, dim=-1):
-    x_d = self.ndim
-    y_d = other.ndim
-    check(
-        x_d == y_d,
-        lambda: "linalg.cross: inputs must have the same number of dimensions.",
-    )
-    check(
-        self.size(dim) == 3 and other.size(dim) == 3,
-        lambda: (
-            f"linalg.cross: inputs dimension {dim} must have length 3. "
-            f"Got {self.size(dim)} and {other.size(dim)}"
-        ),
-    )
-    out_shape = _broadcast_shapes(self.shape, other.shape)
-    return self.new_empty(out_shape)
-
-
 @register_meta(
     [aten.cummax.default, aten.cummax.out, aten.cummin.default, aten.cummin.out]
 )
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
@@ -15,6 +15,7 @@
 from torch._prims_common import (
     elementwise_dtypes,
     ELEMENTWISE_TYPE_PROMOTION_KIND,
+    is_boolean_dtype,
     is_float_dtype,
     is_integer_dtype,
 )
@@ -428,6 +429,8 @@ def local_scalar_dense(fake_mode, func, arg):
         return fake_mode.shape_env.create_unbacked_symfloat()
     elif is_integer_dtype(arg.dtype):
         return fake_mode.shape_env.create_unbacked_symint()
+    elif is_boolean_dtype(arg.dtype):
+        return fake_mode.shape_env.create_unbacked_symbool()
     else:
         raise NotImplementedError(f"local_scalar_dense/item NYI for {arg.dtype}")
 
diff --git a/torch/ao/quantization/backend_config/executorch.py b/torch/ao/quantization/backend_config/executorch.py
@@ -112,6 +112,13 @@ def _get_linear_configs() -> List[BackendPatternConfig]:
             .set_root_module(torch.nn.Linear)
             .set_reference_quantized_module(nnqr.Linear)
             .set_qat_module(nnqat.Linear))
+    # linear qat module
+    linear_configs.append(
+        BackendPatternConfig(nnqat.Linear)
+            .set_observation_type(observation_type)  # noqa: E131
+            .set_dtype_configs(dtype_configs)
+            .set_root_module(torch.nn.Linear)
+            .set_reference_quantized_module(nnqr.Linear))
     # functional linear
     linear_configs.append(
         BackendPatternConfig(torch.nn.functional.linear)
@@ -195,6 +202,33 @@ def _get_conv_configs() -> List[BackendPatternConfig]:
                 .set_fuser_method(fuse_conv_bn_relu)
                 .set_fused_module(convs.fused_conv_bn_relu))
         # TODO: we can add fusion for torch.relu as well
+        # 3.2 conv + bn (+ relu) fused module configs
+        # fused conv bn
+        conv_configs.append(
+            BackendPatternConfig(convs.fused_conv_bn)
+                .set_dtype_configs(dtype_configs)  # noqa: E131
+                .set_qat_module(convs.bn_qat))
+
+        # fused conv bn relu
+        conv_configs.append(
+            BackendPatternConfig(convs.fused_conv_bn_relu)
+                .set_dtype_configs(dtype_configs)  # noqa: E131
+                .set_qat_module(convs.bn_relu_qat))
+
+        # conv bn, qat fused module
+        conv_configs.append(
+            BackendPatternConfig(convs.bn_qat)
+                .set_observation_type(observation_type)  # noqa: E131
+                .set_dtype_configs(dtype_configs)
+                .set_root_module(convs.root)
+                .set_reference_quantized_module(convs.reference))
+        # conv bn relu, qat fused module
+        conv_configs.append(
+            BackendPatternConfig(convs.bn_relu_qat)
+                .set_observation_type(observation_type)  # noqa: E131
+                .set_dtype_configs(dtype_configs)
+                .set_root_module(convs.root)
+                .set_reference_quantized_module(convs.reference))
     return conv_configs
 
 def _get_binary_ops_configs() -> List[BackendPatternConfig]:
@@ -242,7 +276,9 @@ def _get_share_qparams_ops_configs() -> List[BackendPatternConfig]:
         F.relu,
         F.relu6,
         torch.nn.AdaptiveAvgPool2d,
+        torch.nn.ReLU6,
         torch.nn.Hardtanh,
+        torch.mean,
         torch.squeeze,
         "mean",
         "permute",
diff --git a/torch/ao/quantization/fx/prepare.py b/torch/ao/quantization/fx/prepare.py
@@ -633,6 +633,9 @@ def _maybe_insert_input_observer_for_arg_or_kwarg(
                 # input etc.
                 arg_as_input_target_is_dynamic and arg is node.args[0]
             )
+        # if node.meta["target_dtype_info"]["weight_obs_or_fq_ctr"] == None
+        # or node.meta["target_dtype_info"]["input_act_obs_or_fq_ctr"] == None
+        needs_obs = needs_obs and (act_post_process_ctr is not None)
 
     else:
         assert qconfig is not None
diff --git a/torch/distributed/_spmd/distribute.py b/torch/distributed/_spmd/distribute.py
diff --git a/torch/distributed/_spmd/graph_optimization.py b/torch/distributed/_spmd/graph_optimization.py
diff --git a/torch/fx/experimental/symbolic_shapes.py b/torch/fx/experimental/symbolic_shapes.py
diff --git a/torch/utils/backend_registration.py b/torch/utils/backend_registration.py

Original file line number	Diff line number	Diff line change
`@@ -2807,7 +2807,7 @@ static std::vector<Tensor> reshape_input_for_column_stack(TensorList tensors) {`
`2807`	`2807`	`auto transform_lambda = [](const Tensor& input) -> Tensor {`
`2808`	`2808`	`// reshape 0D or 1D tensor t into (t.numel(), 1)`
`2809`	`2809`	`if (input.dim() <= 1) {`
`2810`		`- return input.reshape_symint({input.sym_numel(), 1});`
	`2810`	`+ return input.reshape({input.numel(), 1});`
`2811`	`2811`	`}`
`2812`	`2812`	`return input;`
`2813`	`2813`	`};`
Original file line number	Diff line number	Diff line change
`@@ -171,7 +171,7 @@ class C10_API Scalar {`
`171`	`171`	`}`
`172`	`172`
`173`	`173`	`bool isSymbolic() const {`
`174`		`- return Tag::HAS_si == tag \|\| Tag::HAS_sd == tag;`
	`174`	`+ return Tag::HAS_si == tag \|\| Tag::HAS_sd == tag \|\| Tag::HAS_sb == tag;`
`175`	`175`	`}`
`176`	`176`
`177`	`177`	`C10_ALWAYS_INLINE Scalar& operator=(Scalar&& other) noexcept {`