gofmm.hpp

#ifndef GOFMM_HPP
#define GOFMM_HPP

#include <future>
#include <thread>
#include <chrono>
#include <set>
#include <vector>
#include <map>
#include <unordered_set>
#include <deque>
#include <assert.h>
#include <typeinfo>
#include <algorithm>
#include <functional>
#include <array>
#include <random>
#include <numeric>
#include <sstream>
#include <iostream>
#include <string>
#include <stdio.h>
#include <omp.h>
#include <time.h>

#include <hmlp.h>
#include <hmlp_base.hpp>
#include <primitives/lowrank.hpp>
#include <primitives/combinatorics.hpp>
#include <primitives/gemm.hpp>
#include <containers/VirtualMatrix.hpp>
#include <containers/SPDMatrix.hpp>
#include <tree.hpp>
#include <igofmm.hpp>
#ifdef HMLP_USE_CUDA
#include <cuda_runtime.h>
#include <gofmm_gpu.hpp>
#endif

using namespace std;
using namespace hmlp;



#define MAX_NRHS 1024

#define GEMM_NB 256


//#define DEBUG_SPDASKIT 1
#define REPORT_ANN_ACCURACY 1
#define REPORT_COMPRESS_STATUS 1
#define REPORT_EVALUATE_STATUS 1


namespace hmlp
{
namespace gofmm
{

class CommandLineHelper
{
  public:

    CommandLineHelper( int argc, char *argv[] )
    {
      sscanf( argv[ 1 ], "%lu", &n );
      sscanf( argv[ 2 ], "%lu", &m );
      sscanf( argv[ 3 ], "%lu", &k );
      sscanf( argv[ 4 ], "%lu", &s );
      sscanf( argv[ 5 ], "%lu", &nrhs );
      sscanf( argv[ 6 ], "%lf", &stol );
      sscanf( argv[ 7 ], "%lf", &budget );
      distance_type = argv[ 8 ];
      if ( !distance_type.compare( "geometry" ) )
      {
        metric = GEOMETRY_DISTANCE;
      }
      else if ( !distance_type.compare( "kernel" ) )
      {
        metric = KERNEL_DISTANCE;
      }
      else if ( !distance_type.compare( "angle" ) )
      {
        metric = ANGLE_DISTANCE;
      }
      else
      {
        printf( "%s is not supported\n", argv[ 8 ] );
        exit( 1 );
      }
      spdmatrix_type = argv[ 9 ];
      if ( !spdmatrix_type.compare( "testsuit" ) )
      {
      }
      else if ( !spdmatrix_type.compare( "userdefine" ) )
      {
      }
      else if ( !spdmatrix_type.compare( "pvfmm" ) )
      {
      }
      else if ( !spdmatrix_type.compare( "dense" ) || !spdmatrix_type.compare( "ooc" ) )
      {
        user_matrix_filename = argv[ 10 ];
        if ( argc > 11 ) 
        {
          user_points_filename = argv[ 11 ];
          sscanf( argv[ 12 ], "%lu", &d );
        }
      }
      else if ( !spdmatrix_type.compare( "mlp" ) )
      {
        hidden_layers = argv[ 10 ];
        user_points_filename = argv[ 11 ];
        sscanf( argv[ 12 ], "%lu", &d );
      }
      else if ( !spdmatrix_type.compare( "cov" ) )
      {
        kernelmatrix_type = argv[ 10 ];
        user_points_filename = argv[ 11 ];
        sscanf( argv[ 12 ], "%lu", &d );
        sscanf( argv[ 13 ], "%lu", &nb );
      }
      else if ( !spdmatrix_type.compare( "kernel" ) )
      {
        kernelmatrix_type = argv[ 10 ];
        user_points_filename = argv[ 11 ];
        sscanf( argv[ 12 ], "%lu", &d );
        if ( argc > 13 ) sscanf( argv[ 13 ], "%lf", &h );
      }
      else
      {
        printf( "%s is not supported\n", argv[ 9 ] );
        exit( 1 );
      }
    }; 
    size_t n, m, k, s, nrhs;
    double stol = 1E-3;
    double budget = 0.0;
    DistanceMetric metric = ANGLE_DISTANCE;

    size_t d, nb; 
    double h = 1.0;
    string distance_type;
    string spdmatrix_type;
    string kernelmatrix_type;
    string hidden_layers;
    string user_matrix_filename;
    string user_points_filename;
}; 
template<typename T>
class Configuration
{
   public:

    Configuration() {};

      Configuration( DistanceMetric metric_type,
        size_t problem_size, size_t leaf_node_size, 
      size_t neighbor_size, size_t maximum_rank, 
         T tolerance, T budget ) 
      {
      Set( metric_type, problem_size, leaf_node_size, 
          neighbor_size, maximum_rank, tolerance, budget );
      };

      void Set( DistanceMetric metric_type,
        size_t problem_size, size_t leaf_node_size, 
      size_t neighbor_size, size_t maximum_rank, 
         T tolerance, T budget ) 
      {
         this->metric_type = metric_type;
         this->problem_size = problem_size;
         this->leaf_node_size = leaf_node_size;
         this->neighbor_size = neighbor_size;
         this->maximum_rank = maximum_rank;
         this->tolerance = tolerance;
         this->budget = budget;
      };

    void CopyFrom( Configuration<T> &config ) { *this = config; };

      DistanceMetric MetricType() { return metric_type; };

      size_t ProblemSize() { return problem_size; };

      size_t LeafNodeSize() { return leaf_node_size; };

      size_t NeighborSize() { return neighbor_size; };

      size_t MaximumRank() { return maximum_rank; };

      T Tolerance() { return tolerance; };

      T Budget() { return budget; };

    bool IsSymmetric() { return is_symmetric; };

    bool UseAdaptiveRanks() { return use_adaptive_ranks; };

    bool SecureAccuracy() { return secure_accuracy; };

   private:

      DistanceMetric metric_type = ANGLE_DISTANCE;

      size_t problem_size = 0;

      size_t leaf_node_size = 64;

      size_t neighbor_size = 32;

      size_t maximum_rank = 64;

      T tolerance = 1E-3;

      T budget = 0.03;

    bool is_symmetric = true;

      bool use_adaptive_ranks = true;

    bool secure_accuracy = false;

}; 
template<typename SPDMATRIX, typename SPLITTER, typename T>
class Setup : public tree::Setup<SPLITTER, T>,
              public Configuration<T>
{
  public:

    Setup() {};

    void FromConfiguration( Configuration<T> &config,
        SPDMATRIX &K, SPLITTER &splitter, Data<pair<T, size_t>> *NN )
    { 
      this->CopyFrom( config ); 
      this->K = &K;
      this->splitter = splitter;
      this->NN = NN;
    };

    SPDMATRIX *K = NULL;

    Data<T> *w = NULL;
    Data<T> *u = NULL;

    Data<T> *input = NULL;
    Data<T> *output = NULL;

    T lambda = 0.0;

    bool do_ulv_factorization = true;

  private:




}; 
template<typename T>
class NodeData : public Factor<T>
{
  public:

    NodeData() {};

    Lock lock;

    bool isskel = false;

    vector<size_t> skels;

    vector<int> jpvt;

    Data<T> proj;

    map<size_t, T> snids; 

    vector<size_t> candidate_rows;
    vector<size_t> candidate_cols;

    Data<T> KIJ; 

    Data<T> w_skel;
    Data<T> u_skel;

    Data<T> w_leaf;
    Data<T> u_leaf[ 20 ];

    View<T> w_view;
    View<T> u_view;

    Data<size_t> Nearbmap;
    Data<T> NearKab;
    Data<T> FarKab;


    Event skeletonize;
    Event updateweight;
    Event skeltoskel;
    Event skeltonode;
    Event s2s;
    Event s2n;

    double knn_acc = 0.0;
    size_t num_acc = 0;

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

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "TreeView" );
      label = to_string( arg->treelist_id );
      cost = 1.0;
    };

    void DependencyAnalysis() { arg->DependOnParent( this ); };

    void Execute( Worker* user_worker )
    {
      //printf( "TreeView %lu\n", node->treelist_id );
      auto *node   = arg;
      auto &data   = node->data;
      auto *setup  = node->setup;

      auto &w = *(setup->u);
      auto &u = *(setup->w);

      auto &U  = data.u_view;
      auto &W  = data.w_view;

      if ( !node->parent ) 
      {
        U.Set( u );
        W.Set( w );
      }

      if ( !node->isleaf )
      {
        auto &UL = node->lchild->data.u_view;
        auto &UR = node->rchild->data.u_view;
        auto &WL = node->lchild->data.w_view;
        auto &WR = node->rchild->data.w_view;
        U.Partition2x1( UL, 
                        UR, node->lchild->n, TOP );
        W.Partition2x1( WL, 
                        WR, node->lchild->n, TOP );
      }
    };
}; 
template<typename NODE>
class Summary
{

  public:

    Summary() {};

    deque<Statistic> rank;

    deque<Statistic> skeletonize;

    deque<Statistic> updateweight;

    deque<Statistic> s2s_kij_t;
    deque<Statistic> s2s_t;
    deque<Statistic> s2s_gfp;

    deque<Statistic> s2n_kij_t;
    deque<Statistic> s2n_t;
    deque<Statistic> s2n_gfp;


    void operator() ( NODE *node )
    {
      if ( rank.size() <= node->l )
      {
        rank.push_back( hmlp::Statistic() );
        skeletonize.push_back( hmlp::Statistic() );
        updateweight.push_back( hmlp::Statistic() );
      }

      rank[ node->l ].Update( (double)node->data.skels.size() );
      skeletonize[ node->l ].Update( node->data.skeletonize.GetDuration() );
      updateweight[ node->l ].Update( node->data.updateweight.GetDuration() );

#ifdef DUMP_ANALYSIS_DATA
      if ( node->parent )
      {
        auto *parent = node->parent;
        printf( "@TREE\n" );
        printf( "#%lu (s%lu), #%lu (s%lu), %lu, %lu\n", 
            node->treelist_id, node->data.skels.size(), 
            parent->treelist_id, parent->data.skels.size(),
            node->data.skels.size(), node->l );
      }
      else
      {
        printf( "@TREE\n" );
        printf( "#%lu (s%lu), , %lu, %lu\n", 
            node->treelist_id, node->data.skels.size(), 
            node->data.skels.size(), node->l );
      }
#endif
    };

    void Print()
    {
      for ( size_t l = 1; l < rank.size(); l ++ )
      {
        printf( "@SUMMARY\n" );
        printf( "level %2lu, ", l ); rank[ l ].Print();
        //printf( "skel_t:     " ); skeletonize[ l ].Print();
        //printf( "... ... ...\n" );
        //printf( "n2s_t:      " ); updateweight[ l ].Print();
        //printf( "s2s_t:      " ); s2s_t[ l ].Print();
        //printf( "s2s_gfp:    " ); s2s_gfp[ l ].Print();
        //printf( "s2n_t:      " ); s2n_t[ l ].Print();
        //printf( "s2n_gfp:    " ); s2n_gfp[ l ].Print();
      }
    };

}; 
template<typename SPDMATRIX, int N_SPLIT, typename T> 
struct centersplit
{
  SPDMATRIX *Kptr = NULL;
  DistanceMetric metric = ANGLE_DISTANCE;
  size_t n_centroid_samples = 5;
 
  centersplit() {};

  centersplit( SPDMATRIX& K ) { this->Kptr = &K; };

  vector<vector<size_t>> operator() ( vector<size_t>& gids ) const 
  {
    assert( N_SPLIT == 2 );
    assert( Kptr );


    SPDMATRIX &K = *Kptr;
    vector<vector<size_t>> split( N_SPLIT );
    size_t n = gids.size();
    vector<T> temp( n, 0.0 );

    auto column_samples = combinatorics::SampleWithoutReplacement( 
        n_centroid_samples, gids );


    auto DIC = K.Distances( this->metric, gids, column_samples );

    for ( auto & it : temp ) it = 0;

    for ( size_t j = 0; j < DIC.col(); j ++ )
      for ( size_t i = 0; i < DIC.row(); i ++ ) 
        temp[ i ] += DIC( i, j );

    auto idf2c = distance( temp.begin(), max_element( temp.begin(), temp.end() ) );

    vector<size_t> P( 1, gids[ idf2c ] );

    auto DIP = K.Distances( this->metric, gids, P );

    auto idf2f = distance( DIP.begin(), max_element( DIP.begin(), DIP.end() ) );

    vector<size_t> Q( 1, gids[ idf2f ] );

    auto DIQ = K.Distances( this->metric, gids, P );

    for ( size_t i = 0; i < temp.size(); i ++ )
      temp[ i ] = DIP[ i ] - DIQ[ i ];

    return combinatorics::MedianSplit( temp );
  };
}; 
template<typename SPDMATRIX, int N_SPLIT, typename T> 
struct randomsplit
{
  SPDMATRIX *Kptr = NULL;

  DistanceMetric metric = ANGLE_DISTANCE;

  randomsplit() {};

  randomsplit( SPDMATRIX& K ) { this->Kptr = &K; };

  inline vector<vector<size_t> > operator() ( vector<size_t>& gids ) const 
  {
    assert( Kptr && ( N_SPLIT == 2 ) );

    SPDMATRIX &K = *Kptr;
    size_t n = gids.size();
    vector<vector<size_t> > split( N_SPLIT );
    vector<T> temp( n, 0.0 );

    size_t idf2c = std::rand() % n;
    size_t idf2f = std::rand() % n;
    while ( idf2c == idf2f ) idf2f = std::rand() % n;


    vector<size_t> P( 1, gids[ idf2c ] );
    vector<size_t> Q( 1, gids[ idf2f ] );


    auto DIP = K.Distances( this->metric, gids, P );
    auto DIQ = K.Distances( this->metric, gids, Q );

    for ( size_t i = 0; i < temp.size(); i ++ )
      temp[ i ] = DIP[ i ] - DIQ[ i ];

    return combinatorics::MedianSplit( temp );

  };
}; 
template<typename NODE>
void FindNeighbors( NODE *node, DistanceMetric metric )
{
  using T = typename NODE::T;
  auto & setup = *(node->setup);
  auto & K = *(setup.K);
  auto & NN = *(setup.NN);
  auto & I = node->gids;
  size_t kappa = NN.row();
  pair<T, size_t> init( numeric_limits<T>::max(), NN.col() );
  auto candidates = K.NeighborSearch( metric, kappa, I, I, init );
   #pragma omp parallel
  {
    vector<pair<T, size_t> > aux( 2 * kappa );
    #pragma omp for
    for ( size_t j = 0; j < I.size(); j ++ )
    {
      MergeNeighbors( kappa, NN.columndata( I[ j ] ), 
          candidates.columndata( j ), aux );
    }
  }
}; 
template<class NODE, typename T>
class NeighborsTask : public Task
{
  public:

    NODE *arg = NULL;
   
     DistanceMetric metric = ANGLE_DISTANCE;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "Neighbors" );
      label = to_string( arg->treelist_id );
      metric = arg->setup->MetricType();

      //--------------------------------------
      double flops, mops;
      auto &gids = arg->gids;
      auto &NN = *arg->setup->NN;
      flops = gids.size();
      flops *= ( 4.0 * gids.size() );
      // Heap select worst case
      mops = (size_t)std::log( NN.row() ) * gids.size();
      mops *= gids.size();
      // Access K
      mops += flops;
      event.Set( name + label, flops, mops );
      //--------------------------------------
         
      // TODO: Need an accurate cost model.
      cost = mops / 1E+9;
    };

    void DependencyAnalysis() { arg->DependOnNoOne( this ); };

    void Execute( Worker* user_worker ) { FindNeighbors( arg, metric ); };

}; 
template<bool DOAPPROXIMATE, bool SORTED, typename T, typename CSCMATRIX>
Data<pair<T, size_t>> SparsePattern( size_t n, size_t k, CSCMATRIX &K )
{
  pair<T, size_t> initNN( numeric_limits<T>::max(), n );
  Data<pair<T, size_t>> NN( k, n, initNN );

  printf( "SparsePattern k %lu n %lu, NN.row %lu NN.col %lu ...", 
      k, n, NN.row(), NN.col() ); fflush( stdout );

  #pragma omp parallel for schedule( dynamic )
  for ( size_t j = 0; j < n; j ++ )
  {
    std::set<size_t> NNset;
    size_t nnz = K.ColPtr( j + 1 ) - K.ColPtr( j );
    if ( DOAPPROXIMATE && nnz > 2 * k ) nnz = 2 * k;

    //printf( "j %lu nnz %lu\n", j, nnz );

    for ( size_t i = 0; i < nnz; i ++ )
    {
      // TODO: this is lid. Need to be gid.
      auto row_ind = K.RowInd( K.ColPtr( j ) + i );
      auto val     = K.Value( K.ColPtr( j ) + i );

      if ( val ) val = 1.0 / std::abs( val );
      else       val = std::numeric_limits<T>::max() - 1.0;

      NNset.insert( row_ind );
      std::pair<T, std::size_t> query( val, row_ind );
      if ( nnz < k ) // not enough candidates
      {
        NN[ j * k + i  ] = query;
      }
      else
      {
        hmlp::HeapSelect( 1, NN.row(), &query, NN.data() + j * NN.row() );
      }
    }

    while ( nnz < k )
    {
      std::size_t row_ind = rand() % n;
      if ( !NNset.count( row_ind ) )
      {
        T val = std::numeric_limits<T>::max() - 1.0;
        std::pair<T, std::size_t> query( val, row_ind );
        NNset.insert( row_ind );
        NN[ j * k + nnz ] = query;
        nnz ++;
      }
    }
  }
  printf( "Done.\n" ); fflush( stdout );

  if ( SORTED )
  {
    printf( "Sorting ... " ); fflush( stdout );
    struct 
    {
      bool operator () ( std::pair<T, size_t> a, std::pair<T, size_t> b )
      {   
        return a.first < b.first;
      }   
    } ANNLess;

    //printf( "SparsePattern k %lu n %lu, NN.row %lu NN.col %lu\n", k, n, NN.row(), NN.col() );

    #pragma omp parallel for
    for ( size_t j = 0; j < NN.col(); j ++ )
    {
      std::sort( NN.data() + j * NN.row(), NN.data() + ( j + 1 ) * NN.row(), ANNLess );
    }
    printf( "Done.\n" ); fflush( stdout );
  }

  return NN;
}; 
/* @brief Helper functions for sorting sampling neighbors. */ 
template<typename TA, typename TB>
pair<TB, TA> flip_pair( const pair<TA, TB> &p )
{
  return pair<TB, TA>( p.second, p.first );
}; 
template<typename TA, typename TB>
multimap<TB, TA> flip_map( const map<TA, TB> &src )
{
  multimap<TB, TA> dst;
  transform( src.begin(), src.end(), inserter( dst, dst.begin() ), 
                 flip_pair<TA, TB> );
  return dst;
}; 
template<typename NODE>
void Interpolate( NODE *node )
{
  using T = typename NODE::T;
  if ( !node ) return;

  auto &K = *node->setup->K;
  auto &data = node->data;
  auto &skels = data.skels;
  auto &proj = data.proj;
  auto &jpvt = data.jpvt;
  auto s = proj.row();
  auto n = proj.col();

  if ( !data.isskel || proj[ 0 ] == 0 ) return;

  assert( s );
  assert( s <= n );
  assert( jpvt.size() == n );

  if ( data.isskel )
  {
    data.w_skel.reserve( skels.size(), MAX_NRHS );
    data.u_skel.reserve( skels.size(), MAX_NRHS );
  }


  Data<T> R1( s, s, 0.0 );

  for ( int j = 0; j < s; j ++ )
  {
    for ( int i = 0; i < s; i ++ )
    {
      if ( i <= j ) R1[ j * s + i ] = proj[ j * s + i ];
    }
  }

  Data<T> tmp = proj;

  xtrsm( "L", "U", "N", "N", s, n, 1.0, R1.data(), s, tmp.data(), s );

  for ( int j = 0; j < n; j ++ )
  {
    for ( int i = 0; i < s; i ++ )
    {
     proj[ jpvt[ j ] * s + i ] = tmp[ j * s + i ];
    }
  }
 

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

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "it" );
      label = to_string( arg->treelist_id );
      // Need an accurate cost model.
      cost = 1.0;
    };

    void DependencyAnalysis() { arg->DependOnNoOne( this ); };

    void Execute( Worker* user_worker ) { Interpolate( arg ); };

}; 
template<bool NNPRUNE, typename NODE>
void RowSamples( NODE *node, size_t nsamples )
{
  using T = typename NODE::T;
  auto &setup = *(node->setup);
  auto &data = node->data;
  auto &K = *(setup.K);

  auto &amap = data.candidate_rows;

  amap.clear();

  if ( setup.NN )
  {
    //printf( "construct snids NN.row() %lu NN.col() %lu\n", 
    //    node->setup->NN->row(), node->setup->NN->col() ); fflush( stdout );
    auto &NN = *(setup.NN);
    auto &gids = node->gids;
    auto &snids = data.snids;
    size_t knum = NN.row();

    if ( node->isleaf )
    {
      snids.clear();

      vector<pair<T, size_t>> tmp( knum * gids.size() );
      for ( size_t j = 0; j < gids.size(); j ++ )
        for ( size_t i = 0; i < knum; i ++ )
          tmp[ j * knum + i ] = NN( i, gids[ j ] );

      sort( tmp.begin(), tmp.end() );
    
      for ( auto it : tmp )
      {
        size_t it_gid = it.second;
        size_t it_morton = setup.morton[ it_gid ];

        if ( snids.size() >= nsamples ) break;

        bool is_near;
        if ( NNPRUNE ) is_near = node->NNNearNodeMortonIDs.count( it_morton );
        else           is_near = (it_morton == node->morton );

        if ( !is_near )
        {
          auto ret = snids.insert( make_pair( it.second, it.first ) );
        }
      }
    }
    else
    {
      auto &lsnids = node->lchild->data.snids;
      auto &rsnids = node->rchild->data.snids;

      snids = lsnids;

      for ( auto it = rsnids.begin(); it != rsnids.end(); it ++ )
      {
        auto ret = snids.insert( *it );
        if ( !ret.second )
        {
          if ( ret.first->second > (*it).first )
            ret.first->second = (*it).first;
        }
      }

      for ( auto gid : gids ) snids.erase( gid );
    }


    if ( nsamples < K.col() - node->n )
    {
      multimap<T, size_t> ordered_snids = flip_map( snids );
      amap.reserve( nsamples );

      for ( auto it : ordered_snids )
      {
        if ( amap.size() >= nsamples ) break;
        amap.push_back( it.second );
      }

      while ( amap.size() < nsamples )
      {
        //size_t sample = rand() % K.col();
        auto important_sample = K.ImportantSample( 0 );
        size_t sample_gid = important_sample.second;
        size_t sample_morton = setup.morton[ sample_gid ];

        if ( !MortonHelper::IsMyParent( sample_morton, node->morton ) )
        {
          amap.push_back( sample_gid );
        }
      }
    }
    else 
    {
      for ( size_t sample = 0; sample < K.col(); sample ++ )
      {
        size_t sample_morton = setup.morton[ sample ];
        if ( !MortonHelper::IsMyParent( sample_morton, node->morton ) )
        {
          amap.push_back( sample );
        }
      }
    }
  } 
}; 
template<bool NNPRUNE, typename NODE>
void SkeletonKIJ( NODE *node )
{
  using T = typename NODE::T;
  auto &K = *(node->setup->K);
  auto &data = node->data;
  auto &candidate_rows = data.candidate_rows;
  auto &candidate_cols = data.candidate_cols;
  auto &KIJ = data.KIJ;
  auto *lchild = node->lchild;
  auto *rchild = node->rchild;

  if ( node->isleaf )
  {
    candidate_cols = node->gids;
  }
  else
  {
    auto &lskels = lchild->data.skels;
    auto &rskels = rchild->data.skels;
    if ( !lskels.size() || !rskels.size() ) return;
    candidate_cols = lskels;
    candidate_cols.insert( candidate_cols.end(), 
        rskels.begin(), rskels.end() );
  }

  size_t nsamples = 2 * candidate_cols.size();

  if ( nsamples < 2 * node->setup->LeafNodeSize() ) 
    nsamples = 2 * node->setup->LeafNodeSize();

  RowSamples<NNPRUNE>( node, nsamples );
  
  KIJ = K( candidate_rows, candidate_cols );



}; 
template<bool NNPRUNE, typename NODE, typename T>
class SkeletonKIJTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "par-gskm" );
      label = to_string( arg->treelist_id );
      cost = 5.0;
      priority = true;
    };

    void DependencyAnalysis() { arg->DependOnChildren( this ); };

    void Execute( Worker* user_worker ) { SkeletonKIJ<NNPRUNE>( arg ); };

}; 
template<typename NODE>
void Skeletonize( NODE *node )
{
  using T = typename NODE::T;
  if ( !node->parent ) return;

  auto &K   = *(node->setup->K);
  auto &NN  = *(node->setup->NN);
  auto maxs = node->setup->MaximumRank();
  auto stol = node->setup->Tolerance();
  bool secure_accuracy = node->setup->SecureAccuracy();
  bool use_adaptive_ranks = node->setup->UseAdaptiveRanks();

  auto &data  = node->data;
  auto &skels = data.skels;
  auto &proj  = data.proj;
  auto &jpvt  = data.jpvt;
  auto &KIJ   = data.KIJ;
  auto &candidate_cols = data.candidate_cols;

  size_t N = K.col();
  size_t m = KIJ.row();
  size_t n = KIJ.col();
  size_t q = node->n;

  if ( secure_accuracy )
  {
    if ( !node->isleaf && ( !node->lchild->data.isskel || !node->rchild->data.isskel ) )
    {
      skels.clear();
      proj.resize( 0, 0 );
      data.isskel = false;
      return;
    }
  }

  T scaled_stol = std::sqrt( (T)n / q ) * std::sqrt( (T)m / (N - q) ) * stol;
  scaled_stol *= std::sqrt( (T)q / N );

  lowrank::id( use_adaptive_ranks, secure_accuracy,
    KIJ.row(), KIJ.col(), maxs, scaled_stol, KIJ, skels, proj, jpvt );

  KIJ.resize( 0, 0 );

  if ( secure_accuracy )
  {
    data.isskel = (skels.size() != 0);
  }
  else
  {
    assert( skels.size() && proj.size() && jpvt.size() );
    data.isskel = true;
  }
  
  for ( size_t i = 0; i < skels.size(); i ++ ) skels[ i ] = candidate_cols[ skels[ i ] ];

}; 
template<typename NODE, typename T>
class SkeletonizeTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "sk" );
      label = to_string( arg->treelist_id );
      cost = 5.0;
      priority = true;
    };

    void GetEventRecord()
    {
      double flops = 0.0, mops = 0.0;

      auto &K = *arg->setup->K;
      size_t n = arg->data.proj.col();
      size_t m = 2 * n;
      size_t k = arg->data.proj.row();

      flops += ( 2.0 / 3.0 ) * n * n * ( 3 * m - n );
      mops += ( 2.0 / 3.0 ) * n * n * ( 3 * m - n );

      /* TRSM */
      flops += k * ( k - 1 ) * ( n + 1 );
      mops  += 2.0 * ( k * k + k * n );

      //flops += ( 2.0 / 3.0 ) * k * k * ( 3 * m - k );
      //mops += 2.0 * m * k;
      //flops += 2.0 * m * n * k;
      //mops += 2.0 * ( m * k + k * n + m * n );
      //flops += ( 1.0 / 3.0 ) * k * k * n;
      //mops += 2.0 * ( k * k + k * n );

      event.Set( label + name, flops, mops );
      arg->data.skeletonize = event;
    };

    void DependencyAnalysis() { arg->DependOnNoOne( this ); };

    void Execute( Worker* user_worker ) { Skeletonize( arg ); };

}; 
template<typename NODE>
void UpdateWeights( NODE *node )
{
  using T = typename NODE::T;
  if ( !node->parent || !node->data.isskel ) return;

  auto &w = *node->setup->w;

  auto &data = node->data;
  auto &proj = data.proj;
  auto &skels = data.skels;
  auto &w_skel = data.w_skel;
  auto &w_leaf = data.w_leaf;
  auto *lchild = node->lchild;
  auto *rchild = node->rchild;


  size_t nrhs = w.col();

  w_skel.resize( skels.size(), nrhs );

  //printf( "%lu UpdateWeight w_skel.num() %lu\n", node->treelist_id, w_skel.num() );

  if ( node->isleaf )
  {
    if ( w_leaf.size() )
    {
      //printf( "%8lu w_leaf allocated [%lu %lu]\n", 
      //    node->morton, w_leaf.row(), w_leaf.col() ); fflush( stdout );
      
      xgemm
      (
        "N", "N",
        w_skel.row(), w_skel.col(), w_leaf.row(),
        1.0, proj.data(),   proj.row(),
             w_leaf.data(), w_leaf.row(),
        0.0, w_skel.data(), w_skel.row()
      );
    }
    else
    {
      View<T> W = data.w_view;
      //printf( "%8lu n2s W[%lu %lu ld %lu]\n", 
      //    node->morton, W.row(), W.col(), W.ld() ); fflush( stdout );
      //for ( int i = 0; i < 10; i ++ )
      //  printf( "%lu W.data() + %d = %E\n", node->gids[ i ], i, *(W.data() + i) );
      xgemm
      (
        "N", "N",
        w_skel.row(), w_skel.col(), W.row(),
        1.0, proj.data(),   proj.row(),
                W.data(),       W.ld(),
        0.0, w_skel.data(), w_skel.row()
      );
    }

    //double update_leaf_time = omp_get_wtime() - beg;
    //printf( "%lu, m %lu n %lu k %lu, total %.3E\n", 
    //  node->treelist_id, 
    //  w_skel.row(), w_skel.col(), w_leaf.col(), update_leaf_time );
  }
  else
  {
    //double beg = omp_get_wtime();
    auto &w_lskel = lchild->data.w_skel;
    auto &w_rskel = rchild->data.w_skel;
    auto &lskel = lchild->data.skels;
    auto &rskel = rchild->data.skels;

    //if ( 1 )
    if ( node->treelist_id > 6 )
    {
      //printf( "%8lu n2s\n", node->morton ); fflush( stdout );
      xgemm
      (
        "N", "N",
        w_skel.row(), w_skel.col(), lskel.size(),
        1.0,    proj.data(),    proj.row(),
             w_lskel.data(), w_lskel.row(),
        0.0,  w_skel.data(),  w_skel.row()
      );
      xgemm
      (
        "N", "N",
        w_skel.row(), w_skel.col(), rskel.size(),
        1.0,    proj.data() + proj.row() * lskel.size(), proj.row(),
             w_rskel.data(), w_rskel.row(),
        1.0,  w_skel.data(),  w_skel.row()
      );
    }
    else
    {
      View<T> P( false,   proj ), PL,
                                  PR;
      View<T> W( false, w_skel ), WL( false, w_lskel ),
                                  WR( false, w_rskel );
      P.Partition1x2( PL, PR, lskel.size(), LEFT );
      gemm::xgemm<GEMM_NB>( (T)1.0, PL, WL, (T)0.0, W );
      W.DependencyCleanUp();
      gemm::xgemm<GEMM_NB>( (T)1.0, PR, WR, (T)1.0, W );
      //W.DependencyCleanUp();
    }
  }
}; 
template<typename NODE, typename T>
class UpdateWeightsTask : public Task
{
  public:

    NODE *arg = NULL;

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

      double flops, mops;
      auto &gids = arg->gids;
      auto &skels = arg->data.skels;
      auto &w = *arg->setup->w;
      if ( arg->isleaf )
      {
        auto m = skels.size();
        auto n = w.col();
        auto k = gids.size();
        flops = 2.0 * m * n * k;
        mops = 2.0 * ( m * n + m * k + k * n );
      }
      else
      {
        auto &lskels = arg->lchild->data.skels;
        auto &rskels = arg->rchild->data.skels;
        auto m = skels.size();
        auto n = w.col();
        auto k = lskels.size() + rskels.size();
        flops = 2.0 * m * n * k;
        mops  = 2.0 * ( m * n + m * k + k * n );
      }

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

    void Prefetch( Worker* user_worker )
    {
      auto &proj = arg->data.proj;
      __builtin_prefetch( proj.data() );
      auto &w_skel = arg->data.w_skel;
      __builtin_prefetch( w_skel.data() );
      if ( arg->isleaf )
      {
        auto &w_leaf = arg->data.w_leaf;
        __builtin_prefetch( w_leaf.data() );
      }
      else
      {
        auto &w_lskel = arg->lchild->data.w_skel;
        __builtin_prefetch( w_lskel.data() );
        auto &w_rskel = arg->rchild->data.w_skel;
        __builtin_prefetch( w_rskel.data() );
      }
#ifdef HMLP_USE_CUDA
      hmlp::Device *device = NULL;
      if ( user_worker ) device = user_worker->GetDevice();
      if ( device ) 
      {
        proj.CacheD( device );
        proj.PrefetchH2D( device, 1 );
        if ( arg->isleaf )
        {
          auto &w_leaf = arg->data.w_leaf;
          w_leaf.CacheD( device );
          w_leaf.PrefetchH2D( device, 1 );
        }
        else
        {
          auto &w_lskel = arg->lchild->data.w_skel;
          w_lskel.CacheD( device );
          w_lskel.PrefetchH2D( device, 1 );
          auto &w_rskel = arg->rchild->data.w_skel;
          w_rskel.CacheD( device );
          w_rskel.PrefetchH2D( device, 1 );
        }
      }
#endif
    };

    void DependencyAnalysis() { arg->DependOnChildren( this ); };

    void Execute( Worker* user_worker )
    {
#ifdef HMLP_USE_CUDA 
      hmlp::Device *device = NULL;
      if ( user_worker ) device = user_worker->GetDevice();
      if ( device ) gpu::UpdateWeights( device, arg );
      else               UpdateWeights<NODE, T>( arg );
#else
      UpdateWeights( arg );
#endif
    };

}; 
template<typename NODE>
void SkeletonsToSkeletons( NODE *node )
{
  using T = typename NODE::T;
  if ( !node->parent || !node->data.isskel ) return;

  double beg, u_skel_time, s2s_time;

  auto *FarNodes = &node->NNFarNodes;

  auto &K = *node->setup->K;
  auto &data = node->data;
  auto &amap = node->data.skels;
  auto &u_skel = node->data.u_skel;
  auto &FarKab = node->data.FarKab;

  size_t nrhs = node->setup->w->col();

  beg = omp_get_wtime();
  u_skel.resize( 0, 0 );
  u_skel.resize( amap.size(), nrhs, 0.0 );
  u_skel_time = omp_get_wtime() - beg;

  size_t offset = 0;


  View<T> FarKab_v( FarKab );

  for ( auto it = FarNodes->begin(); it != FarNodes->end(); it ++ )
  {
    auto &bmap = (*it)->data.skels;
    auto &w_skel = (*it)->data.w_skel;
    assert( w_skel.col() == nrhs );
    assert( w_skel.row() == bmap.size() );
    assert( w_skel.size() == nrhs * bmap.size() );

    if ( FarKab.size() ) 
    {
      //if ( node->treelist_id > 6 )
      if ( 1 )
      {
        assert( FarKab.row() == amap.size() );
        assert( u_skel.row() * offset <= FarKab.size() );

        //printf( "%8lu s2s %8lu w_skel[%lu %lu]\n", 
        //    node->morton, (*it)->morton, w_skel.row(), w_skel.col() );
        //fflush( stdout );
        xgemm
        (
          "N", "N",
          u_skel.row(), u_skel.col(), w_skel.row(),
          1.0, FarKab.data() + u_skel.row() * offset, FarKab.row(),
               w_skel.data(),          w_skel.row(),
          1.0, u_skel.data(),          u_skel.row()
        );

      }
      else
      {
        View<T> U( false, u_skel );
        View<T> W( false, w_skel );
        View<T> Kab;
        assert( FarKab.col() >= W.row() + offset );
        Kab.Set( FarKab.row(), W.row(), 0, offset, &FarKab_v );
        gemm::xgemm<GEMM_NB>( (T)1.0, Kab, W, (T)1.0, U );
      }

      offset += w_skel.row();
    }
    else
    {
      printf( "Far Kab not cached treelist_id %lu, l %lu\n\n",
               node->treelist_id, node->l ); fflush( stdout );

      auto Kab = K( amap, bmap );
      xgemm( "N", "N", u_skel.row(), u_skel.col(), w_skel.row(),
        1.0, Kab.data(),       Kab.row(),
             w_skel.data(), w_skel.row(),
        1.0, u_skel.data(), u_skel.row() );
    }
  }

}; 
template<bool NNPRUNE, typename NODE, typename T>
class SkeletonsToSkeletonsTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "s2s" );
      {
        //label = std::to_string( arg->treelist_id );
        ostringstream ss;
        ss << arg->treelist_id;
        label = ss.str();
      }

      double flops = 0.0, mops = 0.0;
      auto &w = *arg->setup->w;
      size_t m = arg->data.skels.size();
      size_t n = w.col();

      std::set<NODE*> *FarNodes;
      if ( NNPRUNE ) FarNodes = &arg->NNFarNodes;
      else           FarNodes = &arg->FarNodes;

      for ( auto it = FarNodes->begin(); it != FarNodes->end(); it ++ )
      {
        size_t k = (*it)->data.skels.size();
        flops += 2.0 * m * n * k;
        mops  += m * k; // cost of Kab
        mops  += 2.0 * ( m * n + n * k + k * n );
      }

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

    void DependencyAnalysis()
    {
      for ( auto it : arg->NNFarNodes ) it->DependencyAnalysis( R, this );
      arg->DependencyAnalysis( RW, this );
      this->TryEnqueue();
    };

    void Execute( Worker* user_worker ) { SkeletonsToSkeletons( arg ); };
}; 
template<typename NODE>
void SkeletonsToNodes( NODE *node )
{
  using T = typename NODE::T;

  auto &K = *node->setup->K;
  auto &w = *node->setup->w;

  auto &gids = node->gids;
  auto &data = node->data;
  auto &proj = data.proj;
  auto &skels = data.skels;
  auto &u_skel = data.u_skel;
  auto *lchild = node->lchild;
  auto *rchild = node->rchild;

  size_t nrhs = w.col();





  if ( node->isleaf )
  {
    View<T> U = data.u_view;

    if ( U.col() == nrhs )
    {
      //printf( "%8lu s2n U[%lu %lu %lu]\n", 
      //    node->morton, U.row(), U.col(), U.ld() ); fflush( stdout );
      xgemm
      (
        "Transpose", "Non-transpose",
        U.row(), U.col(), u_skel.row(),
        1.0,   proj.data(),   proj.row(),
             u_skel.data(), u_skel.row(),
        1.0,      U.data(),       U.ld()
      );
    }
    else
    {
      //printf( "%8lu use u_leaf u_view [%lu %lu ld %lu]\n", 
      //    node->morton, U.row(), U.col(), U.ld()  ); fflush( stdout );

      auto &u_leaf = node->data.u_leaf[ 0 ];

      u_leaf.resize( 0, 0 );
      u_leaf.resize( gids.size(), nrhs, 0.0 );

      if ( data.isskel )
      {
        xgemm
        (
          "T", "N",
          u_leaf.row(), u_leaf.col(), u_skel.row(),
          1.0,   proj.data(),   proj.row(),
               u_skel.data(), u_skel.row(),
          1.0, u_leaf.data(), u_leaf.row()
        );
      }
    }
  }
  else
  {
    if ( !node->parent || !node->data.isskel ) return;

    auto &u_lskel = lchild->data.u_skel;
    auto &u_rskel = rchild->data.u_skel;
    auto &lskel = lchild->data.skels;
    auto &rskel = rchild->data.skels;

    //if ( 1 )
    if ( node->treelist_id > 6 )
    {
      //printf( "%8lu s2n\n", node->morton ); fflush( stdout );
      xgemm
      (
        "Transpose", "No transpose",
        u_lskel.row(), u_lskel.col(), proj.row(),
        1.0, proj.data(),    proj.row(),
             u_skel.data(),  u_skel.row(),
        1.0, u_lskel.data(), u_lskel.row()
      );
      xgemm
      (
        "Transpose", "No transpose",
        u_rskel.row(), u_rskel.col(), proj.row(),
        1.0, proj.data() + proj.row() * lskel.size(), proj.row(),
             u_skel.data(), u_skel.row(),
        1.0, u_rskel.data(), u_rskel.row()
      );
    }
    else
    {
      View<T> P(  true,   proj ), PL,
                                  PR;
      View<T> U( false, u_skel ), UL( false, u_lskel ),
                                  UR( false, u_rskel );
      P.Partition2x1( PL,
                      PR, lskel.size(), TOP );
      gemm::xgemm<GEMM_NB>( (T)1.0, PL, U, (T)1.0, UL );
      gemm::xgemm<GEMM_NB>( (T)1.0, PR, U, (T)1.0, UR );
    }
  }
  //printf( "\n" );

}; 
template<bool NNPRUNE, typename NODE, typename T>
class SkeletonsToNodesTask : public Task
{
  public:

    NODE *arg = NULL;

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

      //--------------------------------------
      double flops = 0.0, mops = 0.0;
      auto &gids = arg->gids;
      auto &data = arg->data;
      auto &proj = data.proj;
      auto &skels = data.skels;
      auto &w = *arg->setup->w;

      if ( arg->isleaf )
      {
        size_t m = proj.col();
        size_t n = w.col();
        size_t k = proj.row();
        flops += 2.0 * m * n * k;
        mops  += 2.0 * ( m * n + n * k + m * k );
      }
      else
      {
        if ( !arg->parent || !arg->data.isskel )
        {
          // No computation.
        }
        else
        {
          size_t m = proj.col();
          size_t n = w.col();
          size_t k = proj.row();
          flops += 2.0 * m * n * k;
          mops  += 2.0 * ( m * n + n * k + m * k );
        }
      }

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

    void Prefetch( Worker* user_worker )
    {
      auto &proj = arg->data.proj;
      __builtin_prefetch( proj.data() );
      auto &u_skel = arg->data.u_skel;
      __builtin_prefetch( u_skel.data() );
      if ( arg->isleaf )
      {
        //__builtin_prefetch( arg->data.u_leaf[ 0 ].data() );
        //__builtin_prefetch( arg->data.u_leaf[ 1 ].data() );
        //__builtin_prefetch( arg->data.u_leaf[ 2 ].data() );
        //__builtin_prefetch( arg->data.u_leaf[ 3 ].data() );
      }
      else
      {
        auto &u_lskel = arg->lchild->data.u_skel;
        __builtin_prefetch( u_lskel.data() );
        auto &u_rskel = arg->rchild->data.u_skel;
        __builtin_prefetch( u_rskel.data() );
      }
#ifdef HMLP_USE_CUDA
      hmlp::Device *device = NULL;
      if ( user_worker ) device = user_worker->GetDevice();
      if ( device )
      {
        int stream_id = arg->treelist_id % 8;
        proj.CacheD( device );
        proj.PrefetchH2D( device, stream_id );
        u_skel.CacheD( device );
        u_skel.PrefetchH2D( device, stream_id );
        if ( arg->isleaf )
        {
        }
        else
        {
          auto &u_lskel = arg->lchild->data.u_skel;
          u_lskel.CacheD( device );
          u_lskel.PrefetchH2D( device, stream_id );
          auto &u_rskel = arg->rchild->data.u_skel;
          u_rskel.CacheD( device );
          u_rskel.PrefetchH2D( device, stream_id );
        }
      }
#endif
    };

    void DependencyAnalysis() { arg->DependOnParent( this ); };

    void Execute( Worker* user_worker )
    {
#ifdef HMLP_USE_CUDA 
      Device *device = NULL;
      if ( user_worker ) device = user_worker->GetDevice();
      if ( device ) gpu::SkeletonsToNodes<NNPRUNE, NODE, T>( device, arg );
      else               SkeletonsToNodes<NNPRUNE, NODE, T>( arg );
#else
    SkeletonsToNodes( arg );
#endif
    };

}; 
template<int SUBTASKID, bool NNPRUNE, typename NODE, typename T>
void LeavesToLeaves( NODE *node, size_t itbeg, size_t itend )
{
  assert( node->isleaf );

  double beg, u_leaf_time, before_writeback_time, after_writeback_time;

  auto &K = *node->setup->K;
  auto &w = *node->setup->w;

  auto &gids = node->gids;
  auto &data = node->data;
  auto &amap = node->gids;
  auto &NearKab = data.NearKab;

  size_t nrhs = w.col();

  set<NODE*> *NearNodes;
  if ( NNPRUNE ) NearNodes = &node->NNNearNodes;
  else           NearNodes = &node->NearNodes;

  auto &u_leaf = data.u_leaf[ SUBTASKID ];
  u_leaf.resize( 0, 0 );

  if ( itbeg == itend ) 
  {
    return;
  }
  else
  {
    u_leaf.resize( gids.size(), nrhs, 0.0 );
  }

  if ( NearKab.size() ) 
  {
    size_t itptr = 0;
    size_t offset = 0;

    for ( auto it = NearNodes->begin(); it != NearNodes->end(); it ++ )
    {
      if ( itbeg <= itptr && itptr < itend )
      {
        //auto wb = w( bmap );
        auto wb = (*it)->data.w_leaf;

        if ( wb.size() )
        {
          xgemm
          (
            "N", "N",
            u_leaf.row(), u_leaf.col(), wb.row(),
            1.0, NearKab.data() + offset * NearKab.row(), NearKab.row(),
                      wb.data(),                               wb.row(),
            1.0,  u_leaf.data(),                           u_leaf.row()
          );
        }
        else
        {
          View<T> W = (*it)->data.w_view;
          xgemm
          (
            "N", "N",
            u_leaf.row(), u_leaf.col(), W.row(),
            1.0, NearKab.data() + offset * NearKab.row(), NearKab.row(),
                       W.data(),                                 W.ld(),
            1.0,  u_leaf.data(),                           u_leaf.row()
          );
        }
      }
      offset += (*it)->gids.size();
      itptr ++;
    }
  }
  else 
  {
    size_t itptr = 0;
    for ( auto it = NearNodes->begin(); it != NearNodes->end(); it ++ )
    {
      if ( itbeg <= itptr && itptr < itend )
      {
        auto &bmap = (*it)->gids;
        auto wb = (*it)->data.w_leaf;

        auto Kab = K( amap, bmap );

        if ( wb.size() )
        {
          xgemm( "N", "N", u_leaf.row(), u_leaf.col(), wb.row(),
            1.0,    Kab.data(),    Kab.row(),
                     wb.data(),     wb.row(),
            1.0, u_leaf.data(), u_leaf.row());
        }
        else
        {
          View<T> W = (*it)->data.w_view;
          xgemm( "N", "N", u_leaf.row(), u_leaf.col(), W.row(),
            1.0,    Kab.data(),    Kab.row(),
                      W.data(),       W.ld(),
            1.0, u_leaf.data(), u_leaf.row() );
        }
      }
      itptr ++;
    }
  }
  before_writeback_time = omp_get_wtime() - beg;

}; 
template<int SUBTASKID, bool NNPRUNE, typename NODE, typename T>
class LeavesToLeavesTask : public Task
{
  public:

    NODE *arg = NULL;

    size_t itbeg;

     size_t itend;

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

      //--------------------------------------
      double flops = 0.0, mops = 0.0;
      auto &gids = arg->gids;
      auto &data = arg->data;
      auto &proj = data.proj;
      auto &skels = data.skels;
      auto &w = *arg->setup->w;
      auto &K = *arg->setup->K;
      auto &NearKab = data.NearKab;

      assert( arg->isleaf );

      size_t m = gids.size();
      size_t n = w.col();

      set<NODE*> *NearNodes;
      if ( NNPRUNE ) NearNodes = &arg->NNNearNodes;
      else           NearNodes = &arg->NearNodes;

      size_t itptr = 0;
      size_t itrange = ( NearNodes->size() + 3 ) / 4;
      if ( itrange < 1 ) itrange = 1;
      itbeg = ( SUBTASKID - 1 ) * itrange;
      itend = ( SUBTASKID + 0 ) * itrange;
      if ( itbeg > NearNodes->size() ) itbeg = NearNodes->size();
      if ( itend > NearNodes->size() ) itend = NearNodes->size();
      if ( SUBTASKID == 4 ) itend = NearNodes->size();

      for ( auto it = NearNodes->begin(); it != NearNodes->end(); it ++ )
      {
        if ( itbeg <= itptr && itptr < itend )
        {
          size_t k = (*it)->gids.size();
          flops += 2.0 * m * n * k;
          mops += m * k;
          mops += 2.0 * ( m * n + n * k + m * k );
        }
        itptr ++;
      }

      event.Set( label + name, flops, mops );

      cost = flops / 1E+9;
    };

    void Prefetch( Worker* user_worker )
    {
      auto &u_leaf = arg->data.u_leaf[ SUBTASKID ];
      __builtin_prefetch( u_leaf.data() );
    };

    void GetEventRecord()
    {
      //arg->data.s2n = event;
    };

    void DependencyAnalysis()
    {
      assert( arg->isleaf );
      this->TryEnqueue();

      //auto &u_leaf = arg->data.u_leaf[ SUBTASKID ];
      //u_leaf.DependencyAnalysis( hmlp::ReadWriteType::W, this );
    };

    void Execute( Worker* user_worker )
    {
      LeavesToLeaves<SUBTASKID, NNPRUNE, NODE, T>( arg, itbeg, itend );
    };

}; 
template<typename NODE>
void PrintSet( set<NODE*> &set )
{
  for ( auto it = set.begin(); it != set.end(); it ++ )
  {
    printf( "%lu, ", (*it)->treelist_id );
  }
  printf( "\n" );
}; 
template<typename NODE>
multimap<size_t, size_t> NearNodeBallots( NODE *node )
{
  assert( node->isleaf );

  auto &setup = *(node->setup);
  auto &NN = *(setup.NN);
  auto &gids = node->gids;

  map<size_t, size_t> ballot;

  size_t HasMissingNeighbors = 0;


  for ( size_t j = 0; j < gids.size(); j ++ )
  {
    for ( size_t i = 0; i < NN.row(); i ++ )
    {
      auto value = NN( i, gids[ j ] ).first;
      size_t neighbor_gid = NN( i, gids[ j ] ).second;
      if ( neighbor_gid >= 0 && neighbor_gid < NN.col() )
      {
        size_t neighbor_morton = setup.morton[ neighbor_gid ];
        size_t weighted_ballot = 1.0 / ( value + 1E-3 );
        //printf( "gid %lu i %lu neighbor_gid %lu morton %lu\n", gids[ j ], i, 
        //    neighbor_gid, neighbor_morton );

        if (  i < NN.row() / 2 )
        {
          if ( ballot.find( neighbor_morton ) != ballot.end() )
          {
            //ballot[ neighbor_morton ] ++;
            ballot[ neighbor_morton ] += weighted_ballot;
          }
          else
          {
            //ballot[ neighbor_morton ] = 1;
            ballot[ neighbor_morton ] = weighted_ballot;
          }
        }
      }
      else
      {
        HasMissingNeighbors ++;
      }
    }
  }

  if ( HasMissingNeighbors )
  {
    printf( "Missing %lu neighbor pairs\n", HasMissingNeighbors ); 
    fflush( stdout );
  }

  return flip_map( ballot );

}; 
template<typename NODE, typename T>
void NearSamples( NODE *node )
{
  auto &setup = *(node->setup);
  auto &NN = *(setup.NN);

  if ( node->isleaf )
  {
    auto &gids = node->gids;
    //double budget = setup.budget;
    double budget = setup.Budget();
    size_t n_nodes = ( 1 << node->l );

    node->NearNodes.insert( node );
    node->NNNearNodes.insert( node );
    node->NNNearNodeMortonIDs.insert( node->morton );

    multimap<size_t, size_t> sorted_ballot = NearNodeBallots( node );

    for ( auto it = sorted_ballot.rbegin(); it != sorted_ballot.rend(); it ++ )
    {
      if ( node->NNNearNodes.size() >= n_nodes * budget ) break;
      auto *target = (*node->morton2node)[ (*it).second ];
      node->NNNearNodeMortonIDs.insert( (*it).second );
      node->NNNearNodes.insert( target );
    }
  }

}; 
template<typename NODE, typename T>
class NearSamplesTask : public Task
{
  public:

    NODE *arg = NULL;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "near" );

      //--------------------------------------
      double flops = 0.0, mops = 0.0;

      event.Set( label + name, flops, mops );
      cost = 1.0;
      priority = true;
    }

    void DependencyAnalysis() { this->TryEnqueue(); };

    void Execute( Worker* user_worker )
    {
      NearSamples<NODE, T>( arg );
    };

}; 
template<typename TREE>
void SymmetrizeNearInteractions( TREE & tree )
{
  int n_nodes = 1 << tree.depth;
  auto level_beg = tree.treelist.begin() + n_nodes - 1;

  for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
  {
    auto *node = *(level_beg + node_ind);
    auto & NearMortonIDs = node->NNNearNodeMortonIDs;
    for ( auto & it : NearMortonIDs )
    {
      auto *target = tree.morton2node[ it ];
      target->NNNearNodes.insert( node );
      target->NNNearNodeMortonIDs.insert( it );
    }
  }
}; 
template<bool NNPRUNE, typename NODE>
class CacheNearNodesTask : public Task
{
  public:

    NODE *arg;

    void Set( NODE *user_arg )
    {
      arg = user_arg;
      name = string( "c-n" );
      label = to_string( arg->treelist_id );
      cost = 1.0;
    };

    void GetEventRecord()
    {
      double flops = 0.0, mops = 0.0;

      NODE *node = arg;
      auto *NearNodes = &node->NearNodes;
      if ( NNPRUNE ) NearNodes = &node->NNNearNodes;
      auto &K = *node->setup->K;

      size_t m = node->gids.size();
      size_t n = 0;
      for ( auto it = NearNodes->begin(); it != NearNodes->end(); it ++ )
      {
        n += (*it)->gids.size();
      }
      event.Set( label + name, flops, mops );
    };

    void DependencyAnalysis() { arg->DependOnNoOne( this ); };

    void Execute( Worker* user_worker )
    {
      //printf( "%lu CacheNearNodes beg\n", arg->treelist_id ); fflush( stdout );

      NODE *node = arg;
      auto *NearNodes = &node->NearNodes;
      if ( NNPRUNE ) NearNodes = &node->NNNearNodes;
      auto &K = *node->setup->K;
      auto &data = node->data;
      auto &amap = node->gids;
      vector<size_t> bmap;
      for ( auto it = NearNodes->begin(); it != NearNodes->end(); it ++ )
      {
        bmap.insert( bmap.end(), (*it)->gids.begin(), (*it)->gids.end() );
      }
      data.NearKab = K( amap, bmap );

      data.Nearbmap.resize( bmap.size(), 1 );
      for ( size_t i = 0; i < bmap.size(); i ++ ) 
        data.Nearbmap[ i ] = bmap[ i ];

#ifdef HMLP_USE_CUDA
      auto *device = hmlp_get_device( 0 );
      node->data.Nearbmap.PrefetchH2D( device, 8 );

      size_t preserve_size = 3000000000;
      //if ( data.NearKab.col() * MAX_NRHS < 1200000000 &&
      //     data.NearKab.size() * 8 + preserve_size < device->get_memory_left() &&
      //     data.NearKab.size() * 8 > 4096 * 4096 * 8 * 4 )
      if ( data.NearKab.col() * MAX_NRHS < 1200000000 &&
           data.NearKab.size() * 8 + preserve_size < device->get_memory_left() )
      {
        data.NearKab.PrefetchH2D( device, 8 );
      }
      else
      {
        printf( "Kab %lu %lu not cache\n", data.NearKab.row(), data.NearKab.col() );
      }
#endif

      //printf( "%lu CacheNearNodesTask end\n", arg->treelist_id ); fflush( stdout );
    };
}; 
template<typename NODE>
void FindFarNodes( NODE *node, NODE *target )
{
  assert( target->isleaf );

  set<NODE*> *NearNodes;
  auto &data = node->data;
  auto *lchild = node->lchild;
  auto *rchild = node->rchild;

  NearNodes = &target->NearNodes;

  if ( !data.isskel || node->ContainAny( *NearNodes ) )
  {
    if ( !node->isleaf )
    {
      FindFarNodes( lchild, target );
      FindFarNodes( rchild, target );
    }
  }
  else
  {
    target->FarNodes.insert( node );
  }

  NearNodes = &target->NNNearNodes;

  if ( !data.isskel || node->ContainAny( *NearNodes ) )
  {
    if ( !node->isleaf )
    {
      FindFarNodes( lchild, target );
      FindFarNodes( rchild, target );
    }
  }
  else
  {
    if ( node->setup->IsSymmetric() && ( node->morton < target->morton ) )
    {
    }
    else
    {
      target->NNFarNodes.insert( node );
    }
  }

}; 
template<typename TREE>
void MergeFarNodes( TREE &tree )
{
  for ( int l = tree.depth; l >= 0; l -- )
  {
    size_t n_nodes = ( 1 << l );
    auto level_beg = tree.treelist.begin() + n_nodes - 1;

    for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
    {
      auto *node = *(level_beg + node_ind);

      if ( !node->data.isskel ) continue;

      if ( node->isleaf )
      {
        FindFarNodes( tree.treelist[ 0 ] , node );
      }
      else
      {
        auto *lchild = node->lchild;
        auto *rchild = node->rchild;

        auto &pFarNodes =   node->FarNodes;
        auto &lFarNodes = lchild->FarNodes;
        auto &rFarNodes = rchild->FarNodes;
        for ( auto it = lFarNodes.begin(); it != lFarNodes.end(); ++ it )
        {
          if ( rFarNodes.count( *it ) ) pFarNodes.insert( *it );
        }
        for ( auto it = pFarNodes.begin(); it != pFarNodes.end(); it ++ )
        {
          lFarNodes.erase( *it ); rFarNodes.erase( *it );
        }


        auto &pNNFarNodes =   node->NNFarNodes;
        auto &lNNFarNodes = lchild->NNFarNodes;
        auto &rNNFarNodes = rchild->NNFarNodes;

        //printf( "node %lu\n", node->treelist_id );
        //PrintSet( pNNFarNodes );
        //PrintSet( lNNFarNodes );
        //PrintSet( rNNFarNodes );


        for ( auto it = lNNFarNodes.begin(); it != lNNFarNodes.end(); ++ it )
        {
          if ( rNNFarNodes.count( *it ) ) pNNFarNodes.insert( *it );
        }
        for ( auto it = pNNFarNodes.begin(); it != pNNFarNodes.end(); it ++ )
        {
          lNNFarNodes.erase( *it ); 
          rNNFarNodes.erase( *it );
        }

        //PrintSet( pNNFarNodes );
        //PrintSet( lNNFarNodes );
        //PrintSet( rNNFarNodes );
      }
    }
  }

  if ( tree.setup.IsSymmetric() )
  {
    for ( int l = tree.depth; l >= 0; l -- )
    {
      std::size_t n_nodes = 1 << l;
      auto level_beg = tree.treelist.begin() + n_nodes - 1;

      for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
      {
        auto *node = *(level_beg + node_ind);
        auto &pFarNodes = node->NNFarNodes;
        for ( auto it = pFarNodes.begin(); it != pFarNodes.end(); it ++ )
        {
          (*it)->NNFarNodes.insert( node );
        }
      }
    }
  }
  
#ifdef DEBUG_SPDASKIT
  for ( int l = tree.depth; l >= 0; l -- )
  {
    std::size_t n_nodes = 1 << l;
    auto level_beg = tree.treelist.begin() + n_nodes - 1;

    for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
    {
      auto *node = *(level_beg + node_ind);
      auto &pFarNodes =   node->NNFarNodes;
      for ( auto it = pFarNodes.begin(); it != pFarNodes.end(); it ++ )
      {
        if ( !( (*it)->NNFarNodes.count( node ) ) )
        {
          printf( "Unsymmetric FarNodes %lu, %lu\n", node->treelist_id, (*it)->treelist_id );
          printf( "Near\n" );
          PrintSet(  node->NNNearNodes );
          PrintSet( (*it)->NNNearNodes );
          printf( "Far\n" );
          PrintSet(  node->NNFarNodes );
          PrintSet( (*it)->NNFarNodes );
          printf( "======\n" );
          break;
        }
      }
      if ( pFarNodes.size() )
      {
        printf( "l %2lu FarNodes(%lu) ", node->l, node->treelist_id );
        PrintSet( pFarNodes );
      }
    }
  }
#endif
};


template<bool NNPRUNE, bool CACHE = true, typename TREE>
void CacheFarNodes( TREE &tree )
{
  #pragma omp parallel for schedule( dynamic )
  for ( size_t i = 0; i < tree.treelist.size(); i ++ )
  {
    auto *node = tree.treelist[ i ];
    if ( node->isleaf )
    {
      node->data.w_leaf.reserve( node->gids.size(), MAX_NRHS );
      node->data.u_leaf[ 0 ].reserve( MAX_NRHS, node->gids.size() );
    }
  }

  if ( CACHE )
  {
    #pragma omp parallel for schedule( dynamic )
    for ( size_t i = 0; i < tree.treelist.size(); i ++ )
    {
      auto *node = tree.treelist[ i ];
      auto *FarNodes = &node->FarNodes;
      if ( NNPRUNE ) FarNodes = &node->NNFarNodes;
      auto &K = *node->setup->K;
      auto &data = node->data;
      auto &amap = data.skels;
      std::vector<size_t> bmap;
      for ( auto it = FarNodes->begin(); it != FarNodes->end(); it ++ )
      {
        bmap.insert( bmap.end(), (*it)->data.skels.begin(), 
                                 (*it)->data.skels.end() );
      }
      data.FarKab = K( amap, bmap );
    }
  }
}; 
template<bool NNPRUNE, typename TREE>
double DrawInteraction( TREE &tree )
{
  double exact_ratio = 0.0;
  FILE * pFile;
  //int n;
  char name[ 100 ];

  pFile = fopen ( "interaction.m", "w" );

  fprintf( pFile, "figure('Position',[100,100,800,800]);" );
  fprintf( pFile, "hold on;" );
  fprintf( pFile, "axis square;" );
  fprintf( pFile, "axis ij;" );

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

    for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
    {
      auto *node = *(level_beg + node_ind);

      if ( NNPRUNE )
      {
        auto &pNearNodes = node->NNNearNodes;
        auto &pFarNodes = node->NNFarNodes;
        for ( auto it = pFarNodes.begin(); it != pFarNodes.end(); it ++ )
        {
          double gb = (double)std::min( node->l, (*it)->l ) / tree.depth;
          //printf( "node->l %lu (*it)->l %lu depth %lu\n", node->l, (*it)->l, tree.depth );
          fprintf( pFile, "rectangle('position',[%lu %lu %lu %lu],'facecolor',[1.0,%lf,%lf]);\n",
              node->offset,      (*it)->offset,
              node->gids.size(), (*it)->gids.size(),
              gb, gb );
        }
        for ( auto it = pNearNodes.begin(); it != pNearNodes.end(); it ++ )
        {
          fprintf( pFile, "rectangle('position',[%lu %lu %lu %lu],'facecolor',[0.2,0.4,1.0]);\n",
              node->offset,      (*it)->offset,
              node->gids.size(), (*it)->gids.size() );

          exact_ratio += node->gids.size() * (*it)->gids.size();
        }  
      }
      else
      {
      }
    }
  }
  fprintf( pFile, "hold off;" );
  fclose( pFile );

  return exact_ratio / ( tree.n * tree.n );
}; 
template<bool SYMBOLIC, bool NNPRUNE, typename NODE, typename T>
void Evaluate
( 
  NODE *node, 
  size_t gid, 
  vector<size_t> &nnandi, // k + 1 non-prunable lists
  Data<T> &potentials 
)
{
  auto &K = *node->setup->K;
  auto &w = *node->setup->w;
  auto &gids = node->gids;
  auto &data = node->data;
  auto *lchild = node->lchild;
  auto *rchild = node->rchild;

  size_t nrhs = w.col();

  auto amap = std::vector<size_t>( 1 );
  amap[ 0 ] = gid;

  if ( !SYMBOLIC ) // No potential evaluation.
  {
    assert( potentials.size() == amap.size() * nrhs );
  }

  if ( !data.isskel || node->ContainAny( nnandi ) )
  {
    if ( node->isleaf )
    {
      if ( SYMBOLIC )
      {
        data.lock.Acquire();
        {
          if ( NNPRUNE ) node->NNNearIDs.insert( gid );
          else           node->NearIDs.insert(   gid );
        }
        data.lock.Release();
      }
      else
      {
#ifdef DEBUG_SPDASKIT
        printf( "level %lu direct evaluation\n", node->l );
#endif

        auto Kab = K( amap, gids ); 

        std::vector<size_t> bmap( nrhs );
        for ( size_t j = 0; j < bmap.size(); j ++ )
          bmap[ j ] = j;

        auto wb  = w( gids, bmap ); 

        xgemm
        (
          "N", "N",
          Kab.row(), wb.col(), wb.row(),
          1.0, Kab.data(),        Kab.row(),
                wb.data(),         wb.row(),
          1.0, potentials.data(), potentials.row()
        );
      }
    }
    else
    {
      Evaluate<SYMBOLIC, NNPRUNE>( lchild, gid, nnandi, potentials );      
      Evaluate<SYMBOLIC, NNPRUNE>( rchild, gid, nnandi, potentials );
    }
  }
  else // need gid's morton and neighbors' mortons
  {
    //printf( "level %lu is prunable\n", node->l );
    if ( SYMBOLIC )
    {
      data.lock.Acquire();
      {
        // Add gid to prunable list.
        if ( NNPRUNE ) node->FarIDs.insert(   gid );
        else           node->NNFarIDs.insert( gid );
      }
      data.lock.Release();
    }
    else
    {
#ifdef DEBUG_SPDASKIT
      printf( "level %lu is prunable\n", node->l );
#endif
      auto Kab = K( amap, node->data.skels );
      auto &w_skel = node->data.w_skel;
      xgemm
      (
        "N", "N",
        Kab.row(), w_skel.col(), w_skel.row(),
        1.0, Kab.data(),        Kab.row(),
          w_skel.data(),     w_skel.row(),
        1.0, potentials.data(), potentials.row()
      );          
    }
  }



}; 
template<bool SYMBOLIC, bool NNPRUNE, typename TREE, typename T>
void Evaluate
( 
  TREE &tree, 
  size_t gid, 
  Data<T> &potentials
)
{
  vector<size_t> nnandi;
  auto &w = *tree.setup.w;

  potentials.clear();
  potentials.resize( 1, w.col(), 0.0 );

  if ( NNPRUNE )
  {
    auto &NN = *tree.setup.NN;
    nnandi.reserve( NN.row() + 1 );
    nnandi.push_back( gid );
    for ( size_t i = 0; i < NN.row(); i ++ )
    {
      nnandi.push_back( NN( i, gid ).second );
    }
#ifdef DEBUG_SPDASKIT
    printf( "nnandi.size() %lu\n", nnandi.size() );
#endif
  }
  else
  {
    nnandi.reserve( 1 );
    nnandi.push_back( gid );
  }

  Evaluate<SYMBOLIC, NNPRUNE>( tree.treelist[ 0 ], gid, nnandi, potentials );

}; 
template<
  bool     USE_RUNTIME = true, 
  bool     USE_OMP_TASK = false, 
  bool     NNPRUNE = true, 
  bool     CACHE = true, 
  typename TREE, 
  typename T>
Data<T> Evaluate
( 
  TREE &tree,
  Data<T> &weights
)
{
  const bool AUTO_DEPENDENCY = true;

  using NODE = typename TREE::NODE;

  double beg, time_ratio, evaluation_time = 0.0;
  double allocate_time, computeall_time;
  double forward_permute_time, backward_permute_time;

  tree.DependencyCleanUp();

  size_t n    = weights.row();
  size_t nrhs = weights.col();

  beg = omp_get_wtime();
  hmlp::Data<T> potentials( n, nrhs, 0.0 );
  tree.setup.w = &weights;
  tree.setup.u = &potentials;
  allocate_time = omp_get_wtime() - beg;

  if ( REPORT_EVALUATE_STATUS )
  {
    printf( "Forward permute ...\n" ); fflush( stdout );
  }
  beg = omp_get_wtime();
  int n_nodes = ( 1 << tree.depth );
  auto level_beg = tree.treelist.begin() + n_nodes - 1;
  #pragma omp parallel for
  for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
  {
    auto *node = *(level_beg + node_ind);


    auto &gids = node->gids;
    auto &w_leaf = node->data.w_leaf;

    if ( w_leaf.row() != gids.size() || w_leaf.col() != weights.col() )
    {
      w_leaf.resize( gids.size(), weights.col() );
    }

    for ( size_t j = 0; j < w_leaf.col(); j ++ )
    {
      for ( size_t i = 0; i < w_leaf.row(); i ++ )
      {
        w_leaf( i, j ) = weights( gids[ i ], j ); 
      }
    };
  }
  forward_permute_time = omp_get_wtime() - beg;



  if ( REPORT_EVALUATE_STATUS )
  {
    printf( "N2S, S2S, S2N, L2L (HMLP Runtime) ...\n" ); fflush( stdout );
  }
  if ( tree.setup.IsSymmetric() )
  {
    beg = omp_get_wtime();
#ifdef HMLP_USE_CUDA
    potentials.AllocateD( hmlp_get_device( 0 ) );
    using LEAFTOLEAFVER2TASK = gpu::LeavesToLeavesVer2Task<CACHE, NNPRUNE, NODE, T>;
    LEAFTOLEAFVER2TASK leaftoleafver2task;
#endif
    using LEAFTOLEAFTASK1 = LeavesToLeavesTask<1, NNPRUNE, NODE, T>;
    using LEAFTOLEAFTASK2 = LeavesToLeavesTask<2, NNPRUNE, NODE, T>;
    using LEAFTOLEAFTASK3 = LeavesToLeavesTask<3, NNPRUNE, NODE, T>;
    using LEAFTOLEAFTASK4 = LeavesToLeavesTask<4, NNPRUNE, NODE, T>;

    using NODETOSKELTASK  = UpdateWeightsTask<NODE, T>;
    using SKELTOSKELTASK  = SkeletonsToSkeletonsTask<NNPRUNE, NODE, T>;
    using SKELTONODETASK  = SkeletonsToNodesTask<NNPRUNE, NODE, T>;

    LEAFTOLEAFTASK1 leaftoleaftask1;
    LEAFTOLEAFTASK2 leaftoleaftask2;
    LEAFTOLEAFTASK3 leaftoleaftask3;
    LEAFTOLEAFTASK4 leaftoleaftask4;

    NODETOSKELTASK  nodetoskeltask;
    SKELTOSKELTASK  skeltoskeltask;
    SKELTONODETASK  skeltonodetask;


//    if ( USE_OMP_TASK )
//    {
//      assert( !USE_RUNTIME );
//      tree.template TraverseLeafs<false, false>( leaftoleaftask1 );
//      tree.template TraverseLeafs<false, false>( leaftoleaftask2 );
//      tree.template TraverseLeafs<false, false>( leaftoleaftask3 );
//      tree.template TraverseLeafs<false, false>( leaftoleaftask4 );
//      tree.template UpDown<true, true, true>( nodetoskeltask, skeltoskeltask, skeltonodetask );
//    }
//    else
//    {
//      assert( !USE_OMP_TASK );
//
//#ifdef HMLP_USE_CUDA
//      tree.template TraverseLeafs<AUTO_DEPENDENCY, USE_RUNTIME>( leaftoleafver2task );
//#else
//      tree.template TraverseLeafs<AUTO_DEPENDENCY, USE_RUNTIME>( leaftoleaftask1 );
//      tree.template TraverseLeafs<AUTO_DEPENDENCY, USE_RUNTIME>( leaftoleaftask2 );
//      tree.template TraverseLeafs<AUTO_DEPENDENCY, USE_RUNTIME>( leaftoleaftask3 );
//      tree.template TraverseLeafs<AUTO_DEPENDENCY, USE_RUNTIME>( leaftoleaftask4 );
//#endif
//
//      /** check scheduler */
//      //hmlp_get_runtime_handle()->scheduler->ReportRemainingTime();
//      tree.template TraverseUp       <AUTO_DEPENDENCY, USE_RUNTIME>( nodetoskeltask );
//      tree.template TraverseUnOrdered<AUTO_DEPENDENCY, USE_RUNTIME>( skeltoskeltask );
//      tree.template TraverseDown     <AUTO_DEPENDENCY, USE_RUNTIME>( skeltonodetask );
//      /** check scheduler */
//      //hmlp_get_runtime_handle()->scheduler->ReportRemainingTime();
//
//      if ( USE_RUNTIME ) hmlp_run();
//
//
//
//#ifdef HMLP_USE_CUDA
//      hmlp::Device *device = hmlp_get_device( 0 );
//      for ( int stream_id = 0; stream_id < 10; stream_id ++ )
//        device->wait( stream_id );
//      //potentials.PrefetchD2H( device, 0 );
//      potentials.FetchD2H( device );
//#endif
//    }




#ifdef HMLP_USE_CUDA
    tree.TraverseLeafs( leaftoleafver2task );
#else
    tree.TraverseLeafs( leaftoleaftask1 );
    tree.TraverseLeafs( leaftoleaftask2 );
    tree.TraverseLeafs( leaftoleaftask3 );
    tree.TraverseLeafs( leaftoleaftask4 );
#endif
    tree.TraverseUp( nodetoskeltask );
    tree.TraverseUnOrdered( skeltoskeltask );
    tree.TraverseDown( skeltonodetask );
    tree.ExecuteAllTasks();
    //if ( USE_RUNTIME ) hmlp_run();


    double d2h_beg_t = omp_get_wtime();
#ifdef HMLP_USE_CUDA
    hmlp::Device *device = hmlp_get_device( 0 );
    for ( int stream_id = 0; stream_id < 10; stream_id ++ )
      device->wait( stream_id );
    //potentials.PrefetchD2H( device, 0 );
    potentials.FetchD2H( device );
#endif
    double d2h_t = omp_get_wtime() - d2h_beg_t;
    printf( "d2h_t %lfs\n", d2h_t );


    double aggregate_beg_t = omp_get_wtime();
    #pragma omp parallel for
    for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
    {
      auto *node = *(level_beg + node_ind);
      auto &u_leaf = node->data.u_leaf[ 0 ];
      for ( size_t p = 1; p < 20; p ++ )
      {
        for ( size_t i = 0; i < node->data.u_leaf[ p ].size(); i ++ )
          u_leaf[ i ] += node->data.u_leaf[ p ][ i ];
      }
    }
    double aggregate_t = omp_get_wtime() - aggregate_beg_t;
    printf( "aggregate_t %lfs\n", d2h_t );
 
#ifdef HMLP_USE_CUDA
    device->wait( 0 );
#endif
    computeall_time = omp_get_wtime() - beg;
  }
  else // TODO: implement unsymmetric prunning
  {
    printf( "Non symmetric ComputeAll is not yet implemented\n" );
    exit( 1 );
  }



  if ( REPORT_EVALUATE_STATUS )
  {
    printf( "Backward permute ...\n" ); fflush( stdout );
  }
  beg = omp_get_wtime();
  #pragma omp parallel for
  for ( int node_ind = 0; node_ind < n_nodes; node_ind ++ )
  {
    auto *node = *(level_beg + node_ind);
    auto &amap = node->gids;
    auto &u_leaf = node->data.u_leaf[ 0 ];



    //for ( size_t j = 0; j < amap.size(); j ++ )
    //  for ( size_t i = 0; i < potentials.row(); i ++ )
    //    potentials[ amap[ j ] * potentials.row() + i ] += u_leaf( j, i );


    for ( size_t j = 0; j < potentials.col(); j ++ )
      for ( size_t i = 0; i < amap.size(); i ++ )
        potentials( amap[ i ], j ) += u_leaf( i, j );


  }
  backward_permute_time = omp_get_wtime() - beg;

  evaluation_time += allocate_time;
  evaluation_time += forward_permute_time;
  evaluation_time += computeall_time;
  evaluation_time += backward_permute_time;
  time_ratio = 100 / evaluation_time;

  if ( REPORT_EVALUATE_STATUS )
  {
    printf( "========================================================\n");
    printf( "GOFMM evaluation phase\n" );
    printf( "========================================================\n");
    printf( "Allocate ------------------------------ %5.2lfs (%5.1lf%%)\n", 
        allocate_time, allocate_time * time_ratio );
    printf( "Forward permute ----------------------- %5.2lfs (%5.1lf%%)\n", 
        forward_permute_time, forward_permute_time * time_ratio );
    printf( "N2S, S2S, S2N, L2L -------------------- %5.2lfs (%5.1lf%%)\n", 
        computeall_time, computeall_time * time_ratio );
    printf( "Backward permute ---------------------- %5.2lfs (%5.1lf%%)\n", 
        backward_permute_time, backward_permute_time * time_ratio );
    printf( "========================================================\n");
    printf( "Evaluate ------------------------------ %5.2lfs (%5.1lf%%)\n", 
        evaluation_time, evaluation_time * time_ratio );
    printf( "========================================================\n\n");
  }


  tree.DependencyCleanUp();

  return potentials;

}; 
template<typename SPLITTER, typename T, typename SPDMATRIX>
Data<pair<T, size_t>> FindNeighbors
(
  SPDMATRIX &K, 
  SPLITTER splitter, 
   Configuration<T> &config,
  size_t n_iter = 10
)
{
  using DATA  = gofmm::NodeData<T>;
  using SETUP = gofmm::Setup<SPDMATRIX, SPLITTER, T>;
  using TREE  = tree::Tree<SETUP, DATA>;
  using NODE  = typename TREE::NODE;
  DistanceMetric metric = config.MetricType();
  size_t n = config.ProblemSize();
   size_t k = config.NeighborSize(); 
  pair<T, size_t> init( numeric_limits<T>::max(), n );
  gofmm::NeighborsTask<NODE, T> NEIGHBORStask;
  TREE rkdt;
  rkdt.setup.FromConfiguration( config, K, splitter, NULL );
  return rkdt.AllNearestNeighbor( n_iter, k, n_iter, init, NEIGHBORStask );
}; 
template<typename SPLITTER, typename RKDTSPLITTER, typename T, typename SPDMATRIX>
tree::Tree< gofmm::Setup<SPDMATRIX, SPLITTER, T>, gofmm::NodeData<T>>
*Compress
( 
  SPDMATRIX &K, 
  Data<pair<T, size_t>> &NN,
  SPLITTER splitter, 
  RKDTSPLITTER rkdtsplitter,
   Configuration<T> &config
)
{
  DistanceMetric metric = config.MetricType();
  size_t n = config.ProblemSize();
   size_t m = config.LeafNodeSize();
   size_t k = config.NeighborSize(); 
   size_t s = config.MaximumRank(); 
  T stol = config.Tolerance();
   T budget = config.Budget(); 

  const bool NNPRUNE   = true;
  const bool CACHE     = true;

  using SETUP = gofmm::Setup<SPDMATRIX, SPLITTER, T>;
  using DATA  = gofmm::NodeData<T>;
  using TREE  = tree::Tree<SETUP, DATA>;
  using NODE  = typename TREE::NODE;


  double beg, omptask45_time, omptask_time, ref_time;
  double time_ratio, compress_time = 0.0, other_time = 0.0;
  double ann_time, tree_time, skel_time, mergefarnodes_time, cachefarnodes_time;
  double nneval_time, nonneval_time, fmm_evaluation_time, symbolic_evaluation_time;

  beg = omp_get_wtime();
  if ( NN.size() != n * k ) 
  {
    NN = gofmm::FindNeighbors( K, rkdtsplitter, config );
  }
  ann_time = omp_get_wtime() - beg;


  auto *tree_ptr = new TREE();
   auto &tree = *tree_ptr;
  tree.setup.FromConfiguration( config, K, splitter, &NN );


  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "TreePartitioning ...\n" ); fflush( stdout );
  }
  beg = omp_get_wtime();
  tree.TreePartition();
  tree_time = omp_get_wtime() - beg;


#ifdef HMLP_AVX512

  assert( omp_get_max_threads() == 68 );
  //mkl_set_dynamic( 0 );
  //mkl_set_num_threads( 4 );
  hmlp_set_num_workers( 17 );
#else
  //if ( omp_get_max_threads() > 8 )
  //{
  //  hmlp_set_num_workers( omp_get_max_threads() / 2 );
  //}
  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "omp_get_max_threads() %d\n", omp_get_max_threads() );
  }
#endif




  NearSamplesTask<NODE, T> NEARSAMPLEStask;
  tree.DependencyCleanUp();
  printf( "Dependency clean up\n" ); fflush( stdout );
  tree.TraverseLeafs( NEARSAMPLEStask );
  tree.ExecuteAllTasks();
  //hmlp_run();
  printf( "Finish NearSamplesTask\n" ); fflush( stdout );
  SymmetrizeNearInteractions( tree );
  printf( "Finish SymmetrizeNearInteractions\n" ); fflush( stdout );



  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "Skeletonization (HMLP Runtime) ...\n" ); fflush( stdout );
  }
  beg = omp_get_wtime();
  gofmm::SkeletonKIJTask<NNPRUNE, NODE, T> GETMTXtask;
  gofmm::SkeletonizeTask<NODE, T> SKELtask;
  gofmm::InterpolateTask<NODE> PROJtask;
  tree.DependencyCleanUp();
  tree.TraverseUp( GETMTXtask, SKELtask );
  tree.TraverseUnOrdered( PROJtask );
  if ( CACHE )
  {
    gofmm::CacheNearNodesTask<NNPRUNE, NODE> KIJtask;
    tree.template TraverseLeafs( KIJtask );
  }
  other_time += omp_get_wtime() - beg;
  hmlp_run();
  skel_time = omp_get_wtime() - beg;




  beg = omp_get_wtime();
  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "MergeFarNodes ...\n" ); fflush( stdout );
  }
  gofmm::MergeFarNodes( tree );
  mergefarnodes_time = omp_get_wtime() - beg;

  beg = omp_get_wtime();
  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "CacheFarNodes ...\n" ); fflush( stdout );
  }
  gofmm::CacheFarNodes<NNPRUNE, CACHE>( tree );
  cachefarnodes_time = omp_get_wtime() - beg;

  auto exact_ratio = hmlp::gofmm::DrawInteraction<true>( tree );

  compress_time += ann_time;
  compress_time += tree_time;
  compress_time += skel_time;
  compress_time += mergefarnodes_time;
  compress_time += cachefarnodes_time;
  time_ratio = 100.0 / compress_time;
  if ( REPORT_COMPRESS_STATUS )
  {
    printf( "========================================================\n");
    printf( "GOFMM compression phase\n" );
    printf( "========================================================\n");
    printf( "NeighborSearch ------------------------ %5.2lfs (%5.1lf%%)\n", ann_time, ann_time * time_ratio );
    printf( "TreePartitioning ---------------------- %5.2lfs (%5.1lf%%)\n", tree_time, tree_time * time_ratio );
    printf( "Skeletonization ----------------------- %5.2lfs (%5.1lf%%)\n", skel_time, skel_time * time_ratio );
    printf( "MergeFarNodes ------------------------- %5.2lfs (%5.1lf%%)\n", mergefarnodes_time, mergefarnodes_time * time_ratio );
    printf( "CacheFarNodes ------------------------- %5.2lfs (%5.1lf%%)\n", cachefarnodes_time, cachefarnodes_time * time_ratio );
    printf( "========================================================\n");
    printf( "Compress (%4.2lf not compressed) -------- %5.2lfs (%5.1lf%%)\n", 
        exact_ratio, compress_time, compress_time * time_ratio );
    printf( "========================================================\n\n");
  }

  tree_ptr->DependencyCleanUp();

  return tree_ptr;

}; 
template<typename T, typename SPDMATRIX>
tree::Tree<
  gofmm::Setup<SPDMATRIX, centersplit<SPDMATRIX, 2, T>, T>, 
  gofmm::NodeData<T>>
*Compress( SPDMATRIX &K, T stol, T budget, size_t m, size_t k, size_t s )
{
  using SPLITTER     = centersplit<SPDMATRIX, 2, T>;
  using RKDTSPLITTER = randomsplit<SPDMATRIX, 2, T>;
  Data<pair<T, size_t>> NN;
  SPLITTER splitter( K );
  splitter.Kptr = &K;
   splitter.metric = ANGLE_DISTANCE;
  RKDTSPLITTER rkdtsplitter( K );
  rkdtsplitter.Kptr = &K;
   rkdtsplitter.metric = ANGLE_DISTANCE;
  size_t n = K.row();

   Configuration<T> config( ANGLE_DISTANCE, n, m, k, s, stol, budget );

  return Compress<SPLITTER, RKDTSPLITTER>
         ( K, NN, //ANGLE_DISTANCE, 
                splitter, rkdtsplitter, //n, m, k, s, stol, budget, 
                config );
}; 
template<typename T, typename SPDMATRIX>
tree::Tree<
  gofmm::Setup<SPDMATRIX, centersplit<SPDMATRIX, 2, T>, T>, 
  gofmm::NodeData<T>>
*Compress( SPDMATRIX &K, T stol, T budget )
{
  using SPLITTER     = centersplit<SPDMATRIX, 2, T>;
  using RKDTSPLITTER = randomsplit<SPDMATRIX, 2, T>;
  Data<pair<T, std::size_t>> NN;
  SPLITTER splitter( K );
  splitter.Kptr = &K;
  splitter.metric = ANGLE_DISTANCE;
  RKDTSPLITTER rkdtsplitter( K );
  rkdtsplitter.Kptr = &K;
  rkdtsplitter.metric = ANGLE_DISTANCE;
  size_t n = K.row();
  size_t m = 128;
  size_t k = 16;
  size_t s = m;

  if ( n >= 16384 )
  {
    m = 128;
    k = 20;
    s = 256;
  }

  if ( n >= 32768 )
  {
    m = 256;
    k = 24;
    s = 384;
  }

  if ( n >= 65536 )
  {
    m = 512;
    k = 32;
    s = 512;
  }

   Configuration<T> config( ANGLE_DISTANCE, n, m, k, s, stol, budget );

  return Compress<SPLITTER, RKDTSPLITTER>
         ( K, NN, //ANGLE_DISTANCE, 
                splitter, rkdtsplitter, config );

}; 
template<typename T>
tree::Tree<
  gofmm::Setup<SPDMatrix<T>, centersplit<SPDMatrix<T>, 2, T>, T>, 
  gofmm::NodeData<T>>
*Compress( SPDMatrix<T> &K, T stol, T budget )
{
   return Compress<T, SPDMatrix<T>>( K, stol, budget );
}; 
template<typename NODE, typename T>
void ComputeError( NODE *node, Data<T> potentials )
{
  auto &K = *node->setup->K;
  auto &w = node->setup->w;

  auto &amap = node->gids;
  std::vector<size_t> bmap = std::vector<size_t>( K.col() );

  for ( size_t j = 0; j < bmap.size(); j ++ ) bmap[ j ] = j;

  auto Kab = K( amap, bmap );

  auto nrm2 = hmlp_norm( potentials.row(), potentials.col(), 
                         potentials.data(), potentials.row() ); 

  xgemm
  (
    "N", "T",
    Kab.row(), w.row(), w.col(),
    -1.0, Kab.data(),        Kab.row(),
          w.data(),          w.row(),
     1.0, potentials.data(), potentials.row()
  );          

  auto err = hmlp_norm( potentials.row(), potentials.col(), 
                        potentials.data(), potentials.row() ); 
  
  printf( "node relative error %E, nrm2 %E\n", err / nrm2, nrm2 );


}; // end void ComputeError()








template<typename TREE, typename T>
T ComputeError( TREE &tree, size_t gid, Data<T> potentials )
{
  auto &K = *tree.setup.K;
  auto &w = *tree.setup.w;

  auto amap = std::vector<size_t>( 1, gid );
  auto bmap = std::vector<size_t>( K.col() );
  for ( size_t j = 0; j < bmap.size(); j ++ ) bmap[ j ] = j;

  auto Kab = K( amap, bmap );
  auto exact = potentials;

  xgemm
  (
    "N", "N",
    Kab.row(), w.col(), w.row(),
    1.0,   Kab.data(),   Kab.row(),
             w.data(),     w.row(),
    0.0, exact.data(), exact.row()
  );          


  auto nrm2 = hmlp_norm( exact.row(),  exact.col(), 
                         exact.data(), exact.row() ); 

  xgemm
  (
    "N", "N",
    Kab.row(), w.col(), w.row(),
    -1.0, Kab.data(),       Kab.row(),
          w.data(),          w.row(),
     1.0, potentials.data(), potentials.row()
  );          

  auto err = hmlp_norm( potentials.row(), potentials.col(), 
                        potentials.data(), potentials.row() ); 

  return err / nrm2;
}; 
template<typename TREE>
void SelfTesting( TREE &tree, size_t ntest, size_t nrhs )
{
  using T = typename TREE::T;
  size_t n = tree.n;
  if ( ntest > n ) ntest = n;
  vector<size_t> all_rhs( nrhs );
  for ( size_t rhs = 0; rhs < nrhs; rhs ++ ) all_rhs[ rhs ] = rhs;

  //auto A = tree.CheckAllInteractions();

  Data<T> w( n, nrhs ); w.rand();
  auto u = Evaluate<true, false, true, true>( tree, w );

  T nnerr_avg = 0.0;
  T nonnerr_avg = 0.0;
  T fmmerr_avg = 0.0;
  printf( "========================================================\n");
  printf( "Accuracy report\n" );
  printf( "========================================================\n");
  for ( size_t i = 0; i < ntest; i ++ )
  {
    size_t tar = i * n / ntest;
    Data<T> potentials;
    Evaluate<false, true>( tree, tar, potentials );
    auto nnerr = ComputeError( tree, tar, potentials );
    Evaluate<false, false>( tree, tar, potentials );
    auto nonnerr = ComputeError( tree, tar, potentials );
    //potentials = u( vector<size_t>( i ), all_rhs );
    for ( size_t p = 0; p < potentials.col(); p ++ )
    {
      potentials[ p ] = u( tar, p );
    }
    auto fmmerr = ComputeError( tree, tar, potentials );

    if ( i < 10 )
    {
      printf( "gid %6lu, ASKIT %3.1E, HODLR %3.1E, GOFMM %3.1E\n", 
          tar, nnerr, nonnerr, fmmerr );
    }
    nnerr_avg += nnerr;
    nonnerr_avg += nonnerr;
    fmmerr_avg += fmmerr;
  }
  printf( "========================================================\n");
  printf( "            ASKIT %3.1E, HODLR %3.1E, GOFMM %3.1E\n", 
      nnerr_avg / ntest , nonnerr_avg / ntest, fmmerr_avg / ntest );
  printf( "========================================================\n");

  if ( !tree.setup.SecureAccuracy() )
  {
    T lambda = 5.0;
    gofmm::Factorize( tree, lambda );
    gofmm::ComputeError( tree, lambda, w, u );
  }

}; 
template<typename SPDMATRIX>
void LaunchHelper( SPDMATRIX &K, CommandLineHelper &cmd )
{
  using T = typename SPDMATRIX::T;

  const int N_CHILDREN = 2;
  using SPLITTER     = gofmm::centersplit<SPDMATRIX, N_CHILDREN, T>;
  using RKDTSPLITTER = gofmm::randomsplit<SPDMATRIX, N_CHILDREN, T>;
  SPLITTER splitter( K );
  splitter.Kptr = &K;
  splitter.metric = cmd.metric;
  RKDTSPLITTER rkdtsplitter( K );
  rkdtsplitter.Kptr = &K;
  rkdtsplitter.metric = cmd.metric;
  gofmm::Configuration<T> config( cmd.metric, 
      cmd.n, cmd.m, cmd.k, cmd.s, cmd.stol, cmd.budget );
  Data<pair<T, size_t>> NN;
  //auto *tree_ptr = gofmm::Compress( X, K, NN, splitter, rkdtsplitter, config );
  auto *tree_ptr = gofmm::Compress( K, NN, splitter, rkdtsplitter, config );
   auto &tree = *tree_ptr;
  gofmm::SelfTesting( tree, 100, cmd.nrhs );


//  //#ifdef DUMP_ANALYSIS_DATA
//  gofmm::Summary<NODE> summary;
//  tree.Summary( summary );
//  summary.Print();

  delete tree_ptr;
}; 
template<typename T, typename SPDMATRIX>
class SimpleGOFMM
{
  public:

    SimpleGOFMM( SPDMATRIX &K, T stol, T budget )
    {
      tree_ptr = Compress( K, stol, budget );
    };

    ~SimpleGOFMM()
    {
      if ( tree_ptr ) delete tree_ptr;
    };

    void Multiply( Data<T> &y, Data<T> &x )
    {
      //hmlp::Data<T> weights( x.col(), x.row() );

      //for ( size_t j = 0; j < x.col(); j ++ )
      //  for ( size_t i = 0; i < x.row(); i ++ )
      //    weights( j, i ) = x( i, j );


      y = gofmm::Evaluate( *tree_ptr, x );
      //auto potentials = hmlp::gofmm::Evaluate( *tree_ptr, weights );

      //for ( size_t j = 0; j < y.col(); j ++ )
      //  for ( size_t i = 0; i < y.row(); i ++ )
      //    y( i, j ) = potentials( j, i );

    };

  private:

    tree::Tree<
      gofmm::Setup<SPDMATRIX, centersplit<SPDMATRIX, 2, T>, T>, 
      gofmm::NodeData<T>> *tree_ptr = NULL; 

}; 
typedef SPDMatrix<double> dSPDMatrix_t;
typedef SPDMatrix<float > sSPDMatrix_t;

typedef hmlp::gofmm::Setup<SPDMatrix<double>, 
    centersplit<SPDMatrix<double>, 2, double>, double> dSetup_t;

typedef hmlp::gofmm::Setup<SPDMatrix<float>, 
    centersplit<SPDMatrix<float >, 2,  float>,  float> sSetup_t;

typedef tree::Tree<dSetup_t, gofmm::NodeData<double>> dTree_t;
typedef tree::Tree<sSetup_t, gofmm::NodeData<float >> sTree_t;





Data<double> Evaluate( dTree_t *tree, Data<double> *weights );
Data<float>  Evaluate( dTree_t *tree, Data<float > *weights );

dTree_t *Compress( dSPDMatrix_t *K, double stol, double budget );
sTree_t *Compress( sSPDMatrix_t *K,  float stol,  float budget );

dTree_t *Compress( dSPDMatrix_t *K, double stol, double budget, 
      size_t m, size_t k, size_t s );
sTree_t *Compress( sSPDMatrix_t *K,  float stol,  float budget, 
      size_t m, size_t k, size_t s );


double ComputeError( dTree_t *tree, size_t gid, hmlp::Data<double> *potentials );
float  ComputeError( sTree_t *tree, size_t gid, hmlp::Data<float>  *potentials );










}; 
}; 
#endif