index_map.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_INDEX_MAP_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_INDEX_MAP_HPP_
 
 
#include <ginkgo/core/base/segmented_array.hpp>
#include <ginkgo/core/distributed/partition.hpp>
 
 
namespace gko {
namespace experimental {
namespace distributed {
 
 
enum class index_space {
    local,      
    non_local,  
    combined    
};
 
 
template <typename LocalIndexType, typename GlobalIndexType = int64>
struct index_map {
    using partition_type = Partition<LocalIndexType, GlobalIndexType>;
 
    array<LocalIndexType> map_to_local(const array<GlobalIndexType>& global_ids,
                                       index_space index_space_v) const;
 
    size_type get_global_size() const;
 
    size_type get_local_size() const;
 
    size_type get_non_local_size() const;
 
    index_map(std::shared_ptr<const Executor> exec,
              std::shared_ptr<const partition_type> partition,
              comm_index_type rank,
              const array<GlobalIndexType>& recv_connections);
 
    index_map(std::shared_ptr<const Executor> exec);
 
    const segmented_array<GlobalIndexType>& get_remote_global_idxs() const;
 
    const segmented_array<LocalIndexType>& get_remote_local_idxs() const;
 
    const array<comm_index_type>& get_remote_target_ids() const;
 
    std::shared_ptr<const Executor> get_executor() const;
 
    index_map(const index_map& other);
 
    index_map(index_map&& other) noexcept;
 
    index_map& operator=(const index_map& other);
 
    index_map& operator=(index_map&& other);
 
private:
    std::shared_ptr<const Executor> exec_;
    std::shared_ptr<const partition_type> partition_;
    comm_index_type rank_;
 
    array<comm_index_type> remote_target_ids_;
    segmented_array<LocalIndexType> remote_local_idxs_;
    segmented_array<GlobalIndexType> remote_global_idxs_;
};
 
 
}  // namespace distributed
}  // namespace experimental
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_INDEX_MAP_HPP_