algorithms.hpp

#ifndef DIY_ALGORITHMS_HPP
#define DIY_ALGORITHMS_HPP

#include <vector>

#include "master.hpp"
#include "assigner.hpp"
#include "reduce.hpp"
#include "reduce-operations.hpp"
#include "partners/swap.hpp"

#include "detail/algorithms/sort.hpp"
#include "detail/algorithms/kdtree.hpp"
#include "detail/algorithms/kdtree-sampling.hpp"

#include "log.hpp"

namespace diy
{
    template<class Block, class T, class Cmp>
    void sort(Master&                   master,               
              const Assigner&           assigner,             
              std::vector<T> Block::*   values,               
              std::vector<T> Block::*   samples,              
              size_t                    num_samples,          
              const Cmp&                cmp,                  
              int                       k   = 2,              
              bool                      samples_only = false) 
    {
        bool immediate = master.immediate();
        master.set_immediate(false);

        // NB: although sorter will go out of scope, its member functions sample()
        //     and exchange() will return functors whose copies get saved inside reduce
        detail::SampleSort<Block,T,Cmp> sorter(values, samples, cmp, num_samples);

        // swap-reduce to all-gather samples
        RegularDecomposer<DiscreteBounds> decomposer(1, interval(0,assigner.nblocks()), assigner.nblocks());
        RegularSwapPartners   partners(decomposer, k);
        reduce(master, assigner, partners, sorter.sample(), detail::SkipIntermediate(partners.rounds()));

        // all_to_all to exchange the values
        if (!samples_only)
            all_to_all(master, assigner, sorter.exchange(), k);

        master.set_immediate(immediate);
    }


    template<class Block, class T>
    void sort(Master&                   master,      
              const Assigner&           assigner,    
              std::vector<T> Block::*   values,      
              std::vector<T> Block::*   samples,     
              size_t                    num_samples, 
              int                       k   = 2)     
    {
        sort(master, assigner, values, samples, num_samples, std::less<T>(), k);
    }

    template<class Block, class Point>
    void kdtree(Master&                         master,      
                const Assigner&                 assigner,    
                int                             dim,         
                const ContinuousBounds&         domain,      
                std::vector<Point>  Block::*    points,      
                size_t                          bins,        
                bool                            wrap = false)
    {
        if (assigner.nblocks() & (assigner.nblocks() - 1))
            throw std::runtime_error(fmt::format("KD-tree requires a number of blocks that's a power of 2, got {}", assigner.nblocks()));

        typedef     diy::RegularContinuousLink      RCLink;

        for (size_t i = 0; i < master.size(); ++i)
        {
            RCLink* link   = static_cast<RCLink*>(master.link(i));
            *link = RCLink(dim, domain, domain);

            if (wrap)       // set up the links to self
            {
                diy::BlockID self = { master.gid(i), master.communicator().rank() };
                for (int j = 0; j < dim; ++j)
                {
                    diy::Direction dir, wrap_dir;

                    // left
                    dir[j] = -1; wrap_dir[j] = -1;
                    link->add_neighbor(self);
                    link->add_bounds(domain);
                    link->add_direction(dir);
                    link->add_wrap(wrap_dir);

                    // right
                    dir[j] = 1; wrap_dir[j] = 1;
                    link->add_neighbor(self);
                    link->add_bounds(domain);
                    link->add_direction(dir);
                    link->add_wrap(wrap_dir);
                }
            }
        }

        detail::KDTreePartition<Block,Point>    kdtree_partition(dim, points, bins);

        detail::KDTreePartners                  partners(dim, assigner.nblocks(), wrap, domain);
        reduce(master, assigner, partners, kdtree_partition);

        // update master.expected to match the links
        int expected = 0;
        for (size_t i = 0; i < master.size(); ++i)
            expected += master.link(i)->size_unique();
        master.set_expected(expected);
    }

    template<class Block, class Point>
    void kdtree_sampling
               (Master&                         master,      
                const Assigner&                 assigner,    
                int                             dim,         
                const ContinuousBounds&         domain,      
                std::vector<Point>  Block::*    points,      
                size_t                          samples,     
                bool                            wrap = false)
    {
        if (assigner.nblocks() & (assigner.nblocks() - 1))
            throw std::runtime_error(fmt::format("KD-tree requires a number of blocks that's a power of 2, got {}", assigner.nblocks()));

        typedef     diy::RegularContinuousLink      RCLink;

        for (size_t i = 0; i < master.size(); ++i)
        {
            RCLink* link   = static_cast<RCLink*>(master.link(i));
            *link = RCLink(dim, domain, domain);

            if (wrap)       // set up the links to self
            {
                diy::BlockID self = { master.gid(i), master.communicator().rank() };
                for (int j = 0; j < dim; ++j)
                {
                    diy::Direction dir, wrap_dir;

                    // left
                    dir[j] = -1; wrap_dir[j] = -1;
                    link->add_neighbor(self);
                    link->add_bounds(domain);
                    link->add_direction(dir);
                    link->add_wrap(wrap_dir);

                    // right
                    dir[j] = 1; wrap_dir[j] = 1;
                    link->add_neighbor(self);
                    link->add_bounds(domain);
                    link->add_direction(dir);
                    link->add_wrap(wrap_dir);
                }
            }
        }

        detail::KDTreeSamplingPartition<Block,Point>    kdtree_partition(dim, points, samples);

        detail::KDTreePartners                          partners(dim, assigner.nblocks(), wrap, domain);
        reduce(master, assigner, partners, kdtree_partition);

        // update master.expected to match the links
        int expected = 0;
        for (size_t i = 0; i < master.size(); ++i)
            expected += master.link(i)->size_unique();
        master.set_expected(expected);
    }
}

#endif