1 /** 2 * Copyright 2021 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef MINDSPORE_CORE_BASE_COMPLEX_STORAGE_H_ 17 #define MINDSPORE_CORE_BASE_COMPLEX_STORAGE_H_ 18 19 #include "base/float16.h" 20 21 namespace mindspore { 22 23 template <typename T> 24 struct alignas(sizeof(T) * 2) ComplexStorage { 25 T real_; 26 T imag_; 27 28 ComplexStorage() = default; 29 ~ComplexStorage() = default; 30 31 ComplexStorage(const ComplexStorage<T> &other) noexcept = default; 32 ComplexStorage(ComplexStorage<T> &&other) noexcept = default; 33 34 ComplexStorage &operator=(const ComplexStorage<T> &other) noexcept = default; 35 ComplexStorage &operator=(ComplexStorage<T> &&other) noexcept = default; 36 real_ComplexStorage37 inline constexpr ComplexStorage(const T &real, const T &imag = T()) : real_(real), imag_(imag) {} 38 ComplexStorageComplexStorage39 inline explicit constexpr ComplexStorage(const float16 &real) : real_(static_cast<T>(real)), imag_(T()) {} 40 41 template <typename U = T> ComplexStorageComplexStorage42 explicit ComplexStorage(const std::enable_if_t<std::is_same<U, float>::value, ComplexStorage<double>> &other) 43 : real_(other.real_), imag_(other.imag_) {} 44 45 template <typename U = T> ComplexStorageComplexStorage46 explicit ComplexStorage(const std::enable_if_t<std::is_same<U, double>::value, ComplexStorage<float>> &other) 47 : real_(other.real_), imag_(other.imag_) {} 48 49 inline explicit operator bool() const { return static_cast<bool>(real_) || static_cast<bool>(imag_); } 50 inline explicit operator signed char() const { return static_cast<signed char>(real_); } 51 inline explicit operator unsigned char() const { return static_cast<unsigned char>(real_); } 52 inline explicit operator double() const { return static_cast<double>(real_); } 53 inline explicit operator float() const { return static_cast<float>(real_); } int16_tComplexStorage54 inline explicit operator int16_t() const { return static_cast<int16_t>(real_); } uint16_tComplexStorage55 inline explicit operator uint16_t() const { return static_cast<uint16_t>(real_); } int32_tComplexStorage56 inline explicit operator int32_t() const { return static_cast<int32_t>(real_); } uint32_tComplexStorage57 inline explicit operator uint32_t() const { return static_cast<uint32_t>(real_); } int64_tComplexStorage58 inline explicit operator int64_t() const { return static_cast<int64_t>(real_); } uint64_tComplexStorage59 inline explicit operator uint64_t() const { return static_cast<uint64_t>(real_); } float16ComplexStorage60 inline explicit operator float16() const { return static_cast<float16>(real_); } 61 }; 62 63 template <typename T> 64 inline bool operator==(const ComplexStorage<T> &lhs, const ComplexStorage<T> &rhs) { 65 return lhs.real_ == rhs.real_ && lhs.imag_ == rhs.imag_; 66 } 67 68 template <typename T> 69 inline std::ostream &operator<<(std::ostream &os, const ComplexStorage<T> &v) { 70 return (os << std::noshowpos << v.real_ << std::showpos << v.imag_ << 'j'); 71 } 72 73 } // namespace mindspore 74 75 #endif // MINDSPORE_CORE_BASE_COMPLEX_STORAGE_H_ 76