GPU: Declare the device std::complex operators before the Core headers libeigen/eigen!2988 Closes #3139 Co-authored-by: Rasmus Munk Larsen <rmlarsen@gmail.com>
diff --git a/Eigen/Core b/Eigen/Core index 9da31c1..5468550 100644 --- a/Eigen/Core +++ b/Eigen/Core
@@ -181,6 +181,11 @@ #endif // IWYU pragma: begin_exports +// Device overloads of the std::complex operators. They must precede every template that applies an operator to a +// dependent Scalar; see the header. +#ifdef EIGEN_GPUCC +#include "src/Core/arch/GPU/Complex.h" +#endif #include "src/Core/util/Constants.h" #include "src/Core/util/Meta.h" #include "src/Core/util/Assert.h" @@ -342,11 +347,6 @@ #include "src/Core/functors/StlFunctors.h" #include "src/Core/functors/AssignmentFunctors.h" -// Specialized functors for GPU. -#ifdef EIGEN_GPUCC -#include "src/Core/arch/GPU/Complex.h" -#endif - // Specializations of vectorized activation functions for NEON. #ifdef EIGEN_VECTORIZE_NEON #include "src/Core/arch/NEON/UnaryFunctors.h"
diff --git a/Eigen/src/Core/arch/GPU/Complex.h b/Eigen/src/Core/arch/GPU/Complex.h index 584a661..d79b569 100644 --- a/Eigen/src/Core/arch/GPU/Complex.h +++ b/Eigen/src/Core/arch/GPU/Complex.h
@@ -16,8 +16,16 @@ // operator/ are not constexpr. Due to this, GCC and older versions of clang do // not treat them as device functions and thus Eigen functors making use of // these operators fail to compile. Here, we manually specialize these -// operators and functors for complex types when building for CUDA to enable -// their use on-device. +// operators for complex types when building for CUDA to enable their use +// on-device. +// +// Eigen/Core includes this header ahead of the Core headers (Meta.h's +// equal_strict, MathFunctions.h, GenericPacketMath.h, the functors), which +// apply these operators to a dependent Scalar: two-phase +// lookup finds non-ADL candidates in the definition context only, and ADL for +// std::complex<T> reaches namespace std alone, so an overload declared later +// leaves those templates bound to the host-only std:: operators. nvcc before +// CUDA 13.3 resolved them at the instantiation point, which masked the ordering. // // NOTES: // - Compound assignment operators +=,-=,*=,/=(Scalar) will not work on device, @@ -30,6 +38,10 @@ // failures. // - Compiling with ICC requires defining _USE_COMPLEX_SPECIALIZATION_ prior // to the first inclusion of <complex>. +// - Device code outside namespace Eigen that applies these operators to a +// dependent Scalar reaches them only through `using namespace Eigen;` or +// using-declarations in scope (see test/gpu_basic.cu); ADL alone finds the +// host-only std:: operators. #if defined(EIGEN_GPUCC) && defined(EIGEN_GPU_COMPILE_PHASE) @@ -60,6 +72,16 @@ namespace Eigen { +namespace internal { +// Defined in MathFunctions.h, which follows this header. +template <typename T> +EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_multiply(const std::complex<T>& a, + const std::complex<T>& b); +template <typename T> +EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_divide(const std::complex<T>& a, + const std::complex<T>& b); +} // namespace internal + // Specialized std::complex overloads. namespace complex_operator_detail { @@ -68,38 +90,42 @@ // since they are already specialized for float/double/long double within // the standard <complex> header. We also do not specialize the stream // operators. +// numext is not declared yet; the std::complex members used below are +// constexpr, hence device-callable under EIGEN_CONSTEXPR_ARE_DEVICE_FUNC +// (nvcc: --expt-relaxed-constexpr), as real_impl<std::complex<T>> +// already assumes. #define EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(T) \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a) { return a; } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a) { \ - return std::complex<T>(-numext::real(a), -numext::imag(a)); \ + return std::complex<T>(-a.real(), -a.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a, \ const std::complex<T>& b) { \ - return std::complex<T>(numext::real(a) + numext::real(b), numext::imag(a) + numext::imag(b)); \ + return std::complex<T>(a.real() + b.real(), a.imag() + b.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a, const T& b) { \ - return std::complex<T>(numext::real(a) + b, numext::imag(a)); \ + return std::complex<T>(a.real() + b, a.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const T& a, const std::complex<T>& b) { \ - return std::complex<T>(a + numext::real(b), numext::imag(b)); \ + return std::complex<T>(a + b.real(), b.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a, \ const std::complex<T>& b) { \ - return std::complex<T>(numext::real(a) - numext::real(b), numext::imag(a) - numext::imag(b)); \ + return std::complex<T>(a.real() - b.real(), a.imag() - b.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a, const T& b) { \ - return std::complex<T>(numext::real(a) - b, numext::imag(a)); \ + return std::complex<T>(a.real() - b, a.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const T& a, const std::complex<T>& b) { \ - return std::complex<T>(a - numext::real(b), -numext::imag(b)); \ + return std::complex<T>(a - b.real(), -b.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const std::complex<T>& a, \ @@ -108,11 +134,11 @@ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const std::complex<T>& a, const T& b) { \ - return std::complex<T>(numext::real(a) * b, numext::imag(a) * b); \ + return std::complex<T>(a.real() * b, a.imag() * b); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const T& a, const std::complex<T>& b) { \ - return std::complex<T>(a * numext::real(b), a * numext::imag(b)); \ + return std::complex<T>(a * b.real(), a * b.imag()); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const std::complex<T>& a, \ @@ -121,7 +147,7 @@ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const std::complex<T>& a, const T& b) { \ - return std::complex<T>(numext::real(a) / b, numext::imag(a) / b); \ + return std::complex<T>(a.real() / b, a.imag() / b); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const T& a, const std::complex<T>& b) { \ @@ -129,14 +155,12 @@ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator+=(std::complex<T>& a, const std::complex<T>& b) { \ - numext::real_ref(a) += numext::real(b); \ - numext::imag_ref(a) += numext::imag(b); \ + a = std::complex<T>(a.real() + b.real(), a.imag() + b.imag()); \ return a; \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator-=(std::complex<T>& a, const std::complex<T>& b) { \ - numext::real_ref(a) -= numext::real(b); \ - numext::imag_ref(a) -= numext::imag(b); \ + a = std::complex<T>(a.real() - b.real(), a.imag() - b.imag()); \ return a; \ } \ \ @@ -151,15 +175,15 @@ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const std::complex<T>& a, const std::complex<T>& b) { \ - return numext::real(a) == numext::real(b) && numext::imag(a) == numext::imag(b); \ + return a.real() == b.real() && a.imag() == b.imag(); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const std::complex<T>& a, const T& b) { \ - return numext::real(a) == b && numext::imag(a) == 0; \ + return a.real() == b && a.imag() == 0; \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const T& a, const std::complex<T>& b) { \ - return a == numext::real(b) && 0 == numext::imag(b); \ + return a == b.real() && 0 == b.imag(); \ } \ \ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator!=(const std::complex<T>& a, const std::complex<T>& b) { \
diff --git a/test/gpu_basic.cu b/test/gpu_basic.cu index 6bb1a42..69497a2 100644 --- a/test/gpu_basic.cu +++ b/test/gpu_basic.cu
@@ -76,6 +76,27 @@ } }; +// Applies the complex operators inside Eigen's own templates, which see the device overloads only through +// Eigen/Core's include order; complex_operators below finds them through its using-directive. +template <typename ComplexType> +struct complex_internal_operators { + EIGEN_DEVICE_FUNC void operator()(int i, const ComplexType* in, ComplexType* out) const { + const int num_operators = 8; + int out_idx = i * num_operators; + const ComplexType a = in[i]; + const ComplexType b = in[i + 1]; + + out[out_idx++] = numext::negate(a); + out[out_idx++] = numext::conj(a); + out[out_idx++] = internal::padd(a, b); + out[out_idx++] = internal::psub(a, b); + out[out_idx++] = internal::pmul(a, b); + out[out_idx++] = internal::pdiv(a, b); + out[out_idx++] = internal::pnegate(a); + out[out_idx++] = internal::pconj(a); + } +}; + template <typename T> struct complex_sqrt { EIGEN_DEVICE_FUNC void operator()(int i, const typename T::Scalar* in, typename T::Scalar* out) const { @@ -599,6 +620,7 @@ // Test std::complex. CALL_SUBTEST(run_and_compare_to_gpu(complex_operators<Vector3cf>(), nthreads, cfin, cfout)); + CALL_SUBTEST(run_and_compare_to_gpu(complex_internal_operators<std::complex<float>>(), nthreads, cfin, cfout)); CALL_SUBTEST(test_with_infs_nans(complex_sqrt<Vector3cf>(), nthreads, cfin, cfout)); // numeric_limits