tree_mpi.hpp

#ifndef MPITREE_HPP
#define MPITREE_HPP

#include <tree.hpp>
//#include <DistData.hpp>
using namespace std;
using namespace hmlp;


namespace hmlp
{
namespace mpitree
{



// *  @brief This is the default ball tree splitter. Given coordinates,
// *         compute the direction from the two most far away points.
// *         Project all points to this line and split into two groups
// *         using a median select.
// *
// *  @para
// *
// *  @TODO  Need to explit the parallelism.
// */ 
//template<int N_SPLIT, typename T>
//struct centersplit
//{
//  // closure
//  Data<T> *Coordinate;
//
//  inline vector<vector<size_t> > operator()
//  ( 
//    vector<size_t>& gids
//  ) const 
//  {
//    assert( N_SPLIT == 2 );
//
//    Data<T> &X = *Coordinate;
//    size_t d = X.row();
//    size_t n = gids.size();
//
//    T rcx0 = 0.0, rx01 = 0.0;
//    size_t x0, x1;
//    vector<vector<size_t> > split( N_SPLIT );
//
//
//    vector<T> centroid = combinatorics::Mean( d, n, X, gids );
//    vector<T> direction( d );
//    vector<T> projection( n, 0.0 );
//
//    //printf( "After Mean\n" );
//
//    // Compute the farest x0 point from the centroid
//    for ( int i = 0; i < n; i ++ )
//    {
//      T rcx = 0.0;
//      for ( int p = 0; p < d; p ++ )
//      {
//        T tmp = X[ gids[ i ] * d + p ] - centroid[ p ];
//        rcx += tmp * tmp;
//      }
//      //printf( "\n" );
//      if ( rcx > rcx0 ) 
//      {
//        rcx0 = rcx;
//        x0 = i;
//      }
//    }
//
//
//    // Compute the farest point x1 from x0
//    for ( int i = 0; i < n; i ++ )
//    {
//      T rxx = 0.0;
//      for ( int p = 0; p < d; p ++ )
//      {
//        T tmp = X[ gids[ i ] * d + p ] - X[ gids[ x0 ] * d + p ];
//        rxx += tmp * tmp;
//      }
//      if ( rxx > rx01 )
//      {
//        rx01 = rxx;
//        x1 = i;
//      }
//    }
//
//    // Compute direction
//    for ( int p = 0; p < d; p ++ )
//    {
//      direction[ p ] = X[ gids[ x1 ] * d + p ] - X[ gids[ x0 ] * d + p ];
//    }
//
//    // Compute projection
//    projection.resize( n, 0.0 );
//    for ( int i = 0; i < n; i ++ )
//      for ( int p = 0; p < d; p ++ )
//        projection[ i ] += X[ gids[ i ] * d + p ] * direction[ p ];
//
//    /** Parallel median search */
//    T median;
//    
//    if ( 1 )
//    {
//      median = hmlp::combinatorics::Select( n, n / 2, projection );
//    }
//    else
//    {
//      auto proj_copy = projection;
//      std::sort( proj_copy.begin(), proj_copy.end() );
//      median = proj_copy[ n / 2 ];
//    }
//
//    split[ 0 ].reserve( n / 2 + 1 );
//    split[ 1 ].reserve( n / 2 + 1 );
//
//
//    /** TODO: Can be parallelized */
//    std::vector<std::size_t> middle;
//    for ( size_t i = 0; i < n; i ++ )
//    {
//      if      ( projection[ i ] < median ) split[ 0 ].push_back( i );
//      else if ( projection[ i ] > median ) split[ 1 ].push_back( i );
//      else                                 middle.push_back( i );
//    }
//
//    for ( size_t i = 0; i < middle.size(); i ++ )
//    {
//      if ( split[ 0 ].size() <= split[ 1 ].size() ) split[ 0 ].push_back( middle[ i ] );
//      else                                          split[ 1 ].push_back( middle[ i ] );
//    }
//
//
//    return split;
//  };
//
//
//  inline std::vector<std::vector<size_t> > operator()
//  ( 
//    std::vector<size_t>& gids,
//    hmlp::mpi::Comm comm
//  ) const 
//  {
//    std::vector<std::vector<size_t> > split( N_SPLIT );
//
//    return split;
//  };
//
//};
//
//
//
//
//
//template<int N_SPLIT, typename T>
//struct randomsplit
//{
//  Data<T> *Coordinate = NULL;
//
//  inline vector<vector<size_t> > operator() ( vector<size_t>& gids ) const 
//  {
//    vector<vector<size_t> > split( N_SPLIT );
//    return split;
//  };
//
//  inline vector<vector<size_t> > operator() ( vector<size_t>& gids, mpi::Comm comm ) const 
//  {
//    vector<vector<size_t> > split( N_SPLIT );
//    return split;
//  };
//};
//





template<typename NODE>
class DistSplitTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "DistSplit" );
      label = to_string( arg->treelist_id );

      double flops = 6.0 * arg->n;
      double  mops = 6.0 * arg->n;

      event.Set( label + name, flops, mops );
      cost = mops / 1E+9;
      priority = true;
    };


    void DependencyAnalysis()
    {
      arg->DependencyAnalysis( R, this );

      if ( !arg->isleaf )
      {
        if ( arg->GetCommSize() > 1 )
        {
               assert( arg->child );
          arg->child->DependencyAnalysis( RW, this );
        }
        else
        {
               assert( arg->lchild && arg->rchild );
          arg->lchild->DependencyAnalysis( RW, this );
          arg->rchild->DependencyAnalysis( RW, this );
        }
      }
      this->TryEnqueue();
      };

    void Execute( Worker* user_worker ) { arg->Split(); };

}; 
template<typename SPLITTER, typename DATATYPE>
class Setup
{
  public:

    typedef DATATYPE T;

    Setup() {};

    ~Setup() {};




      vector<size_t> ContainAny( vector<size_t> &queries, size_t target )
    {
         vector<size_t> validation( queries.size(), 0 );

      if ( !morton.size() )
      {
        printf( "Morton id was not initialized.\n" );
        exit( 1 );
      }

      for ( size_t i = 0; i < queries.size(); i ++ )
      {
        //auto it = this->setup->morton.find( queries[ i ] );

            //if ( it != this->setup->morton.end() )
            //{
        //  if ( tree::IsMyParent( *it, this->morton ) ) validation[ i ] = 1;
            //}


       if ( MortonHelper::IsMyParent( morton[ queries[ i ] ], target ) ) 
             validation[ i ] = 1;

      }
      return validation;

    }; 
    size_t m;
    
    size_t max_depth = 15;

    //DistData<STAR, CBLK, T> *X_cblk = NULL;
    //DistData<STAR, CIDS, T> *X      = NULL;

    DistData<STAR, CBLK, pair<T, size_t>> *NN_cblk = NULL;
    DistData<STAR, CIDS, pair<T, size_t>> *NN      = NULL;

    vector<size_t> morton;

    SPLITTER splitter;

}; 
template<typename NODE>
class DistIndexPermuteTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      name = std::string( "Permutation" );
      arg = user_arg;
      // Need an accurate cost model.
      cost = 1.0;
    };

    void DependencyAnalysis() { arg->DependOnChildren( this ); };
    //{
    //  arg->DependencyAnalysis( hmlp::ReadWriteType::RW, this );
    //  if ( !arg->isleaf && !arg->child )
    //  {
    //    arg->lchild->DependencyAnalysis( hmlp::ReadWriteType::R, this );
    //    arg->rchild->DependencyAnalysis( hmlp::ReadWriteType::R, this );
    //  }
    //  this->TryEnqueue();
    //};


    void Execute( Worker* user_worker )
    {
      if ( !arg->isleaf && !arg->child )
      {
        auto &gids = arg->gids; 
        auto &lgids = arg->lchild->gids;
        auto &rgids = arg->rchild->gids;
        gids = lgids;
        gids.insert( gids.end(), rgids.begin(), rgids.end() );
      }
    };

}; 
//template<typename SETUP, int N_CHILDREN, typename NODEDATA>
template<typename SETUP, typename NODEDATA>
class Node : public tree::Node<SETUP, NODEDATA>
{
  public:

    typedef typename SETUP::T T;

    static const int N_CHILDREN = 2;

    Node( SETUP *setup, size_t n, size_t l, 
        Node *parent,
        unordered_map<size_t, tree::Node<SETUP, NODEDATA>*> *morton2node,
        Lock *treelock, mpi::Comm comm ) 
      : tree::Node<SETUP, NODEDATA>( setup, n, l, 
          parent, morton2node, treelock ) 
    {
      this->comm = comm;
      mpi::Comm_size( comm, &size );
      mpi::Comm_rank( comm, &rank );
    };

    Node( SETUP *setup, size_t n, size_t l, vector<size_t> &gids, 
        Node *parent, 
        unordered_map<size_t, tree::Node<SETUP, NODEDATA>*> *morton2node,
        Lock *treelock, mpi::Comm comm ) 
      : Node<SETUP, NODEDATA>( setup, n, l, parent, 
          morton2node, treelock, comm ) 
    {
      this->gids = gids;
    };

    Node( size_t morton ) : tree::Node<SETUP, NODEDATA>( morton )
    { 
    };


    //void SetupChild( class Node *child )
    //{
    //  this->kids[ 0 ] = child;
    //  this->child = child;
    //};



    void Split()
    {
      int num_points_total = 0;
      int num_points_owned = (this->gids).size();
      mpi::Allreduce( &num_points_owned, &num_points_total, 1, MPI_SUM, comm );
      this->n = num_points_total;

      if ( child )
      {
        assert( size > 1 );

        auto split = this->setup->splitter( this->gids, comm );

        int partner_rank = 0;
        int sent_size = 0; 
        int recv_size = 0;
        vector<size_t> &kept_gids = child->gids;
        vector<int>     sent_gids;
        vector<int>     recv_gids;

        if ( rank < size / 2 ) 
        {
          partner_rank = rank + size / 2;
          kept_gids.resize( split[ 0 ].size() );
          for ( size_t i = 0; i < kept_gids.size(); i ++ )
            kept_gids[ i ] = this->gids[ split[ 0 ][ i ] ];
          sent_gids.resize( split[ 1 ].size() );
          sent_size = sent_gids.size();
          for ( size_t i = 0; i < sent_gids.size(); i ++ )
            sent_gids[ i ] = this->gids[ split[ 1 ][ i ] ];
        }
        else       
        {
          partner_rank = rank - size / 2;
          kept_gids.resize( split[ 1 ].size() );
          for ( size_t i = 0; i < kept_gids.size(); i ++ )
            kept_gids[ i ] = this->gids[ split[ 1 ][ i ] ];
          sent_gids.resize( split[ 0 ].size() );
          sent_size = sent_gids.size();
          for ( size_t i = 0; i < sent_gids.size(); i ++ )
            sent_gids[ i ] = this->gids[ split[ 0 ][ i ] ];
        }
        assert( partner_rank >= 0 );





        //mpi::Sendrecv( &sent_size, 1, partner_rank, 10,
        //               &recv_size, 1, partner_rank, 10, comm, &status );

        //printf( "rank %d kept_size %lu sent_size %d recv_size %d\n", 
        //    rank, kept_gids.size(), sent_size, recv_size ); fflush( stdout );

        //recv_gids.resize( recv_size );

        //mpi::Sendrecv( 
        //    sent_gids.data(), sent_size, MPI_INT, partner_rank, 20,
        //    recv_gids.data(), recv_size, MPI_INT, partner_rank, 20,
        //    comm, &status );

        mpi::ExchangeVector( sent_gids, partner_rank, 20,
                             recv_gids, partner_rank, 20, comm, &status );
        for ( auto it : recv_gids ) kept_gids.push_back( it );
        //kept_gids.reserve( kept_gids.size() + recv_gids.size() );
        //for ( size_t i = 0; i < recv_gids.size(); i ++ )
        //  kept_gids.push_back( recv_gids[ i ] );
        

      }
         else
         {
        tree::Node<SETUP, NODEDATA>::Split();
        }
      mpi::Barrier( comm );
    }; 
    void DependOnChildren( Task *task )
    {
      this->DependencyAnalysis( RW, task );
      if ( size < 2 )
      {
        if ( this->lchild ) this->lchild->DependencyAnalysis( R, task );
        if ( this->rchild ) this->rchild->DependencyAnalysis( R, task );
      }
      else
      {
        if ( child ) child->DependencyAnalysis( R, task );
      }
      task->TryEnqueue();
    }; 
    void DependOnParent( Task *task )
    {
      this->DependencyAnalysis( R, task );
      if ( size < 2 )
      {
        if ( this->lchild ) this->lchild->DependencyAnalysis( RW, task );
        if ( this->rchild ) this->rchild->DependencyAnalysis( RW, task );
      }
      else
      {
        if ( child ) child->DependencyAnalysis( RW, task );
      }
      task->TryEnqueue();
    }; 
    void Print() {};

    mpi::Comm GetComm() { return comm; };
    int GetCommSize() { return size; };
    int GetCommRank() { return rank; };

    Node *child = NULL;

  private:

    mpi::Comm comm = MPI_COMM_WORLD;
    mpi::Status status;
    int size = 1;
    int rank = 0;

}; 
template<class SETUP, class NODEDATA>
class Tree : public tree::Tree<SETUP, NODEDATA>,
             public mpi::MPIObject
{
  public:

    typedef typename SETUP::T T;

    typedef tree::Node<SETUP, NODEDATA> NODE;

    typedef Node<SETUP, NODEDATA> MPINODE;

    vector<MPINODE*> mpitreelists;

    
    Tree( mpi::Comm comm ) : tree::Tree<SETUP, NODEDATA>::Tree(),
                             mpi::MPIObject( comm )
    {
      //this->comm = comm;
      //mpi::Comm_size( comm, &size );
      //mpi::Comm_rank( comm, &rank );
      //NearRecvFrom.resize( size );
      NearRecvFrom.resize( this->GetCommSize() );
      //FarRecvFrom.resize( size );
      FarRecvFrom.resize( this->GetCommSize() );
    };

    ~Tree()
    {
      //printf( "~Tree() distributed, mpitreelists.size() %lu\n",
      //    mpitreelists.size() ); fflush( stdout );
      if ( mpitreelists.size() )
      {
        for ( size_t i = 0; i < mpitreelists.size() - 1; i ++ )
          if ( mpitreelists[ i ] ) delete mpitreelists[ i ];
        mpitreelists.clear();
      }
      //printf( "end ~Tree() distributed\n" ); fflush( stdout );
    };


    void CleanUp()
    {
      if ( this->treelist.size() )
      {
        for ( size_t i = 0; i < this->treelist.size(); i ++ )
          if ( this->treelist[ i ] ) delete this->treelist[ i ];
      }
      this->treelist.clear();

      if ( mpitreelists.size() )
      {
        for ( size_t i = 0; i < mpitreelists.size() - 1; i ++ )
          if ( mpitreelists[ i ] ) delete mpitreelists[ i ];
      }
      mpitreelists.clear();

    }; 
    void AllocateNodes( vector<size_t> &gids )
    {
      //auto mycomm  = comm;
      auto mycomm  = this->GetComm();
      //int mysize  = size;
      int mysize  = this->GetCommSize();
      //int myrank  = rank;
      int myrank  = this->GetCommRank();
      int mycolor = 0;
      size_t mylevel = 0;

      auto *root = new MPINODE( &(this->setup), 
          this->n, mylevel, gids, NULL, 
          &(this->morton2node), &(this->lock), mycomm );

      mpitreelists.push_back( root );

      while ( mysize > 1 )
      {
        mpi::Comm childcomm;

        mylevel += 1;
        mycolor = ( myrank < mysize / 2 ) ? 0 : 1;
        ierr = mpi::Comm_split( mycomm, mycolor, myrank, &(childcomm) );
        mycomm = childcomm;
        mpi::Comm_size( mycomm, &mysize );
        mpi::Comm_rank( mycomm, &myrank );

        auto *parent = mpitreelists.back();
        auto *child  = new MPINODE( &(this->setup), 
            (size_t)0, mylevel, parent, 
            &(this->morton2node), &(this->lock), mycomm );

        child->sibling = new NODE( (size_t)0 ); // Node morton is computed later.
        //parent->SetupChild( child );
        parent->kids[ 0 ] = child;
        parent->child = child;
        mpitreelists.push_back( child );
      }
      this->Barrier();

      auto *local_tree_root = mpitreelists.back();
      tree::Tree<SETUP, NODEDATA>::AllocateNodes( local_tree_root );

    }; 
    vector<size_t> GetPermutation()
    {
      vector<size_t> perm_loc, perm_glb;
      perm_loc = tree::Tree<SETUP, NODEDATA>::GetPermutation();
      mpi::GatherVector( perm_loc, perm_glb, 0, this->GetComm() );

      //if ( rank == 0 )
      //{
      //  /** Sanity check using an 0:N-1 table. */
      //  vector<bool> Table( this->n, false );
      //  for ( size_t i = 0; i < perm_glb.size(); i ++ )
      //    Table[ perm_glb[ i ] ] = true;
      //  for ( size_t i = 0; i < Table.size(); i ++ ) assert( Table[ i ] );
      //}

      return perm_glb;
    }; 
    template<typename KNNTASK>
    DistData<STAR, CBLK, pair<T, size_t>>
    AllNearestNeighbor( size_t n_tree, size_t n, size_t k,
      pair<T, size_t> initNN, KNNTASK &dummy )
    {
      mpi::PrintProgress( "[BEG] NeighborSearch ...", this->GetComm() );

      this->n = n;
      DistData<STAR, CBLK, pair<T, size_t>> NN( k, n, initNN, this->GetComm() );
      this->setup.m = 4 * k;
      if ( this->setup.m < 512 ) this->setup.m = 512;
      this->m = this->setup.m;


      //num_points_owned = ( n - 1 ) / this->GetCommSize() + 1;

      //if ( n % this->GetCommSize() )
      //{
      //  //if ( rank >= ( n % size ) ) num_points_owned -= 1;
      //  if ( this->GetCommRank() >= ( n % this->GetCommSize() ) ) 
      //    num_points_owned -= 1;
      //}

      num_points_owned = n / this->GetCommSize();
      if ( this->GetCommRank() < ( n % this->GetCommSize() ) ) 
         num_points_owned += 1;



      vector<size_t> gids( num_points_owned, 0 );
      for ( size_t i = 0; i < num_points_owned; i ++ )
        //gids[ i ] = i * size + rank;
        gids[ i ] = i * this->GetCommSize() + this->GetCommRank();

      AllocateNodes( gids );


      DistSplitTask<MPINODE> mpisplittask;
      tree::SplitTask<NODE>  seqsplittask;
      for ( size_t t = 0; t < n_tree; t ++ )
      {
        DistTraverseDown( mpisplittask );
        LocaTraverseDown( seqsplittask );
        ExecuteAllTasks();
        
        DistData<STAR, CIDS, pair<T, size_t>> Q_cids( k, this->n, this->treelist[ 0 ]->gids, initNN, this->GetComm() );
        this->setup.NN = &Q_cids;
        LocaTraverseLeafs( dummy );
        ExecuteAllTasks();

        DistData<STAR, CBLK, pair<T, size_t>> Q_cblk( k, this->n, this->GetComm() );
        Q_cblk = Q_cids;
        assert( Q_cblk.col_owned() == NN.col_owned() );
        MergeNeighbors( k, NN.col_owned(), NN, Q_cblk );
      }







//      /** Metric tree partitioning. */
//      DistSplitTask<MPINODE> mpisplittask;
//      tree::SplitTask<NODE>  seqsplittask;
//      DependencyCleanUp();
//      DistTraverseDown( mpisplittask );
//      LocaTraverseDown( seqsplittask );
//      ExecuteAllTasks();
//
//
//      for ( size_t t = 0; t < n_tree; t ++ )
//      {
//        this->Barrier();
//        //if ( this->GetCommRank() == 0 ) printf( "Iteration #%lu\n", t );
//
//        /** Query neighbors computed in CIDS distribution.  */
//        DistData<STAR, CIDS, pair<T, size_t>> Q_cids( k, this->n, 
//            this->treelist[ 0 ]->gids, initNN, this->GetComm() );
//        /** Pass in neighbor pointer. */
//        this->setup.NN = &Q_cids;
//        /** Overlap */
//        if ( t != n_tree - 1 )
//        {
//          //DependencyCleanUp();
//          DistTraverseDown( mpisplittask );
//          ExecuteAllTasks();
//        }
//        mpi::PrintProgress( "[MID] Here ...", this->GetComm() );
//        DependencyCleanUp();
//        LocaTraverseLeafs( dummy );
//        LocaTraverseDown( seqsplittask );
//        ExecuteAllTasks();
//        mpi::PrintProgress( "[MID] Here 22...", this->GetComm() );
//
//        if ( t == 0 )
//        {
//          /** Redistribute from CIDS to CBLK */
//          NN = Q_cids; 
//        }
//        else
//        {
//          /** Queries computed in CBLK distribution */
//          DistData<STAR, CBLK, pair<T, size_t>> Q_cblk( k, this->n, this->GetComm() );
//          /** Redistribute from CIDS to CBLK */
//          Q_cblk = Q_cids;
//          /** Merge Q_cblk into NN (sort and remove duplication) */
//          assert( Q_cblk.col_owned() == NN.col_owned() );
//          MergeNeighbors( k, NN.col_owned(), NN, Q_cblk );
//        }
//
//        //double mer_time = omp_get_wtime() - beg;
//
//        //if ( rank == 0 )
//        //printf( "%lfs %lfs %lfs\n", mpi_time, seq_time, mer_time ); fflush( stdout );
//      }

      for ( auto &neig : NN )
      {
        if ( neig.second < 0 || neig.second >= NN.col() )
        {
          printf( "Illegle neighbor gid %lu\n", neig.second );
          break;
        }
      }

      mpi::PrintProgress( "[END] NeighborSearch ...", this->GetComm() );
      return NN;
    }; 
    void TreePartition() 
    {
      mpi::PrintProgress( "[BEG] TreePartitioning ...", this->GetComm() );

      this->n = this->setup.ProblemSize();
      this->m = this->setup.LeafNodeSize();

      //for ( size_t i = rank; i < this->n; i += size ) 
      for ( size_t i = this->GetCommRank(); i < this->n; i += this->GetCommSize() ) 
        this->global_indices.push_back( i );
      num_points_owned = this->global_indices.size();
      AllocateNodes( this->global_indices );



      DependencyCleanUp();



      
      DistSplitTask<MPINODE> mpiSPLITtask;
      tree::SplitTask<NODE> seqSPLITtask;
      DistTraverseDown( mpiSPLITtask );
      LocaTraverseDown( seqSPLITtask );
      ExecuteAllTasks();



      tree::IndexPermuteTask<NODE> seqINDXtask;
      LocaTraverseUp( seqINDXtask );
      DistIndexPermuteTask<MPINODE> mpiINDXtask;
      DistTraverseUp( mpiINDXtask );
      ExecuteAllTasks();

      //printf( "rank %d finish split\n", rank ); fflush( stdout );


      (this->setup).morton.resize( this->n );

      RecursiveMorton( mpitreelists[ 0 ], MortonHelper::Root() );

      this->morton2node.clear();

      for ( auto node : this->treelist ) this->morton2node[ node->morton ] = node;

      for ( auto node : this->mpitreelists ) 
      {
        this->morton2node[ node->morton ] = node;
        auto *sibling = node->sibling;
        if ( node->l ) this->morton2node[ sibling->morton ] = sibling;
      }

      this->Barrier();
      mpi::PrintProgress( "[END] TreePartitioning ...", this->GetComm() ); 
    }; 
    void RecursiveMorton( MPINODE *node, MortonHelper::Recursor r ) 
    {
      int comm_size = this->GetCommSize();
      int comm_rank = this->GetCommRank();
      int node_size = node->GetCommSize(); 
      int node_rank = node->GetCommRank(); 

      node->morton = MortonHelper::MortonID( r );
      if ( node->sibling )
        node->sibling->morton = MortonHelper::SiblingMortonID( r );

      if ( node_size < 2 )
      {
        tree::Tree<SETUP, NODEDATA>::RecursiveMorton( node, r );
        auto &gids = this->treelist[ 0 ]->gids;
        vector<int> recv_size( comm_size, 0 );
        vector<int> recv_disp( comm_size, 0 );
        vector<pair<size_t, size_t>> send_pairs;
        vector<pair<size_t, size_t>> recv_pairs( this->n );

        for ( auto it : gids ) 
        {
          send_pairs.push_back( 
              pair<size_t, size_t>( it, this->setup.morton[ it ]) );
        }

        int send_size = send_pairs.size();
        mpi::Allgather( &send_size, 1, recv_size.data(), 1, this->GetComm() );
        for ( size_t p = 1; p < comm_size; p ++ )
        {
          recv_disp[ p ] = recv_disp[ p - 1 ] + recv_size[ p - 1 ];
        }
        size_t total_gids = 0;
        for ( size_t p = 0; p < comm_size; p ++ ) 
        {
          total_gids += recv_size[ p ];
        }
        assert( total_gids == this->n );
        mpi::Allgatherv( send_pairs.data(), send_size, 
            recv_pairs.data(), recv_size.data(), recv_disp.data(), this->GetComm() );
        for ( auto it : recv_pairs ) this->setup.morton[ it.first ] = it.second;
      }
      else
      {
        if ( node_rank < node_size / 2 )
        {
          RecursiveMorton( node->child, MortonHelper::RecurLeft( r ) );
        }
        else                   
        {
          RecursiveMorton( node->child, MortonHelper::RecurRight( r ) );
        }
      }
    }; 
    Data<int> CheckAllInteractions()
    {
      int total_depth = this->treelist.back()->l;
      int num_leafs = 1 << total_depth;
      Data<int> A( num_leafs, num_leafs, 0 );
      Data<int> B( num_leafs, num_leafs, 0 );
      for ( int t = 1; t < this->treelist.size(); t ++ )
      {
        auto *node = this->treelist[ t ];
        //for ( auto it : node->NNNearNodeMortonIDs )
        //{
        //  auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
        //  auto J = MortonHelper::Morton2Offsets(   it, total_depth );
        //  for ( auto i : I ) for ( auto j : J ) A( i, j ) += 1; 
        //}
        //for ( auto it : node->NNFarNodeMortonIDs )
        //{
        //  auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
        //  auto J = MortonHelper::Morton2Offsets(   it, total_depth );
        //  for ( auto i : I ) for ( auto j : J ) A( i, j ) += 1; 
        //}

        for ( int p = 0; p < this->GetCommSize(); p ++ )
        {
          if ( node->isleaf )
          {
            for ( auto & it : node->DistNear[ p ] )
            {
              auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
              auto J = MortonHelper::Morton2Offsets(   it.first, total_depth );
              for ( auto i : I ) for ( auto j : J ) 
              {
                assert( i < num_leafs && j < num_leafs );
                A( i, j ) += 1; 
              }
            }
          }
          for ( auto & it : node->DistFar[ p ] )
          {
            auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
            auto J = MortonHelper::Morton2Offsets(   it.first, total_depth );
            for ( auto i : I ) for ( auto j : J ) 
            {
              assert( i < num_leafs && j < num_leafs );
              A( i, j ) += 1; 
            }
          }
        }
      }

      for ( auto *node : mpitreelists )
      {
        //for ( auto it : node->NNNearNodeMortonIDs )
        //{
        //  auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
        //  auto J = MortonHelper::Morton2Offsets(   it, total_depth );
        //  for ( auto i : I ) for ( auto j : J ) A( i, j ) += 1; 
        //}
        //for ( auto it : node->NNFarNodeMortonIDs )
        //{
        //  auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
        //  auto J = MortonHelper::Morton2Offsets(   it, total_depth );
        //  for ( auto i : I ) for ( auto j : J ) A( i, j ) += 1; 
        //}
        for ( int p = 0; p < this->GetCommSize(); p ++ )
        {
          if ( node->isleaf )
          {
          for ( auto & it : node->DistNear[ p ] )
          {
            auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
            auto J = MortonHelper::Morton2Offsets(   it.first, total_depth );
            for ( auto i : I ) for ( auto j : J ) 
            {
              assert( i < num_leafs && j < num_leafs );
              A( i, j ) += 1; 
            }
          }
          }
          for ( auto & it : node->DistFar[ p ] )
          {
            auto I = MortonHelper::Morton2Offsets( node->morton, total_depth );
            auto J = MortonHelper::Morton2Offsets(   it.first, total_depth );
            for ( auto i : I ) for ( auto j : J ) 
            {
              assert( i < num_leafs && j < num_leafs );
              A( i, j ) += 1; 
            }
          }
        }
      }

      mpi::Reduce( A.data(), B.data(), A.size(), MPI_SUM, 0, this->GetComm() );

      if ( this->GetCommRank() == 0 )
      {
        for ( size_t i = 0; i < num_leafs; i ++ )
        {
          for ( size_t j = 0; j < num_leafs; j ++ ) printf( "%d", B( i, j ) );
          printf( "\n" );
        }
      }

      return B;
    }; 
    int Morton2Rank( size_t it )
    {
      return MortonHelper::Morton2Rank( it, this->GetCommSize() );
    }; 
    int Index2Rank( size_t gid )
    {
       return Morton2Rank( this->setup.morton[ gid ] );
    }; 
    template<typename TASK, typename... Args>
    void LocaTraverseUp( TASK &dummy, Args&... args )
    {
      assert( this->treelist.size() );

         //printf( "depth %lu\n", this->depth ); fflush( stdout );

      for ( int l = this->depth; l >= 1; l -- )
      {
        size_t n_nodes = 1 << l;
        auto level_beg = this->treelist.begin() + n_nodes - 1;

        for ( size_t node_ind = 0; node_ind < n_nodes; node_ind ++ )
        {
          auto *node = *(level_beg + node_ind);
          RecuTaskSubmit( node, dummy, args... );
        }
      }
    }; 
    template<typename TASK, typename... Args>
    void DistTraverseUp( TASK &dummy, Args&... args )
    {
      MPINODE *node = mpitreelists.back();
      while ( node )
      {
        if ( this->DoOutOfOrder() ) RecuTaskSubmit(  node, dummy, args... );
        else                        RecuTaskExecute( node, dummy, args... );
        node = (MPINODE*)node->parent;
      }
    }; 
    template<typename TASK, typename... Args>
    void LocaTraverseDown( TASK &dummy, Args&... args )
    {
      assert( this->treelist.size() );

      for ( int l = 1; l <= this->depth; l ++ )
      {
        size_t n_nodes = 1 << l;
        auto level_beg = this->treelist.begin() + n_nodes - 1;

        for ( size_t node_ind = 0; node_ind < n_nodes; node_ind ++ )
        {
          auto *node = *(level_beg + node_ind);
          RecuTaskSubmit( node, dummy, args... );
        }
      }
    }; 
    template<typename TASK, typename... Args>
    void DistTraverseDown( TASK &dummy, Args&... args )
    {
      auto *node = mpitreelists.front();
      while ( node )
      {
            //printf( "now at level %lu\n", node->l ); fflush( stdout );
        if ( this->DoOutOfOrder() ) RecuTaskSubmit(  node, dummy, args... );
        else                        RecuTaskExecute( node, dummy, args... );
            //printf( "RecuTaskSubmit at level %lu\n", node->l ); fflush( stdout );

        node = node->child;
      }
    }; 
    template<typename TASK, typename... Args>
    void LocaTraverseLeafs( TASK &dummy, Args&... args )
    {
      assert( this->treelist.size() );

      int n_nodes = 1 << this->depth;
      auto level_beg = this->treelist.begin() + n_nodes - 1;

      for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
      {
        auto *node = *(level_beg + node_ind);
        RecuTaskSubmit( node, dummy, args... );
      }
    }; 
    template<typename TASK, typename... Args>
    void LocaTraverseUnOrdered( TASK &dummy, Args&... args )
    {
      LocaTraverseDown( dummy, args... );
    }; 
    template<typename TASK, typename... Args>
    void DistTraverseUnOrdered( TASK &dummy, Args&... args )
    {
      DistTraverseDown( dummy, args... );
    }; 
    void DependencyCleanUp()
    {
      for ( auto node : mpitreelists ) node->DependencyCleanUp();
      //for ( size_t i = 0; i < mpitreelists.size(); i ++ )
      //{
      //  mpitreelists[ i ]->DependencyCleanUp();
      //}

      tree::Tree<SETUP, NODEDATA>::DependencyCleanUp();






      for ( auto p : NearRecvFrom ) p.DependencyCleanUp();
      for ( auto p :  FarRecvFrom ) p.DependencyCleanUp();

    }; 
    void ExecuteAllTasks()
    {
      hmlp_run();
      this->Barrier();
      DependencyCleanUp();
    }; 
    void DependOnNearInteractions( int p, Task *task )
    {
      for ( auto it : NearSentToRank[ p ] )
      {
        auto *node = this->morton2node[ it ];
        node->DependencyAnalysis( R, task );
      }
      task->TryEnqueue();
    }; 
    void DependOnFarInteractions( int p, Task *task )
    {
      for ( auto it : FarSentToRank[ p ] )
      {
        auto *node = this->morton2node[ it ];
        node->DependencyAnalysis( R, task );
      }
      task->TryEnqueue();
    }; 
    vector<vector<size_t>>   NearSentToRank;
    vector<map<size_t, int>> NearRecvFromRank;
    vector<vector<size_t>>   FarSentToRank;
    vector<map<size_t, int>> FarRecvFromRank;

    vector<ReadWrite> NearRecvFrom;
    vector<ReadWrite> FarRecvFrom;
    
  private:

    int ierr = 0;
    size_t num_points_owned = 0;

}; 
}; 
}; 
#endif