Support Half type in randperm.

xuhdev · facebook-github-bot · commit 0f7c3710ddf5 · 2019-07-10T12:23:25.000-07:00
Summary: Pull Request resolved: #22102 Test Plan: Imported from OSS Differential Revision: D16153586 Pulled By: li-roy fbshipit-source-id: d58e3dbc5da893005f4eaf521a28b0d752274eff
diff --git a/aten/src/ATen/native/TensorFactories.cpp b/aten/src/ATen/native/TensorFactories.cpp
@@ -525,11 +525,12 @@ Tensor& randperm_out(Tensor& result, int64_t n) {
 
 Tensor& randperm_out_cpu(Tensor& result, int64_t n, Generator* generator) {
   TORCH_CHECK(n >= 0, "n must be non-negative, got", n);
+  check_supported_max_int_with_precision(n, result);
   result.resize_({n});
   auto gen = get_generator_or_default<CPUGenerator>(generator, detail::getDefaultCPUGenerator());
   // See Note [Acquire lock when using random generators]
   std::lock_guard<std::mutex> lock(gen->mutex_);
-  AT_DISPATCH_ALL_TYPES(result.scalar_type(), "randperm", [&]() -> void {
+  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Half, result.scalar_type(), "randperm", [&]() -> void {
     randperm_cpu<scalar_t>(result, n, gen);
   });
 
diff --git a/aten/src/ATen/native/TensorFactories.h b/aten/src/ATen/native/TensorFactories.h
@@ -64,5 +64,25 @@ inline void check_size_nonnegative(IntArrayRef size) {
     TORCH_CHECK(x >= 0, "Trying to create tensor with negative dimension ", x, ": ", size);
   }
 }
+
+inline void check_supported_max_int_with_precision(int64_t n, const Tensor& tensor) {
+  TORCH_CHECK(at::scalar_tensor(n, tensor.options()).defined(),
+              "n is too large for result tensor type: '", tensor.type().toString(), "'");
+
+  // Ensure sufficient precision for floating point representation.
+  switch (tensor.scalar_type()) {
+    case at::ScalarType::Half:
+      TORCH_CHECK(n <= (int64_t(1) << 11) + 1, "n cannot be greater than 2049 for Half type.");
+      break;
+    case at::ScalarType::Float:
+      TORCH_CHECK(n <= (int64_t(1) << 24) + 1, "n cannot be greater than 2^24+1 for Float type.");
+      break;
+    case at::ScalarType::Double:  // Unlikely to happen, but doesn't hurt to check
+      TORCH_CHECK(n <= (int64_t(1) << 53) + 1, "n cannot be greater than 2^53+1 for Double type.");
+      break;
+    default:
+      break;
+  }
+}
 } // namespace native
 } // namespace at
diff --git a/aten/src/ATen/native/cuda/TensorFactories.cu b/aten/src/ATen/native/cuda/TensorFactories.cu
@@ -78,8 +78,7 @@ Tensor empty_strided_cuda(IntArrayRef size, IntArrayRef stride, const TensorOpti
 
 Tensor& randperm_out_cuda(Tensor& result, int64_t n, Generator* generator) {
   TORCH_CHECK(n >= 0, "n must be non-negative, got", n);
-  TORCH_CHECK(at::scalar_tensor(n, result.options()).defined(),
-  "n is too large for result tensor type: '", result.type().toString(), "'");
+  check_supported_max_int_with_precision(n, result);
 
   result.resize_({n});
 
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -2790,8 +2790,8 @@ def test_randperm_cuda(self):
         self.assertEqual(res1, res2, 0)
 
         with torch.random.fork_rng(devices=[0]):
-            res1 = torch.randperm(50000, dtype=torch.half, device=cuda)
-        res2 = torch.cuda.HalfTensor()
+            res1 = torch.randperm(50000, dtype=torch.float, device=cuda)
+        res2 = torch.cuda.FloatTensor()
         torch.randperm(50000, out=res2, device=cuda)
         self.assertEqual(res1, res2, 0)
 
@@ -2802,6 +2802,14 @@ def test_randperm_cuda(self):
         self.assertEqual(res1.numel(), 0)
         self.assertEqual(res2.numel(), 0)
 
+        # Test exceptions when n is too large for a floating point type
+        for res, small_n, large_n in ((torch.cuda.HalfTensor(), 2**11 + 1, 2**11 + 2),
+                                      (torch.cuda.FloatTensor(), 2**24 + 1, 2**24 + 2),
+                                      (torch.cuda.DoubleTensor(), 2**25,  # 2**53 + 1 is too large to run
+                                       2**53 + 2)):
+            torch.randperm(small_n, out=res)  # No exception expected
+            self.assertRaises(RuntimeError, lambda: torch.randperm(large_n, out=res))
+
     def test_random_neg_values(self):
         _TestTorchMixin._test_random_neg_values(self, use_cuda=True)
 
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -4512,6 +4512,14 @@ def test_randperm(self):
         self.assertEqual(res1.numel(), 0)
         self.assertEqual(res2.numel(), 0)
 
+        # Test exceptions when n is too large for a floating point type
+        for res, small_n, large_n in ((torch.HalfTensor(), 2**11 + 1, 2**11 + 2),
+                                      (torch.FloatTensor(), 2**24 + 1, 2**24 + 2),
+                                      (torch.DoubleTensor(), 2**25,  # 2**53 + 1 is too large to run
+                                       2**53 + 2)):
+            torch.randperm(small_n, out=res)  # No exception expected
+            self.assertRaises(RuntimeError, lambda: torch.randperm(large_n, out=res))
+
     def test_random(self):
         # This test is flaky with p<=(2/(ub-lb))^200=6e-36
         t = torch.FloatTensor(200)
