review updates:

- documentation
- renaming

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 remove_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_

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_;

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"
 
@@ -29,35 +30,12 @@ void build_mapping(
  array<GlobalIndexType>& remote_global_idxs, array<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];
@@ -127,11 +105,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,
  device_segmented_array<const GlobalIndexType> remote_global_idxs,
  experimental::distributed::comm_index_type rank,
  const array<GlobalIndexType>& global_ids,
@@ -162,14 +140,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>();
  }
 
@@ -220,7 +198,8 @@ 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