Added support for non-deterministic random number generation on GPU

1 files changed, 75 insertions, 23 deletions

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h b/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
index bfd01ad6e..b5adb4041 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
@@ -185,21 +185,25 @@ template <typename T> struct ProdReducer
 
 // Random number generation
 namespace {
+#ifdef __CUDA_ARCH__
+__device__ int get_random_seed() {
+    return clock();
+}
+#else
 int get_random_seed() {
-#if defined _WIN32
+#ifdef _WIN32
     SYSTEMTIME st;
     GetSystemTime(&st);
     return st.wSecond + 1000 * st.wMilliseconds;
 #elif defined __APPLE__
     return mach_absolute_time();
-#elif defined __CUDA_ARCH__
-    return clock();
 #else
     timespec ts;
     clock_gettime(CLOCK_REALTIME, &ts);
-    return static_cast<int>(ts.tv_nsec);
+    return ts.tv_nsec;
 #endif
 }
+#endif
 }
 
 #if !defined (EIGEN_USE_GPU) || !defined(__CUDACC__) || !defined(__CUDA_ARCH__)
@@ -209,11 +213,14 @@ template <typename T> class UniformRandomGenerator {
  public:
   static const bool PacketAccess = true;
 
-  UniformRandomGenerator(bool deterministic = true) {
+  UniformRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     if (!deterministic) {
       srand(get_random_seed());
     }
   }
+  UniformRandomGenerator(const UniformRandomGenerator& other) {
+    m_deterministic = other.m_deterministic;
+  }
 
   template<typename Index>
   T operator()(Index, Index = 0) const {
@@ -228,6 +235,9 @@ template <typename T> class UniformRandomGenerator {
     }
     return internal::pload<typename internal::packet_traits<T>::type>(values);
   }
+
+ private:
+  bool m_deterministic;
 };
 
 #if __cplusplus > 199711
@@ -235,13 +245,14 @@ template <> class UniformRandomGenerator<float> {
  public:
   static const bool PacketAccess = true;
 
-  UniformRandomGenerator(bool deterministic = true) {
+  UniformRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     if (!deterministic) {
       m_generator.seed(get_random_seed());
     }
   }
   UniformRandomGenerator(const UniformRandomGenerator<float>& other) {
     m_generator.seed(other(0, 0) * UINT_MAX);
+    m_deterministic = other.m_deterministic;
   }
 
   template<typename Index>
@@ -260,6 +271,9 @@ template <> class UniformRandomGenerator<float> {
 
  private:
   UniformRandomGenerator& operator = (const UniformRandomGenerator&);
+  // Make sure m_deterministic comes first to match the layout of the cpu
+  // version of the code.
+  bool m_deterministic;
   mutable std::mt19937 m_generator;
   mutable std::uniform_real_distribution<float> m_distribution;
 };
@@ -268,13 +282,14 @@ template <> class UniformRandomGenerator<double> {
  public:
   static const bool PacketAccess = true;
 
-  UniformRandomGenerator(bool deterministic = true) {
+  UniformRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     if (!deterministic) {
       m_generator.seed(get_random_seed());
     }
   }
   UniformRandomGenerator(const UniformRandomGenerator<double>& other) {
     m_generator.seed(other(0, 0) * UINT_MAX);
+    m_deterministic = other.m_deterministic;
   }
 
   template<typename Index>
@@ -293,6 +308,9 @@ template <> class UniformRandomGenerator<double> {
 
  private:
   UniformRandomGenerator& operator = (const UniformRandomGenerator&);
+  // Make sure m_deterministic comes first to match the layout of the cpu
+  // version of the code.
+  bool m_deterministic;
   mutable std::mt19937 m_generator;
   mutable std::uniform_real_distribution<double> m_distribution;
 };
@@ -307,22 +325,30 @@ template <> class UniformRandomGenerator<float> {
  public:
   static const bool PacketAccess = true;
 
-  EIGEN_DEVICE_FUNC UniformRandomGenerator(bool deterministic = true) {
+  __device__ UniformRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     const int tid = blockIdx.x * blockDim.x + threadIdx.x;
     const int seed = deterministic ? 0 : get_random_seed();
     curand_init(seed, tid, 0, &m_state);
   }
 
+  __device__ UniformRandomGenerator(const UniformRandomGenerator& other) {
+    m_deterministic = other.m_deterministic;
+    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int seed = m_deterministic ? 0 : get_random_seed();
+     curand_init(seed, tid, 0, &m_state);
+  }
+
   template<typename Index>
-  EIGEN_DEVICE_FUNC float operator()(Index, Index = 0) const {
+  __device__ float operator()(Index, Index = 0) const {
     return curand_uniform(&m_state);
   }
   template<typename Index>
-  EIGEN_DEVICE_FUNC float4 packetOp(Index, Index = 0) const {
+  __device__ float4 packetOp(Index, Index = 0) const {
     return curand_uniform4(&m_state);
   }
 
  private:
+  bool m_deterministic;
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
@@ -330,21 +356,28 @@ template <> class UniformRandomGenerator<double> {
  public:
   static const bool PacketAccess = true;
 
-  EIGEN_DEVICE_FUNC UniformRandomGenerator(bool deterministic = true) {
+  __device__ UniformRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     const int tid = blockIdx.x * blockDim.x + threadIdx.x;
     const int seed = deterministic ? 0 : get_random_seed();
     curand_init(seed, tid, 0, &m_state);
   }
+  __device__ UniformRandomGenerator(const UniformRandomGenerator& other) {
+    m_deterministic = other.m_deterministic;
+    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int seed = m_deterministic ? 0 : get_random_seed();
+    curand_init(seed, tid, 0, &m_state);
+  }
   template<typename Index>
-  EIGEN_DEVICE_FUNC double operator()(Index, Index = 0) const {
+  __device__ double operator()(Index, Index = 0) const {
     return curand_uniform_double(&m_state);
   }
   template<typename Index>
-  EIGEN_DEVICE_FUNC double2 packetOp(Index, Index = 0) const {
+  __device__ double2 packetOp(Index, Index = 0) const {
     return curand_uniform2_double(&m_state);
   }
 
  private:
+  bool m_deterministic;
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
@@ -357,12 +390,13 @@ template <typename T> class NormalRandomGenerator {
  public:
   static const bool PacketAccess = true;
 
-  NormalRandomGenerator(bool deterministic = true) : m_distribution(0, 1) {
+  NormalRandomGenerator(bool deterministic = true) : m_deterministic(deterministic), m_distribution(0, 1) {
     if (!deterministic) {
       m_generator.seed(get_random_seed());
     }
   }
-  NormalRandomGenerator(const NormalRandomGenerator& other) : m_distribution(other.m_distribution) {
+  NormalRandomGenerator(const NormalRandomGenerator& other)
+      : m_deterministic(other.m_deterministic), m_distribution(other.m_distribution) {
     m_generator.seed(other(0, 0) * UINT_MAX);
   }
 
@@ -380,6 +414,8 @@ template <typename T> class NormalRandomGenerator {
     return internal::pload<typename internal::packet_traits<T>::type>(values);
   }
 
+ private:
+  bool m_deterministic;
   mutable std::normal_distribution<T> m_distribution;
   mutable std::mt19937 m_generator;
 };
@@ -393,22 +429,28 @@ template <> class NormalRandomGenerator<float> {
  public:
   static const bool PacketAccess = true;
 
-  EIGEN_DEVICE_FUNC NormalRandomGenerator(bool deterministic = true) {
+  __device__ NormalRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     const int tid = blockIdx.x * blockDim.x + threadIdx.x;
     const int seed = deterministic ? 0 : get_random_seed();
     curand_init(seed, tid, 0, &m_state);
   }
-
+  __device__ NormalRandomGenerator(const NormalRandomGenerator<float>& other) {
+    m_deterministic = other.m_deterministic;
+    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int seed = m_deterministic ? 0 : get_random_seed();
+    curand_init(seed, tid, 0, &m_state);
+  }
   template<typename Index>
-  EIGEN_DEVICE_FUNC float operator()(Index, Index = 0) const {
+   __device__ float operator()(Index, Index = 0) const {
     return curand_normal(&m_state);
   }
   template<typename Index>
-  EIGEN_DEVICE_FUNC float4 packetOp(Index, Index = 0) const {
+   __device__ float4 packetOp(Index, Index = 0) const {
     return curand_normal4(&m_state);
   }
 
  private:
+  bool m_deterministic;
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
@@ -416,21 +458,28 @@ template <> class NormalRandomGenerator<double> {
  public:
   static const bool PacketAccess = true;
 
-  EIGEN_DEVICE_FUNC NormalRandomGenerator(bool deterministic = true) {
+  __device__ NormalRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {
     const int tid = blockIdx.x * blockDim.x + threadIdx.x;
     const int seed = deterministic ? 0 : get_random_seed();
     curand_init(seed, tid, 0, &m_state);
   }
+  __device__ NormalRandomGenerator(const NormalRandomGenerator<double>& other) {
+    m_deterministic = other.m_deterministic;
+    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    const int seed = m_deterministic ? 0 : get_random_seed();
+    curand_init(seed, tid, 0, &m_state);
+  }
   template<typename Index>
-  EIGEN_DEVICE_FUNC double operator()(Index, Index = 0) const {
+  __device__ double operator()(Index, Index = 0) const {
     return curand_normal_double(&m_state);
   }
   template<typename Index>
-  EIGEN_DEVICE_FUNC double2 packetOp(Index, Index = 0) const {
+  __device__ double2 packetOp(Index, Index = 0) const {
     return curand_normal2_double(&m_state);
   }
 
  private:
+  bool m_deterministic;
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
@@ -438,7 +487,10 @@ template <> class NormalRandomGenerator<double> {
 
 template <typename T> class NormalRandomGenerator {
  public:
-  NormalRandomGenerator(bool = true) {}
+  NormalRandomGenerator(bool deterministic = true) : m_deterministic(deterministic) {}
+
+ private:
+  bool m_deterministic;
 };
 
 #endif