Repository: pytorch/pytorch
Pull request: Add mutex to THC random number generator
This fixes https://github.com/pytorch/pytorch/issues/4777 and https://github.com/pytorch/pytorch/issues/6288 by adding `std::mutex` to protect access to `THCGenerator` in multithread scenario, using the same design as https://github.com/pytorch/pytorch/blob/master/aten/src/TH/THGenerator.h.

Problem statement:
Thread-save random number generator
The generator works for one GPU, but when I'm trying `DataParallel` which creates multiple threads for GPUs, I get the segmentation error. After a few exploration in gdb, I found that this line of code may be the culprit. It's highly likely that the generator is shared among threads and race conditions happen here.

```C
  y = *(_generator->state + (_generator->next)++);
```


#### env
- Pytorch 0.3.0 (pip)
- Ubuntu 16.04
- python 3.6

#### minimal reproducing snippet:

```python
from threading import Thread
import torch
import time

def _worker():
    while True:
        probs = torch.Tensor(500, 500).uniform_(0, 1)
        random = probs
        # random = torch.geometric(probs)
        print(random[0][0])
        time.sleep(1)

threads = [Thread(target=_worker) for i in range(8)]

for thread in threads:
    thread.start()
for thread in threads:
    thread.join()
```

Hints from issue discussion:
And I think the 64 bit version may be a better alternative than the current hack approach.
http://www.math.sci.hiroshima-u.ac.jp/~m-mat/MT/VERSIONS/C-LANG/mt19937-64.c
Do you mean CPU generators aren't thread-safe? That issue is tracked at https://github.com/pytorch/pytorch/issues/3794

Commit message: Build test_rng only if CUDA exists

Patch to review:
diff --git a/.jenkins/pytorch/test.sh b/.jenkins/pytorch/test.sh
index 248f0346facd..92d5de7e6721 100755
--- a/.jenkins/pytorch/test.sh
+++ b/.jenkins/pytorch/test.sh
@@ -43,6 +43,7 @@ if [[ "$BUILD_TEST_LIBTORCH" == "1" ]]; then
    CPP_BUILD="$PWD/../cpp-build"
    if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
      "$CPP_BUILD"/libtorch/bin/test_jit
+     "$CPP_BUILD"/libtorch/bin/test_rng
    else
      "$CPP_BUILD"/libtorch/bin/test_jit "[cpu]"
    fi
diff --git a/aten/src/ATen/native/cuda/Distributions.cu b/aten/src/ATen/native/cuda/Distributions.cu
index 50d162f14fce..85d28f566e50 100644
--- a/aten/src/ATen/native/cuda/Distributions.cu
+++ b/aten/src/ATen/native/cuda/Distributions.cu
@@ -10,6 +10,7 @@
 
 #include <THC/THCGeneral.h>
 #include <THC/THCTensorRandom.h>
+#include <THC/THCGenerator.h>
 #include <THC/THCApply.cuh>
 #include <THC/THCNumerics.cuh>
 
@@ -21,8 +22,8 @@ THCGenerator* THCRandom_getGenerator(THCState* state);
 namespace {
 std::pair<uint64_t, uint64_t> next_philox_seed(at::Generator* gen) {
   auto gen_ = THCRandom_getGenerator(at::globalContext().thc_state);
-  uint64_t offset = THAtomicAddLong(&gen_->philox_seed_offset, 1);
-  return std::make_pair(gen_->initial_seed, offset);
+  uint64_t offset = THAtomicAddLong(&gen_->state.philox_seed_offset, 1);
+  return std::make_pair(gen_->state.initial_seed, offset);
 }
 
 template <typename scalar_t>
diff --git a/aten/src/THC/THCGenerator.h b/aten/src/THC/THCGenerator.h
new file mode 100644
index 000000000000..0eeb5f64c50e
--- /dev/null
+++ b/aten/src/THC/THCGenerator.h
@@ -0,0 +1,19 @@
+#ifndef THC_GENERATOR_INC
+#define THC_GENERATOR_INC
+
+#include <mutex>
+
+typedef struct THCGeneratorState {
+  struct curandStateMtgp32* gen_states;
+  struct mtgp32_kernel_params *kernel_params;
+  int initf;
+  uint64_t initial_seed;
+  int64_t philox_seed_offset;
+} THCGeneratorState;
+
+struct THCGenerator {
+  std::mutex mutex; /* mutex for using this generator */
+  THCGeneratorState state;
+};
+
+#endif
diff --git a/aten/src/THC/THCTensorRandom.cpp b/aten/src/THC/THCTensorRandom.cpp
index ddccb7c5a80d..404b5ff10e8b 100644
--- a/aten/src/THC/THCTensorRandom.cpp
+++ b/aten/src/THC/THCTensorRandom.cpp
@@ -1,4 +1,5 @@
 #include "THCTensorRandom.h"
+#include "THCGenerator.h"
 
 #include <random>
 #include <curand.h>
@@ -11,15 +12,16 @@ void createGeneratorState(THCGenerator* gen, uint64_t seed);
 /* Frees memory allocated during setup. */
 void destroyGenerator(THCState *state, THCGenerator* gen)
 {
-  if (gen->gen_states)
+  std::lock_guard<std::mutex> lock(gen->mutex);
+  if (gen->state.gen_states)
   {
-    THCudaCheck(THCudaFree(state, gen->gen_states));
-    gen->gen_states = NULL;
+    THCudaCheck(THCudaFree(state, gen->state.gen_states));
+    gen->state.gen_states = NULL;
   }
-  if (gen->kernel_params)
+  if (gen->state.kernel_params)
   {
-    THCudaCheck(THCudaFree(state, gen->kernel_params));
-    gen->kernel_params = NULL;
+    THCudaCheck(THCudaFree(state, gen->state.kernel_params));
+    gen->state.kernel_params = NULL;
   }
 }
 
@@ -39,11 +41,12 @@ void THCRandom_init(THCState* state, int devices, int current_device)
   std::random_device rd;
   for (int i = 0; i < rng_state->num_devices; ++i)
   {
-    rng_state->gen[i].initf = 0;
-    rng_state->gen[i].initial_seed = createSeed(rd);
-    rng_state->gen[i].philox_seed_offset = 0;
-    rng_state->gen[i].gen_states = NULL;
-    rng_state->gen[i].kernel_params = NULL;
+    new (&rng_state->gen[i].mutex) std::mutex();
+    rng_state->gen[i].state.initf = 0;
+    rng_state->gen[i].state.initial_seed = createSeed(rd);
+    rng_state->gen[i].state.philox_seed_offset = 0;
+    rng_state->gen[i].state.gen_states = NULL;
+    rng_state->gen[i].state.kernel_params = NULL;
   }
 }
 
@@ -74,18 +77,20 @@ static THCGenerator* THCRandom_rawGenerator(THCState* state)
 THCGenerator* THCRandom_getGenerator(THCState* state)
 {
   THCGenerator* gen = THCRandom_rawGenerator(state);
-  if (gen->initf == 0)
+  std::lock_guard<std::mutex> lock(gen->mutex);
+  if (gen->state.initf == 0)
   {
     initializeGenerator(state, gen);
-    createGeneratorState(gen, gen->initial_seed);
-    gen->initf = 1;
+    createGeneratorState(gen, gen->state.initial_seed);
+    gen->state.initf = 1;
   }
   return gen;
 }
 
 struct curandStateMtgp32* THCRandom_generatorStates(struct THCState* state)
 {
-  return THCRandom_getGenerator(state)->gen_states;
+  THCGenerator* gen = THCRandom_getGenerator(state);
+  return gen->state.gen_states;
 }
 
 /* Random seed */
@@ -109,8 +114,9 @@ uint64_t THCRandom_seedAll(THCState* state)
 void THCRandom_manualSeed(THCState* state, uint64_t seed)
 {
   THCGenerator* gen = THCRandom_rawGenerator(state);
-  gen->initial_seed = seed;
-  if (gen->initf) {
+  std::lock_guard<std::mutex> lock(gen->mutex);
+  gen->state.initial_seed = seed;
+  if (gen->state.initf) {
     createGeneratorState(gen, seed);
   }
 }
@@ -130,5 +136,6 @@ void THCRandom_manualSeedAll(THCState* state, uint64_t seed)
 /* Get the initial seed */
 uint64_t THCRandom_initialSeed(THCState* state)
 {
-  return THCRandom_getGenerator(state)->initial_seed;
+  THCGenerator* gen = THCRandom_getGenerator(state);
+  return gen->state.initial_seed;
 }
diff --git a/aten/src/THC/THCTensorRandom.cu b/aten/src/THC/THCTensorRandom.cu
index a179213cd52e..51d282abe202 100644
--- a/aten/src/THC/THCTensorRandom.cu
+++ b/aten/src/THC/THCTensorRandom.cu
@@ -5,6 +5,7 @@
 #include "THCTensorMath.h"
 #include "THCReduceApplyUtils.cuh"
 #include "THCTensorRandom.cuh"
+#include "THCGenerator.h"
 
 #include <thrust/functional.h>
 #include <curand.h>
@@ -18,22 +19,22 @@
 
 THCGenerator* THCRandom_getGenerator(THCState* state);
 
-/* Sets up generator. Allocates but does not create the generator states. */
+/* Sets up generator. Allocates but does not create the generator states. Not thread-safe. */
 __host__ void initializeGenerator(THCState *state, THCGenerator* gen)
 {
-  THCudaCheck(THCudaMalloc(state, (void**)&gen->gen_states, MAX_NUM_BLOCKS * sizeof(curandStateMtgp32)));
-  THCudaCheck(THCudaMalloc(state, (void**)&gen->kernel_params, sizeof(mtgp32_kernel_params)));
+  THCudaCheck(THCudaMalloc(state, (void**)&gen->state.gen_states, MAX_NUM_BLOCKS * sizeof(curandStateMtgp32)));
+  THCudaCheck(THCudaMalloc(state, (void**)&gen->state.kernel_params, sizeof(mtgp32_kernel_params)));
 }
 
-/* Creates a new generator state given the seed. */
+/* Creates a new generator state given the seed. Not thread-safe. */
 __host__ void createGeneratorState(THCGenerator* gen, uint64_t seed)
 {
-  if (curandMakeMTGP32Constants(mtgp32dc_params_fast_11213, gen->kernel_params) != CURAND_STATUS_SUCCESS)
+  if (curandMakeMTGP32Constants(mtgp32dc_params_fast_11213, gen->state.kernel_params) != CURAND_STATUS_SUCCESS)
   {
     THError("Creating MTGP constants failed.");
   }
-  if (curandMakeMTGP32KernelState(gen->gen_states, mtgp32dc_params_fast_11213,
-                                  gen->kernel_params, MAX_NUM_BLOCKS, seed) != CURAND_STATUS_SUCCESS)
+  if (curandMakeMTGP32KernelState(gen->state.gen_states, mtgp32dc_params_fast_11213,
+                                  gen->state.kernel_params, MAX_NUM_BLOCKS, seed) != CURAND_STATUS_SUCCESS)
   {
     THError("Creating MTGP kernel state failed.");
   }
@@ -42,19 +43,20 @@ __host__ void createGeneratorState(THCGenerator* gen, uint64_t seed)
 __host__ void THCRandom_getRNGState(THCState* state, THByteTensor *rng_state)
 {
   THCGenerator* gen = THCRandom_getGenerator(state);
+  std::lock_guard<std::mutex> lock(gen->mutex);
 
   // The RNG state comprises the MTPG32 states, the seed, and an offset used for Philox
   static const size_t states_size = MAX_NUM_BLOCKS * sizeof(curandStateMtgp32);
-  static const size_t seed_size = sizeof(gen->initial_seed);
-  static const size_t offset_size = sizeof(gen->philox_seed_offset);
+  static const size_t seed_size = sizeof(gen->state.initial_seed);
+  static const size_t offset_size = sizeof(gen->state.philox_seed_offset);
   static const size_t total_size = states_size + seed_size + offset_size;
   THByteTensor_resize1d(rng_state, total_size);
   THArgCheck(THByteTensor_nElement(rng_state) == total_size, 1, "RNG state is wrong size");
   THArgCheck(THByteTensor_isContiguous(rng_state), 1, "RNG state must be contiguous");
-  THCudaCheck(cudaMemcpy(THByteTensor_data(rng_state), gen->gen_states,
+  THCudaCheck(cudaMemcpy(THByteTensor_data(rng_state), gen->state.gen_states,
                          states_size, cudaMemcpyDeviceToHost));
-  memcpy(THByteTensor_data(rng_state) + states_size, &gen->initial_seed, seed_size);
-  memcpy(THByteTensor_data(rng_state) + states_size + seed_size, &gen->philox_seed_offset, offset_size);
+  memcpy(THByteTensor_data(rng_state) + states_size, &gen->state.initial_seed, seed_size);
+  memcpy(THByteTensor_data(rng_state) + states_size + seed_size, &gen->state.philox_seed_offset, offset_size);
 }
 
 __global__ void set_rngstate_kernel(curandStateMtgp32 *state, mtgp32_kernel_params *kernel)
@@ -65,10 +67,11 @@ __global__ void set_rngstate_kernel(curandStateMtgp32 *state, mtgp32_kernel_para
 __host__ void THCRandom_setRNGState(THCState* state, THByteTensor *rng_state)
 {
   THCGenerator* gen = THCRandom_getGenerator(state);
+  std::lock_guard<std::mutex> lock(gen->mutex);
 
   static const size_t states_size = MAX_NUM_BLOCKS * sizeof(curandStateMtgp32);
-  static const size_t seed_size = sizeof(gen->initial_seed);
-  static const size_t offset_size = sizeof(gen->philox_seed_offset);
+  static const size_t seed_size = sizeof(gen->state.initial_seed);
+  static const size_t offset_size = sizeof(gen->state.philox_seed_offset);
   static const size_t total_size = states_size + seed_size + offset_size;
   bool no_philox_seed = false;
   if (THByteTensor_nElement(rng_state) == total_size - offset_size) {
@@ -79,16 +82,16 @@ __host__ void THCRandom_setRNGState(THCState* state, THByteTensor *rng_state)
   }
   THArgCheck(THByteTensor_isContiguous(rng_state), 1, "RNG state must be contiguous");
 
-  THCudaCheck(cudaMemcpy(gen->gen_states, THByteTensor_data(rng_state),
+  THCudaCheck(cudaMemcpy(gen->state.gen_states, THByteTensor_data(rng_state),
                          states_size, cudaMemcpyHostToDevice));
   set_rngstate_kernel<<<1, MAX_NUM_BLOCKS, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, gen->kernel_params);
-  memcpy(&gen->initial_seed, THByteTensor_data(rng_state) + states_size, seed_size);
+      gen->state.gen_states, gen->state.kernel_params);
+  memcpy(&gen->state.initial_seed, THByteTensor_data(rng_state) + states_size, seed_size);
   if (!no_philox_seed) {
-    memcpy(&gen->philox_seed_offset, THByteTensor_data(rng_state) + states_size + seed_size, offset_size);
+    memcpy(&gen->state.philox_seed_offset, THByteTensor_data(rng_state) + states_size + seed_size, offset_size);
   }
   else {
-    gen->philox_seed_offset = 0;
+    gen->state.philox_seed_offset = 0;
   }
 }
 
diff --git a/aten/src/THC/THCTensorRandom.h b/aten/src/THC/THCTensorRandom.h
index 21fe6d94256e..5203df28c78e 100644
--- a/aten/src/THC/THCTensorRandom.h
+++ b/aten/src/THC/THCTensorRandom.h
@@ -6,14 +6,7 @@
 #include "generic/THCTensorRandom.h"
 #include "THCGenerateAllTypes.h"
 
-/* Generator */
-typedef struct _Generator {
-  struct curandStateMtgp32* gen_states;
-  struct mtgp32_kernel_params *kernel_params;
-  int initf;
-  uint64_t initial_seed;
-  int64_t philox_seed_offset;
-} THCGenerator;
+typedef struct THCGenerator THCGenerator;
 
 typedef struct THCRNGState {
   /* One generator per GPU */
diff --git a/aten/src/THC/generic/THCTensorRandom.cu b/aten/src/THC/generic/THCTensorRandom.cu
index a577a53c19bd..6e67e9763f2e 100644
--- a/aten/src/THC/generic/THCTensorRandom.cu
+++ b/aten/src/THC/generic/THCTensorRandom.cu
@@ -16,7 +16,7 @@ THC_API void THCTensor_(uniform)(THCState* state, THCTensor *self_, double a, do
   real *data = THCTensor_(data)(state, self);
 
   generate_uniform<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, a, b);
+      gen->state.gen_states, size, data, a, b);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -31,7 +31,7 @@ THC_API void THCTensor_(normal)(THCState* state, THCTensor *self_, double mean,
   real *data = THCTensor_(data)(state, self);
 
   generate_normal<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, mean, stdv);
+      gen->state.gen_states, size, data, mean, stdv);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -70,7 +70,7 @@ THC_API void THCTensor_(logNormal)(THCState* state, THCTensor *self_, double mea
   real *data = THCTensor_(data)(state, self);
 
   generateLogNormal<real><<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, mean, stdv);
+      gen->state.gen_states, size, data, mean, stdv);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -86,7 +86,7 @@ THC_API void THCTensor_(exponential)(THCState* state, THCTensor *self_, double l
   real *data = THCTensor_(data)(state, self);
 
   generate_exponential<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, lambda);
+      gen->state.gen_states, size, data, lambda);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -102,7 +102,7 @@ THC_API void THCTensor_(cauchy)(THCState* state, THCTensor *self_, double median
   real *data = THCTensor_(data)(state, self);
 
   generate_cauchy<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, median, sigma);
+      gen->state.gen_states, size, data, median, sigma);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -242,7 +242,7 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
 
       sampleMultinomialWithReplacement
         <<<grid, block, 0, THCState_getCurrentStream(state)>>>(
-          gen->gen_states,
+          gen->state.gen_states,
           n_sample,
           THCudaLongTensor_data(state, self),
           numDist, numCategories,
@@ -275,7 +275,7 @@ THC_API void THCTensor_(multinomial)(struct THCState *state,
         // recalculate our distribution
         sampleMultinomialWithoutReplacement
           <<<grid, block, 0, THCState_getCurrentStream(state)>>>(
-            gen->gen_states,
+            gen->state.gen_states,
             n_sample,
             sample,
             THCudaLongTensor_data(state, self),
@@ -412,7 +412,7 @@ THC_API void THCTensor_(bernoulli)(THCState* state, THCTensor *self_, double p)
   real *data = THCTensor_(data)(state, self);
 
   generate_bernoulli<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, p);
+      gen->state.gen_states, size, data, p);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -443,7 +443,7 @@ THC_API void THCTensor_(NAME)(THCState* state,                                 \
   THArgCheck(size == prob_size, 3, "inconsistent tensor size");                \
                                                                                \
   generate_bernoulli_tensor<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>( \
-      gen->gen_states, size, result_data, probs_data);                         \
+      gen->state.gen_states, size, result_data, probs_data);                         \
                                                                                \
   PROB_TYPE##_free(state, probs);                                              \
   THCTensor_(freeCopyTo)(state, self, self_);                                  \
@@ -479,7 +479,7 @@ THC_API void THCTensor_(geometric)(THCState* state, THCTensor *self_, double p)
   real *data = THCTensor_(data)(state, self);
 
   generate_geometric<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, p);
+      gen->state.gen_states, size, data, p);
 
   THCTensor_(freeCopyTo)(state, self, self_);
 };
@@ -500,11 +500,11 @@ THC_API void THCTensor_(clampedRandom)(THCState* state, THCTensor *self_, int64_
 #if defined(THC_REAL_IS_LONG) || defined(THC_REAL_IS_DOUBLE) || defined(THC_REAL_IS_FLOAT)
   if (range > 1ULL << 32) {
     generate_random_64<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-        gen->gen_states, size, data, min_val, range);
+        gen->state.gen_states, size, data, min_val, range);
   } else {
 #endif
     generate_random<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-        gen->gen_states, size, data, min_val, range);
+        gen->state.gen_states, size, data, min_val, range);
 #if defined(THC_REAL_IS_LONG) || defined(THC_REAL_IS_DOUBLE) || defined(THC_REAL_IS_FLOAT)
   }
 #endif
@@ -530,19 +530,19 @@ THC_API void THCTensor_(random)(THCState* state, THCTensor *self_)
 
 #if defined(THC_REAL_IS_HALF)
   generate_random<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, 0UL, (1UL << HLF_MANT_DIG) + 1);
+      gen->state.gen_states, size, data, 0UL, (1UL << HLF_MANT_DIG) + 1);
 #elif defined(THC_REAL_IS_FLOAT)
   generate_random<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, 0UL, (1UL << FLT_MANT_DIG) + 1);
+      gen->state.gen_states, size, data, 0UL, (1UL << FLT_MANT_DIG) + 1);
 #elif defined(THC_REAL_IS_DOUBLE)
   generate_random_64<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, 0ULL, (1ULL << DBL_MANT_DIG) + 1);
+      gen->state.gen_states, size, data, 0ULL, (1ULL << DBL_MANT_DIG) + 1);
 #elif defined(THC_REAL_IS_LONG)
   generate_random_64<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, 0ULL, static_cast<uint64_t>(std::numeric_limits<real>::max()) + 1);
+      gen->state.gen_states, size, data, 0ULL, static_cast<uint64_t>(std::numeric_limits<real>::max()) + 1);
 #else
   generate_random<<<NUM_BLOCKS, BLOCK_SIZE, 0, THCState_getCurrentStream(state)>>>(
-      gen->gen_states, size, data, 0UL, static_cast<uint32_t>(std::numeric_limits<real>::max()) + 1);
+      gen->state.gen_states, size, data, 0UL, static_cast<uint32_t>(std::numeric_limits<real>::max()) + 1);
 #endif
 
   THCTensor_(freeCopyTo)(state, self, self_);
diff --git a/tools/cpp_build/libtorch/CMakeLists.txt b/tools/cpp_build/libtorch/CMakeLists.txt
index a555be39158c..ec7bf3975b4d 100644
--- a/tools/cpp_build/libtorch/CMakeLists.txt
+++ b/tools/cpp_build/libtorch/CMakeLists.txt
@@ -201,3 +201,21 @@ INSTALL(TARGETS test_jit
   RUNTIME DESTINATION "${TORCH_INSTALL_BIN_DIR}"
   LIBRARY DESTINATION "${TORCH_INSTALL_LIB_DIR}"
   ARCHIVE DESTINATION "${TORCH_INSTALL_LIB_DIR}")
+
+IF(NOT NO_CUDA)
+  SET(TORCH_TEST_SRCS ${TORCH_SRC_DIR}/csrc/cuda/test_rng.cpp)
+
+  ADD_EXECUTABLE(test_rng ${TORCH_TEST_SRCS})
+
+  TARGET_LINK_LIBRARIES(test_rng torch)
+
+  TARGET_INCLUDE_DIRECTORIES(test_rng PUBLIC
+    "${TORCH_SRC_DIR}/../"
+    "${TORCH_SRC_DIR}/../aten/src/ATen/utils/catch/single_include"
+    "${COMMON_INCLUDES}")
+
+  INSTALL(TARGETS test_rng
+    RUNTIME DESTINATION "${TORCH_INSTALL_BIN_DIR}"
+    LIBRARY DESTINATION "${TORCH_INSTALL_LIB_DIR}"
+    ARCHIVE DESTINATION "${TORCH_INSTALL_LIB_DIR}")
+ENDIF(NOT NO_CUDA)
diff --git a/torch/csrc/cuda/test_rng.cpp b/torch/csrc/cuda/test_rng.cpp
new file mode 100644
index 000000000000..536598eeba24
--- /dev/null
+++ b/torch/csrc/cuda/test_rng.cpp
@@ -0,0 +1,27 @@
+#define CATCH_CONFIG_MAIN
+#include "catch.hpp"
+
+#include "ATen/ATen.h"
+#include "cuda.h"
+#include "cuda_runtime.h"
+#include <thread>
+
+void makeRandomNumber() {
+  cudaSetDevice(std::rand() % 2);
+  auto x = at::CUDA(at::kFloat).randn({1000});
+}
+
+void testCudaRNGMultithread() {
+  auto threads = std::vector<std::thread>();
+  for (auto i = 0; i < 1000; i++) {
+    threads.emplace_back(makeRandomNumber);
+  }
+  for (auto& t : threads) {
+    t.join();
+  }
+};
+
+TEST_CASE( "CUDA RNG test", "[cuda]" ) {
+  SECTION( "multithread" )
+    testCudaRNGMultithread();
+}
\ No newline at end of file


Code review comments:
diff_hunk: @@ -201,3 +201,21 @@ INSTALL(TARGETS test_jit
   RUNTIME DESTINATION "${TORCH_INSTALL_BIN_DIR}"
   LIBRARY DESTINATION "${TORCH_INSTALL_LIB_DIR}"
   ARCHIVE DESTINATION "${TORCH_INSTALL_LIB_DIR}")
+
+IF(NOT NO_CUDA)
+  SET(TORCH_TEST_SRCS ${TORCH_SRC_DIR}/csrc/cuda/test_rng.cpp)
+
+  ADD_EXECUTABLE(test_rng ${TORCH_TEST_SRCS})
+
+  TARGET_LINK_LIBRARIES(test_rng torch)
+
+  TARGET_INCLUDE_DIRECTORIES(test_rng PUBLIC
+    "${TORCH_SRC_DIR}/../"
+    "${TORCH_SRC_DIR}/../aten/src/ATen/utils/catch/single_include"
+    "${COMMON_INCLUDES}")
+
+  INSTALL(TARGETS test_rng
+    RUNTIME DESTINATION "${TORCH_INSTALL_BIN_DIR}"
+    LIBRARY DESTINATION "${TORCH_INSTALL_LIB_DIR}"
+    ARCHIVE DESTINATION "${TORCH_INSTALL_LIB_DIR}")
path: tools/cpp_build/libtorch/CMakeLists.txt
text: @user1:
That's much better, but:
1. Can you please move `test_rng.cpp` out of `torch/csrc` and into `test/cpp`? I know that's where `test_jit.cpp` is, but that's for legacy reasons. We don't want to keep tests right next to our code.
2. Please remove the `INSTALL` line. No one wants to install the tests.

@user2:
Will is going to move it to `aten/src/ATen/test`