#pragma once
#include <complex>
#include <c10/macros/Macros.h>
#if defined(__CUDACC__) || defined(__HIPCC__)
#include <thrust/complex.h>
#endif
C10_CLANG_DIAGNOSTIC_PUSH()
#if C10_CLANG_HAS_WARNING("-Wimplicit-float-conversion")
C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-float-conversion")
#endif
#if C10_CLANG_HAS_WARNING("-Wfloat-conversion")
C10_CLANG_DIAGNOSTIC_IGNORE("-Wfloat-conversion")
#endif
namespace c10 {
// c10::complex is an implementation of complex numbers that aims
// to work on all devices supported by PyTorch
//
// Most of the APIs duplicates std::complex
// Reference: https://en.cppreference.com/w/cpp/numeric/complex
//
// [NOTE: Complex Operator Unification]
// Operators currently use a mix of std::complex, thrust::complex, and
// c10::complex internally. The end state is that all operators will use
// c10::complex internally. Until then, there may be some hacks to support all
// variants.
//
//
// [Note on Constructors]
//
// The APIs of constructors are mostly copied from C++ standard:
// https://en.cppreference.com/w/cpp/numeric/complex/complex
//
// Since C++14, all constructors are constexpr in std::complex
//
// There are three types of constructors:
// - initializing from real and imag:
// `constexpr complex( const T& re = T(), const T& im = T() );`
// - implicitly-declared copy constructor
// - converting constructors
//
// Converting constructors:
// - std::complex defines converting constructor between float/double/long
// double,
// while we define converting constructor between float/double.
// - For these converting constructors, upcasting is implicit, downcasting is
// explicit.
// - We also define explicit casting from std::complex/thrust::complex
// - Note that the conversion from thrust is not constexpr, because
// thrust does not define them as constexpr ????
//
//
// [Operator =]
//
// The APIs of operator = are mostly copied from C++ standard:
// https://en.cppreference.com/w/cpp/numeric/complex/operator%3D
//
// Since C++20, all operator= are constexpr. Although we are not building with
// C++20, we also obey this behavior.
//
// There are three types of assign operator:
// - Assign a real value from the same scalar type
// - In std, this is templated as complex& operator=(const T& x)
// with specialization `complex& operator=(T x)` for float/double/long
// double Since we only support float and double, on will use `complex&
// operator=(T x)`
// - Copy assignment operator and converting assignment operator
// - There is no specialization of converting assignment operators, which type
// is
// convertible is solely dependent on whether the scalar type is convertible
//
// In addition to the standard assignment, we also provide assignment operators
// with std and thrust
//
//
// [Casting operators]
//
// std::complex does not have casting operators. We define casting operators
// casting to std::complex and thrust::complex
//
//
// [Operator ""]
//
// std::complex has custom literals `i`, `if` and `il` defined in namespace
// `std::literals::complex_literals`. We define our own custom literals in the
// namespace `c10::complex_literals`. Our custom literals does not follow the
// same behavior as in std::complex, instead, we define _if, _id to construct
// float/double complex literals.
//
//
// [real() and imag()]
//
// In C++20, there are two overload of these functions, one it to return the
// real/imag, another is to set real/imag, they are both constexpr. We follow
// this design.
//
//
// [Operator +=,-=,*=,/=]
//
// Since C++20, these operators become constexpr. In our implementation, they
// are also constexpr.
//
// There are two types of such operators: operating with a real number, or
// operating with another complex number. For the operating with a real number,
// the generic template form has argument type `const T &`, while the overload
// for float/double/long double has `T`. We will follow the same type as
// float/double/long double in std.
//
// [Unary operator +-]
//
// Since C++20, they are constexpr. We also make them expr
//
// [Binary operators +-*/]
//
// Each operator has three versions (taking + as example):
// - complex + complex
// - complex + real
// - real + complex
//
// [Operator ==, !=]
//
// Each operator has three versions (taking == as example):
// - complex == complex
// - complex == real
// - real == complex
//
// Some of them are removed on C++20, but we decide to keep them
//
// [Operator <<, >>]
//
// These are implemented by casting to std::complex
//
//
//
// TODO(@zasdfgbnm): c10::complex<c10::Half> is not currently supported,
// because:
// - lots of members and functions of c10::Half are not constexpr
// - thrust::complex only support float and double
template <typename T>
struct alignas(sizeof(T) * 2) complex {
using value_type = T;
T real_ = T(0);
T imag_ = T(0);
constexpr complex() = default;
C10_HOST_DEVICE constexpr complex(const T& re, const T& im = T())
: real_(re), imag_(im) {}
template <typename U>
explicit constexpr complex(const std::complex<U>& other)
: complex(other.real(), other.imag()) {}
#if defined(__CUDACC__) || defined(__HIPCC__)
template <typename U>
explicit C10_HOST_DEVICE complex(const thrust::complex<U>& other)
: real_(other.real()), imag_(other.imag()) {}
// NOTE can not be implemented as follow due to ROCm bug:
// explicit C10_HOST_DEVICE complex(const thrust::complex<U> &other):
// complex(other.real(), other.imag()) {}
#endif
// Use SFINAE to specialize casting constructor for c10::complex<float> and
// c10::complex<double>
template <typename U = T>
C10_HOST_DEVICE explicit constexpr complex(
const std::enable_if_t<std::is_same<U, float>::value, complex<double>>&
other)
: real_(other.real_), imag_(other.imag_) {}
template <typename U = T>
C10_HOST_DEVICE constexpr complex(
const std::enable_if_t<std::is_same<U, double>::value, complex<float>>&
other)
: real_(other.real_), imag_(other.imag_) {}
constexpr complex<T>& operator=(T re) {
real_ = re;
imag_ = 0;
return *this;
}
constexpr complex<T>& operator+=(T re) {
real_ += re;
return *this;
}
constexpr complex<T>& operator-=(T re) {
real_ -= re;
return *this;
}
constexpr complex<T>& operator*=(T re) {
real_ *= re;
imag_ *= re;
return *this;
}
constexpr complex<T>& operator/=(T re) {
real_ /= re;
imag_ /= re;
return *this;
}
template <typename U>
constexpr complex<T>& operator=(const complex<U>& rhs) {
real_ = rhs.real();
imag_ = rhs.imag();
return *this;
}
template <typename U>
constexpr complex<T>& operator+=(const complex<U>& rhs) {
real_ += rhs.real();
imag_ += rhs.imag();
return *this;
}
template <typename U>
constexpr complex<T>& operator-=(const complex<U>& rhs) {
real_ -= rhs.real();
imag_ -= rhs.imag();
return *this;
}
template <typename U>
constexpr complex<T>& operator*=(const complex<U>& rhs) {
// (a + bi) * (c + di) = (a*c - b*d) + (a * d + b * c) i
T a = real_;
T b = imag_;
U c = rhs.real();
U d = rhs.imag();
real_ = a * c - b * d;
imag_ = a * d + b * c;
return *this;
}
#ifdef __APPLE__
#define FORCE_INLINE_APPLE __attribute__((always_inline))
#else
#define FORCE_INLINE_APPLE
#endif
template <typename U>
constexpr FORCE_INLINE_APPLE complex<T>& operator/=(const complex<U>& rhs)
__ubsan_ignore_float_divide_by_zero__ {
// (a + bi) / (c + di) = (ac + bd)/(c^2 + d^2) + (bc - ad)/(c^2 + d^2) i
// the calculation below follows numpy's complex division
T ar = real_;
T ai = imag_;
U br = rhs.real();
U bi = rhs.imag();
#if defined(__GNUC__) && !defined(__clang__)
// std::abs is already constexpr by gcc
auto abs_br = std::abs(br);
auto abs_bi = std::abs(bi);
#else
auto abs_br = br < 0 ? -br : br;
auto abs_bi = bi < 0 ? -bi : bi;
#endif
if (abs_br >= abs_bi) {
if (abs_br == 0 && abs_bi == 0) {
/* divide by zeros should yield a complex inf or nan */
real_ = ar / abs_br;
imag_ = ai / abs_bi;
} else {
auto rat = bi / br;
auto scl = 1.0 / (br + bi * rat);
real_ = (ar + ai * rat) * scl;
imag_ = (ai - ar * rat) * scl;
}
} else {
auto rat = br / bi;
auto scl = 1.0 / (bi + br * rat);
real_ = (ar * rat + ai) * scl;
imag_ = (ai * rat - ar) * scl;
}
return *this;
}
#undef FORCE_INLINE_APPLE
template <typename U>
constexpr complex<T>& operator=(const std::complex<U>& rhs) {
real_ = rhs.real();
imag_ = rhs.imag();
return *this;
}
#if defined(__CUDACC__) || defined(__HIPCC__)
template <typename U>
C10_HOST_DEVICE complex<T>& operator=(const thrust::complex<U>& rhs) {
real_ = rhs.real();
imag_ = rhs.imag();
return *this;
}
#endif
template <typename U>
explicit constexpr operator std::complex<U>() const {
return std::complex<U>(std::complex<T>(real(), imag()));
}
#if defined(__CUDACC__) || defined(__HIPCC__)
template <typename U>
C10_HOST_DEVICE explicit operator thrust::complex<U>() const {
return static_cast<thrust::complex<U>>(thrust::complex<T>(real(), imag()));
}
#endif
// consistent with NumPy behavior
explicit constexpr operator bool() const {
return real() || imag();
}
C10_HOST_DEVICE constexpr T real() const {
return real_;
}
constexpr void real(T value) {
real_ = value;
}
constexpr T imag() const {
return imag_;
}
constexpr void imag(T value) {
imag_ = value;
}
};
namespace complex_literals {
constexpr complex<float> operator"" _if(long double imag) {
return complex<float>(0.0f, static_cast<float>(imag));
}
constexpr complex<double> operator"" _id(long double imag) {
return complex<double>(0.0, static_cast<double>(imag));
}
constexpr complex<float> operator"" _if(unsigned long long imag) {
return complex<float>(0.0f, static_cast<float>(imag));
}
Loading ...