review updates:

- documentation
- renaming
- use gko::vector
- test move-out state

Co-authored-by: Pratik Nayak <pratik.nayak@kit.edu>
Co-authored-by: Tobias Ribizel <mail@ribizel.de>
Co-authored-by: Gregor Olenik <gregor.olenik@kit.edu>

common/unified/distributed/index_map_kernels.cpp

@@ -30,7 +30,7 @@ GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
 
 
 template <typename LocalIndexType, typename GlobalIndexType>
-void get_local(
+void map_to_local(
     std::shared_ptr<const DefaultExecutor> exec,
     const experimental::distributed::Partition<LocalIndexType, GlobalIndexType>*
         partition,
@@ -42,7 +42,7 @@ void get_local(
     array<LocalIndexType>& local_ids) GKO_NOT_IMPLEMENTED;
 
 GKO_INSTANTIATE_FOR_EACH_LOCAL_GLOBAL_INDEX_TYPE(
-    GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY);
+    GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL);
 
 
 }  // namespace index_map

core/device_hooks/common_kernels.inc.cpp

@@ -256,7 +256,7 @@ 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);
+GKO_STUB_LOCAL_GLOBAL_TYPE(GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL);
 
 
 }  // namespace index_map

core/distributed/index_map.cpp

@@ -13,7 +13,8 @@ namespace index_map_kernels {
 
 
 GKO_REGISTER_OPERATION(build_mapping, index_map::build_mapping);
-GKO_REGISTER_OPERATION(get_local, index_map::get_local);
+GKO_REGISTER_OPERATION(map_to_local, index_map::map_to_local);
+
 
 }  // namespace index_map_kernels
 
@@ -76,25 +77,26 @@ size_type index_map<LocalIndexType, GlobalIndexType>::get_global_size() const
 
 
 template <typename LocalIndexType, typename GlobalIndexType>
-array<LocalIndexType> index_map<LocalIndexType, GlobalIndexType>::get_local(
-    const array<GlobalIndexType>& global_ids, index_space is) const
+array<LocalIndexType> index_map<LocalIndexType, GlobalIndexType>::map_to_local(
+    const array<GlobalIndexType>& global_ids, index_space index_space_v) const
 {
     array<LocalIndexType> local_ids(exec_);
 
-    exec_->run(index_map_kernels::make_get_local(
+    exec_->run(index_map_kernels::make_map_to_local(
         partition_.get(), remote_target_ids_, to_device(remote_global_idxs_),
-        rank_, global_ids, is, local_ids));
+        rank_, global_ids, index_space_v, 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)
+    std::shared_ptr<const Executor> exec,
+    std::shared_ptr<const partition_type> partition, comm_index_type rank,
+    const array<GlobalIndexType>& recv_connections)
     : exec_(std::move(exec)),
-      partition_(std::move(part)),
+      partition_(std::move(partition)),
       rank_(rank),
       remote_target_ids_(exec_),
       remote_local_idxs_(exec_),

core/distributed/index_map_kernels.hpp

@@ -2,13 +2,15 @@
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
-#ifndef INDEX_MAP_KERNELS_HPP
-#define INDEX_MAP_KERNELS_HPP
+#ifndef GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_
+#define GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_
+
+
+#include <ginkgo/core/distributed/index_map.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>
 
 
@@ -20,6 +22,23 @@ namespace gko {
 namespace kernels {
 
 
+/**
+ * This kernel creates an index map from a partition and global remote indices.
+ *
+ * The index map is defined by the output parameters remote_local_idxs,
+ * remote_global_idxs, and remote_sizes. After this functions:
+ *
+ * - remote_global_idxs contains the unique indices from recv_connections,
+ *   sorted first by the owning part (as defined in the partition) and then by
+ *   the global index
+ * - remote_local_idxs contains the indices of remote_global_idxs but mapped
+ *   into the local index spaces of their owning parts
+ * - remote_sizes contains the number of remote indices (either in
+ *   remote_global_idxs, or remote_local_indices) per owning part.
+ *
+ * The sizes array is used to create segmented arrays for both output index
+ * arrays.
+ */
 #define GKO_DECLARE_INDEX_MAP_BUILD_MAPPING(_ltype, _gtype)                  \
     void build_mapping(                                                      \
         std::shared_ptr<const DefaultExecutor> exec,                         \
@@ -30,12 +49,26 @@ namespace kernels {
         array<int64>& remote_sizes)
 
 
-#define GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(_ltype, _gtype)      \
-    void get_local(                                                            \
+/**
+ * This kernel maps global indices to local indices.
+ *
+ * The global indices in remote_global_idxs are mapped into the local index
+ * space defined by is. The resulting indices are stored in local_ids.
+ * The index map is defined by the input parameters:
+ *
+ * - partition:  the global partition
+ * - remote_target_ids: the owning part ids of each segment of
+ *   remote_global_idxs
+ * - remote_global_idxs: the remote global indices, segmented by the owning part
+ *   ids
+ * - rank: the part id of this process
+ */
+#define GKO_DECLARE_INDEX_MAP_MAP_TO_LOCAL(_ltype, _gtype)                     \
+    void map_to_local(                                                         \
         std::shared_ptr<const DefaultExecutor> exec,                           \
         const experimental::distributed::Partition<_ltype, _gtype>* partition, \
         const array<experimental::distributed::comm_index_type>&               \
-            remote_targed_ids,                                                 \
+            remote_target_ids,                                                 \
         device_segmented_array<const _gtype> remote_global_idxs,               \
         experimental::distributed::comm_index_type rank,                       \
         const array<_gtype>& global_ids,                                       \
@@ -46,8 +79,7 @@ namespace kernels {
     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_INDEX_MAP_MAP_TO_LOCAL(LocalIndexType, GlobalIndexType)
 
 
 GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(index_map,
@@ -60,4 +92,4 @@ GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(index_map,
 }  // namespace kernels
 }  // namespace gko
 
-#endif  // INDEX_MAP_KERNELS_HPP
+#endif  // GKO_CORE_DISTRIBUTED_INDEX_MAP_KERNELS_HPP_

core/test/distributed/index_map.cpp

@@ -89,6 +89,7 @@ TYPED_TEST(IndexMap, CanMoveConstruct)
 {
     using map_type = typename TestFixture::map_type;
     using global_index_type = typename TestFixture::global_index_type;
+    using local_index_type = typename TestFixture::local_index_type;
     gko::array<global_index_type> connections(this->exec, {4, 3, 3, 4, 2});
     auto imap = map_type(this->exec, this->part, 0, connections);
     auto copy = imap;
@@ -112,6 +113,13 @@ TYPED_TEST(IndexMap, CanMoveConstruct)
               imap_remote_local_it);
     ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_data(),
               imap_remote_global_it);
+    ASSERT_EQ(imap.get_global_size(), 0);
+    ASSERT_EQ(imap.get_local_size(), 0);
+    GKO_ASSERT_ARRAY_EQ(imap.get_remote_target_ids(), {});
+    assert_collection_eq(imap.get_remote_global_idxs(),
+                         gko::segmented_array<global_index_type>{this->exec});
+    assert_collection_eq(imap.get_remote_local_idxs(),
+                         gko::segmented_array<local_index_type>{this->exec});
 }
 
 
@@ -138,6 +146,7 @@ TYPED_TEST(IndexMap, CanMoveAssign)
 {
     using map_type = typename TestFixture::map_type;
     using global_index_type = typename TestFixture::global_index_type;
+    using local_index_type = typename TestFixture::local_index_type;
     gko::array<global_index_type> connections(this->exec, {4, 3, 3, 4, 2});
     auto imap = map_type(this->exec, this->part, 0, connections);
     auto copy = imap;
@@ -162,4 +171,11 @@ TYPED_TEST(IndexMap, CanMoveAssign)
               imap_remote_local_it);
     ASSERT_EQ(move.get_remote_global_idxs().get_const_flat_data(),
               imap_remote_global_it);
+    ASSERT_EQ(imap.get_global_size(), 0);
+    ASSERT_EQ(imap.get_local_size(), 0);
+    GKO_ASSERT_ARRAY_EQ(imap.get_remote_target_ids(), {});
+    assert_collection_eq(imap.get_remote_global_idxs(),
+                         gko::segmented_array<global_index_type>{this->exec});
+    assert_collection_eq(imap.get_remote_local_idxs(),
+                         gko::segmented_array<local_index_type>{this->exec});
 }

include/ginkgo/core/distributed/index_map.hpp

@@ -17,6 +17,8 @@ namespace distributed {
 
 /**
  * \brief Index space classification for the locally stored indices.
+ *
+ * The definitions of the enum values is clarified in @ref index_map.
  */
 enum class index_space {
     local,      //!< indices that are locally owned
@@ -42,17 +44,20 @@ enum class index_space {
  * (index_space::non_local). The locally owned indices are defined as
  * $[0, \dots, |I_k|)$, and the non-locally owned as $[0, \dots, |R_k|)$.
  * With these index sets, the following mappings are defined:
+ *
  * - $c_k : \hat{I}_k \mapsto [0, \dots, |\hat{I}_k|)$ which maps global indices
- *   into the combined/full local index space,
+ *   into the combined/full local index space (denoted as
+ *   index_space::combined),
  * - $l_k: I_k \mapsto [0, \dots, |I_k|)$ which maps global indices into the
- *   locally owned index space,
+ *   locally owned index space (denoted as index_space::local),
  * - $r_k: R_k \mapsto [0, \dots, |R_k|)$ which maps global indices into the
- *   non-locally owned index space.
+ *   non-locally owned index space (denoted as index_space::non_local).
+ *
  * The required map can be selected by passing the appropriate type of an
  * index_space.
  *
  * The index map for $I_k$ has no knowledge about any other index maps for
- * $I_l, l \neq k$. In particular, any global index passed to the `get_local`
+ * $I_l, l \neq k$. In particular, any global index passed to the `map_to_local`
  * map that is not part of the specified index space, will be mapped to an
  * invalid_index.
  *
@@ -61,18 +66,19 @@ enum class index_space {
  */
 template <typename LocalIndexType, typename GlobalIndexType = int64>
 struct index_map {
-    using part_type = Partition<LocalIndexType, GlobalIndexType>;
+    using partition_type = Partition<LocalIndexType, GlobalIndexType>;
 
     /**
      * \brief Maps global indices to local indices
      *
-     * \param global_ids the global indices to map
-     * \param is the index space in which the returned local indices are defined
+     * \param global_ids  the global indices to map
+     * \param index_space_v  the index space in which the returned local indices
+     *                       are defined
      *
-     * \return the mapped local indices. Any global index that is not in the
-     *         specified index space is mapped to invalid_index.
+     * \return  the mapped local indices. Any global index that is not in the
+     *          specified index space is mapped to invalid_index.
      */
-    array<LocalIndexType> get_local(
+    array<LocalIndexType> map_to_local(
         const array<GlobalIndexType>& global_ids,
         index_space index_space_v = index_space::combined) const;
 
@@ -98,13 +104,14 @@ struct index_map {
      * filtered out.
      *
      * \param exec  the executor
-     * \param part  the partition of the global index set
+     * \param partition  the partition of the global index set
      * \param rank  the id of the global index space subset
      * \param recv_connections  the global indices that are not owned by this
      *                          rank, but accessed by it
      */
     index_map(std::shared_ptr<const Executor> exec,
-              std::shared_ptr<const part_type> part, comm_index_type rank,
+              std::shared_ptr<const partition_type> partition,
+              comm_index_type rank,
               const array<GlobalIndexType>& recv_connections);
 
     /**
@@ -123,7 +130,8 @@ struct index_map {
      * \brief get the index set $R_k$, but mapped to their respective local
      *        index space.
      *
-     * The indices are ordered by their owning rank and global index.
+     * The indices are grouped by their owning rank and sorted according to
+     * their global index within each group.
      *
      * The set $R_k = \hat{I}_k \setminus I_k$ can also be written as the union
      * of the intersection of $\hat{I}_k$ with other disjoint sets
@@ -157,7 +165,7 @@ struct index_map {
 
 private:
     std::shared_ptr<const Executor> exec_;
-    std::shared_ptr<const part_type> partition_;
+    std::shared_ptr<const partition_type> partition_;
     comm_index_type rank_;
 
     array<comm_index_type> remote_target_ids_;