GPU/Complex.h
Go to the documentation of this file.
1 // This file is part of Eigen, a lightweight C++ template library
2 // for linear algebra.
3 //
4 // Copyright (C) 2014 Benoit Steiner <benoit.steiner.goog@gmail.com>
5 // Copyright (C) 2021 C. Antonio Sanchez <cantonios@google.com>
6 //
7 // This Source Code Form is subject to the terms of the Mozilla
8 // Public License v. 2.0. If a copy of the MPL was not distributed
9 // with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
10 
11 #ifndef EIGEN_COMPLEX_GPU_H
12 #define EIGEN_COMPLEX_GPU_H
13 
14 // Many std::complex methods such as operator+, operator-, operator* and
15 // operator/ are not constexpr. Due to this, GCC and older versions of clang do
16 // not treat them as device functions and thus Eigen functors making use of
17 // these operators fail to compile. Here, we manually specialize these
18 // operators and functors for complex types when building for CUDA to enable
19 // their use on-device.
20 //
21 // NOTES:
22 // - Compound assignment operators +=,-=,*=,/=(Scalar) will not work on device,
23 // since they are already specialized in the standard. Using them will result
24 // in silent kernel failures.
25 // - Compiling with MSVC and using +=,-=,*=,/=(std::complex<Scalar>) will lead
26 // to duplicate definition errors, since these are already specialized in
27 // Visual Studio's <complex> header (contrary to the standard). This is
28 // preferable to removing such definitions, which will lead to silent kernel
29 // failures.
30 // - Compiling with ICC requires defining _USE_COMPLEX_SPECIALIZATION_ prior
31 // to the first inclusion of <complex>.
32 
33 #if defined(EIGEN_GPUCC) && defined(EIGEN_GPU_COMPILE_PHASE)
34 
35 // ICC already specializes std::complex<float> and std::complex<double>
36 // operators, preventing us from making them device functions here.
37 // This will lead to silent runtime errors if the operators are used on device.
38 //
39 // To allow std::complex operator use on device, define _OVERRIDE_COMPLEX_SPECIALIZATION_
40 // prior to first inclusion of <complex>. This prevents ICC from adding
41 // its own specializations, so our custom ones below can be used instead.
42 #if !(EIGEN_COMP_ICC && defined(_USE_COMPLEX_SPECIALIZATION_))
43 
44 // Import Eigen's internal operator specializations.
45 #define EIGEN_USING_STD_COMPLEX_OPERATORS \
46  using Eigen::complex_operator_detail::operator+; \
47  using Eigen::complex_operator_detail::operator-; \
48  using Eigen::complex_operator_detail::operator*; \
49  using Eigen::complex_operator_detail::operator/; \
50  using Eigen::complex_operator_detail::operator+=; \
51  using Eigen::complex_operator_detail::operator-=; \
52  using Eigen::complex_operator_detail::operator*=; \
53  using Eigen::complex_operator_detail::operator/=; \
54  using Eigen::complex_operator_detail::operator==; \
55  using Eigen::complex_operator_detail::operator!=;
56 
57 // IWYU pragma: private
58 #include "../../InternalHeaderCheck.h"
59 
60 namespace Eigen {
61 
62 // Specialized std::complex overloads.
63 namespace complex_operator_detail {
64 
65 template <typename T>
66 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_multiply(const std::complex<T>& a,
67  const std::complex<T>& b) {
68  const T a_real = numext::real(a);
69  const T a_imag = numext::imag(a);
70  const T b_real = numext::real(b);
71  const T b_imag = numext::imag(b);
72  return std::complex<T>(a_real * b_real - a_imag * b_imag, a_imag * b_real + a_real * b_imag);
73 }
74 
75 template <typename T>
76 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_divide_fast(const std::complex<T>& a,
77  const std::complex<T>& b) {
78  const T a_real = numext::real(a);
79  const T a_imag = numext::imag(a);
80  const T b_real = numext::real(b);
81  const T b_imag = numext::imag(b);
82  const T norm = (b_real * b_real + b_imag * b_imag);
83  return std::complex<T>((a_real * b_real + a_imag * b_imag) / norm, (a_imag * b_real - a_real * b_imag) / norm);
84 }
85 
86 template <typename T>
87 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_divide_stable(const std::complex<T>& a,
88  const std::complex<T>& b) {
89  const T a_real = numext::real(a);
90  const T a_imag = numext::imag(a);
91  const T b_real = numext::real(b);
92  const T b_imag = numext::imag(b);
93  // Smith's complex division (https://arxiv.org/pdf/1210.4539.pdf),
94  // guards against over/under-flow.
95  const bool scale_imag = numext::abs(b_imag) <= numext::abs(b_real);
96  const T rscale = scale_imag ? T(1) : b_real / b_imag;
97  const T iscale = scale_imag ? b_imag / b_real : T(1);
98  const T denominator = b_real * rscale + b_imag * iscale;
99  return std::complex<T>((a_real * rscale + a_imag * iscale) / denominator,
100  (a_imag * rscale - a_real * iscale) / denominator);
101 }
102 
103 template <typename T>
104 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> complex_divide(const std::complex<T>& a,
105  const std::complex<T>& b) {
106 #if EIGEN_FAST_MATH
107  return complex_divide_fast(a, b);
108 #else
109  return complex_divide_stable(a, b);
110 #endif
111 }
112 
113 // NOTE: We cannot specialize compound assignment operators with Scalar T,
114 // (i.e. operator@=(const T&), for @=+,-,*,/)
115 // since they are already specialized for float/double/long double within
116 // the standard <complex> header. We also do not specialize the stream
117 // operators.
118 #define EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(T) \
119  \
120  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a) { return a; } \
121  \
122  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a) { \
123  return std::complex<T>(-numext::real(a), -numext::imag(a)); \
124  } \
125  \
126  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a, \
127  const std::complex<T>& b) { \
128  return std::complex<T>(numext::real(a) + numext::real(b), numext::imag(a) + numext::imag(b)); \
129  } \
130  \
131  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const std::complex<T>& a, const T& b) { \
132  return std::complex<T>(numext::real(a) + b, numext::imag(a)); \
133  } \
134  \
135  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator+(const T& a, const std::complex<T>& b) { \
136  return std::complex<T>(a + numext::real(b), numext::imag(b)); \
137  } \
138  \
139  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a, \
140  const std::complex<T>& b) { \
141  return std::complex<T>(numext::real(a) - numext::real(b), numext::imag(a) - numext::imag(b)); \
142  } \
143  \
144  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const std::complex<T>& a, const T& b) { \
145  return std::complex<T>(numext::real(a) - b, numext::imag(a)); \
146  } \
147  \
148  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator-(const T& a, const std::complex<T>& b) { \
149  return std::complex<T>(a - numext::real(b), -numext::imag(b)); \
150  } \
151  \
152  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const std::complex<T>& a, \
153  const std::complex<T>& b) { \
154  return complex_multiply(a, b); \
155  } \
156  \
157  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const std::complex<T>& a, const T& b) { \
158  return std::complex<T>(numext::real(a) * b, numext::imag(a) * b); \
159  } \
160  \
161  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator*(const T& a, const std::complex<T>& b) { \
162  return std::complex<T>(a * numext::real(b), a * numext::imag(b)); \
163  } \
164  \
165  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const std::complex<T>& a, \
166  const std::complex<T>& b) { \
167  return complex_divide(a, b); \
168  } \
169  \
170  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const std::complex<T>& a, const T& b) { \
171  return std::complex<T>(numext::real(a) / b, numext::imag(a) / b); \
172  } \
173  \
174  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T> operator/(const T& a, const std::complex<T>& b) { \
175  return complex_divide(std::complex<T>(a, 0), b); \
176  } \
177  \
178  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator+=(std::complex<T>& a, const std::complex<T>& b) { \
179  numext::real_ref(a) += numext::real(b); \
180  numext::imag_ref(a) += numext::imag(b); \
181  return a; \
182  } \
183  \
184  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator-=(std::complex<T>& a, const std::complex<T>& b) { \
185  numext::real_ref(a) -= numext::real(b); \
186  numext::imag_ref(a) -= numext::imag(b); \
187  return a; \
188  } \
189  \
190  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator*=(std::complex<T>& a, const std::complex<T>& b) { \
191  a = complex_multiply(a, b); \
192  return a; \
193  } \
194  \
195  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE std::complex<T>& operator/=(std::complex<T>& a, const std::complex<T>& b) { \
196  a = complex_divide(a, b); \
197  return a; \
198  } \
199  \
200  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const std::complex<T>& a, const std::complex<T>& b) { \
201  return numext::real(a) == numext::real(b) && numext::imag(a) == numext::imag(b); \
202  } \
203  \
204  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const std::complex<T>& a, const T& b) { \
205  return numext::real(a) == b && numext::imag(a) == 0; \
206  } \
207  \
208  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator==(const T& a, const std::complex<T>& b) { \
209  return a == numext::real(b) && 0 == numext::imag(b); \
210  } \
211  \
212  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator!=(const std::complex<T>& a, const std::complex<T>& b) { \
213  return !(a == b); \
214  } \
215  \
216  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator!=(const std::complex<T>& a, const T& b) { return !(a == b); } \
217  \
218  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool operator!=(const T& a, const std::complex<T>& b) { return !(a == b); }
219 
220 // Do not specialize for long double, since that reduces to double on device.
221 EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(float)
222 EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(double)
223 
224 #undef EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS
225 
226 } // namespace complex_operator_detail
227 
228 EIGEN_USING_STD_COMPLEX_OPERATORS
229 
230 namespace numext {
231 EIGEN_USING_STD_COMPLEX_OPERATORS
232 } // namespace numext
233 
234 namespace internal {
235 EIGEN_USING_STD_COMPLEX_OPERATORS
236 
237 } // namespace internal
238 } // namespace Eigen
239 
240 #endif // !(EIGEN_COMP_ICC && _USE_COMPLEX_SPECIALIZATION_)
241 
242 #endif // EIGEN_GPUCC && EIGEN_GPU_COMPILE_PHASE
243 
244 #endif // EIGEN_COMPLEX_GPU_H
imag
AnnoyingScalar imag(const AnnoyingScalar &)
Definition: AnnoyingScalar.h:132
T
Eigen::Triplet< double > T
Definition: EigenUnitTest.cpp:11
EIGEN_DEVICE_FUNC
#define EIGEN_DEVICE_FUNC
Definition: Macros.h:892
EIGEN_STRONG_INLINE
#define EIGEN_STRONG_INLINE
Definition: Macros.h:834
b
Scalar * b
Definition: benchVecAdd.cpp:17
real
float real
Definition: datatypes.h:10
a
const Scalar * a
Definition: level2_cplx_impl.h:32
Eigen::numext::abs
EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE std::enable_if_t< NumTraits< T >::IsSigned||NumTraits< T >::IsComplex, typename NumTraits< T >::Real > abs(const T &x)
Definition: MathFunctions.h:1355
Eigen
Namespace containing all symbols from the Eigen library.
Definition: bench_norm.cpp:70
internal
Definition: Eigen_Colamd.h:49