ginkgo-project · upsj · Dec 21, 2019 · Dec 19, 2019 · Dec 19, 2019 · Dec 21, 2019
diff --git a/cuda/test/matrix/csr_kernels.cpp b/cuda/test/matrix/csr_kernels.cpp
@@ -131,7 +131,7 @@ class Csr : public ::testing::Test {
 
     matrix_pair gen_unsorted_mtx()
     {
-        constexpr int min_nnz_per_row = 2;  // Must be larger/equal than 2
+        constexpr int min_nnz_per_row = 2;  // Must be at least 2
         auto local_mtx_ref =
             gen_mtx<Mtx>(mtx_size[0], mtx_size[1], min_nnz_per_row);
         for (size_t row = 0; row < mtx_size[0]; ++row) {

diff --git a/hip/base/hipsparse_bindings.hip.hpp b/hip/base/hipsparse_bindings.hip.hpp
@@ -602,6 +602,84 @@ inline void destroy_spgemm_info(csrgemm2Info_t info)
 }
 
 
+template <typename IndexType>
+void create_identity_permutation(hipsparseHandle_t handle, IndexType size,
+                                 IndexType *permutation) GKO_NOT_IMPLEMENTED;
+
+template <>
+inline void create_identity_permutation<int32>(hipsparseHandle_t handle,
+                                               int32 size, int32 *permutation)
+{
+    GKO_ASSERT_NO_HIPSPARSE_ERRORS(
+        hipsparseCreateIdentityPermutation(handle, size, permutation));
+}
+
+
+template <typename IndexType>
+void csrsort_buffer_size(hipsparseHandle_t handle, IndexType m, IndexType n,
+                         IndexType nnz, const IndexType *row_ptrs,
+                         const IndexType *col_idxs,
+                         size_type &buffer_size) GKO_NOT_IMPLEMENTED;
+
+template <>
+inline void csrsort_buffer_size<int32>(hipsparseHandle_t handle, int32 m,
+                                       int32 n, int32 nnz,
+                                       const int32 *row_ptrs,
+                                       const int32 *col_idxs,
+                                       size_type &buffer_size)
+{
+    GKO_ASSERT_NO_HIPSPARSE_ERRORS(hipsparseXcsrsort_bufferSizeExt(
+        handle, m, n, nnz, row_ptrs, col_idxs, &buffer_size));
+}
+
+
+template <typename IndexType>
+void csrsort(hipsparseHandle_t handle, IndexType m, IndexType n, IndexType nnz,
+             const hipsparseMatDescr_t descr, const IndexType *row_ptrs,
+             IndexType *col_idxs, IndexType *permutation,
+             void *buffer) GKO_NOT_IMPLEMENTED;
+
+template <>
+inline void csrsort<int32>(hipsparseHandle_t handle, int32 m, int32 n,
+                           int32 nnz, const hipsparseMatDescr_t descr,
+                           const int32 *row_ptrs, int32 *col_idxs,
+                           int32 *permutation, void *buffer)
+{
+    GKO_ASSERT_NO_HIPSPARSE_ERRORS(hipsparseXcsrsort(
+        handle, m, n, nnz, descr, row_ptrs, col_idxs, permutation, buffer));
+}
+
+
+template <typename IndexType, typename ValueType>
+void gather(hipsparseHandle_t handle, IndexType nnz, const ValueType *in,
+            ValueType *out, const IndexType *permutation) GKO_NOT_IMPLEMENTED;
+
+#define GKO_BIND_HIPSPARSE_GATHER(ValueType, HipsparseName)                   \
+    template <>                                                               \
+    inline void gather<int32, ValueType>(hipsparseHandle_t handle, int32 nnz, \
+                                         const ValueType *in, ValueType *out, \
+                                         const int32 *permutation)            \
+    {                                                                         \
+        GKO_ASSERT_NO_HIPSPARSE_ERRORS(HipsparseName(                         \
+            handle, nnz, as_hiplibs_type(in), as_hiplibs_type(out),           \
+            permutation, HIPSPARSE_INDEX_BASE_ZERO));                         \
+    }                                                                         \
+    static_assert(true,                                                       \
+                  "This assert is used to counter the false positive extra "  \
+                  "semi-colon warnings")
+
+GKO_BIND_HIPSPARSE_GATHER(float, hipsparseSgthr);
+GKO_BIND_HIPSPARSE_GATHER(double, hipsparseDgthr);
+#if defined(hipsparseVersionMajor) && defined(hipsparseVersionMinor) && \
+    ((hipsparseVersionMajor > 1) ||                                     \
+     (hipsparseVersionMajor == 1 && hipsparseVersionMinor >= 4))
+GKO_BIND_HIPSPARSE_GATHER(std::complex<float>, hipsparseCgthr);
+GKO_BIND_HIPSPARSE_GATHER(std::complex<double>, hipsparseZgthr);
+#endif  // hipsparse version >= 1.4
+
+#undef GKO_BIND_HIPSPARSE_GATHER
+
+
 }  // namespace hipsparse
 }  // namespace hip
 }  // namespace kernels

diff --git a/hip/matrix/csr_kernels.hip.cpp b/hip/matrix/csr_kernels.hip.cpp
@@ -1086,7 +1086,46 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 template <typename ValueType, typename IndexType>
 void sort_by_column_index(std::shared_ptr<const HipExecutor> exec,
                           matrix::Csr<ValueType, IndexType> *to_sort)
-    GKO_NOT_IMPLEMENTED;
+{
+    if (hipsparse::is_supported<ValueType, IndexType>::value) {
+        auto handle = exec->get_hipsparse_handle();
+        auto descr = hipsparse::create_mat_descr();
+        auto m = IndexType(to_sort->get_size()[0]);
+        auto n = IndexType(to_sort->get_size()[1]);
+        auto nnz = IndexType(to_sort->get_num_stored_elements());
+        auto row_ptrs = to_sort->get_const_row_ptrs();
+        auto col_idxs = to_sort->get_col_idxs();
+        auto vals = to_sort->get_values();
+
+        // copy values
+        Array<ValueType> tmp_vals_array(exec, nnz);
+        exec->copy_from(exec.get(), nnz, vals, tmp_vals_array.get_data());
+        auto tmp_vals = tmp_vals_array.get_const_data();
+
+        // init identity permutation
+        Array<IndexType> permutation_array(exec, nnz);
+        auto permutation = permutation_array.get_data();
+        hipsparse::create_identity_permutation(handle, nnz, permutation);
+
+        // allocate buffer
+        size_type buffer_size{};
+        hipsparse::csrsort_buffer_size(handle, m, n, nnz, row_ptrs, col_idxs,
+                                       buffer_size);
+        Array<char> buffer_array{exec, buffer_size};
+        auto buffer = buffer_array.get_data();
+
+        // sort column indices
+        hipsparse::csrsort(handle, m, n, nnz, descr, row_ptrs, col_idxs,
+                           permutation, buffer);
+
+        // sort values
+        hipsparse::gather(handle, nnz, tmp_vals, vals, permutation);
+
+        hipsparse::destroy(descr);
+    } else {
+        GKO_NOT_IMPLEMENTED;
+    }
+}
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_SORT_BY_COLUMN_INDEX);

diff --git a/hip/test/matrix/csr_kernels.hip.cpp b/hip/test/matrix/csr_kernels.hip.cpp
@@ -112,6 +112,38 @@ class Csr : public ::testing::Test {
         dbeta->copy_from(beta.get());
     }
 
+    struct matrix_pair {
+        std::unique_ptr<Mtx> ref;
+        std::unique_ptr<Mtx> hip;
+    };
+
+    matrix_pair gen_unsorted_mtx()
+    {
+        constexpr int min_nnz_per_row = 2;  // Must be at least 2
+        auto local_mtx_ref =
+            gen_mtx<Mtx>(mtx_size[0], mtx_size[1], min_nnz_per_row);
+        for (size_t row = 0; row < mtx_size[0]; ++row) {
+            const auto row_ptrs = local_mtx_ref->get_const_row_ptrs();
+            const auto start_row = row_ptrs[row];
+            auto col_idx = local_mtx_ref->get_col_idxs() + start_row;
+            auto vals = local_mtx_ref->get_values() + start_row;
+            const auto nnz_in_this_row = row_ptrs[row + 1] - row_ptrs[row];
+            auto swap_idx_dist =
+                std::uniform_int_distribution<>(0, nnz_in_this_row - 1);
+            // shuffle `nnz_in_this_row / 2` times
+            for (size_t perm = 0; perm < nnz_in_this_row; perm += 2) {
+                const auto idx1 = swap_idx_dist(rand_engine);
+                const auto idx2 = swap_idx_dist(rand_engine);
+                std::swap(col_idx[idx1], col_idx[idx2]);
+                std::swap(vals[idx1], vals[idx2]);
+            }
+        }
+        auto local_mtx_hip = Mtx::create(hip);
+        local_mtx_hip->copy_from(local_mtx_ref.get());
+
+        return {std::move(local_mtx_ref), std::move(local_mtx_hip)};
+    }
+
     std::shared_ptr<gko::ReferenceExecutor> ref;
     std::shared_ptr<const gko::HipExecutor> hip;
 
@@ -552,4 +584,28 @@ TEST_F(Csr, MoveToHybridIsEquivalentToRef)
 }
 
 
+TEST_F(Csr, SortSortedMatrixIsEquivalentToRef)
+{
+    set_up_apply_data(std::make_shared<Mtx::automatical>(hip));
+
+    mtx->sort_by_column_index();
+
6D47
    dmtx->sort_by_column_index();
+
+    // Values must be unchanged, therefore, tolerance is `0`
+    GKO_ASSERT_MTX_NEAR(mtx, dmtx, 0);
+}
+
+
+TEST_F(Csr, SortUnsortedMatrixIsEquivalentToRef)
+{
+    auto uns_mtx = gen_unsorted_mtx();
+
+    uns_mtx.ref->sort_by_column_index();
+    uns_mtx.hip->sort_by_column_index();
+
+    // Values must be unchanged, therefore, tolerance is `0`
+    GKO_ASSERT_MTX_NEAR(uns_mtx.ref, uns_mtx.hip, 0);
+}
+
+
 }  // namespace
diff --git a/omp/test/matrix/csr_kernels.cpp b/omp/test/matrix/csr_kernels.cpp
@@ -141,7 +141,7 @@ class Csr : public ::testing::Test {
 
     matrix_pair gen_unsorted_mtx()
     {
-        constexpr int min_nnz_per_row = 2;  // Must be larger/equal than 2
+        constexpr int min_nnz_per_row = 2;  // Must be at least 2
         auto local_mtx_ref =
             gen_mtx<Mtx>(mtx_size[0], mtx_size[1], min_nnz_per_row);
         for (size_t row = 0; row < mtx_size[0]; ++row) {

diff --git a/omp/test/matrix/sparsity_csr_kernels.cpp b/omp/test/matrix/sparsity_csr_kernels.cpp
@@ -117,7 +117,7 @@ class SparsityCsr : public ::testing::Test {
 
     matrix_pair gen_unsorted_mtx()
     {
-        constexpr int min_nnz_per_row = 2;  // Must be larger/equal than 2
+        constexpr int min_nnz_per_row = 2;  // Must be at least 2
         auto local_mtx_ref =
             gen_mtx<Mtx>(mtx_size[0], mtx_size[1], min_nnz_per_row);
         for (size_t row = 0; row < mtx_size[0]; ++row) {