mirror/eigen
2
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2025-03-07 18:27:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

Updated the tensor sum and mean reducer to enable them to process complex numbers on cuda gpus.

Browse Source
This commit is contained in:
Benoit Steiner 2016-09-28 17:08:41 -07:00
parent f3a00dd2b5
commit 2bda1b0d93
2 changed files with 41 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
unsupported/Eigen/CXX11/src/Tensor/TensorFunctors.h
View File

@ -99,7 +99,8 @@ template <typename T> struct SumReducer
static const bool IsStateful = false;
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) const {
(*accum) += t;
internal::scalar_sum_op<T> sum_op;
*accum = sum_op(*accum, t);
}
template <typename Packet>
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reducePacket(const Packet& p, Packet* accum) const {
@ -145,7 +146,8 @@ template <typename T> struct MeanReducer
MeanReducer() : scalarCount_(0), packetCount_(0) { }
EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void reduce(const T t, T* accum) {
(*accum) += t;
internal::scalar_sum_op<T> sum_op;
*accum = sum_op(*accum, t);
scalarCount_++;
}
template <typename Packet>

37
unsupported/test/cxx11_tensor_complex_cuda.cu
View File

@ -71,8 +71,45 @@ void test_cuda_nullary() {
}
static void test_cuda_sum_reductions() {
Eigen::CudaStreamDevice stream;
Eigen::GpuDevice gpu_device(&stream);
const int num_rows = internal::random<int>(1024, 5*1024);
const int num_cols = internal::random<int>(1024, 5*1024);
Tensor<std::complex<float>, 2> in(num_rows, num_cols);
in.setRandom();
Tensor<std::complex<float>, 0> full_redux;
full_redux = in.sum();
std::size_t in_bytes = in.size() * sizeof(std::complex<float>);
std::size_t out_bytes = full_redux.size() * sizeof(std::complex<float>);
std::complex<float>* gpu_in_ptr = static_cast<std::complex<float>*>(gpu_device.allocate(in_bytes));
std::complex<float>* gpu_out_ptr = static_cast<std::complex<float>*>(gpu_device.allocate(out_bytes));
gpu_device.memcpyHostToDevice(gpu_in_ptr, in.data(), in_bytes);
TensorMap<Tensor<std::complex<float>, 2> > in_gpu(gpu_in_ptr, num_rows, num_cols);
TensorMap<Tensor<std::complex<float>, 0> > out_gpu(gpu_out_ptr);
out_gpu.device(gpu_device) = in_gpu.sum();
Tensor<std::complex<float>, 0> full_redux_gpu;
gpu_device.memcpyDeviceToHost(full_redux_gpu.data(), gpu_out_ptr, out_bytes);
gpu_device.synchronize();
// Check that the CPU and GPU reductions return the same result.
VERIFY_IS_APPROX(full_redux(), full_redux_gpu());
gpu_device.deallocate(gpu_in_ptr);
gpu_device.deallocate(gpu_out_ptr);
}
void test_cxx11_tensor_complex()
{
CALL_SUBTEST(test_cuda_nullary());
CALL_SUBTEST(test_cuda_sum_reductions());
}