Revert D26815021: Revert D26744062: Add assert_async

Test Plan: revert-hammer

Differential Revision:
D26815021

Original commit changeset: 972eaafcdf14

fbshipit-source-id: e528260e1aa91df1873c73af00aa57addd671607

Author: ezyang

aten/src/ATen/native/TensorCompare.cpp

@@ -193,6 +193,10 @@ bool is_nonzero(const Tensor& self) {
   TORCH_INTERNAL_ASSERT(false, "Expected non-Tensor backend scalar");
 }
 
+void assert_async_cpu(const Tensor& self) {
+  TORCH_CHECK(native::is_nonzero(self), "Expected Tensor with single nonzero value, but got zero");
+}
+
 namespace {
 
 // DO NOT USE THIS -- it's just an implementation detail of wrapped_scalar tensor below.

aten/src/ATen/native/cuda/TensorCompare.cu

@@ -59,4 +59,26 @@ REGISTER_DISPATCH(where_kernel, &where_kernel_impl);
 REGISTER_DISPATCH(isposinf_stub, &isposinf_kernel_impl);
 REGISTER_DISPATCH(isneginf_stub, &isneginf_kernel_impl);
 
+template <typename scalar_t>
+__global__ void assert_async_cuda_kernel(scalar_t* input) {
+  CUDA_KERNEL_ASSERT(input[0] != 0);
+}
+
+__global__ void assert_async_cuda_kernel(c10::complex<float>* input) {
+  CUDA_KERNEL_ASSERT(input[0] != c10::complex<float>(0, 0));
+}
+__global__ void assert_async_cuda_kernel(c10::complex<double>* input) {
+  CUDA_KERNEL_ASSERT(input[0] != c10::complex<double>(0, 0));
+}
+
+void assert_async_cuda(const Tensor& self) {
+  auto n = self.numel();
+  TORCH_CHECK(n != 0, "Boolean value of Tensor with no values is ambiguous");
+  TORCH_CHECK(n < 2, "Boolean value of Tensor with more than one value is ambiguous");
+  auto stream = at::cuda::getCurrentCUDAStream();
+  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(at::ScalarType::Half, at::ScalarType::Bool, at::ScalarType::BFloat16, self.scalar_type(), "assert_async_cuda", [&] {
+    assert_async_cuda_kernel<<<1, 1, 0, stream>>>(self.data_ptr<scalar_t>());
+  });
+}
+
 }} // namespace at::native

aten/src/ATen/native/native_functions.yaml

@@ -118,6 +118,14 @@
 
 - func: align_tensors(Tensor[] tensors) -> Tensor[]
 
+# Not assert because it's a keyword; not Assert because FX already
+# took that syntax
+# TODO: need to specify this is side-effectful somehow
+- func: assert_async(Tensor self) -> ()
+  dispatch:
+    CPU: assert_async_cpu
+    CUDA: assert_async_cuda
+
 - func: refine_names(Tensor(a) self, Dimname[] names) -> Tensor(a)
   variants: method
 

test/test_cuda.py

@@ -1825,6 +1825,37 @@ def worker(rank):
 t2.start()
 """])
 
+    def test_cuda_assert_async(self):
+        with self.assertRaisesRegex(RuntimeError, "Boolean value of Tensor with no values is ambiguous"):
+            torch.assert_async(torch.tensor([], device="cuda"))
+        with self.assertRaisesRegex(RuntimeError, "Boolean value of Tensor with more than one value is ambiguous"):
+            torch.assert_async(torch.tensor([0, 0], device="cuda"))
+
+        torch.assert_async(torch.tensor(1, device="cuda"))
+        torch.assert_async(torch.tensor(0.1, device="cuda"))
+        torch.assert_async(torch.tensor(-0.1, device="cuda"))
+        torch.assert_async(torch.tensor(True, device="cuda"))
+        torch.assert_async(torch.tensor(0 + 0.1j, device="cuda"))
+
+        fail_stmts = [
+            "torch.assert_async(torch.tensor(0, device='cuda'))",
+            "torch.assert_async(torch.tensor(0.0, device='cuda'))",
+            "torch.assert_async(torch.tensor(False, device='cuda'))",
+            "torch.assert_async(torch.tensor(0+ 0 j, device='cuda'))",
+        ]
+
+        import subprocess
+        for stmt in fail_stmts:
+            with self.subTest(stmt=stmt):
+                r = subprocess.call([sys.executable, '-c', f"""\
+import torch
+
+{stmt}
+torch.cuda.synchronize()
+"""])
+                self.assertTrue(r != 0)
+
+
     def test_grad_scaling_unscale(self, dtype=torch.float):
         inv_scale = torch.full((1,), 0.25, dtype=torch.float, device="cuda:0")
         found_inf = torch.full((1,), 0.0, dtype=torch.float, device="cuda:0")

test/test_torch.py

@@ -2356,6 +2356,32 @@ def test_is_nonzero(self):
             self.assertTrue(torch.tensor([1]).is_nonzero())
             self.assertFalse(torch.tensor([[0]]).is_nonzero())
             self.assertTrue(torch.tensor([[1]]).is_nonzero())
+            self.assertTrue(torch.tensor(0.1).is_nonzero())
+            self.assertTrue(torch.tensor(-0.1).is_nonzero())
+            self.assertFalse(torch.tensor(0.0).is_nonzero())
+            self.assertTrue(torch.tensor(True).is_nonzero())
+            self.assertFalse(torch.tensor(False).is_nonzero())
+            self.assertFalse(torch.tensor(0 + 0j).is_nonzero())
+            self.assertTrue(torch.tensor(0 + 0.1j).is_nonzero())
+
+        def test_assert_async(self):
+            with self.assertRaisesRegex(RuntimeError, "Boolean value of Tensor with no values is ambiguous"):
+                torch.assert_async(torch.tensor([]))
+            with self.assertRaisesRegex(RuntimeError, "Boolean value of Tensor with more than one value is ambiguous"):
+                torch.assert_async(torch.tensor([0, 0]))
+            with self.assertRaisesRegex(RuntimeError, "Expected Tensor with single nonzero value, but got zero"):
+                torch.assert_async(torch.tensor(0))
+            torch.assert_async(torch.tensor(1))
+            torch.assert_async(torch.tensor(0.1))
+            torch.assert_async(torch.tensor(-0.1))
+            with self.assertRaisesRegex(RuntimeError, "Expected Tensor with single nonzero value, but got zero"):
+                torch.assert_async(torch.tensor(0.0))
+            torch.assert_async(torch.tensor(True))
+            with self.assertRaisesRegex(RuntimeError, "Expected Tensor with single nonzero value, but got zero"):
+                torch.assert_async(torch.tensor(False))
+            torch.assert_async(torch.tensor(0 + 0.1j))
+            with self.assertRaisesRegex(RuntimeError, "Expected Tensor with single nonzero value, but got zero"):
+                torch.assert_async(torch.tensor(0 + 0j))
 
         # NB: we must not be built with CUDA; if we are built with CUDA but no CUDA
         # is available, we get a different error.

torch/_torch_docs.py

@@ -10565,6 +10565,24 @@ def merge_dicts(*dicts):
     device(type='cpu')
 """)
 
+add_docstr(torch.assert_async,
+           r"""
+assert_async(tensor) -> void
+
+Asynchronously assert that the contents of tensor are nonzero.  For CPU tensors,
+this is equivalent to ``assert tensor`` or ``assert tensor.is_nonzero()``; for
+CUDA tensors, we DO NOT synchronize and you may only find out the assertion
+failed at a later CUDA kernel launch.  Asynchronous assertion can be helpful for
+testing invariants in CUDA tensors without giving up performance.  This function
+is NOT intended to be used for regular error checking, as it will trash your CUDA
+context if the assert fails (forcing you to restart your PyTorch process.)
+
+Args:
+    tensor (Tensor): a one element tensor to test to see if it is nonzero.  Zero
+        elements (including False for boolean tensors) cause an assertion failure
+        to be raised.
+""")
+
 add_docstr(torch.searchsorted,
            r"""
 searchsorted(sorted_sequence, values, *, out_int32=False, right=False, out=None) -> Tensor

torch/overrides.py

@@ -272,6 +272,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]:
         torch.argmin: lambda input: -1,
         torch.argsort: lambda input, dim=None: -1,
         torch.asin: lambda input, out=None: -1,
+        torch.assert_async: lambda input: -1,
         torch.arcsin: lambda input, out=None: -1,
         torch.asinh: lambda input, out=None: -1,
         torch.arcsinh: lambda input, out=None: -1,