balance_tree.hpp 9.42 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193
#ifndef _INRIA_BALANCE_TREE_
#define _INRIA_BALANCE_TREE_

#include "weight_traits.hpp"
#include "distributed_regions_to_linear_tree.hpp"
#include "gather_octant_weights.hpp"

#include "inria/algorithm/distributed/sort.hpp"
#include "inria/algorithm/distributed/distribute.hpp"

namespace inria {
namespace linear_tree {
namespace details {

/**
 * \brief Node info with added weigh information
 */
template<class P>
struct weighted_node_info : node::info<P::position_t::Dim> {
    /// Weight attribute type
    using weight_t = decltype(inria::meta::get_weight(std::declval<P>()));
    /// Node weight
    weight_t weight = 0;

    // Use parent constructor
    using node::info<P::position_t::Dim>::info;

    /**
     * \brief Output stream operator
     *
     * \param os Output stream
     * \param n  Weighted node info object to stream
     */
    friend std::ostream& operator<<(std::ostream& os, const weighted_node_info& n) {
        using node::level;
        using node::morton_index;
        return os << '(' << morton_index(n) << ", " << level(n) << ", " << n.weight << ')';
    }
};

/**
 * \brief Function object to access weight
 */
struct weight_accessor {
    template<class T>
    auto operator()(const T& e) const noexcept(noexcept(inria::meta::get_weight(e)))
        -> decltype(inria::meta::get_weight(e))
    {
        return inria::meta::get_weight(e);
    }
};


/**
 * \brief Helper type aliases for create_balanced_linear_tree
 */
namespace cblt {

/// Alias to extract the node::info<Dim> type from a particle iterator
template<class T>
using node_info_from_it = node::info<std::iterator_traits<T>::value_type::position_t::Dim>;

/// Alias to extract the node::info<Dim> type from a particle range
template<class T>
using node_info_from_range = node::info<T::value_type::position_t::Dim>;

} // close namespace [inria::linear_tree::details]::cblts

} // close namespace [inria::linear_tree]::details


/**
 * \brief Create a distributed linear tree from a sorted particle list
 *
 * \warning The particle list must be sorted accross the processes.
 *
 * \param conf   MPI configuration
 * \param level_ Maximum tree depth
 * \param box    The space bounding box
 * \param sorted_particles Distributed list of particles sorted by morton index
 *
 * \note The is box expected to be a cube.
 *
 * \tparam Range Particle list, must define begin and end functions and a
 *               value_type type alias.
 * \tparam Box   Must define `Box::width(int axis)` and `Box::corner(int morton_index)`.
 *
 * \return A holding the local leaf information (inria::linear_tree::node::info)
 *         of the tree.
 */
template<class Range, class Box>
std::vector<details::cblt::node_info_from_range<Range>>
create_balanced_linear_tree(
    inria::mpi_config conf,
    std::size_t level_,
    Box box,
    Range& sorted_particles)
{
    using particle_t = typename Range::value_type;
    using node_info_t = details::weighted_node_info<particle_t>;

    // Compute minimal and maximal octant for local range.
    node_info_t min_oct{}, max_oct{};
    if(! sorted_particles.empty()) {
        auto min_idx = get_morton_index(sorted_particles.front().position(), box, level_);
        min_oct = node_info_t{min_idx, level_};
        auto max_idx = get_morton_index(sorted_particles.back().position(), box, level_);
        max_oct = node_info_t{max_idx, level_};
    }

    // Complete the local region based on sorted particles
    std::vector<node_info_t> local_region;
    complete_region(min_oct, max_oct, local_region);

    // Keep only the coarsest octants
    std::size_t min_level = std::accumulate(begin(local_region), end(local_region), level_,
                                            [](std::size_t l, const node_info_t& n) {
                                                return std::min(l, level(n));
                                            });
    local_region.erase(std::remove_if(begin(local_region), end(local_region),
                                      [&](const node_info_t& n) {
                                          return level(n) > min_level;
                                      }),
                       end(local_region));
    // Complete distributed octree
    std::vector<node_info_t> local_tree;
    distributed_regions_to_linear_tree(conf, begin(local_region), end(local_region), local_tree);

    // Compute the weight of each octant to redistribute them among processes
    gather_octants_weight(conf, begin(local_tree), end(local_tree),
                         begin(sorted_particles), end(sorted_particles), box);

    distribute(conf, local_tree, uniform_distribution<details::weight_accessor>{conf, local_tree});

    coarsen_region(local_tree);

    return {begin(local_tree), end(local_tree)};
}


/**
 * \brief Create a distributed linear tree from a particle list
 *
 * \warning A copy of the particle list is made then sorted.
 *
 * \param conf   MPI configuration
 * \param level_ Maximum tree depth
 * \param box    The space bounding box
 * \param first  Particle list beginning
 * \param last   Particle list end
 * \param comp   Comparison function object used to sort the particles
 *
 * \note The is box expected to be a cube.
 *
 * \tparam ParticleForwardIt Forward iterator
 * \tparam Box   Must define `Box::width(int axis)` and `Box::corner(int morton_index)`.
 *
 * \return A holding the local leaf information (inria::linear_tree::node::info)
 *         of the tree.
 */
template<class ParticleForwardIt, class Comp, class Box>
std::vector<details::cblt::node_info_from_it<ParticleForwardIt>>
create_balanced_linear_tree(
    inria::mpi_config conf,
    std::size_t level_,
    Box box,
    ParticleForwardIt first,
    ParticleForwardIt last,
    Comp&& comp)
{
    auto sorted_particles = inria::sort(conf, first, last, comp);
    return create_balanced_linear_tree(conf, level_, box, sorted_particles);
}


namespace details {

/**
 * \brief Distribution object using pivots to choose target process
 *
 * This class is to be used with the inria::distribute function. When called, it
 * returns the rank of the process that should hold the given element after the
 * sort.
 *
 * \tparam Comp   The comparison callable type.
 * \tparam Pivots The list of pivots type.
 */
template<class Comp, class Pivots>
struct interval_distribution {
    Comp comp;
    Pivots pivots;
    template<class U>
    std::uint64_t operator()(const U& e) {
194
        auto p_it = std::lower_bound(std::begin(pivots), std::end(pivots), e, comp);
195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275
        std::uint64_t p_idx = p_it - std::begin(pivots);
        if(p_idx > pivots.size()) {
            throw std::runtime_error("Element distributed to non existing "
                                     "process " + std::to_string(p_idx));
        }
        return p_idx;
    }
};

template<class Comp, class Pivots>
interval_distribution<Comp, Pivots> make_interval_distribution(Comp&& comp, Pivots&& pivots)
{
    return {std::forward<Comp>(comp), std::forward<Pivots>(pivots)};
}


}

/**
 * \brief Redistribute particles using the linear tree structure as a criteria
 *
 * \param conf        MPI configuration
 * \param linear_tree Local linear tree segment
 * \param particles   Particle list
 */
template<class AssignableRange, class LinearTree,
         class IsRange = inria::enable_if_t<inria::is_range<AssignableRange>::value>,
         class IsAssignable = inria::enable_if_t<
             inria::is_assignable<AssignableRange,
                                  typename AssignableRange::iterator,
                                  typename AssignableRange::iterator
                                  >::value>
         >
void redistribute_particles(inria::mpi_config conf, const LinearTree& linear_tree, AssignableRange& particles) {
    using node::morton_index;
    using morton_index_t = typename std::remove_reference<decltype(morton_index(particles.front()))>::type;
    using particle_t = typename AssignableRange::value_type;

    // Theoretical tree max depth, assumes sizeof(morton_index_t) is the exact
    // storage size (i.e. morton_index_t does not store any other information)
    constexpr auto max_depth = 8 * sizeof(morton_index_t) / 3;
    // Setup MPI datatype for future sends
    inria::mpi::datatype_commit_guard MortonType{inria::mpi::get_datatype<morton_index_t>()};
    // Setup pivots for distribution, initial value is 0b111....111
    std::vector<morton_index_t> pivots(conf.comm.size(), ~morton_index_t{0});
    // Initialise current process pivot
    if(! linear_tree.empty()) {
        auto& n = linear_tree.back();
        pivots[conf.comm.rank()] = morton_index(last_descendant(n, max_depth - n.level));
    }
    // Share pivots
    conf.comm.allgather(pivots.data(), 1, MortonType.datatype);

    // Correct pivots, if a process holds no elements, copy the previous process
    // pivot. Note: this breaks if process 0 does not hold elements
    for(std::size_t i = 1u; i < pivots.size(); ++i) {
        if(pivots[i] == ~morton_index_t{0}){
            pivots[i] = pivots[i-1];
        }
    }

    // Object for comparison between morton index and particle
    struct {
        bool operator()(const morton_index_t& m, const particle_t& p) {
            return m < morton_index(p);
        };
        bool operator()(const particle_t& p, const morton_index_t& m) {
            return morton_index(p) < m;
        }
    } comp;
    // Setup an run distribution algorithm
    auto dist = details::make_interval_distribution(comp, std::move(pivots));
    inria::distribute(conf, particles, dist);
}


}} // close namespace inria::linear_tree



#endif /* _INRIA_BALANCE_TREE_ */