Add Coo Transpose

common/cuda_hip/base/device_matrix_data_kernels.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -118,13 +118,12 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType, typename IndexType>
 void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
-                    device_matrix_data<ValueType, IndexType>& data)
+                    size_type num_elems, IndexType* row_idxs,
+                    IndexType* col_idxs, ValueType* vals)
 {
-    auto it = thrust::make_zip_iterator(
-        thrust::make_tuple(data.get_row_idxs(), data.get_col_idxs()));
-    auto vals = as_device_type(data.get_values());
-    thrust::sort_by_key(thrust_policy(exec), it,
-                        it + data.get_num_stored_elements(), vals);
+    auto it = thrust::make_zip_iterator(thrust::make_tuple(row_idxs, col_idxs));
+    auto vals_it = as_device_type(vals);
+    thrust::sort_by_key(thrust_policy(exec), it, it + num_elems, vals_it);
 }
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(

common/cuda_hip/multigrid/pgm_kernels.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -81,22 +81,6 @@ void sort_agg(std::shared_ptr<const DefaultExecutor> exec, IndexType num,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PGM_SORT_AGG_KERNEL);
 
 
-template <typename ValueType, typename IndexType>
-void sort_row_major(std::shared_ptr<const DefaultExecutor> exec, size_type nnz,
-                    IndexType* row_idxs, IndexType* col_idxs, ValueType* vals)
-{
-    auto vals_it = as_device_type(vals);
-    auto it = thrust::make_zip_iterator(thrust::make_tuple(row_idxs, col_idxs));
-    // Because reduce_by_key is not deterministic, so we do not need
-    // stable_sort_by_key
-    // TODO: If we have deterministic reduce_by_key, it should be
-    // stable_sort_by_key
-    thrust::sort_by_key(thrust_policy(exec), it, it + nnz, vals_it);
-}
-
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM_SORT_ROW_MAJOR);
-
-
 template <typename ValueType, typename IndexType>
 void compute_coarse_coo(std::shared_ptr<const DefaultExecutor> exec,
                         size_type fine_nnz, const IndexType* row_idxs,

common/unified/matrix/coo_kernels.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -62,6 +62,21 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_COO_FILL_IN_DENSE_KERNEL);
 
 
+template <typename ValueType>
+void conj_array(std::shared_ptr<const DefaultExecutor> exec,
+                size_type num_elems, ValueType* values)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto tidx, auto values) {
+            values[tidx] = conj(values[tidx]);
+        },
+        num_elems, values);
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_TYPE(GKO_DECLARE_COO_CONJ_ARRAY_KERNEL);
+
+
 }  // namespace coo
 }  // namespace GKO_DEVICE_NAMESPACE
 }  // namespace kernels

core/base/device_matrix_data.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -105,7 +105,9 @@ void device_matrix_data<ValueType, IndexType>::fill_zero()
 template <typename ValueType, typename IndexType>
 void device_matrix_data<ValueType, IndexType>::sort_row_major()
 {
-    this->values_.get_executor()->run(components::make_sort_row_major(*this));
+    this->values_.get_executor()->run(components::make_sort_row_major(
+        this->get_num_stored_elements(), this->get_row_idxs(),
+        this->get_col_idxs(), this->get_values()));
 }
 
 

core/base/device_matrix_data_kernels.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -45,7 +45,8 @@ namespace kernels {
 #define GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL(ValueType, \
                                                              IndexType) \
     void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,    \
-                        device_matrix_data<ValueType, IndexType>& data)
+                        size_type num_elems, IndexType* row_idxs,       \
+                        IndexType* col_idxs, ValueType* values)
 
 
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                          \

core/device_hooks/common_kernels.inc.cpp

@@ -810,6 +810,7 @@
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COO_ADVANCED_SPMV2_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COO_FILL_IN_DENSE_KERNEL);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_COO_EXTRACT_DIAGONAL_KERNEL);
+GKO_STUB_VALUE_TYPE(GKO_DECLARE_COO_CONJ_ARRAY_KERNEL);
 
 
 }  // namespace coo
@@ -1096,7 +1097,6 @@
 GKO_STUB_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_PGM_FIND_STRONGEST_NEIGHBOR);
 GKO_STUB_NON_COMPLEX_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM_ASSIGN_TO_EXIST_AGG);
-GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM_SORT_ROW_MAJOR);
 GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM_COMPUTE_COARSE_COO);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_PGM_GATHER_INDEX);
 

core/matrix/coo.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -42,6 +42,8 @@ GKO_REGISTER_OPERATION(inplace_absolute_array,
 GKO_REGISTER_OPERATION(outplace_absolute_array,
                        components::outplace_absolute_array);
 GKO_REGISTER_OPERATION(aos_to_soa, components::aos_to_soa);
+GKO_REGISTER_OPERATION(sort_row_major, components::sort_row_major);
+GKO_REGISTER_OPERATION(conj_array, coo::conj_array);
 
 
 }  // anonymous namespace
@@ -378,6 +380,27 @@ void Coo<ValueType, IndexType>::write(mat_data& data) const
     }
 }
 
+template <typename ValueType, typename IndexType>
+std::unique_ptr<LinOp> Coo<ValueType, IndexType>::transpose() const
+{
+    auto coo = this->clone();
+    std::swap(coo->row_idxs_, coo->col_idxs_);
+    auto size = this->get_size();
+    coo->set_size(dim<2>{size[1], size[0]});
+    coo->get_executor()->run(coo::make_sort_row_major(
+        coo->get_num_stored_elements(), coo->get_row_idxs(),
+        coo->get_col_idxs(), coo->get_values()));
+    return coo;
+}
+
+template <typename ValueType, typename IndexType>
+std::unique_ptr<LinOp> Coo<ValueType, IndexType>::conj_transpose() const
+{
+    auto coo = as<transposed_type>(this->transpose());
+    coo->get_executor()->run(coo::make_conj_array(
+        coo->get_num_stored_elements(), coo->get_values()));
+    return coo;
+}
 
 template <typename ValueType, typename IndexType>
 std::unique_ptr<Diagonal<ValueType>>

core/matrix/coo_kernels.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -54,19 +54,26 @@ namespace kernels {
                           const matrix::Coo<ValueType, IndexType>* orig, \
                           matrix::Diagonal<ValueType>* diag)
 
-#define GKO_DECLARE_ALL_AS_TEMPLATES                             \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_SPMV_KERNEL(ValueType, IndexType);           \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_ADVANCED_SPMV_KERNEL(ValueType, IndexType);  \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_SPMV2_KERNEL(ValueType, IndexType);          \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_ADVANCED_SPMV2_KERNEL(ValueType, IndexType); \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_FILL_IN_DENSE_KERNEL(ValueType, IndexType);  \
-    template <typename ValueType, typename IndexType>            \
-    GKO_DECLARE_COO_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType)
+#define GKO_DECLARE_COO_CONJ_ARRAY_KERNEL(ValueType)             \
+    void conj_array(std::shared_ptr<const DefaultExecutor> exec, \
+                    size_type num_elems, ValueType* values)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                               \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_SPMV_KERNEL(ValueType, IndexType);             \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_ADVANCED_SPMV_KERNEL(ValueType, IndexType);    \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_SPMV2_KERNEL(ValueType, IndexType);            \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_ADVANCED_SPMV2_KERNEL(ValueType, IndexType);   \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_FILL_IN_DENSE_KERNEL(ValueType, IndexType);    \
+    template <typename ValueType, typename IndexType>              \
+    GKO_DECLARE_COO_EXTRACT_DIAGONAL_KERNEL(ValueType, IndexType); \
+    template <typename ValueType>                                  \
+    GKO_DECLARE_COO_CONJ_ARRAY_KERNEL(ValueType)
 
 
 GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(coo, GKO_DECLARE_ALL_AS_TEMPLATES);

core/multigrid/pgm.cpp

@@ -23,6 +23,7 @@
 #include <ginkgo/core/matrix/row_gatherer.hpp>
 #include <ginkgo/core/matrix/sparsity_csr.hpp>
 
+#include "core/base/device_matrix_data_kernels.hpp"
 #include "core/base/dispatch_helper.hpp"
 #include "core/base/iterator_factory.hpp"
 #include "core/base/utils.hpp"
@@ -48,7 +49,7 @@ GKO_REGISTER_OPERATION(assign_to_exist_agg, pgm::assign_to_exist_agg);
 GKO_REGISTER_OPERATION(sort_agg, pgm::sort_agg);
 GKO_REGISTER_OPERATION(map_row, pgm::map_row);
 GKO_REGISTER_OPERATION(map_col, pgm::map_col);
-GKO_REGISTER_OPERATION(sort_row_major, pgm::sort_row_major);
+GKO_REGISTER_OPERATION(sort_row_major, components::sort_row_major);
 GKO_REGISTER_OPERATION(count_unrepeated_nnz, pgm::count_unrepeated_nnz);
 GKO_REGISTER_OPERATION(compute_coarse_coo, pgm::compute_coarse_coo);
 GKO_REGISTER_OPERATION(fill_array, components::fill_array);
@@ -115,6 +116,10 @@ std::shared_ptr<matrix::Csr<ValueType, IndexType>> generate_coarse(
                                 non_local_agg.get_const_data(),
                                 col_idxs.get_data()));
     // sort by row, col
+    // Because reduce_by_key is not deterministic, so we do not need
+    // stable_sort_by_key
+    // TODO: If we have deterministic reduce_by_key, we might consider
+    // stable_sort_by_key
     exec->run(pgm::make_sort_row_major(nnz, row_idxs.get_data(),
                                        col_idxs.get_data(), vals.get_data()));
     // compute the total nnz and create the fine csr

core/multigrid/pgm_kernels.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -69,11 +69,6 @@ namespace pgm {
         const matrix::Diagonal<ValueType>* diag, array<IndexType>& agg, \
         array<IndexType>& intermediate_agg)
 
-#define GKO_DECLARE_PGM_SORT_ROW_MAJOR(ValueType, IndexType)         \
-    void sort_row_major(std::shared_ptr<const DefaultExecutor> exec, \
-                        size_type nnz, IndexType* row_idxs,          \
-                        IndexType* col_idxs, ValueType* vals)
-
 #define GKO_DECLARE_PGM_COMPUTE_COARSE_COO(ValueType, IndexType)              \
     void compute_coarse_coo(std::shared_ptr<const DefaultExecutor> exec,      \
                             size_type fine_nnz, const IndexType* row_idxs,    \
@@ -106,8 +101,6 @@ namespace pgm {
     template <typename ValueType, typename IndexType>              \
     GKO_DECLARE_PGM_ASSIGN_TO_EXIST_AGG(ValueType, IndexType);     \
     template <typename ValueType, typename IndexType>              \
-    GKO_DECLARE_PGM_SORT_ROW_MAJOR(ValueType, IndexType);          \
-    template <typename ValueType, typename IndexType>              \
     GKO_DECLARE_PGM_COMPUTE_COARSE_COO(ValueType, IndexType);      \
     template <typename IndexType>                                  \
     GKO_DECLARE_PGM_GATHER_INDEX(IndexType)

dpcpp/base/device_matrix_data_kernels.dp.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -103,16 +103,15 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
 
 template <typename ValueType, typename IndexType>
 void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
-                    device_matrix_data<ValueType, IndexType>& data)
+                    size_type num_elems, IndexType* row_idxs,
+                    IndexType* col_idxs, ValueType* vals)
 {
     auto policy = onedpl_policy(exec);
-    auto input_it = oneapi::dpl::make_zip_iterator(
-        data.get_row_idxs(), data.get_col_idxs(), data.get_values());
-    std::sort(policy, input_it, input_it + data.get_num_stored_elements(),
-              [](auto a, auto b) {
-                  return std::tie(std::get<0>(a), std::get<1>(a)) <
-                         std::tie(std::get<0>(b), std::get<1>(b));
-              });
+    auto input_it = oneapi::dpl::make_zip_iterator(row_idxs, col_idxs, vals);
+    std::sort(policy, input_it, input_it + num_elems, [](auto a, auto b) {
+        return std::tie(std::get<0>(a), std::get<1>(a)) <
+               std::tie(std::get<0>(b), std::get<1>(b));
+    });
 }
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(

dpcpp/multigrid/pgm_kernels.dp.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -70,25 +70,6 @@ void sort_agg(std::shared_ptr<const DefaultExecutor> exec, IndexType num,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_PGM_SORT_AGG_KERNEL);
 
 
-template <typename ValueType, typename IndexType>
-void sort_row_major(std::shared_ptr<const DefaultExecutor> exec, size_type nnz,
-                    IndexType* row_idxs, IndexType* col_idxs, ValueType* vals)
-{
-    auto policy = onedpl_policy(exec);
-    auto it = oneapi::dpl::make_zip_iterator(row_idxs, col_idxs, vals);
-    // Because reduce_by_segment is not deterministic, so we do not need
-    // stable_sort
-    // TODO: If we have deterministic reduce_by_segment, it should be
-    // stable_sort
-    std::sort(policy, it, it + nnz, [](auto a, auto b) {
-        return std::tie(std::get<0>(a), std::get<1>(a)) <
-               std::tie(std::get<0>(b), std::get<1>(b));
-    });
-}
-
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_PGM_SORT_ROW_MAJOR);
-
-
 template <typename ValueType, typename IndexType>
 class coarse_coo_policy {};
 

include/ginkgo/core/matrix/coo.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -58,6 +58,7 @@ class Coo : public EnableLinOp<Coo<ValueType, IndexType>>,
             public DiagonalExtractable<ValueType>,
             public ReadableFromMatrixData<ValueType, IndexType>,
             public WritableToMatrixData<ValueType, IndexType>,
+            public Transposable,
             public EnableAbsoluteComputation<
                 remove_complex<Coo<ValueType, IndexType>>> {
     friend class EnablePolymorphicObject<Coo, LinOp>;
@@ -80,6 +81,7 @@ class Coo : public EnableLinOp<Coo<ValueType, IndexType>>,
 
     using value_type = ValueType;
     using index_type = IndexType;
+    using transposed_type = Coo<ValueType, IndexType>;
     using mat_data = matrix_data<ValueType, IndexType>;
     using device_mat_data = device_matrix_data<ValueType, IndexType>;
     using absolute_type = remove_complex<Coo>;
@@ -122,6 +124,10 @@ class Coo : public EnableLinOp<Coo<ValueType, IndexType>>,
 
     void write(mat_data& data) const override;
 
+    std::unique_ptr<LinOp> transpose() const override;
+
+    std::unique_ptr<LinOp> conj_transpose() const override;
+
     std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
     std::unique_ptr<absolute_type> compute_absolute() const override;