add csr absolute and store weight explicitly

Co-authored-by: Tobias Ribizel <[email protected]>

common/matrix/csr_kernels.hpp.inc

@@ -901,6 +901,18 @@ __global__ __launch_bounds__(default_block_size) void check_unsorted(
 }
 
 
+template <typename ValueType>
+__global__ __launch_bounds__(default_block_size) void absolute_kernel(
+    size_type num, const ValueType *__restrict__ source,
+    remove_complex<ValueType> *__restrict__ result)
+{
+    const auto tidx = thread::get_thread_id_flat();
+    if (tidx < num) {
+        result[tidx] = gko::abs(source[tidx]);
+    }
+}
+
+
 }  // namespace kernel
 
 

core/device_hooks/common_kernels.inc.cpp

@@ -615,6 +615,11 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_IS_SORTED_BY_COLUMN_INDEX);
 
 
+template <typename ValueType, typename IndexType>
+GKO_DECLARE_CSR_ABSOLUTE(ValueType, IndexType)
+GKO_NOT_COMPILED(GKO_HOOK_MODULE);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ABSOLUTE);
+
 }  // namespace csr
 
 
@@ -995,19 +1000,19 @@ GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_AMGX_PGM_RENUMBER_KERNEL);
 template <typename ValueType, typename IndexType>
 GKO_DECLARE_AMGX_PGM_EXTRACT_DIAG(ValueType, IndexType)
 GKO_NOT_COMPILED(GKO_HOOK_MODULE);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_EXTRACT_DIAG);
 
 template <typename ValueType, typename IndexType>
 GKO_DECLARE_AMGX_PGM_FIND_STRONGEST_NEIGHBOR(ValueType, IndexType)
 GKO_NOT_COMPILED(GKO_HOOK_MODULE);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_FIND_STRONGEST_NEIGHBOR);
 
 template <typename ValueType, typename IndexType>
 GKO_DECLARE_AMGX_PGM_ASSIGN_TO_EXIST_AGG(ValueType, IndexType)
 GKO_NOT_COMPILED(GKO_HOOK_MODULE);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_ASSIGN_TO_EXIST_AGG);
 
 template <typename ValueType, typename IndexType>

core/matrix/csr.cpp

@@ -78,7 +78,7 @@ GKO_REGISTER_OPERATION(calculate_nonzeros_per_row,
 GKO_REGISTER_OPERATION(sort_by_column_index, csr::sort_by_column_index);
 GKO_REGISTER_OPERATION(is_sorted_by_column_index,
                        csr::is_sorted_by_column_index);
-
+GKO_REGISTER_OPERATION(absolute, csr::absolute);
 
 }  // namespace csr
 
@@ -464,6 +464,22 @@ bool Csr<ValueType, IndexType>::is_sorted_by_column_index() const
 }
 
 
+template <typename ValueType, typename IndexType>
+std::unique_ptr<LinOp> Csr<ValueType, IndexType>::absolute() const
+{
+    using abs_type = remove_complex<ValueType>;
+    using abs_csr = Csr<abs_type, IndexType>;
+    auto exec = this->get_executor();
+    auto abs_cpy = abs_csr::create(exec, this->get_size(),
+                                   this->get_num_stored_elements());
+    abs_cpy->col_idxs_ = this->col_idxs_;
+    abs_cpy->row_ptrs_ = this->row_ptrs_;
+    exec->run(csr::make_absolute(this, lend(abs_cpy)));
+    this->convert_strategy_helper(lend(abs_cpy));
+    return std::move(abs_cpy);
+}
+
+
 #define GKO_DECLARE_CSR_MATRIX(ValueType, IndexType) \
     class Csr<ValueType, IndexType>
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_MATRIX);

core/matrix/csr_kernels.hpp

@@ -174,6 +174,11 @@ namespace kernels {
         std::shared_ptr<const DefaultExecutor> exec,                    \
         const matrix::Csr<ValueType, IndexType> *to_check, bool *is_sorted)
 
+#define GKO_DECLARE_CSR_ABSOLUTE(ValueType, IndexType)             \
+    void absolute(std::shared_ptr<const DefaultExecutor> exec,     \
+                  const matrix::Csr<ValueType, IndexType> *source, \
+                  matrix::Csr<remove_complex<ValueType>, IndexType> *result)
+
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                         \
     template <typename ValueType, typename IndexType>                        \
     GKO_DECLARE_CSR_SPMV_KERNEL(ValueType, IndexType);                       \
@@ -216,7 +221,9 @@ namespace kernels {
     template <typename ValueType, typename IndexType>                        \
     GKO_DECLARE_CSR_SORT_BY_COLUMN_INDEX(ValueType, IndexType);              \
     template <typename ValueType, typename IndexType>                        \
-    GKO_DECLARE_CSR_IS_SORTED_BY_COLUMN_INDEX(ValueType, IndexType)
+    GKO_DECLARE_CSR_IS_SORTED_BY_COLUMN_INDEX(ValueType, IndexType);         \
+    template <typename ValueType, typename IndexType>                        \
+    GKO_DECLARE_CSR_ABSOLUTE(ValueType, IndexType)
 
 
 namespace omp {

core/multigrid/amgx_pgm.cpp

@@ -41,6 +41,7 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/base/utils.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/identity.hpp>
 
 
 #include "core/matrix/csr_builder.hpp"
@@ -92,32 +93,43 @@ template <typename ValueType, typename IndexType>
 void AmgxPgm<ValueType, IndexType>::generate()
 {
     using matrix_type = matrix::Csr<ValueType, IndexType>;
+    using rmc_value_type = remove_complex<ValueType>;
+    using weight_matrix_type = matrix::Csr<rmc_value_type, IndexType>;
     auto exec = this->get_executor();
     const auto num = this->system_matrix_->get_size()[0];
-    Array<ValueType> diag(this->get_executor(), num);
+    Array<rmc_value_type> diag(this->get_executor(), num);
     Array<IndexType> strongest_neighbor(this->get_executor(), num);
     Array<IndexType> intermediate_agg(this->get_executor(),
                                       parameters_.deterministic * num);
     const auto amgxpgm_op = gko::as<matrix_type>(this->system_matrix_.get());
     // Initial agg = -1
     exec->run(amgx_pgm::make_initial(agg_));
-    // Extract the diagonal value of matrix
-    exec->run(amgx_pgm::make_extract_diag(amgxpgm_op, diag));
     size_type num_unagg{0};
     size_type num_unagg_prev{0};
+    // TODO: if mtx is a hermitian matrix, weight_mtx = abs(mtx)
+    // compute weight_mtx = (abs(mtx) + abs(mtx'))/2;
+    auto abs_mtx = gko::as<weight_matrix_type>(amgxpgm_op->absolute());
+    // abs_mtx is already real valuetype, so transpose is enough
+    auto weight_mtx = gko::as<weight_matrix_type>(abs_mtx->transpose());
+    auto half_scaler = initialize<matrix::Dense<rmc_value_type>>({0.5}, exec);
+    auto identity = matrix::Identity<rmc_value_type>::create(exec, num);
+    // W = (abs_mtx + transpose(abs_mtx))/2
+    abs_mtx->apply(lend(half_scaler), lend(identity), lend(half_scaler),
+                   lend(weight_mtx));
+    // Extract the diagonal value of matrix
+    exec->run(amgx_pgm::make_extract_diag(weight_mtx.get(), diag));
     for (int i = 0; i < parameters_.max_iterations; i++) {
         // Find the strongest neighbor of each row
-        exec->run(amgx_pgm::make_find_strongest_neighbor(amgxpgm_op, diag, agg_,
-                                                         strongest_neighbor));
+        exec->run(amgx_pgm::make_find_strongest_neighbor(
+            weight_mtx.get(), diag, agg_, strongest_neighbor));
         // Match edges
         exec->run(amgx_pgm::make_match_edge(strongest_neighbor, agg_));
         // Get the num_unagg
         exec->run(amgx_pgm::make_count_unagg(agg_, &num_unagg));
         // no new match, all match, or the ratio of num_unagg/num is lower
         // than parameter.max_unassigned_percentage
         if (num_unagg == 0 || num_unagg == num_unagg_prev ||
-            static_cast<double>(num_unagg) / num <
-                parameters_.max_unassigned_percentage) {
+            num_unagg < parameters_.max_unassigned_percentage * num) {
             break;
         }
         num_unagg_prev = num_unagg;
@@ -128,8 +140,8 @@ void AmgxPgm<ValueType, IndexType>::generate()
         intermediate_agg = agg_;
     }
     while (num_unagg != 0) {
-        exec->run(amgx_pgm::make_assign_to_exist_agg(amgxpgm_op, diag, agg_,
-                                                     intermediate_agg));
+        exec->run(amgx_pgm::make_assign_to_exist_agg(weight_mtx.get(), diag,
+                                                     agg_, intermediate_agg));
         exec->run(amgx_pgm::make_count_unagg(agg_, &num_unagg));
     }
     size_type num_agg;

core/multigrid/amgx_pgm_kernels.hpp

@@ -85,16 +85,16 @@ namespace amgx_pgm {
 #define GKO_DECLARE_AMGX_PGM_FIND_STRONGEST_NEIGHBOR(ValueType, IndexType) \
     void find_strongest_neighbor(                                          \
         std::shared_ptr<const DefaultExecutor> exec,                       \
-        const matrix::Csr<ValueType, IndexType> *source,                   \
+        const matrix::Csr<ValueType, IndexType> *weight_mtx,               \
         const Array<ValueType> &diag, Array<IndexType> &agg,               \
         Array<IndexType> &strongest_neighbor)
 
-#define GKO_DECLARE_AMGX_PGM_ASSIGN_TO_EXIST_AGG(ValueType, IndexType)        \
-    void assign_to_exist_agg(std::shared_ptr<const DefaultExecutor> exec,     \
-                             const matrix::Csr<ValueType, IndexType> *source, \
-                             const Array<ValueType> &diag,                    \
-                             Array<IndexType> &agg,                           \
-                             Array<IndexType> &intermediate_agg)
+#define GKO_DECLARE_AMGX_PGM_ASSIGN_TO_EXIST_AGG(ValueType, IndexType) \
+    void assign_to_exist_agg(                                          \
+        std::shared_ptr<const DefaultExecutor> exec,                   \
+        const matrix::Csr<ValueType, IndexType> *weight_mtx,           \
+        const Array<ValueType> &diag, Array<IndexType> &agg,           \
+        Array<IndexType> &intermediate_agg)
 
 #define GKO_DECLARE_AMGX_PGM_GENERATE(ValueType, IndexType)                 \
     void amgx_pgm_generate(std::shared_ptr<const DefaultExecutor> exec,     \

cuda/matrix/csr_kernels.cu

@@ -635,16 +635,17 @@ void spgeam(std::shared_ptr<const DefaultExecutor> exec,
     auto total_nnz =
         a->get_num_stored_elements() + b->get_num_stored_elements();
     auto nnz_per_row = total_nnz / a->get_size()[0];
-    select_spgeam(spgeam_kernels(),
-                  [&](int compiled_subwarp_size) {
-                      return compiled_subwarp_size >= nnz_per_row ||
-                             compiled_subwarp_size == config::warp_size;
-                  },
-                  syn::value_list<int>(), syn::type_list<>(), exec,
-                  alpha->get_const_values(), a->get_const_row_ptrs(),
-                  a->get_const_col_idxs(), a->get_const_values(),
-                  beta->get_const_values(), b->get_const_row_ptrs(),
-                  b->get_const_col_idxs(), b->get_const_values(), c);
+    select_spgeam(
+        spgeam_kernels(),
+        [&](int compiled_subwarp_size) {
+            return compiled_subwarp_size >= nnz_per_row ||
+                   compiled_subwarp_size == config::warp_size;
+        },
+        syn::value_list<int>(), syn::type_list<>(), exec,
+        alpha->get_const_values(), a->get_const_row_ptrs(),
+        a->get_const_col_idxs(), a->get_const_values(),
+        beta->get_const_values(), b->get_const_row_ptrs(),
+        b->get_const_col_idxs(), b->get_const_values(), c);
 }
 
 GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_SPGEAM_KERNEL);
@@ -1101,6 +1102,23 @@ GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_CSR_IS_SORTED_BY_COLUMN_INDEX);
 
 
+template <typename ValueType, typename IndexType>
+void absolute(std::shared_ptr<const CudaExecutor> exec,
+              const matrix::Csr<ValueType, IndexType> *source,
+              matrix::Csr<remove_complex<ValueType>, IndexType> *result)
+{
+    auto result_val = result->get_values();
+    auto source_val = source->get_const_values();
+    const auto num = source->get_num_stored_elements();
+    const dim3 grid(ceildiv(num, default_block_size));
+
+    kernel::absolute_kernel<<<grid, default_block_size>>>(num, source_val,
+                                                          result_val);
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ABSOLUTE);
+
+
 }  // namespace csr
 }  // namespace cuda
 }  // namespace kernels

cuda/multigrid/amgx_pgm_kernels.cu

@@ -120,34 +120,34 @@ void extract_diag(std::shared_ptr<const CudaExecutor> exec,
     GKO_NOT_IMPLEMENTED;
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_EXTRACT_DIAG);
 
 
 template <typename ValueType, typename IndexType>
-void find_strongest_neighbor(std::shared_ptr<const CudaExecutor> exec,
-                             const matrix::Csr<ValueType, IndexType> *source,
-                             const Array<ValueType> &diag,
-                             Array<IndexType> &agg,
-                             Array<IndexType> &strongest_neighbor)
+void find_strongest_neighbor(
+    std::shared_ptr<const CudaExecutor> exec,
+    const matrix::Csr<ValueType, IndexType> *weight_mtx,
+    const Array<ValueType> &diag, Array<IndexType> &agg,
+    Array<IndexType> &strongest_neighbor)
 {
     GKO_NOT_IMPLEMENTED;
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_FIND_STRONGEST_NEIGHBOR);
 
 
 template <typename ValueType, typename IndexType>
 void assign_to_exist_agg(std::shared_ptr<const CudaExecutor> exec,
-                         const matrix::Csr<ValueType, IndexType> *source,
+                         const matrix::Csr<ValueType, IndexType> *weight_mtx,
                          const Array<ValueType> &diag, Array<IndexType> &agg,
                          Array<IndexType> &intermediate_agg)
 {
     GKO_NOT_IMPLEMENTED;
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_NON_COMPLEX_VALUE_AND_INDEX_TYPE(
     GKO_DECLARE_AMGX_PGM_ASSIGN_TO_EXIST_AGG);
 
 

cuda/test/matrix/csr_kernels.cpp

@@ -60,6 +60,7 @@ namespace {
 class Csr : public ::testing::Test {
 protected:
     using Mtx = gko::matrix::Csr<>;
+    using AbsMtx = gko::matrix::Csr<>;
     using Vec = gko::matrix::Dense<>;
     using ComplexVec = gko::matrix::Dense<std::complex<double>>;
     using ComplexMtx = gko::matrix::Csr<std::complex<double>>;
@@ -725,4 +726,16 @@ TEST_F(Csr, OneAutomaticalWorksWithDifferentMatrices)
 }
 
 
+TEST_F(Csr, AbsoluteMatrixIsEquivalentToRef)
+{
+    set_up_apply_data(std::make_shared<Mtx::classical>());
+
+    auto abs_mtx = mtx->absolute();
+    auto dabs_mtx = dmtx->absolute();
+
+    GKO_ASSERT_MTX_NEAR(static_cast<AbsMtx *>(abs_mtx.get()),
+                        static_cast<AbsMtx *>(dabs_mtx.get()), 1e-14);
+}
+
+
 }  // namespace