devtools/eigen
1
0
Fork 0
mirror of https://gitlab.com/libeigen/eigen.git synced 2026-04-10 11:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add CUDA complex sqrt.

Browse Source 
This is to support scalar `sqrt` of complex numbers `std::complex<T>` on
device, requested by Tensorflow folks.

Technically `std::complex` is not supported by NVCC on device
(though it is by clang), so the default `sqrt(std::complex<T>)` function only
works on the host. Here we create an overload to add back the
functionality.

Also modified the CMake file to add `--relaxed-constexpr` (or
equivalent) flag for NVCC to allow calling constexpr functions from
device functions, and added support for specifying compute architecture for
NVCC (was already available for clang).
This commit is contained in:
Antonio Sanchez
2020-12-22 22:49:06 -08:00
parent fdf2ee62c5
commit 070d303d56
7 changed files with 217 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
Eigen/src/Core/MathFunctions.h
View File

@@ -323,6 +323,27 @@ struct abs2_retval
typedef typename NumTraits<Scalar>::Real type;
};
/****************************************************************************
* Implementation of sqrt *
****************************************************************************/
template<typename Scalar>
struct sqrt_impl
{
EIGEN_DEVICE_FUNC
static EIGEN_ALWAYS_INLINE Scalar run(const Scalar& x)
{
EIGEN_USING_STD(sqrt);
return sqrt(x);
}
};
template<typename Scalar>
struct sqrt_retval
{
typedef Scalar type;
};
/****************************************************************************
* Implementation of norm1 *
****************************************************************************/
@@ -1368,12 +1389,11 @@ inline int log2(int x)
*
* It's usage is justified in performance critical functions, like norm/normalize.
*/
template<typename T>
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE
T sqrt(const T &x)
template<typename Scalar>
EIGEN_DEVICE_FUNC
EIGEN_ALWAYS_INLINE EIGEN_MATHFUNC_RETVAL(sqrt, Scalar) sqrt(const Scalar& x)
{
EIGEN_USING_STD(sqrt);
return sqrt(x);
return EIGEN_MATHFUNC_IMPL(sqrt, Scalar)::run(x);
}
// Boolean specialization, avoids implicit float to bool conversion (-Wimplicit-conversion-floating-point-to-bool).