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 #include "../../InternalHeaderCheck.h"
58 
59 namespace Eigen {
60 
61 // Specialized std::complex overloads.
62 namespace complex_operator_detail {
63 
64 template<typename T>
65 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
66 std::complex<T> complex_multiply(const std::complex<T>& a, const std::complex<T>& b) {
67  const T a_real = numext::real(a);
68  const T a_imag = numext::imag(a);
69  const T b_real = numext::real(b);
70  const T b_imag = numext::imag(b);
71  return std::complex<T>(
72  a_real * b_real - a_imag * b_imag,
73  a_imag * b_real + a_real * b_imag);
74 }
75 
76 template<typename T>
77 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
78 std::complex<T> complex_divide_fast(const std::complex<T>& a, const std::complex<T>& b) {
79  const T a_real = numext::real(a);
80  const T a_imag = numext::imag(a);
81  const T b_real = numext::real(b);
82  const T b_imag = numext::imag(b);
83  const T norm = (b_real * b_real + b_imag * b_imag);
84  return std::complex<T>((a_real * b_real + a_imag * b_imag) / norm,
85  (a_imag * b_real - a_real * b_imag) / norm);
86 }
87 
88 template<typename T>
89 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
90 std::complex<T> complex_divide_stable(const std::complex<T>& a, const std::complex<T>& b) {
91  const T a_real = numext::real(a);
92  const T a_imag = numext::imag(a);
93  const T b_real = numext::real(b);
94  const T b_imag = numext::imag(b);
95  // Smith's complex division (https://arxiv.org/pdf/1210.4539.pdf),
96  // guards against over/under-flow.
97  const bool scale_imag = numext::abs(b_imag) <= numext::abs(b_real);
98  const T rscale = scale_imag ? T(1) : b_real / b_imag;
99  const T iscale = scale_imag ? b_imag / b_real : T(1);
100  const T denominator = b_real * rscale + b_imag * iscale;
101  return std::complex<T>((a_real * rscale + a_imag * iscale) / denominator,
102  (a_imag * rscale - a_real * iscale) / denominator);
103 }
104 
105 template<typename T>
106 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE
107 std::complex<T> complex_divide(const std::complex<T>& a, const std::complex<T>& b) {
108 #if EIGEN_FAST_MATH
109  return complex_divide_fast(a, b);
110 #else
111  return complex_divide_stable(a, b);
112 #endif
113 }
114 
115 // NOTE: We cannot specialize compound assignment operators with Scalar T,
116 // (i.e. operator@=(const T&), for @=+,-,*,/)
117 // since they are already specialized for float/double/long double within
118 // the standard <complex> header. We also do not specialize the stream
119 // operators.
120 #define EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(T) \
121  \
122 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
123 std::complex<T> operator+(const std::complex<T>& a) { return a; } \
124  \
125 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
126 std::complex<T> operator-(const std::complex<T>& a) { \
127  return std::complex<T>(-numext::real(a), -numext::imag(a)); \
128 } \
129  \
130 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
131 std::complex<T> operator+(const std::complex<T>& a, const std::complex<T>& b) { \
132  return std::complex<T>(numext::real(a) + numext::real(b), numext::imag(a) + numext::imag(b)); \
133 } \
134  \
135 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
136 std::complex<T> operator+(const std::complex<T>& a, const T& b) { \
137  return std::complex<T>(numext::real(a) + b, numext::imag(a)); \
138 } \
139  \
140 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
141 std::complex<T> operator+(const T& a, const std::complex<T>& b) { \
142  return std::complex<T>(a + numext::real(b), numext::imag(b)); \
143 } \
144  \
145 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
146 std::complex<T> operator-(const std::complex<T>& a, const std::complex<T>& b) { \
147  return std::complex<T>(numext::real(a) - numext::real(b), numext::imag(a) - numext::imag(b)); \
148 } \
149  \
150 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
151 std::complex<T> operator-(const std::complex<T>& a, const T& b) { \
152  return std::complex<T>(numext::real(a) - b, numext::imag(a)); \
153 } \
154  \
155 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
156 std::complex<T> operator-(const T& a, const std::complex<T>& b) { \
157  return std::complex<T>(a - numext::real(b), -numext::imag(b)); \
158 } \
159  \
160 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
161 std::complex<T> operator*(const std::complex<T>& a, const std::complex<T>& b) { \
162  return complex_multiply(a, b); \
163 } \
164  \
165 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
166 std::complex<T> operator*(const std::complex<T>& a, const T& b) { \
167  return std::complex<T>(numext::real(a) * b, numext::imag(a) * b); \
168 } \
169  \
170 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
171 std::complex<T> operator*(const T& a, const std::complex<T>& b) { \
172  return std::complex<T>(a * numext::real(b), a * numext::imag(b)); \
173 } \
174  \
175 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
176 std::complex<T> operator/(const std::complex<T>& a, const std::complex<T>& b) { \
177  return complex_divide(a, b); \
178 } \
179  \
180 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
181 std::complex<T> operator/(const std::complex<T>& a, const T& b) { \
182  return std::complex<T>(numext::real(a) / b, numext::imag(a) / b); \
183 } \
184  \
185 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
186 std::complex<T> operator/(const T& a, const std::complex<T>& b) { \
187  return complex_divide(std::complex<T>(a, 0), b); \
188 } \
189  \
190 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
191 std::complex<T>& operator+=(std::complex<T>& a, const std::complex<T>& b) { \
192  numext::real_ref(a) += numext::real(b); \
193  numext::imag_ref(a) += numext::imag(b); \
194  return a; \
195 } \
196  \
197 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
198 std::complex<T>& operator-=(std::complex<T>& a, const std::complex<T>& b) { \
199  numext::real_ref(a) -= numext::real(b); \
200  numext::imag_ref(a) -= numext::imag(b); \
201  return a; \
202 } \
203  \
204 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
205 std::complex<T>& operator*=(std::complex<T>& a, const std::complex<T>& b) { \
206  a = complex_multiply(a, b); \
207  return a; \
208 } \
209  \
210 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
211 std::complex<T>& operator/=(std::complex<T>& a, const std::complex<T>& b) { \
212  a = complex_divide(a, b); \
213  return a; \
214 } \
215  \
216 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
217 bool operator==(const std::complex<T>& a, const std::complex<T>& b) { \
218  return numext::real(a) == numext::real(b) && numext::imag(a) == numext::imag(b); \
219 } \
220  \
221 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
222 bool operator==(const std::complex<T>& a, const T& b) { \
223  return numext::real(a) == b && numext::imag(a) == 0; \
224 } \
225  \
226 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
227 bool operator==(const T& a, const std::complex<T>& b) { \
228  return a == numext::real(b) && 0 == numext::imag(b); \
229 } \
230  \
231 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
232 bool operator!=(const std::complex<T>& a, const std::complex<T>& b) { \
233  return !(a == b); \
234 } \
235  \
236 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
237 bool operator!=(const std::complex<T>& a, const T& b) { \
238  return !(a == b); \
239 } \
240  \
241 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
242 bool operator!=(const T& a, const std::complex<T>& b) { \
243  return !(a == b); \
244 }
245 
246 // Do not specialize for long double, since that reduces to double on device.
247 EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(float)
248 EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(double)
249 
250 #undef EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS
251 
252 
253 } // namespace complex_operator_detail
254 
255 EIGEN_USING_STD_COMPLEX_OPERATORS
256 
257 namespace numext {
258 EIGEN_USING_STD_COMPLEX_OPERATORS
259 } // namespace numext
260 
261 namespace internal {
262 EIGEN_USING_STD_COMPLEX_OPERATORS
263 
264 } // namespace internal
265 } // namespace Eigen
266 
267 #endif // !(EIGEN_COMP_ICC && _USE_COMPLEX_SPECIALIZATION_)
268 
269 #endif // EIGEN_GPUCC && EIGEN_GPU_COMPILE_PHASE
270 
271 #endif // EIGEN_COMPLEX_GPU_H
Array< int, 3, 1 > b
const ImagReturnType imag() const
RealReturnType real() const
#define EIGEN_DEVICE_FUNC
Definition: Macros.h:883
Eigen::Triplet< double > T
EIGEN_ALWAYS_INLINE std::enable_if_t< NumTraits< T >::IsSigned||NumTraits< T >::IsComplex, typename NumTraits< T >::Real > abs(const T &x)
: InteropHeaders
Definition: Core:139