add dpcpp transpose kernels

dpcpp/components/atomic.dp.hpp

@@ -0,0 +1,214 @@
+/*******************************<GINKGO LICENSE>******************************
+Copyright (c) 2017-2021, the Ginkgo authors
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions
+are met:
+
+1. Redistributions of source code must retain the above copyright
+notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+contributors may be used to endorse or promote products derived from
+this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
+IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
+TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
+PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+******************************<GINKGO LICENSE>*******************************/
+
+#ifndef GKO_DPCPP_COMPONENTS_ATOMIC_DP_HPP_
+#define GKO_DPCPP_COMPONENTS_ATOMIC_DP_HPP_
+
+
+#include <type_traits>
+
+
+#include <CL/sycl.hpp>
+
+
+#include "dpcpp/base/dpct.hpp"
+
+// #include "dpct/atomic.hpp"
+
+namespace gko {
+namespace kernels {
+namespace dpcpp {
+namespace atomic {
+constexpr cl::sycl::access::address_space local_space =
+    cl::sycl::access::address_space::local_space;
+constexpr cl::sycl::access::address_space global_space =
+    cl::sycl::access::address_space::global_space;
+}  // namespace atomic
+
+namespace {
+
+
+template <cl::sycl::access::address_space addressSpace = atomic::global_space,
+          typename T>
+T atomic_compare_exchange_strong(
+    cl::sycl::multi_ptr<T, addressSpace> addr, T expected, T desired,
+    cl::sycl::memory_order success = cl::sycl::memory_order::relaxed,
+    cl::sycl::memory_order fail = cl::sycl::memory_order::relaxed)
+{
+    cl::sycl::atomic<T, addressSpace> obj(addr);
+    obj.compare_exchange_strong(expected, desired, success, fail);
+    return expected;
+}
+
+template <cl::sycl::access::address_space addressSpace = atomic::global_space,
+          typename T>
+T atomic_compare_exchange_strong(
+    T *addr, T expected, T desired,
+    cl::sycl::memory_order success = cl::sycl::memory_order::relaxed,
+    cl::sycl::memory_order fail = cl::sycl::memory_order::relaxed)
+{
+    return atomic_compare_exchange_strong(
+        cl::sycl::multi_ptr<T, addressSpace>(addr), expected, desired, success,
+        fail);
+}
+
+
+template <cl::sycl::access::address_space addressSpace = atomic::global_space,
+          typename T>
+inline T atomic_fetch_add(
+    T *addr, T operand,
+    cl::sycl::memory_order memoryOrder = cl::sycl::memory_order::relaxed)
+{
+    cl::sycl::atomic<T, addressSpace> obj(
+        (cl::sycl::multi_ptr<T, addressSpace>(addr)));
+    return cl::sycl::atomic_fetch_add(obj, operand, memoryOrder);
+}
+
+template <typename T>
+struct fake_complex {
+    T x;
+    T y;
+};
+
+}  // namespace
+
+
+namespace detail {
+
+
+template <cl::sycl::access::address_space addressSpace, typename ValueType,
+          typename = void>
+struct atomic_helper {
+    __dpct_inline__ static ValueType atomic_add(ValueType *, ValueType)
+    {
+        static_assert(sizeof(ValueType) == 0,
+                      "This default function is not implemented, only the "
+                      "specializations are.");
+        // TODO: add proper implementation of generic atomic add
+    }
+};
+
+
+template <typename ResultType, typename ValueType>
+__dpct_inline__ ResultType reinterpret(ValueType val)
+{
+    static_assert(sizeof(ValueType) == sizeof(ResultType),
+                  "The type to reinterpret to must be of the same size as the "
+                  "original type.");
+    return reinterpret_cast<ResultType &>(val);
+}
+
+
+#define GKO_BIND_ATOMIC_HELPER_STRUCTURE(CONVERTER_TYPE)                   \
+    template <cl::sycl::access::address_space addressSpace,                \
+              typename ValueType>                                          \
+    struct atomic_helper<                                                  \
+        addressSpace, ValueType,                                           \
+        std::enable_if_t<(sizeof(ValueType) == sizeof(CONVERTER_TYPE))>> { \
+        __dpct_inline__ static ValueType atomic_add(                       \
+            ValueType *__restrict__ addr, ValueType val)                   \
+        {                                                                  \
+            CONVERTER_TYPE *address_as_converter =                         \
+                reinterpret_cast<CONVERTER_TYPE *>(addr);                  \
+            CONVERTER_TYPE old = *address_as_converter;                    \
+            CONVERTER_TYPE assumed;                                        \
+            do {                                                           \
+                assumed = old;                                             \
+                old = atomic_compare_exchange_strong<addressSpace>(        \
+                    address_as_converter, assumed,                         \
+                    reinterpret<CONVERTER_TYPE>(                           \
+                        val + reinterpret<ValueType>(assumed)));           \
+            } while (assumed != old);                                      \
+            return reinterpret<ValueType>(old);                            \
+        }                                                                  \
+    };
+
+// Support 64-bit ATOMIC_ADD
+GKO_BIND_ATOMIC_HELPER_STRUCTURE(unsigned long long int);
+// Support 32-bit ATOMIC_ADD
+GKO_BIND_ATOMIC_HELPER_STRUCTURE(unsigned int);
+
+
+#undef GKO_BIND_ATOMIC_HELPER_STRUCTURE
+
+#define GKO_BIND_ATOMIC_HELPER_ValueType(ValueType)                         \
+    template <cl::sycl::access::address_space addressSpace>                 \
+    struct atomic_helper<addressSpace, ValueType, std::enable_if_t<true>> { \
+        __dpct_inline__ static ValueType atomic_add(                        \
+            ValueType *__restrict__ addr, ValueType val)                    \
+        {                                                                   \
+            return atomic_fetch_add<addressSpace>(addr, val);               \
+        }                                                                   \
+    };
+
+GKO_BIND_ATOMIC_HELPER_ValueType(int);
+GKO_BIND_ATOMIC_HELPER_ValueType(unsigned int);
+GKO_BIND_ATOMIC_HELPER_ValueType(unsigned long long int);
+
+
+template <cl::sycl::access::address_space addressSpace, typename ValueType>
+struct atomic_helper<
+    addressSpace, ValueType,
+    std::enable_if_t<is_complex<ValueType>() && sizeof(ValueType) >= 16>> {
+    __dpct_inline__ static ValueType atomic_add(ValueType *__restrict__ addr,
+                                                ValueType val)
+    {
+        using real_type = remove_complex<ValueType>;
+        fake_complex<real_type> *fake_addr =
+            reinterpret_cast<fake_complex<real_type> *>(addr);
+        // Separate to real part and imag part
+        auto real = atomic_helper<addressSpace, real_type>::atomic_add(
+            &(fake_addr->x), val.real());
+        auto imag = atomic_helper<addressSpace, real_type>::atomic_add(
+            &(fake_addr->y), val.imag());
+        return {real, imag};
+    }
+};
+
+
+}  // namespace detail
+
+
+template <cl::sycl::access::address_space addressSpace = atomic::global_space,
+          typename T>
+__dpct_inline__ T atomic_add(T *__restrict__ addr, T val)
+{
+    return detail::atomic_helper<addressSpace, T>::atomic_add(addr, val);
+}
+
+
+}  // namespace dpcpp
+}  // namespace kernels
+}  // namespace gko
+
+
+#endif  // GKO_DPCPP_COMPONENTS_ATOMIC_DP_HPP_

dpcpp/matrix/csr_kernels.dp.cpp

@@ -52,8 +52,10 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include "core/base/allocator.hpp"
 #include "core/base/iterator_factory.hpp"
 #include "core/base/utils.hpp"
+#include "core/components/fill_array.hpp"
 #include "core/components/prefix_sum.hpp"
 #include "core/matrix/csr_builder.hpp"
+#include "dpcpp/components/atomic.dp.hpp"
 #include "dpcpp/components/format_conversion.dp.hpp"
 
 
@@ -629,10 +631,64 @@ void transpose_and_transform(std::shared_ptr<const DpcppExecutor> exec,
                              UnaryOperator op) GKO_NOT_IMPLEMENTED;
 
 
+template <bool conjugate, typename ValueType, typename IndexType>
+void generic_transpose(std::shared_ptr<const DpcppExecutor> exec,
+                       const matrix::Csr<ValueType, IndexType> *orig,
+                       matrix::Csr<ValueType, IndexType> *trans)
+{
+    const auto num_rows = orig->get_size()[0];
+    const auto num_cols = orig->get_size()[1];
+    auto queue = exec->get_queue();
+    const auto row_ptrs = orig->get_const_row_ptrs();
+    const auto cols = orig->get_const_col_idxs();
+    const auto vals = orig->get_const_values();
+
+    Array<IndexType> counts{exec, num_cols + 1};
+    auto tmp_counts = counts.get_data();
+    auto out_row_ptrs = trans->get_row_ptrs();
+    auto out_cols = trans->get_col_idxs();
+    auto out_vals = trans->get_values();
+    components::fill_array(exec, tmp_counts, num_cols, IndexType{});
+
+    queue->submit([&](sycl::handler &cgh) {
+        cgh.parallel_for(sycl::range<1>{num_rows}, [=](sycl::id<1> idx) {
+            const auto row = static_cast<size_type>(idx[0]);
+            const auto begin = row_ptrs[row];
+            const auto end = row_ptrs[row + 1];
+            for (auto i = begin; i < end; i++) {
+                atomic_fetch_add(tmp_counts + cols[i], IndexType{1});
+            }
+        });
+    });
+
+    components::prefix_sum(exec, tmp_counts, num_cols + 1);
+    exec->copy(num_cols + 1, tmp_counts, out_row_ptrs);
+
+    queue->submit([&](sycl::handler &cgh) {
+        cgh.parallel_for(sycl::range<1>{num_rows}, [=](sycl::id<1> idx) {
+            const auto row = static_cast<size_type>(idx[0]);
+            const auto begin = row_ptrs[row];
+            const auto end = row_ptrs[row + 1];
+            for (auto i = begin; i < end; i++) {
+                auto out_nz =
+                    atomic_fetch_add(tmp_counts + cols[i], IndexType{1});
+                out_cols[out_nz] = row;
+                out_vals[out_nz] = conjugate ? conj(vals[i]) : vals[i];
+            }
+        });
+    });
+
+    sort_by_column_index(exec, trans);
+}
+
+
 template <typename ValueType, typename IndexType>
 void transpose(std::shared_ptr<const DpcppExecutor> exec,
                const matrix::Csr<ValueType, IndexType> *orig,
-               matrix::Csr<ValueType, IndexType> *trans) GKO_NOT_IMPLEMENTED;
+               matrix::Csr<ValueType, IndexType> *trans)
+{
+    generic_transpose<false>(exec, orig, trans);
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_TRANSPOSE_KERNEL);
 
@@ -641,7 +697,9 @@ template <typename ValueType, typename IndexType>
 void conj_transpose(std::shared_ptr<const DpcppExecutor> exec,
                     const matrix::Csr<ValueType, IndexType> *orig,
                     matrix::Csr<ValueType, IndexType> *trans)
-    GKO_NOT_IMPLEMENTED;
+{
+    generic_transpose<true>(exec, orig, trans);
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_CONJ_TRANSPOSE_KERNEL);

dpcpp/test/matrix/csr_kernels.cpp

@@ -65,6 +65,8 @@ class Csr : public ::testing::Test {
     using Arr = gko::Array<int>;
     using Mtx = gko::matrix::Csr<value_type>;
     using Vec = gko::matrix::Dense<value_type>;
+    using ComplexVec = gko::matrix::Dense<std::complex<value_type>>;
+    using ComplexMtx = gko::matrix::Csr<std::complex<value_type>>;
 
     Csr() : mtx_size(532, 231), rand_engine(42) {}
 
@@ -102,12 +104,17 @@ class Csr : public ::testing::Test {
     {
         mtx = Mtx::create(ref);
         mtx->copy_from(gen_mtx<Vec>(mtx_size[0], mtx_size[1], 1));
+        complex_mtx = ComplexMtx::create(ref);
+        complex_mtx->copy_from(
+            gen_mtx<ComplexVec>(mtx_size[0], mtx_size[1], 1));
         square_mtx = Mtx::create(ref);
         square_mtx->copy_from(gen_mtx<Vec>(mtx_size[0], mtx_size[0], 1));
         alpha = gko::initialize<Vec>({2.0}, ref);
         beta = gko::initialize<Vec>({-1.0}, ref);
         dmtx = Mtx::create(dpcpp);
         dmtx->copy_from(mtx.get());
+        complex_dmtx = ComplexMtx::create(dpcpp);
+        complex_dmtx->copy_from(complex_mtx.get());
         square_dmtx = Mtx::create(dpcpp);
         square_dmtx->copy_from(square_mtx.get());
         dalpha = Vec::create(dpcpp);
@@ -141,11 +148,13 @@ class Csr : public ::testing::Test {
     std::ranlux48 rand_engine;
 
     std::unique_ptr<Mtx> mtx;
+    std::unique_ptr<ComplexMtx> complex_mtx;
     std::unique_ptr<Mtx> square_mtx;
     std::unique_ptr<Vec> alpha;
     std::unique_ptr<Vec> beta;
 
     std::unique_ptr<Mtx> dmtx;
+    std::unique_ptr<ComplexMtx> complex_dmtx;
     std::unique_ptr<Mtx> square_dmtx;
     std::unique_ptr<Vec> dalpha;
     std::unique_ptr<Vec> dbeta;
@@ -258,6 +267,30 @@ TEST_F(Csr, AdvancedApplyToIdentityMatrixIsEquivalentToRef)
 }
 
 
+TEST_F(Csr, TransposeIsEquivalentToRef)
+{
+    set_up_apply_data();
+
+    auto trans = gko::as<Mtx>(mtx->transpose());
+    auto d_trans = gko::as<Mtx>(dmtx->transpose());
+
+    GKO_ASSERT_MTX_NEAR(d_trans, trans, 0.0);
+    ASSERT_TRUE(d_trans->is_sorted_by_column_index());
+}
+
+
+TEST_F(Csr, ConjugateTransposeIsEquivalentToRef)
+{
+    set_up_apply_data();
+
+    auto trans = gko::as<ComplexMtx>(complex_mtx->conj_transpose());
+    auto d_trans = gko::as<ComplexMtx>(complex_dmtx->conj_transpose());
+
+    GKO_ASSERT_MTX_NEAR(d_trans, trans, 0.0);
+    ASSERT_TRUE(d_trans->is_sorted_by_column_index());
+}
+
+
 TEST_F(Csr, RecognizeSortedMatrixIsEquivalentToRef)
 {
     set_up_apply_data();