review updates:

- documentation - renaming Co-authored-by: Pratik Nayak <[email protected]> Co-authored-by: Tobias Ribizel <[email protected]>
ginkgo-project · May 6, 2024 · 8cb01b4 · 8cb01b4
1 parent 9164009
commit 8cb01b4
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 73 deletions.
diff --git a/common/unified/distributed/index_map_kernels.cpp b/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

diff --git a/core/device_hooks/common_kernels.inc.cpp b/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

diff --git a/core/distributed/index_map.cpp b/core/distributed/index_map.cpp
@@ -13,7 +13,7 @@ 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 +76,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_, remote_global_idxs_, rank_,
-        global_ids, is, local_ids));
+        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_),

diff --git a/core/distributed/index_map_kernels.hpp b/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>
 
 
@@ -29,12 +31,12 @@ namespace kernels {
         std::vector<int64>& remote_sizes)
 
 
-#define GKO_DECLARE_INDEX_MAP_GET_LOCAL_FROM_GLOBAL_ARRAY(_ltype, _gtype)      \
-    void get_local(                                                            \
+#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,                                                 \
         const segmented_array<_gtype>& remote_global_idxs,                     \
         experimental::distributed::comm_index_type rank,                       \
         const array<_gtype>& global_ids,                                       \
@@ -45,8 +47,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,
@@ -59,4 +60,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_
diff --git a/include/ginkgo/core/distributed/index_map.hpp b/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);
 
     /**
@@ -157,7 +164,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_;

diff --git a/reference/distributed/index_map_kernels.cpp b/reference/distributed/index_map_kernels.cpp
@@ -8,6 +8,7 @@
 #include <ginkgo/core/distributed/index_map.hpp>
 
 
+#include "core/base/allocator.hpp"
 #include "core/base/segmented_array.hpp"
 #include "reference/distributed/partition_helpers.hpp"
 
@@ -30,35 +31,12 @@ void build_mapping(
     std::vector<int64>& remote_sizes)
 {
     using experimental::distributed::comm_index_type;
-    using partition_type =
-        experimental::distributed::Partition<LocalIndexType, GlobalIndexType>;
     auto part_ids = part->get_part_ids();
 
-    std::vector<GlobalIndexType> unique_indices(recv_connections.get_size());
+    vector<GlobalIndexType> unique_indices(recv_connections.get_size(), {exec});
     std::copy_n(recv_connections.get_const_data(), recv_connections.get_size(),
                 unique_indices.begin());
 
-    auto find_range = [](GlobalIndexType idx, const partition_type* partition,
-                         size_type hint) {
-        auto range_bounds = partition->get_range_bounds();
-        auto num_ranges = partition->get_num_ranges();
-        if (range_bounds[hint] <= idx && idx < range_bounds[hint + 1]) {
-            return hint;
-        } else {
-            auto it = std::upper_bound(range_bounds + 1,
-                                       range_bounds + num_ranges + 1, idx);
-            return static_cast<size_type>(std::distance(range_bounds + 1, it));
-        }
-    };
-
-    auto map_to_local = [](GlobalIndexType idx, const partition_type* partition,
-                           size_type range_id) {
-        auto range_bounds = partition->get_range_bounds();
-        auto range_starting_indices = partition->get_range_starting_indices();
-        return static_cast<LocalIndexType>(idx - range_bounds[range_id]) +
-               range_starting_indices[range_id];
-    };
-
     auto find_part = [&](GlobalIndexType idx) {
         auto range_id = find_range(idx, part, 0);
         return part_ids[range_id];
@@ -132,11 +110,11 @@ 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,
-    const array<experimental::distributed::comm_index_type>& remote_targed_ids,
+    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,
@@ -167,14 +145,14 @@ void get_local(
         // global index. As a result, the array is not sorted wrt.
         // the global indexing. So find the part-id that corresponds
         // to the global index first
-        auto set_id = std::distance(
-            remote_targed_ids.get_const_data(),
-            std::lower_bound(remote_targed_ids.get_const_data(),
-                             remote_targed_ids.get_const_data() +
-                                 remote_targed_ids.get_size(),
-                             part_id));
-
-        if (set_id == remote_targed_ids.get_size()) {
+        auto set_id =
+            std::distance(remote_target_ids.get_const_data(),
+                          std::lower_bound(remote_target_ids.get_const_data(),
+                                           remote_target_ids.get_const_data() +
+                                               remote_target_ids.get_size(),
+                                           part_id));
+
+        if (set_id == remote_target_ids.get_size()) {
             return invalid_index<LocalIndexType>();
         }
 
@@ -225,7 +203,7 @@ void get_local(
 }
 
 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
 }  // namespace reference

diff --git a/reference/test/distributed/index_map_kernels.cpp b/reference/test/distributed/index_map_kernels.cpp
@@ -20,8 +20,6 @@
 #include "core/distributed/index_map_kernels.hpp"
 #include "core/test/utils.hpp"
 
-namespace {
-
 
 using comm_index_type = gko::experimental::distributed::comm_index_type;
 
@@ -195,6 +193,3 @@ TEST_F(IndexMap, CanGetLocalWithCombinedISWithInvalid)
     gko::array<local_index_type> expected(ref, {2, 3, 0, 1, 2, 4, -1, 1});
     GKO_ASSERT_ARRAY_EQ(local_ids, expected);
 }
-
-
-}  // namespace