adds distributed index map core & reference

common/unified/CMakeLists.txt

@@ -6,6 +6,7 @@ set(UNIFIED_SOURCES
     components/format_conversion_kernels.cpp
     components/precision_conversion_kernels.cpp
     components/reduce_array_kernels.cpp
+    distributed/index_map_kernels.cpp
     distributed/partition_helpers_kernels.cpp
     distributed/partition_kernels.cpp
     matrix/coo_kernels.cpp
@@ -30,4 +31,4 @@ set(UNIFIED_SOURCES
     solver/ir_kernels.cpp
     )
 list(TRANSFORM UNIFIED_SOURCES PREPEND ${CMAKE_CURRENT_SOURCE_DIR}/)
-set(GKO_UNIFIED_COMMON_SOURCES ${UNIFIED_SOURCES} PARENT_SCOPE)
+set(GKO_UNIFIED_COMMON_SOURCES ${UNIFIED_SOURCES} PARENT_SCOPE)

common/unified/distributed/index_map_kernels.cpp

@@ -0,0 +1,51 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "core/distributed/index_map_kernels.hpp"
+
+
+#include "common/unified/base/kernel_launch.hpp"
+
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace index_map {
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void build_mapping(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const experimental::distributed::Partition<LocalIndexType, GlobalIndexType>*
+        part,
+    const array<GlobalIndexType>& recv_connections,
+    array<experimental::distributed::comm_index_type>& remote_part_ids,
+    array<LocalIndexType>& remote_local_idxs,
+    array<GlobalIndexType>& remote_global_idxs,
+    array<int64>& remote_sizes) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_INDEX_MAP_BUILD_MAPPING);
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+void get_local(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const experimental::distributed::Partition<LocalIndexType, GlobalIndexType>*
+        partition,
+    const array<experimental::distributed::comm_index_type>& remote_target_ids,
+    const segmented_array<GlobalIndexType>& remote_global_idxs,
+    experimental::distributed::comm_index_type rank,
+    const array<GlobalIndexType>& global_ids,
+    experimental::distributed::index_space is,
+    array<LocalIndexType>& local_ids) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
+    GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY);
+
+
+}  // namespace index_map
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

core/CMakeLists.txt

@@ -20,6 +20,7 @@ target_sources(ginkgo
     base/timer.cpp
     base/version.cpp
     config/property_tree.cpp
+    distributed/index_map.cpp
     distributed/partition.cpp
     factorization/cholesky.cpp
     factorization/elimination_forest.cpp

core/device_hooks/common_kernels.inc.cpp

@@ -16,6 +16,7 @@
 #include "core/components/precision_conversion_kernels.hpp"
 #include "core/components/prefix_sum_kernels.hpp"
 #include "core/components/reduce_array_kernels.hpp"
+#include "core/distributed/index_map_kernels.hpp"
 #include "core/distributed/matrix_kernels.hpp"
 #include "core/distributed/partition_helpers_kernels.hpp"
 #include "core/distributed/partition_kernels.hpp"
@@ -251,6 +252,16 @@ GKO_STUB_INDEX_TYPE(GKO_DECLARE_PARTITION_HELPERS_COMPRESS_RANGES);
 }  // namespace partition_helpers
 
 
+namespace index_map {
+
+
+GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_BUILD_MAPPING);
+GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY);
+
+
+}  // namespace index_map
+
+
 namespace distributed_vector {
 
 

core/distributed/index_map.cpp

@@ -0,0 +1,185 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/distributed/index_map.hpp>
+
+
+#include "core/distributed/index_map_kernels.hpp"
+
+
+namespace gko {
+namespace index_map_kernels {
+
+
+GKO_REGISTER_OPERATION(build_mapping, index_map::build_mapping);
+GKO_REGISTER_OPERATION(get_local, index_map::get_local);
+
+}  // namespace index_map_kernels
+
+
+namespace experimental {
+namespace distributed {
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+size_type index_map<LocalIndexType, GlobalIndexType>::get_local_size() const
+{
+    return partition_ ? partition_->get_part_size(rank_) : 0;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+size_type index_map<LocalIndexType, GlobalIndexType>::get_non_local_size() const
+{
+    return remote_global_idxs_.get_size();
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+const segmented_array<GlobalIndexType>&
+index_map<LocalIndexType, GlobalIndexType>::get_remote_global_idxs() const
+{
+    return remote_global_idxs_;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+const segmented_array<LocalIndexType>&
+index_map<LocalIndexType, GlobalIndexType>::get_remote_local_idxs() const
+{
+    return remote_local_idxs_;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+const array<comm_index_type>&
+index_map<LocalIndexType, GlobalIndexType>::get_remote_target_ids() const
+{
+    return remote_target_ids_;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+std::shared_ptr<const Executor>
+index_map<LocalIndexType, GlobalIndexType>::get_executor() const
+{
+    return exec_;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+size_type index_map<LocalIndexType, GlobalIndexType>::get_global_size() const
+{
+    return partition_ ? partition_->get_size() : 0;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+array<LocalIndexType> index_map<LocalIndexType, GlobalIndexType>::get_local(
+    const array<GlobalIndexType>& global_ids, index_space is) const
+{
+    array<LocalIndexType> local_ids(exec_);
+
+    exec_->run(index_map_kernels::make_get_local(
+        partition_.get(), remote_target_ids_, remote_global_idxs_, rank_,
+        global_ids, is, local_ids));
+
+    return local_ids;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>::index_map(
+    std::shared_ptr<const Executor> exec, std::shared_ptr<const part_type> part,
+    comm_index_type rank, const array<GlobalIndexType>& recv_connections)
+    : exec_(std::move(exec)),
+      partition_(std::move(part)),
+      rank_(rank),
+      remote_target_ids_(exec_),
+      remote_local_idxs_(exec_),
+      remote_global_idxs_(exec_)
+{
+    array<LocalIndexType> flat_remote_local_idxs(exec_);
+    array<GlobalIndexType> flat_remote_global_idxs(exec_);
+    array<int64> remote_sizes(exec_);
+    exec_->run(index_map_kernels::make_build_mapping(
+        partition_.get(), recv_connections, remote_target_ids_,
+        flat_remote_local_idxs, flat_remote_global_idxs, remote_sizes));
+    remote_local_idxs_ = segmented_array<LocalIndexType>::create_from_sizes(
+        std::move(flat_remote_local_idxs), remote_sizes);
+    remote_global_idxs_ = segmented_array<GlobalIndexType>::create_from_sizes(
+        std::move(flat_remote_global_idxs), remote_sizes);
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>::index_map(
+    std::shared_ptr<const Executor> exec)
+    : exec_(exec),
+      partition_(nullptr),
+      remote_target_ids_(exec),
+      remote_local_idxs_(exec),
+      remote_global_idxs_(exec)
+{}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>&
+index_map<LocalIndexType, GlobalIndexType>::operator=(const index_map& other)
+{
+    if (this != &other) {
+        partition_ = other.partition_;
+        rank_ = other.rank_;
+        remote_target_ids_ = other.remote_target_ids_;
+        remote_local_idxs_ = other.remote_local_idxs_;
+        remote_global_idxs_ = other.remote_global_idxs_;
+    }
+    return *this;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>&
+index_map<LocalIndexType, GlobalIndexType>::operator=(index_map&& other)
+{
+    if (this != &other) {
+        partition_ = std::move(other.partition_);
+        rank_ = other.rank_;
+        remote_target_ids_ = std::move(other.remote_target_ids_);
+        remote_local_idxs_ = std::move(other.remote_local_idxs_);
+        remote_global_idxs_ = std::move(other.remote_global_idxs_);
+    }
+    return *this;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>::index_map(const index_map& other)
+    : exec_(other.get_executor()),
+      remote_local_idxs_(other.get_executor()),
+      remote_global_idxs_(other.get_executor())
+{
+    *this = other;
+}
+
+
+template <typename LocalIndexType, typename GlobalIndexType>
+index_map<LocalIndexType, GlobalIndexType>::index_map(
+    index_map&& other) noexcept
+    : exec_(other.exec_),
+      remote_local_idxs_(other.get_executor()),
+      remote_global_idxs_(other.get_executor())
+{
+    *this = std::move(other);
+}
+
+
+#define GKO_DECLARE_INDEX_MAP(_ltype, _gtype) class index_map<_ltype, _gtype>
+
+GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(GKO_DECLARE_INDEX_MAP);
+
+
+}  // namespace distributed
+}  // namespace experimental
+}  // namespace gko

core/distributed/index_map_kernels.hpp

@@ -0,0 +1,62 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef INDEX_MAP_KERNELS_HPP
+#define INDEX_MAP_KERNELS_HPP
+
+
+#include <ginkgo/core/base/array.hpp>
+#include <ginkgo/core/base/segmented_array.hpp>
+#include <ginkgo/core/distributed/index_map.hpp>
+#include <ginkgo/core/distributed/partition.hpp>
+
+
+#include "core/base/kernel_declaration.hpp"
+
+
+namespace gko {
+namespace kernels {
+
+
+#define GKO_DECLARE_INDEX_MAP_BUILD_MAPPING(_ltype, _gtype)                  \
+    void build_mapping(                                                      \
+        std::shared_ptr<const DefaultExecutor> exec,                         \
+        const experimental::distributed::Partition<_ltype, _gtype>* part,    \
+        const array<_gtype>& recv_connections,                               \
+        array<experimental::distributed::comm_index_type>& part_ids,         \
+        array<_ltype>& remote_local_idxs, array<_gtype>& remote_global_idxs, \
+        array<int64>& remote_sizes)
+
+
+#define GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(_ltype, _gtype)      \
+    void get_local(                                                            \
+        std::shared_ptr<const DefaultExecutor> exec,                           \
+        const experimental::distributed::Partition<_ltype, _gtype>* partition, \
+        const array<experimental::distributed::comm_index_type>&               \
+            remote_targed_ids,                                                 \
+        const segmented_array<_gtype>& remote_global_idxs,                     \
+        experimental::distributed::comm_index_type rank,                       \
+        const array<_gtype>& global_ids,                                       \
+        experimental::distributed::index_space is, array<_ltype>& local_ids)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                                      \
+    template <typename LocalIndexType, typename GlobalIndexType>          \
+    GKO_DECLARE_INDEX_MAP_BUILD_MAPPING(LocalIndexType, GlobalIndexType); \
+    template <typename LocalIndexType, typename GlobalIndexType>          \
+    GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(LocalIndexType,     \
+                                                      GlobalIndexType)
+
+
+GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(index_map,
+                                        GKO_DECLARE_ALL_AS_TEMPLATES);
+
+
+#undef GKO_DECLARE_ALL_AS_TEMPLATES
+
+
+}  // namespace kernels
+}  // namespace gko
+
+#endif  // INDEX_MAP_KERNELS_HPP

reference/CMakeLists.txt

@@ -12,6 +12,7 @@ target_sources(ginkgo_reference
     components/reduce_array_kernels.cpp
     components/precision_conversion_kernels.cpp
     components/prefix_sum_kernels.cpp
+    distributed/index_map_kernels.cpp
     distributed/matrix_kernels.cpp
     distributed/partition_helpers_kernels.cpp
     distributed/partition_kernels.cpp