wip

core/distributed/sparse_communicator.cpp

@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017-2023 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -129,6 +129,15 @@ sparse_communicator::sparse_communicator(
     };
     fill_size_offsets(recv_sizes_, recv_offsets_, part->get_recv_indices());
     fill_size_offsets(send_sizes_, send_offsets_, part->get_send_indices());
+
+    send_idxs_ = array<IndexType>(exec, send_offsets_.back());
+    for (int i = 0; i < part->get_send_indices().get_num_groups(); ++i) {
+        make_array_view(exec, send_sizes_[i],
+                        std::get<array<IndexType>>(send_idxs_).get_data() +
+                            send_offsets_[i]) =
+            part->get_send_indices().get_indices(i);
+    }
+
     part_ = std::move(part);
 }
 
@@ -160,20 +169,14 @@ mpi::request sparse_communicator::communicate_impl_(
 
     send_buffer.init(exec,
                      {send_idxs.get_num_elems(), local_vector->get_size()[1]});
-    size_type offset = 0;
-    for (int i = 0; i < send_idxs.get_num_groups(); ++i) {
-        // need direct support for index_set
-        auto full_idxs = send_idxs.get_indices(i);
-        local_vector->row_gather(
-            &full_idxs, send_buffer.get()->create_submatrix(
-                            {offset, offset + full_idxs.get_num_elems()},
-                            {0, local_vector->get_size()[1]}));
-        offset += full_idxs.get_num_elems();
-    }
+
+    auto& full_send_idxs = std::get<array<IndexType>>(send_idxs_);
+    local_vector->row_gather(&full_send_idxs, send_buffer.get());
 
     auto recv_ptr = recv_buffer->get_values();
     auto send_ptr = send_buffer->get_values();
 
+    exec->synchronize();
     mpi::contiguous_type type(local_vector->get_size()[1],
                               mpi::type_impl<ValueType>::get_type());
     mpi::request req;

include/ginkgo/core/distributed/sparse_communicator.hpp

@@ -1,44 +1,20 @@
-/*******************************<GINKGO LICENSE>******************************
-Copyright (c) 2017-2023, 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>*******************************/
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
 
 #ifndef GKO_PUBLIC_CORE_DISTRIBUTED_SPARSE_COMMUNICATOR_HPP_
 #define GKO_PUBLIC_CORE_DISTRIBUTED_SPARSE_COMMUNICATOR_HPP_
 
+
 #include <ginkgo/config.hpp>
 
 
 #if GINKGO_BUILD_MPI
 
 
+#include <variant>
+
+
 #include <ginkgo/core/base/dense_cache.hpp>
 #include <ginkgo/core/base/index_set.hpp>
 #include <ginkgo/core/base/mpi.hpp>
@@ -48,9 +24,6 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 #include <ginkgo/core/matrix/dense.hpp>
 
 
-#include <variant>
-
-
 namespace gko::experimental::distributed {
 
 
@@ -169,6 +142,7 @@ class sparse_communicator
     std::vector<comm_index_type> send_offsets_;
     std::vector<comm_index_type> recv_sizes_;
     std::vector<comm_index_type> recv_offsets_;
+    std::variant<array<int32>, array<int64>> send_idxs_;
 };
 
 

include/ginkgo/ginkgo.hpp

@@ -58,10 +58,10 @@
 #include <ginkgo/core/distributed/partition.hpp>
 #include <ginkgo/core/distributed/partition_helpers.hpp>
 #include <ginkgo/core/distributed/polymorphic_object.hpp>
-#include <ginkgo/core/distributed/sparse_communicator.hpp>
 
 #include <ginkgo/core/distributed/preconditioner/schwarz.hpp>
 
+#include <ginkgo/core/distributed/sparse_communicator.hpp>
 #include <ginkgo/core/distributed/vector.hpp>
 
 #include <ginkgo/core/factorization/cholesky.hpp>