move to gko::complex<sycl::half> alone

yhmtsai · yhmtsai · commit a5ad85c33db5 · 2024-12-03T10:33:17.000+01:00
diff --git a/accessor/sycl_helper.hpp b/accessor/sycl_helper.hpp
@@ -33,6 +33,10 @@ namespace gko {
 class half;
 
 
+template <typename V>
+class complex;
+
+
 namespace acc {
 namespace detail {
 
@@ -81,6 +85,12 @@ struct sycl_type<std::complex<T>> {
 };
 
 
+template <>
+struct sycl_type<std::complex<gko::half>> {
+    using type = gko::complex<typename sycl_type<gko::half>::type>;
+};
+
+
 }  // namespace detail
 
 
diff --git a/cmake/sycl.cmake b/cmake/sycl.cmake
@@ -22,10 +22,6 @@ function(gko_add_sycl_to_target)
         "${one_value_args}"
         "${multi_value_args}"
         ${ARGN})
-    # trick for complex header chain
-    if("${GINKGO_DPCPP_MAJOR_VERSION}.${GINKGO_DPCPP_MINOR_VERSION}" VERSION_GREATER_EQUAL 7.1)
-        target_include_directories(${SYCL_TARGET} PRIVATE "${PROJECT_BINARY_DIR}/dpcpp/base")
-    endif()
     if(COMMAND add_sycl_to_target)
         add_sycl_to_target(${ARGN})
         return()
diff --git a/dpcpp/CMakeLists.txt b/dpcpp/CMakeLists.txt
@@ -4,11 +4,6 @@ find_package(oneDPL REQUIRED HINTS "$ENV{DPL_ROOT}" "$ENV{DPLROOT}")
 set(GINKGO_MKL_ROOT "${MKL_DIR}" PARENT_SCOPE)
 set(GINKGO_DPL_ROOT "${oneDPL_DIR}" PARENT_SCOPE)
 
-# trick for complex header chain
-if("${GINKGO_DPCPP_MAJOR_VERSION}.${GINKGO_DPCPP_MINOR_VERSION}" VERSION_GREATER_EQUAL 7.1)
-    configure_file(base/complex.hpp ${CMAKE_CURRENT_BINARY_DIR}/base/complex)
-endif()
-
 include(${PROJECT_SOURCE_DIR}/cmake/template_instantiation.cmake)
 add_instantiation_files(${PROJECT_SOURCE_DIR}/common/unified matrix/dense_kernels.instantiate.cpp DENSE_INSTANTIATE)
 add_instantiation_files(. solver/batch_bicgstab_launch.instantiate.dp.cpp BATCH_BICGSTAB_INSTANTIATE)
diff --git a/dpcpp/base/complex.hpp b/dpcpp/base/complex.hpp
@@ -5,29 +5,19 @@
 #ifndef GKO_DPCPP_BASE_COMPLEX_HPP_
 #define GKO_DPCPP_BASE_COMPLEX_HPP_
 
+#include <complex>
+
 #include <sycl/half_type.hpp>
 
 #include <ginkgo/config.hpp>
 
-// this file is to workaround for the intel sycl complex different loading.
-// intel sycl provides complex and the corresponding searching path. When users
-// load complex with -fsycl, the compiler will load intel's <complex> header
-// first and then load usual <complex> header. However, it implicitly
-// instantiates and uses std::complex<sycl::half>, so we need to provide the
-// implementation before that. In ginkgo, we will definitely load <complex> in
-// the public interface, which is before sycl backend, so we have no normal way
-// to provide the std::complex<sycl::half> implementation in sycl.
-// We apply the same trick to load this file first and then load their header
-// later. We will also configure this file as <complex> and provide the search
-// path in sycl module.
-// They start to do this from LIBSYCL 7.1.0.
-
-namespace std {
+
+namespace gko {
 
 template <typename>
 class complex;
 
-// implement std::complex<sycl::half> before knowing std::complex<float>
+
 template <>
 class complex<sycl::half> {
 public:
@@ -53,7 +43,7 @@ class complex<sycl::half> {
     {}
 
     template <typename T, typename = std::enable_if_t<std::is_scalar<T>::value>>
-    complex(const complex<T>& other)
+    complex(const std::complex<T>& other)
         : real_(static_cast<value_type>(other.real())),
           imag_(static_cast<value_type>(other.imag()))
     {}
@@ -62,7 +52,18 @@ class complex<sycl::half> {
 
     value_type imag() const noexcept { return imag_; }
 
-    inline operator std::complex<float>() const noexcept;
+    operator std::complex<float>() const noexcept
+    {
+        return std::complex<float>(static_cast<float>(real_),
+                                   static_cast<float>(imag_));
+    }
+
+    bool operator!=(const complex& r) const { return !this->operator==(r); }
+
+    bool operator==(const complex& r) const
+    {
+        return real_ == r.real() && imag_ == r.imag();
+    }
 
     template <typename V>
     complex& operator=(const V& val)
@@ -107,37 +108,83 @@ class complex<sycl::half> {
     }
 
     template <typename T>
-    complex& operator+=(const complex<T>& val)
+    complex& operator+=(const std::complex<T>& val)
     {
         real_ += val.real();
         imag_ += val.imag();
         return *this;
     }
 
     template <typename T>
-    complex& operator-=(const complex<T>& val)
+    complex& operator-=(const std::complex<T>& val)
     {
         real_ -= val.real();
         imag_ -= val.imag();
         return *this;
     }
 
     template <typename T>
-    inline complex& operator*=(const complex<T>& val);
+    complex& operator*=(const std::complex<T>& val)
+    {
+        auto val_f = static_cast<std::complex<float>>(val);
+        auto result_f = static_cast<std::complex<float>>(*this);
+        result_f *= val_f;
+        real_ = result_f.real();
+        imag_ = result_f.imag();
+        return *this;
+    }
 
     template <typename T>
-    inline complex& operator/=(const complex<T>& val);
+    complex& operator/=(const std::complex<T>& val)
+    {
+        auto val_f = static_cast<std::complex<float>>(val);
+        auto result_f = static_cast<std::complex<float>>(*this);
+        result_f /= val_f;
+        real_ = result_f.real();
+        imag_ = result_f.imag();
+        return *this;
+    }
+
+    complex& operator+=(const complex& val)
+    {
+        real_ += val.real();
+        imag_ += val.imag();
+        return *this;
+    }
+
+    complex& operator-=(const complex& val)
+    {
+        real_ -= val.real();
+        imag_ -= val.imag();
+        return *this;
+    }
 
-// It's for MacOS.
-// TODO: check whether mac compiler always use complex version even when real
-// half
-#define COMPLEX_HALF_OPERATOR(_op, _opeq)                                  \
-    friend complex<sycl::half> operator _op(const complex<sycl::half> lhf, \
-                                            const complex<sycl::half> rhf) \
-    {                                                                      \
-        auto a = lhf;                                                      \
-        a _opeq rhf;                                                       \
-        return a;                                                          \
+    complex& operator*=(const complex& val)
+    {
+        auto val_f = static_cast<std::complex<float>>(val);
+        auto result_f = static_cast<std::complex<float>>(*this);
+        result_f *= val_f;
+        real_ = result_f.real();
+        imag_ = result_f.imag();
+        return *this;
+    }
+
+    complex& operator/=(const complex& val)
+    {
+        auto val_f = static_cast<std::complex<float>>(val);
+        auto result_f = static_cast<std::complex<float>>(*this);
+        result_f /= val_f;
+        real_ = result_f.real();
+        imag_ = result_f.imag();
+        return *this;
+    }
+
+#define COMPLEX_HALF_OPERATOR(_op, _opeq)                               \
+    friend complex operator _op(const complex& lhf, const complex& rhf) \
+    {                                                                   \
+        auto a = lhf;                                                   \
+        a _opeq rhf;                                                    \
+        return a;                                                       \
     }
 
     COMPLEX_HALF_OPERATOR(+, +=)
@@ -147,77 +194,15 @@ class complex<sycl::half> {
 
 #undef COMPLEX_HALF_OPERATOR
 
+    complex operator-() const { return complex(-real_, -imag_); }
+
 private:
     value_type real_;
     value_type imag_;
 };
 
-}  // namespace std
-
-
-// after providing std::complex<sycl::half>, we can load their <complex> to
-// complete the header chain.
-
-#if GINKGO_DPCPP_MAJOR_VERSION > 7 || \
-    (GINKGO_DPCPP_MAJOR_VERSION == 7 && GINKGO_DPCPP_MINOR_VERSION >= 1)
-
-#if defined(__has_include_next)
-// GCC/clang support go through this path.
-#include_next <complex>
-#else
-// MSVC doesn't support "#include_next", so we take the same workaround in
-// stl_wrappers/complex.
-#include <../stl_wrappers/complex>
-#endif
-
-#else
-
-
-#include <complex>
-
-
-#endif
-
-
-// we know the complex<float> now, so we implement those functions requiring
-// complex<float>
-namespace std {
-
-
-inline complex<sycl::half>::operator complex<float>() const noexcept
-{
-    return std::complex<float>(static_cast<float>(real_),
-                               static_cast<float>(imag_));
-}
-
-
-template <typename T>
-inline complex<sycl::half>& complex<sycl::half>::operator*=(
-    const complex<T>& val)
-{
-    auto val_f = static_cast<std::complex<float>>(val);
-    auto result_f = static_cast<std::complex<float>>(*this);
-    result_f *= val_f;
-    real_ = result_f.real();
-    imag_ = result_f.imag();
-    return *this;
-}
-
-
-template <typename T>
-inline complex<sycl::half>& complex<sycl::half>::operator/=(
-    const complex<T>& val)
-{
-    auto val_f = static_cast<std::complex<float>>(val);
-    auto result_f = static_cast<std::complex<float>>(*this);
-    result_f /= val_f;
-    real_ = result_f.real();
-    imag_ = result_f.imag();
-    return *this;
-}
-
 
-}  // namespace std
+}  // namespace gko
 
 
 #endif  // GKO_DPCPP_BASE_COMPLEX_HPP_
diff --git a/dpcpp/base/math.hpp b/dpcpp/base/math.hpp
@@ -32,6 +32,41 @@ struct basic_float_traits<sycl::half> {
 template <>
 struct is_complex_or_scalar_impl<sycl::half> : public std::true_type {};
 
+template <typename ValueType>
+struct complex_helper {
+    using type = std::complex<ValueType>;
+};
+
+template <>
+struct complex_helper<sycl::half> {
+    using type = gko::complex<sycl::half>;
+};
+
+
+template <typename T>
+struct type_size_impl<gko::complex<T>> {
+    static constexpr auto value = sizeof(T) * byte_size;
+};
+
+
+template <typename T>
+struct remove_complex_impl<gko::complex<T>> {
+    using type = T;
+};
+
+
+template <typename T>
+struct truncate_type_impl<gko::complex<T>> {
+    using type =
+        typename complex_helper<typename truncate_type_impl<T>::type>::type;
+};
+
+template <typename T>
+struct is_complex_impl<gko::complex<T>> : public std::true_type {};
+
+template <typename T>
+struct is_complex_or_scalar_impl<gko::complex<T>>
+    : public is_complex_or_scalar_impl<T> {};
 
 }  // namespace detail
 
@@ -41,7 +76,7 @@ bool __dpct_inline__ is_nan(const sycl::half& val)
     return std::isnan(static_cast<float>(val));
 }
 
-bool __dpct_inline__ is_nan(const std::complex<sycl::half>& val)
+bool __dpct_inline__ is_nan(const gko::complex<sycl::half>& val)
 {
     return is_nan(val.real()) || is_nan(val.imag());
 }
@@ -52,7 +87,7 @@ sycl::half __dpct_inline__ abs(const sycl::half& val)
     return abs(static_cast<float>(val));
 }
 
-sycl::half __dpct_inline__ abs(const std::complex<sycl::half>& val)
+sycl::half __dpct_inline__ abs(const gko::complex<sycl::half>& val)
 {
     return abs(static_cast<std::complex<float>>(val));
 }
@@ -62,8 +97,8 @@ sycl::half __dpct_inline__ sqrt(const sycl::half& val)
     return sqrt(static_cast<float>(val));
 }
 
-std::complex<sycl::half> __dpct_inline__
-sqrt(const std::complex<sycl::half>& val)
+gko::complex<sycl::half> __dpct_inline__
+sqrt(const gko::complex<sycl::half>& val)
 {
     return sqrt(static_cast<std::complex<float>>(val));
 }
@@ -74,7 +109,7 @@ bool __dpct_inline__ is_finite(const sycl::half& value)
     return abs(value) < std::numeric_limits<sycl::half>::infinity();
 }
 
-bool __dpct_inline__ is_finite(const std::complex<sycl::half>& value)
+bool __dpct_inline__ is_finite(const gko::complex<sycl::half>& value)
 {
     return is_finite(value.real()) && is_finite(value.imag());
 }
diff --git a/dpcpp/base/types.hpp b/dpcpp/base/types.hpp
@@ -14,6 +14,8 @@
 #include <ginkgo/core/base/matrix_data.hpp>
 #include <ginkgo/core/base/types.hpp>
 
+#include "dpcpp/base/complex.hpp"
+
 
 namespace gko {
 namespace kernels {
@@ -56,6 +58,11 @@ struct sycl_type_impl<std::complex<T>> {
     using type = std::complex<typename sycl_type_impl<T>::type>;
 };
 
+template <>
+struct sycl_type_impl<std::complex<gko::half>> {
+    using type = gko::complex<typename sycl_type_impl<gko::half>::type>;
+};
+
 template <typename ValueType, typename IndexType>
 struct sycl_type_impl<matrix_data_entry<ValueType, IndexType>> {
     using type =
diff --git a/dpcpp/preconditioner/batch_block_jacobi.hpp b/dpcpp/preconditioner/batch_block_jacobi.hpp
@@ -131,7 +131,7 @@ class BlockJacobi final {
 
             // reduction (it does not support complex<half>)
             if constexpr (std::is_same_v<value_type,
-                                         std::complex<sycl::half>>) {
+                                         gko::complex<sycl::half>>) {
                 for (int i = sg_size / 2; i > 0; i /= 2) {
                     sum += sycl::shift_group_left(sg, sum, i);
                 }

Original file line number	Diff line number	Diff line change
`@@ -131,7 +131,7 @@ class BlockJacobi final {`
`131`	`131`
`132`	`132`	`// reduction (it does not support complex<half>)`
`133`	`133`	`if constexpr (std::is_same_v<value_type,`
`134`		`- std::complex<sycl::half>>) {`
	`134`	`+ gko::complex<sycl::half>>) {`
`135`	`135`	`for (int i = sg_size / 2; i > 0; i /= 2) {`
`136`	`136`	`sum += sycl::shift_group_left(sg, sum, i);`
`137`	`137`	`}`