Added support for generation of random complex numbers on CUDA devices

3 files changed, 140 insertions, 0 deletions

diff --git a/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h b/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
index 33e8c01c2..14ffd5c93 100644
--- a/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
+++ b/unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
@@ -387,6 +387,58 @@ template <> class UniformRandomGenerator<double> {
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
+template <> class UniformRandomGenerator<std::complex<float> > {
+ public:
+  static const bool PacketAccess = false;
+
+  __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>
+  __device__ std::complex<float> operator()(Index, Index = 0) const {
+    float4 vals = curand_uniform4(&m_state);
+    return std::complex<float>(vals.x, vals.y);
+  }
+
+ private:
+  bool m_deterministic;
+  mutable curandStatePhilox4_32_10_t m_state;
+};
+
+template <> class UniformRandomGenerator<std::complex<double> > {
+ public:
+  static const bool PacketAccess = false;
+
+  __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>
+  __device__ std::complex<double> operator()(Index, Index = 0) const {
+    double2 vals = curand_uniform2_double(&m_state);
+    return std::complex<double>(vals.x, vals.y);
+  }
+
+ private:
+  bool m_deterministic;
+  mutable curandStatePhilox4_32_10_t m_state;
+};
+
 #endif
 
 
@@ -489,6 +541,58 @@ template <> class NormalRandomGenerator<double> {
   mutable curandStatePhilox4_32_10_t m_state;
 };
 
+template <> class NormalRandomGenerator<std::complex<float> > {
+ public:
+  static const bool PacketAccess = false;
+
+  __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& 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>
+  __device__ std::complex<float> operator()(Index, Index = 0) const {
+    float4 vals = curand_normal4(&m_state);
+    return std::complex<float>(vals.x, vals.y);
+  }
+
+ private:
+  bool m_deterministic;
+  mutable curandStatePhilox4_32_10_t m_state;
+};
+
+template <> class NormalRandomGenerator<std::complex<double> > {
+ public:
+  static const bool PacketAccess = false;
+
+  __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& 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>
+  __device__ std::complex<double> operator()(Index, Index = 0) const {
+    double2 vals = curand_normal2_double(&m_state);
+    return std::complex<double>(vals.x, vals.y);
+  }
+
+ private:
+  bool m_deterministic;
+  mutable curandStatePhilox4_32_10_t m_state;
+};
+
 #else
 
 template <typename T> class NormalRandomGenerator {
diff --git a/unsupported/test/CMakeLists.txt b/unsupported/test/CMakeLists.txt
index 155bfcd76..845cda8ce 100644
--- a/unsupported/test/CMakeLists.txt
+++ b/unsupported/test/CMakeLists.txt
@@ -144,5 +144,6 @@ if(EIGEN_TEST_CXX11)
 #  ei_add_test(cxx11_tensor_cuda "-std=c++0x")
 #  ei_add_test(cxx11_tensor_contract_cuda "-std=c++0x")
 #  ei_add_test(cxx11_tensor_reduction_cuda "-std=c++0x")
+#  ei_add_test(cxx11_tensor_random_cuda "-std=c++0x")
 
 endif()
diff --git a/unsupported/test/cxx11_tensor_random_cuda.cpp b/unsupported/test/cxx11_tensor_random_cuda.cpp
new file mode 100644
index 000000000..5d091de15
--- /dev/null
+++ b/unsupported/test/cxx11_tensor_random_cuda.cpp
@@ -0,0 +1,35 @@
+// This file is part of Eigen, a lightweight C++ template library
+// for linear algebra.
+//
+// Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
+//
+// This Source Code Form is subject to the terms of the Mozilla
+// Public License v. 2.0. If a copy of the MPL was not distributed
+// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#define EIGEN_TEST_NO_LONGDOUBLE
+#define EIGEN_TEST_NO_COMPLEX
+#define EIGEN_TEST_FUNC cxx11_tensor_random_cuda
+#define EIGEN_DEFAULT_DENSE_INDEX_TYPE int
+#define EIGEN_USE_GPU
+
+#include "main.h"
+#include <Eigen/CXX11/Tensor>
+
+static void test_default()
+{
+  Tensor<std::complex<float>, 1> vec(6);
+  vec.setRandom();
+
+  // Fixme: we should check that the generated numbers follow a uniform
+  // distribution instead.
+  for (int i = 1; i < 6; ++i) {
+    VERIFY_IS_NOT_EQUAL(vec(i), vec(i-1));
+  }
+}
+
+
+void test_cxx11_tensor_random_cuda()
+{
+  CALL_SUBTEST(test_default());
+}