Eigen  3.4.90 (git rev 5a9f66fb35d03a4da9ef8976e67a61b30aa16dcf)
 
Loading...
Searching...
No Matches
GPU/Complex.h
1// This file is part of Eigen, a lightweight C++ template library
2// for linear algebra.
3//
4// Copyright (C) 2014 Benoit Steiner <[email protected]>
5// Copyright (C) 2021 C. Antonio Sanchez <[email protected]>
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
60namespace Eigen {
61
62// Specialized std::complex overloads.
63namespace complex_operator_detail {
64
65template <typename T>
66EIGEN_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
75template <typename T>
76EIGEN_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
86template <typename T>
87EIGEN_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
103template <typename T>
104EIGEN_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.
221EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(float)
222EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS(double)
223
224#undef EIGEN_CREATE_STD_COMPLEX_OPERATOR_SPECIALIZATIONS
225
226} // namespace complex_operator_detail
227
228EIGEN_USING_STD_COMPLEX_OPERATORS
229
230namespace numext {
231EIGEN_USING_STD_COMPLEX_OPERATORS
232} // namespace numext
233
234namespace internal {
235EIGEN_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
Eigen
Namespace containing all symbols from the Eigen library.
Definition Core:137