Expose alias multinomial methods to ATen (#17904)

Summary:
This PR exposes the multinomialAliasSetup and multinomialAliasDraw methods.

cc: neerajprad
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17904

Differential Revision: D14700205

Pulled By: ezyang

fbshipit-source-id: 16462fb1f1ef1d560fd586632ea356b23e966ee3

diff --git a/aten/src/ATen/Declarations.cwrap b/aten/src/ATen/Declarations.cwrap
index c67a719..8b8ffd3 100644 (file)
--- a/aten/src/ATen/Declarations.cwrap
+++ b/aten/src/ATen/Declarations.cwrap
@@ -2486,6 +2486,45 @@
         - int64_t to
 ]]
 [[
+  name: _th_multinomial_alias_setup
+  cname: multinomialAliasSetup
+  variants:
+    - function
+  types:
+    - floating_point
+  backends:
+    - CPU
+    - CUDA
+  return: argument 1,2
+  arguments:
+    - arg: THTensor* probs
+    - arg: THIndexTensor* J
+      output: True
+    - arg: THTensor* q
+      output: True
+]]
+[[
+  name: _th_multinomial_alias_draw
+  cname: multinomialAliasDraw
+  types:
+    - floating_point
+  backends:
+    - CPU
+    - CUDA
+  variants:
+    - function
+  return: argument 0
+  arguments:
+    - arg: THIndexTensor* result
+      output: True
+    - arg: THGenerator* generator
+      default: nullptr
+      kwarg_only: True
+    - THTensor* q
+    - THIndexTensor* J
+    - long num_samples
+]]
+[[
   name: _th_multinomial
   cname: multinomial
   types:

diff --git a/aten/src/ATen/native/LegacyDefinitions.cpp b/aten/src/ATen/native/LegacyDefinitions.cpp
index 42efd87..79ee5aa 100644 (file)
--- a/aten/src/ATen/native/LegacyDefinitions.cpp
+++ b/aten/src/ATen/native/LegacyDefinitions.cpp
@@ -516,6 +516,14 @@ Tensor ormqr(const Tensor & self, const Tensor & input2, const Tensor & input3,
   return at::legacy::th::_th_ormqr(self, input2, input3, left, transpose);
 }
 
+std::tuple<Tensor,Tensor> _multinomial_alias_setup(const Tensor & probs) {
+  return at::legacy::th::_th_multinomial_alias_setup(probs);
+}
+
+Tensor _multinomial_alias_draw(const Tensor & q, const Tensor & J, int64_t num_samples, Generator * generator) {
+  return at::legacy::th::_th_multinomial_alias_draw(q, J, num_samples, generator);
+}
+
 Tensor & btrisolve_out(Tensor & result, const Tensor & self, const Tensor & LU_data, const Tensor & LU_pivots) {
   return at::legacy::th::_th_btrisolve_out(result, self, LU_data, LU_pivots);
 }

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
index ba49a53..f8cd135 100644 (file)
--- a/aten/src/ATen/native/native_functions.yaml
+++ b/aten/src/ATen/native/native_functions.yaml
@@ -3837,6 +3837,14 @@
   matches_jit_signature: True
   variants: method, function
 
+- func: _multinomial_alias_setup(Tensor probs) -> (Tensor, Tensor)
+  matches_jit_signature: True
+  variants: function
+
+- func: _multinomial_alias_draw(Tensor J, Tensor q, int num_samples, *, Generator? generator=None) -> Tensor
+  matches_jit_signature: True
+  variants: function
+
 - func: lgamma(Tensor self, *, Tensor(a!) out) -> Tensor(a!)
   matches_jit_signature: True
 

diff --git a/aten/src/TH/generic/THTensorRandom.cpp b/aten/src/TH/generic/THTensorRandom.cpp
index 43c68d6..ea0561c 100644 (file)
--- a/aten/src/TH/generic/THTensorRandom.cpp
+++ b/aten/src/TH/generic/THTensorRandom.cpp
@@ -132,6 +132,9 @@ void THTensor_(logNormal)(THTensor *self, THGenerator *_generator, double mean,
 void THTensor_(multinomialAliasSetup)(THTensor *probs, THLongTensor *J, THTensor *q)
 {
   int64_t inputsize = THTensor_(nElement)(probs);
+  THArgCheck(probs->dim() == 1, 1,
+             "expected 1-D probability tensor, got %d-D probability tensor instead",
+             probs->dim());
   int64_t i = 0;
   THLongTensor *smaller = THLongTensor_newWithSize1d(inputsize);
   THLongTensor *larger = THLongTensor_newWithSize1d(inputsize);
@@ -216,16 +219,23 @@ void THTensor_(multinomialAliasSetup)(THTensor *probs, THLongTensor *J, THTensor
   THLongTensor_free(smaller);
   THLongTensor_free(larger);
 }
-void THTensor_(multinomialAliasDraw)(THLongTensor *self, THGenerator *_generator, THLongTensor *J, THTensor *q)
+void THTensor_(multinomialAliasDraw)(THLongTensor *self, THGenerator *_generator, THTensor *q, THLongTensor *J, int n_sample)
 {
   std::lock_guard<std::mutex> lock(_generator->mutex);
+  THArgCheck(q->dim() == 1, 1,
+             "expected 1-D probability table, got %d-D probability table instead",
+             q->dim());
+  THArgCheck(J->dim() == 1, 2,
+             "expected 1-D alias table, got %d-D alias table instead",
+             J->dim());
+  THArgCheck(n_sample > 0, 3, "cannot sample <= 0 samples");
   int64_t K = THLongTensor_nElement(J);
-  int64_t output_nelem = THLongTensor_nElement(self);
   int64_t i = 0, _mask=0;
   scalar_t _q;
+  THLongTensor_resize1d(self, n_sample);
   int64_t rand_ind, sample_idx, J_sample;
 
-  for (i=0; i < output_nelem; i++)
+  for (i=0; i < n_sample; i++)
     {
       rand_ind = THRandom_uniform(_generator, 0, K);
 

diff --git a/aten/src/TH/generic/THTensorRandom.h b/aten/src/TH/generic/THTensorRandom.h
index d52f34a..947a4d8 100644 (file)
--- a/aten/src/TH/generic/THTensorRandom.h
+++ b/aten/src/TH/generic/THTensorRandom.h
@@ -18,7 +18,7 @@ TH_API void THTensor_(cauchy)(THTensor *self, THGenerator *_generator, double me
 TH_API void THTensor_(logNormal)(THTensor *self, THGenerator *_generator, double mean, double stdv);
 TH_API void THTensor_(multinomial)(THLongTensor *self, THGenerator *_generator, THTensor *prob_dist, int n_sample, int with_replacement);
 TH_API void THTensor_(multinomialAliasSetup)(THTensor *prob_dist, THLongTensor *J, THTensor *q);
-TH_API void THTensor_(multinomialAliasDraw)(THLongTensor *self, THGenerator *_generator, THLongTensor *J, THTensor *q);
+TH_API void THTensor_(multinomialAliasDraw)(THLongTensor *self, THGenerator *_generator, THTensor *q, THLongTensor *J, int n_sample);
 #endif
 
 #if defined(TH_REAL_IS_BYTE)

diff --git a/aten/src/THC/generic/THCTensorRandom.cu b/aten/src/THC/generic/THCTensorRandom.cu
index c18bbfe..842259b 100644 (file)
--- a/aten/src/THC/generic/THCTensorRandom.cu
+++ b/aten/src/THC/generic/THCTensorRandom.cu
@@ -303,10 +303,14 @@ void THCTensor_(multinomial)(struct THCState *state,
 }
 
 void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_probs, THCudaLongTensor *_J, THCTensor *_q){
+  THArgCheck(_probs->dim() == 1, 1,
+             "expected 1-D probability tensor, got %d-D probability tensor instead",
+             _probs->dim());
   THAssert(THCTensor_(isContiguous)(state, _q));
   THAssert(THCudaLongTensor_isContiguous(state, _J));
-  THAssert(THCTensor_(isContiguous)(state, _probs));
-  int64_t inputsize = THCTensor_(nElement)(state, _probs);
+  THCTensor *probs = THCTensor_(newContiguous)(state, _probs);
+  THAssert(THCTensor_(isContiguous)(state, probs));
+  int64_t inputsize = THCTensor_(nElement)(state, probs);
   THCudaLongTensor *smaller = THCudaLongTensor_newWithSize1d(state, inputsize);
   THCudaLongTensor *larger = THCudaLongTensor_newWithSize1d(state, inputsize);
   THCudaLongTensor *smaller_short = THCudaLongTensor_newWithSize1d(state, inputsize);
@@ -320,7 +324,7 @@ void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_probs, THCud
   aliasMultinomialFilter
     <<<inputBlockDim, BLOCK_SIZE, 0, THCState_getCurrentStream(state) >>>(
                      THCTensor_(data)(state, _q),
-                     THCTensor_(data)(state, _probs),
+                     THCTensor_(data)(state, probs),
                      THCudaLongTensor_data(state, smaller),
                      THCudaLongTensor_data(state, larger),
                      THCudaLongTensor_data(state, _J),
@@ -355,24 +359,32 @@ void THCTensor_(multinomialAliasSetup)(THCState *state, THCTensor *_probs, THCud
   THCudaLongTensor_free(state, larger);
   THCudaLongTensor_free(state, smaller_short);
   THCudaLongTensor_free(state, larger_short);
+  THCTensor_free(state, probs);
 }
 
-void THCTensor_(multinomialAliasDraw)(THCState *state, THCudaLongTensor *self, THCudaLongTensor *_J, THCTensor *_q){
+void THCTensor_(multinomialAliasDraw)(THCState *state, THCudaLongTensor *self, THCTensor *_q, THCudaLongTensor *_J, int n_sample){
+  THArgCheck(_q->dim() == 1, 1,
+             "expected 1-D probability table, got %d-D probability table instead",
+             _q->dim());
+  THArgCheck(_J->dim() == 1, 2,
+             "expected 1-D alias table, got %d-D alias table instead",
+             _J->dim());
+  THArgCheck(n_sample > 0, 3, "cannot sample <= 0 samples");
   THAssert(THCTensor_(isContiguous)(state, _q));
   THAssert(THCudaLongTensor_isContiguous(state, _J));
   THCGenerator* gen = THCRandom_getGenerator(state);
   int64_t K = THCudaLongTensor_nElement(state, _J);
-  int64_t output_nelem = THCudaLongTensor_nElement(state, self);
+  THCudaLongTensor_resize1d(state, self, n_sample);
   ptrdiff_t size = THCudaLongTensor_nElement(state, self);
 
-  THCTensor *uniform = THCTensor_(newWithSize1d)(state, output_nelem);
-  THCTensor *bernoulli = THCTensor_(newWithSize1d)(state, output_nelem);
+  THCTensor *uniform = THCTensor_(newWithSize1d)(state, n_sample);
+  THCTensor *bernoulli = THCTensor_(newWithSize1d)(state, n_sample);
 
   THCTensor_(uniform)(state, uniform, 0, K);
   THCTensor_(uniform)(state, bernoulli, 0, 1);
 
   multinomialAliasDrawKernel
-    <<<THCCeilDiv((int)output_nelem+BLOCK_SIZE-1, BLOCK_SIZE), BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
+    <<<THCCeilDiv((int)n_sample+BLOCK_SIZE-1, BLOCK_SIZE), BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
           size,
           THCudaLongTensor_data(state, self),
           THCudaLongTensor_data(state, _J),
@@ -381,6 +393,8 @@ void THCTensor_(multinomialAliasDraw)(THCState *state, THCudaLongTensor *self, T
           THCTensor_(data)(state, uniform),
           THCTensor_(data)(state, bernoulli)
           );
+  THCTensor_(free)(state, uniform);
+  THCTensor_(free)(state, bernoulli);
 }
 
 #endif

diff --git a/aten/src/THC/generic/THCTensorRandom.h b/aten/src/THC/generic/THCTensorRandom.h
index b1e4f55..552207d 100644 (file)
--- a/aten/src/THC/generic/THCTensorRandom.h
+++ b/aten/src/THC/generic/THCTensorRandom.h
@@ -14,7 +14,7 @@ THC_API void THCTensor_(exponential)(struct THCState *state, THCTensor *self, do
 THC_API void THCTensor_(cauchy)(struct THCState *state, THCTensor *self, double median, double sigma);
 THC_API void THCTensor_(multinomial)(struct THCState *state, THCudaLongTensor *self, THCTensor *prob_dist, int n_sample, int with_replacement);
 THC_API void THCTensor_(multinomialAliasSetup)(struct THCState *state, THCTensor *probs, THCudaLongTensor *J, THCTensor *q);
-THC_API void THCTensor_(multinomialAliasDraw)(THCState *state, THCudaLongTensor *self, THCudaLongTensor *_J, THCTensor *_q);
+THC_API void THCTensor_(multinomialAliasDraw)(THCState *state, THCudaLongTensor *self, THCTensor *_q, THCudaLongTensor *_J, int n_sample);
 
 #endif
 

diff --git a/test/test_cuda.py b/test/test_cuda.py
index 0665421..e7a7f01 100644 (file)
--- a/test/test_cuda.py
+++ b/test/test_cuda.py
@@ -2249,6 +2249,9 @@ class TestCuda(TestCase):
         samples = probs.multinomial(1000000, replacement=True)
         self.assertGreater(probs[samples].min().item(), 0)
 
+    def test_multinomial_alias(self):
+        _TestTorchMixin._test_multinomial_alias(self, lambda t: t.cuda())
+
     @staticmethod
     def mute():
         os.dup2(os.open(os.devnull, os.O_WRONLY), sys.stderr.fileno())

diff --git a/test/test_torch.py b/test/test_torch.py
index 2abde8e..855e8f9 100644 (file)
--- a/test/test_torch.py
+++ b/test/test_torch.py
@@ -3292,6 +3292,40 @@ class _TestTorchMixin(object):
     def test_multinomial(self):
         self._test_multinomial(self, torch.FloatTensor)
 
+    @staticmethod
+    def _test_multinomial_alias(self, cast):
+        # Get probs vector to use in setup
+        def get_probs(length, is_contiguous):
+            probs = torch.softmax(torch.randn(length), 0)
+            if not is_contiguous:
+                probs = torch.softmax(torch.randn(length, 2), 0)[:, 1]
+            assert not (is_contiguous ^ probs.is_contiguous()), "contiguity requirement not met"
+            return cast(probs)
+
+        for is_contiguous in [True, False]:
+            probs = get_probs(4, is_contiguous)
+            alias_table, prob_table = torch._multinomial_alias_setup(probs)
+            for n_samples in [-1, 1, 10]:
+                if n_samples > 0:
+                    samples = torch._multinomial_alias_draw(prob_table, alias_table, n_samples)
+                    self.assertEqual(prob_table.size(), torch.Size([4]), "size mismatch: probability table")
+                    self.assertEqual(alias_table.size(), torch.Size([4]), "size mismatch: alias table")
+                    self.assertEqual(samples.size(), torch.Size([n_samples]), "wrong number of samples")
+                else:
+                    with self.assertRaisesRegex(RuntimeError, "cannot sample <= 0 samples"):
+                        torch._multinomial_alias_draw(prob_table, alias_table, n_samples)
+
+            with self.assertRaisesRegex(RuntimeError, "expected 1-D"):
+                probs = probs.view(2, 2)
+                torch._multinomial_alias_setup(probs)
+
+            with self.assertRaisesRegex(RuntimeError, "expected 1-D"):
+                a_t, p_t = torch._multinomial_alias_setup(probs)
+                torch._multinomial_alias_draw(p_t.view(2, 2), a_t.view(2, 2))
+
+    def test_multinomial_alias(self):
+        self._test_multinomial_alias(self, lambda t: t)
+
     def _spawn_method(self, method, arg):
         try:
             mp.set_start_method('spawn')