Move geometric_() to DistributionTemplates (pytorch#37418)

Summary:
Pull Request resolved: pytorch#37418

Fixes pytorch#37369

Test Plan: Imported from OSS

Differential Revision: D21290757

Pulled By: pbelevich

fbshipit-source-id: 42133f35edcbe716a07987bef2e68a4cdc27236a

Author: pbelevich

aten/src/ATen/native/DistributionTemplates.h

@@ -298,6 +298,16 @@ at::Tensor& log_normal_impl_(at::Tensor& self, double mean, double std, c10::opt
   return self;
 }
 
+// =================================================== Geometric ======================================================
+
+template<template<typename> class geometric_kernel, typename RNG>
+Tensor& geometric_impl_(Tensor& self, double p, c10::optional<Generator> gen) {
+  TORCH_CHECK(0 < p && p < 1, "geometric_ expects p to be in (0, 1), but got p=", p);
+  auto iter = TensorIterator::nullary_op(self);
+  geometric_kernel<RNG>()(iter, p, gen);
+  return self;
+}
+
 #undef CHECK_OUT_OF_BOUNDS_AND_SHOW_WARNING
 
 }}}

aten/src/ATen/native/Distributions.cpp

@@ -221,11 +221,17 @@ Tensor& exponential_(Tensor& self, double lambda, c10::optional<Generator> gen)
   return self;
 }
 
+// =================================================== Geometric ======================================================
+
+template<typename RNG>
+struct GeometricStub {
+  void operator()(TensorIterator& iter, double p, c10::optional<Generator> gen) {
+    geometric_stub(iter.device_type(), iter, p, gen);
+  }
+};
+
 Tensor& geometric_(Tensor& self, double p, c10::optional<Generator> gen) {
-  TORCH_CHECK(0 < p && p < 1, "geometric_ expects p to be in (0, 1), but got p=", p);
-  auto iter = TensorIterator::nullary_op(self);
-  geometric_stub(iter.device_type(), iter, p, gen);
-  return self;
+  return at::native::templates::geometric_impl_<GeometricStub, Generator>(self, p, gen);
 }
 
 // ==================================================== Uniform =======================================================

aten/src/ATen/native/cpu/DistributionTemplates.h

@@ -256,4 +256,24 @@ struct LogNormalKernel {
   }
 };
 
+// =================================================== Geometric ======================================================
+
+template<typename RNG>
+void geometric_kernel(TensorIterator& iter, double p, RNG generator) {
+  AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "geometric_cpu", [&]() {
+    std::lock_guard<std::mutex> lock(generator->mutex_);
+    cpu_serial_kernel(iter, [p, generator]() -> scalar_t {
+      at::geometric_distribution<double> geometric(p);
+      return (scalar_t)geometric(generator);
+    });
+  });
+}
+
+template<typename RNG>
+struct GeometricKernel {
+  void operator()(TensorIterator& iter, double p, c10::optional<Generator> gen) {
+    geometric_kernel(iter, p, check_generator<RNG>(gen));
+  }
+};
+
 }}}}}

aten/src/ATen/native/cpu/UnaryOpsKernel.cpp

@@ -329,14 +329,8 @@ static void exponential_kernel(TensorIterator& iter, double lambda, c10::optiona
 }
 
 static void geometric_kernel(TensorIterator& iter, double p, c10::optional<Generator> gen) {
-  AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "geometric_cpu", [&]() {
-    CPUGeneratorImpl* generator = get_generator_or_default<CPUGeneratorImpl>(gen, detail::getDefaultCPUGenerator());
-    std::lock_guard<std::mutex> lock(generator->mutex_);
-    cpu_serial_kernel(iter, [p, generator]() -> scalar_t {
-      at::geometric_distribution<double> geometric(p);
-      return (scalar_t)geometric(generator);
-    });
-  });
+  CPUGeneratorImpl* generator = get_generator_or_default<CPUGeneratorImpl>(gen, detail::getDefaultCPUGenerator());
+  templates::cpu::geometric_kernel(iter, p, generator);
 }
 
 static void log_normal_kernel(TensorIterator& iter, double mean, double std, c10::optional<Generator> gen) {

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

@@ -29,32 +29,11 @@
 
 namespace at { namespace native {
 
-void geometric_kernel_cuda(TensorIterator& iter, double p_, c10::optional<Generator> gen_) {
-  auto gen = get_generator_or_default<CUDAGeneratorImpl>(gen_, cuda::detail::getDefaultCUDAGenerator());
-  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "geometric_cuda", [&] {
-    if (std::is_same<scalar_t, double>::value) {
-      // define lambda for geometric transformation
-      auto geometric_func = [p_] __device__ (double rand) {
-        return static_cast<scalar_t>(::ceil(::log(rand) / ::log(static_cast<double>(1.0)-p_)));
-      };
-      distribution_nullary_kernel<scalar_t, double, curand4_engine_calls/2>(iter,
-        gen,
-        [] __device__ (curandStatePhilox4_32_10_t* state) { return curand_uniform2_double(state); },
-        geometric_func);
-    } else {
-      auto p = static_cast<float>(p_);
-      auto geometric_func = [p] __device__ (float rand) {
-        // use __logf fast approximation for peak bandwidth
-        return static_cast<scalar_t>(::ceil(__logf(rand) / __logf(static_cast<float>(1.0)-p)));
-      };
-      distribution_nullary_kernel<scalar_t, float, curand4_engine_calls>(iter,
-        gen,
-        [] __device__ (curandStatePhilox4_32_10_t* state) { return curand_uniform4(state); },
-        geometric_func);
-    }
-   });
+void geometric_kernel(TensorIterator& iter, double p_, c10::optional<Generator> gen) {
+  auto generator = get_generator_or_default<CUDAGeneratorImpl>(gen, cuda::detail::getDefaultCUDAGenerator());
+  at::native::templates::cuda::geometric_kernel(iter, p_, generator);
 }
 
-REGISTER_DISPATCH(geometric_stub, &geometric_kernel_cuda);
+REGISTER_DISPATCH(geometric_stub, &geometric_kernel);
 
 }} // namespace at::native

aten/src/ATen/native/cuda/DistributionTemplates.h

@@ -518,4 +518,39 @@ struct LogNormalKernel {
   }
 };
 
+// =================================================== Geometric ======================================================
+
+template<typename RNG>
+void geometric_kernel(TensorIterator& iter, double p_, RNG gen) {
+  AT_DISPATCH_ALL_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "geometric_cuda", [&] {
+    if (std::is_same<scalar_t, double>::value) {
+      // define lambda for geometric transformation
+      auto geometric_func = [p_] __device__ (double rand) {
+        return static_cast<scalar_t>(::ceil(::log(rand) / ::log(static_cast<double>(1.0)-p_)));
+      };
+      distribution_nullary_kernel<scalar_t, double, curand4_engine_calls/2>(iter,
+        gen,
+        [] __device__ (curandStatePhilox4_32_10_t* state) { return curand_uniform2_double(state); },
+        geometric_func);
+    } else {
+      auto p = static_cast<float>(p_);
+      auto geometric_func = [p] __device__ (float rand) {
+        // use __logf fast approximation for peak bandwidth
+        return static_cast<scalar_t>(::ceil(__logf(rand) / __logf(static_cast<float>(1.0)-p)));
+      };
+      distribution_nullary_kernel<scalar_t, float, curand4_engine_calls>(iter,
+        gen,
+        [] __device__ (curandStatePhilox4_32_10_t* state) { return curand_uniform4(state); },
+        geometric_func);
+    }
+  });
+}
+
+template<typename RNG>
+struct GeometricKernel {
+  void operator()(TensorIterator& iter, double p, c10::optional<Generator> gen) {
+    geometric_kernel(iter, p, check_generator<RNG>(gen));
+  }
+};
+
 }}}}

aten/src/ATen/test/cpu_rng_test.cpp

@@ -101,6 +101,12 @@ Tensor& log_normal_(Tensor& self, double mean, double std, c10::optional<Generat
   return at::native::templates::log_normal_impl_<native::templates::cpu::LogNormalKernel, TestCPUGenerator>(self, mean, std, gen);
 }
 
+// ================================================== Geometric =======================================================
+
+Tensor& geometric_(Tensor& self, double p, c10::optional<Generator> gen) {
+  return at::native::templates::geometric_impl_<native::templates::cpu::GeometricKernel, TestCPUGenerator>(self, p, gen);
+}
+
 TORCH_LIBRARY_IMPL(aten, CustomRNGKeyId, m) {
   // Random
   m.impl_UNBOXED("random_.from",             random_from_to);
@@ -119,6 +125,8 @@ TORCH_LIBRARY_IMPL(aten, CustomRNGKeyId, m) {
   m.impl_UNBOXED("cauchy_",                  custom_rng_cauchy_);
   // LogNormal
   m.impl_UNBOXED("log_normal_",              log_normal_);
+  // Geometric
+  m.impl_UNBOXED("geometric_",               geometric_);
 }
 
 class RNGTest : public ::testing::Test {
@@ -307,4 +315,20 @@ TEST_F(RNGTest, LogNormal) {
   ASSERT_TRUE(torch::allclose(actual, expected));
 }
 
+// ================================================== Geometric =======================================================
+
+TEST_F(RNGTest, Geometric) {
+  const auto p = 0.42;
+  auto gen = at::make_generator<TestCPUGenerator>(42.0);
+
+  auto actual = torch::empty({3, 3});
+  actual.geometric_(p, gen);
+
+  auto expected = torch::empty_like(actual);
+  auto iter = TensorIterator::nullary_op(expected);
+  native::templates::cpu::geometric_kernel(iter, p, check_generator<TestCPUGenerator>(gen));
+
+  ASSERT_TRUE(torch::allclose(actual, expected));
+}
+
 }

test/test_torch.py

@@ -10224,6 +10224,12 @@ def test_log_normal(self, device, dtype):
         self.assertEqual(a.dtype, dtype)
         self.assertEqual(a.size(), torch.Size([1]))
 
+    @dtypes(torch.float, torch.double)
+    def test_geometric(self, device, dtype):
+        a = torch.tensor([10], dtype=dtype, device=device).geometric_(0.5)
+        self.assertEqual(a.dtype, dtype)
+        self.assertEqual(a.size(), torch.Size([1]))
+
     def test_empty_strided(self, device):
         for shape in [(2, 3, 4), (0, 2, 0)]:
             # some of these cases are pretty strange, just verifying that if as_strided
@@ -10361,11 +10367,6 @@ def test_logical_all(self, device):
         y[-1][-1][-1] = 0
         self.assertEqual(y, x.all(2, keepdim=True))
 
-    def test_geometric(self, device):
-        a = torch.tensor([10], dtype=torch.float, device=device).geometric_(0.5)
-        self.assertEqual(a.dtype, torch.float)
-        self.assertEqual(a.size(), torch.Size([1]))
-
     @dtypes(torch.float32)
     def test_exponential(self, device, dtype):
         a = torch.tensor([10], dtype=torch.float, device=device).exponential_(0.5)