allow pgm to generate weight without symmetrization

common/cuda_hip/base/kernel_launch_reduction.hpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -294,26 +294,31 @@
     const auto block = group::this_thread_block();
     const auto warp = group::tiled_partition<warp_size>(block);
     const auto warp_rank = warp.thread_rank();
-    const auto col = warp_rank + static_cast<int64>(blockIdx.y) * warp_size;
-    auto partial = identity;
-    // accumulate within a thread
-    if (col < cols) {
-        for (auto row = warp_id; row < rows; row += warp_num) {
-            partial = op(partial, fn(row, col, args...));
+
+    for (auto block_col = static_cast<int64>(blockIdx.y) * warp_size;
+         block_col < cols;
+         block_col += static_cast<int64>(gridDim.y) * warp_size) {
+        const auto col = warp_rank + block_col;
+        auto partial = identity;
+        // accumulate within a thread
+        if (col < cols) {
+            for (auto row = warp_id; row < rows; row += warp_num) {
+                partial = op(partial, fn(row, col, args...));
+            }
         }
-    }
-    block_partial[threadIdx.x] = partial;
-    block.sync();
-    // in a single warp: accumulate the results
-    if (threadIdx.x < warp_size) {
-        partial = identity;
-        // accumulate the partial results within a thread
+        block_partial[threadIdx.x] = partial;
+        block.sync();
+        // in a single warp: accumulate the results
+        if (threadIdx.x < warp_size) {
+            partial = identity;
+            // accumulate the partial results within a thread
 #pragma unroll
-        for (int i = 0; i < default_block_size; i += warp_size) {
-            partial = op(partial, block_partial[i + warp_rank]);
-        }
-        if (col < cols) {
-            result[col + blockIdx.x * cols] = finalize(partial);
+            for (int i = 0; i < default_block_size; i += warp_size) {
+                partial = op(partial, block_partial[i + warp_rank]);
+            }
+            if (col < cols) {
+                result[col + blockIdx.x * cols] = finalize(partial);
+            }
         }
     }
 }
@@ -488,7 +493,10 @@
             syn::value_list<int>(), syn::type_list<>(), max_blocks, exec, fn,
             op, finalize, identity, result, size, tmp, map_to_device(args)...);
     } else {
-        const auto col_blocks = ceildiv(cols, config::warp_size);
+        // cuda only accept up to 65545 for grid's y-axix
+        constexpr int64 max_grid_y = 65535;
+        const auto col_blocks =
+            std::min(ceildiv(cols, config::warp_size), max_grid_y);
         const auto row_blocks =
             ceildiv(std::min<int64>(
                         ceildiv(rows * config::warp_size, default_block_size),

core/multigrid/pgm.cpp

@@ -176,6 +176,9 @@ Pgm<ValueType, IndexType>::parse(const config::pnode& config,
     if (auto& obj = config_check.get("skip_sorting")) {
         params.with_skip_sorting(config::get_value<bool>(obj));
     }
+    if (auto& obj = config_check.get("weight_symmetrization")) {
+        params.with_weight_symmetrization(config::get_value<bool>(obj));
+    }
 
     return params;
 }
@@ -202,15 +205,19 @@ Pgm<ValueType, IndexType>::generate_local(
                                    -one<IndexType>()));
     IndexType num_unagg = num_rows;
     IndexType num_unagg_prev = num_rows;
-    // TODO: if mtx is a hermitian matrix, weight_mtx = abs(mtx)
-    // compute weight_mtx = (abs(mtx) + abs(mtx'))/2;
-    auto abs_mtx = local_matrix->compute_absolute();
-    // abs_mtx is already real valuetype, so transpose is enough
-    auto weight_mtx = gko::as<weight_csr_type>(abs_mtx->transpose());
-    auto half_scalar = initialize<matrix::Dense<real_type>>({0.5}, exec);
-    auto identity = matrix::Identity<real_type>::create(exec, num_rows);
-    // W = (abs_mtx + transpose(abs_mtx))/2
-    abs_mtx->apply(half_scalar, identity, half_scalar, weight_mtx);
+    std::shared_ptr<weight_csr_type> weight_mtx = nullptr;
+    if (parameters_.weight_symmetrization) {
+        // compute weight_mtx = (abs(mtx) + abs(mtx'))/2;
+        auto abs_mtx = local_matrix->compute_absolute();
+        // abs_mtx is already real valuetype, so transpose is enough
+        weight_mtx = gko::as<weight_csr_type>(abs_mtx->transpose());
+        auto half_scalar = initialize<matrix::Dense<real_type>>({0.5}, exec);
+        auto identity = matrix::Identity<real_type>::create(exec, num_rows);
+        // W = (abs_mtx + transpose(abs_mtx))/2
+        abs_mtx->apply(half_scalar, identity, half_scalar, weight_mtx);
+    } else {
+        weight_mtx = local_matrix->compute_absolute();
+    }
     // Extract the diagonal value of matrix
     auto diag = weight_mtx->extract_diagonal();
     for (int i = 0; i < parameters_.max_iterations; i++) {

core/test/config/multigrid.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -54,6 +54,8 @@ struct Pgm : MultigridLevelConfigTest<gko::multigrid::Pgm<float, int>,
         param.with_deterministic(true);
         config_map["skip_sorting"] = pnode{true};
         param.with_skip_sorting(true);
+        config_map["weight_symmetrization"] = pnode{false};
+        param.with_weight_symmetrization(false);
     }
 
     template <typename AnswerType>
@@ -67,6 +69,8 @@ struct Pgm : MultigridLevelConfigTest<gko::multigrid::Pgm<float, int>,
                   ans_param.max_unassigned_ratio);
         ASSERT_EQ(res_param.deterministic, ans_param.deterministic);
         ASSERT_EQ(res_param.skip_sorting, ans_param.skip_sorting);
+        ASSERT_EQ(res_param.weight_symmetrization,
+                  ans_param.weight_symmetrization);
     }
 };
 

include/ginkgo/core/multigrid/pgm.hpp

@@ -125,6 +125,13 @@ class Pgm : public LinOp, public EnableMultigridLevel<ValueType> {
          * incorrect.
          */
         bool GKO_FACTORY_PARAMETER_SCALAR(skip_sorting, false);
+
+        /**
+         * `weight_symmetrization` decides whether to symmetrize the absolute
+         * weight matrix. When it is true, the weight matrix is (abs(mtx) +
+         * abs(mtx)')/2. Otherwise, the weight matrix is abs(mtx).
+         */
+        bool GKO_FACTORY_PARAMETER_SCALAR(weight_symmetrization, true);
     };
     GKO_ENABLE_LIN_OP_FACTORY(Pgm, parameters, Factory);
     GKO_ENABLE_BUILD_METHOD(Factory);

reference/test/multigrid/pgm_kernels.cpp

@@ -506,4 +506,62 @@ TYPED_TEST(Pgm, GenerateMgLevelOnUnsortedMatrix)
 }
 
 
+TYPED_TEST(Pgm, GenerateMgLevelWithoutSymmetrization)
+{
+    using value_type = typename TestFixture::value_type;
+    using index_type = typename TestFixture::index_type;
+    using Mtx = typename TestFixture::Mtx;
+    using SparsityCsr = typename TestFixture::SparsityCsr;
+    using MgLevel = typename TestFixture::MgLevel;
+    using RowGatherer = typename TestFixture::RowGatherer;
+    auto mglevel = MgLevel::build()
+                       .with_max_iterations(2u)
+                       .with_max_unassigned_ratio(0.1)
+                       .on(this->exec);
+    /* this matrix will generate the same weight matrix without symmetrization
+     * as this->fine, so the prolong and restriction will be the same.
+     *
+     *  5 -3   -3    0    0
+     * -3  5    0   -2.5 -1.5
+     * -3  0    5    0   -1.5
+     *  0 -2.5  0    5    0
+     *  0 -1.5 -1.5  0    5
+     */
+    auto matrix = gko::share(Mtx::create(this->exec));
+    matrix->read({{5, 5},
+                  {{0, 0, 5},
+                   {0, 1, -3},
+                   {0, 2, -3},
+                   {1, 0, -3},
+                   {1, 1, 5},
+                   {1, 3, -2.5},
+                   {1, 4, -1.5},
+                   {2, 0, -3},
+                   {2, 2, 5},
+                   {2, 4, -1.5},
+                   {3, 1, -2.5},
+                   {3, 3, 5},
+                   {4, 1, -1.5},
+                   {4, 2, -1.5},
+                   {4, 4, 5}}});
+    auto prolong_op = gko::share(Mtx::create(this->exec, gko::dim<2>{5, 2}, 0));
+    // 0-2-4, 1-3
+    prolong_op->read(
+        {{5, 2}, {{0, 0, 1}, {1, 1, 1}, {2, 0, 1}, {3, 1, 1}, {4, 0, 1}}});
+    auto restrict_op = gko::share(gko::as<Mtx>(prolong_op->transpose()));
+
+    auto coarse_fine = mglevel->generate(matrix);
+    auto row_gatherer = gko::as<RowGatherer>(coarse_fine->get_prolong_op());
+    auto row_gather_view = gko::array<index_type>::const_view(
+        this->exec, row_gatherer->get_size()[0],
+        row_gatherer->get_const_row_idxs());
+    auto expected_row_gather =
+        gko::array<index_type>(this->exec, {0, 1, 0, 1, 0});
+
+    GKO_ASSERT_MTX_NEAR(gko::as<SparsityCsr>(coarse_fine->get_restrict_op()),
+                        restrict_op, r<value_type>::value);
+    GKO_ASSERT_ARRAY_EQ(row_gather_view, expected_row_gather);
+}
+
+
 }  // namespace

test/multigrid/pgm_kernels.cpp

@@ -336,3 +336,42 @@ TEST_F(Pgm, GenerateMgLevelIsEquivalentToRefOnUnsortedMatrix)
                         r<value_type>::value);
     GKO_ASSERT_ARRAY_EQ(d_row_gather_view, row_gather_view);
 }
+
+
+TEST_F(Pgm, GenerateMgLevelWithoutSymmetrizationIsEquivalentOnHpd)
+{
+    initialize_data();
+    auto d_mg_level_factory = gko::multigrid::Pgm<value_type, int>::build()
+                                  .with_deterministic(true)
+                                  .with_skip_sorting(true)
+                                  .with_weight_symmetrization(true)
+                                  .on(exec);
+    auto d_mg_level_factory_wo_sym =
+        gko::multigrid::Pgm<value_type, int>::build()
+            .with_deterministic(true)
+            .with_skip_sorting(true)
+            .with_weight_symmetrization(false)
+            .on(exec);
+
+    auto d_mg_level = d_mg_level_factory->generate(d_system_mtx->clone());
+    auto d_mg_level_wo_sym = d_mg_level_factory_wo_sym->generate(d_system_mtx);
+    auto d_row_gatherer = gko::as<RowGatherer>(d_mg_level->get_prolong_op());
+    auto d_row_gatherer_wo_sym =
+        gko::as<RowGatherer>(d_mg_level_wo_sym->get_prolong_op());
+    auto d_row_gather_view = gko::array<index_type>::const_view(
+        d_row_gatherer->get_executor(), d_row_gatherer->get_size()[0],
+        d_row_gatherer->get_const_row_idxs());
+    auto d_row_gather_view_wo_sym = gko::array<index_type>::const_view(
+        d_row_gatherer_wo_sym->get_executor(),
+        d_row_gatherer_wo_sym->get_size()[0],
+        d_row_gatherer_wo_sym->get_const_row_idxs());
+
+    GKO_ASSERT_MTX_NEAR(
+        gko::as<SparsityCsr>(d_mg_level_wo_sym->get_restrict_op()),
+        gko::as<SparsityCsr>(d_mg_level->get_restrict_op()),
+        r<value_type>::value);
+    GKO_ASSERT_MTX_NEAR(gko::as<Csr>(d_mg_level_wo_sym->get_coarse_op()),
+                        gko::as<Csr>(d_mg_level->get_coarse_op()),
+                        r<value_type>::value);
+    GKO_ASSERT_ARRAY_EQ(d_row_gather_view, d_row_gather_view_wo_sym);
+}