337 lines
11 KiB
C++
337 lines
11 KiB
C++
/* Template implementations for kdtree.h. */
|
|
|
|
namespace thor {
|
|
|
|
/**
|
|
* A generalized partition function for cached arrays.
|
|
*/
|
|
template<typename PartitionCondition, typename PointCache, typename Bound>
|
|
index_t Partition(
|
|
PartitionCondition splitcond,
|
|
index_t begin, index_t count,
|
|
PointCache* points,
|
|
Bound* left_bound, Bound* right_bound);
|
|
|
|
};
|
|
|
|
template<typename PartitionCondition, typename PointCache, typename Bound>
|
|
index_t thor::Partition(
|
|
PartitionCondition splitcond,
|
|
index_t begin, index_t count,
|
|
PointCache* points,
|
|
Bound* left_bound, Bound* right_bound) {
|
|
index_t left_i = begin;
|
|
index_t right_i = begin + count - 1;
|
|
|
|
/* At any point:
|
|
* every thing that strictly precedes left_i is correct
|
|
* every thing that strictly succeeds right_i is correct
|
|
*/
|
|
for (;;) {
|
|
for (;;) {
|
|
if (unlikely(left_i > right_i)) return left_i;
|
|
CacheRead<typename PointCache::Element> left_v(points, left_i);
|
|
if (!splitcond.is_left(left_v->vec())) {
|
|
*right_bound |= left_v->vec();
|
|
break;
|
|
}
|
|
*left_bound |= left_v->vec();
|
|
left_i++;
|
|
}
|
|
|
|
for (;;) {
|
|
if (unlikely(left_i > right_i)) return left_i;
|
|
CacheRead<typename PointCache::Element> right_v(points, right_i);
|
|
if (splitcond.is_left(right_v->vec())) {
|
|
*left_bound |= right_v->vec();
|
|
break;
|
|
}
|
|
*right_bound |= right_v->vec();
|
|
right_i--;
|
|
}
|
|
|
|
points->Swap(left_i, right_i);
|
|
|
|
DEBUG_ASSERT(left_i <= right_i);
|
|
right_i--;
|
|
}
|
|
}
|
|
|
|
template<typename TPoint, typename TNode, typename TParam>
|
|
void KdTreeHybridBuilder<TPoint, TNode, TParam>::Doit(
|
|
struct datanode* module, const Param* param_in,
|
|
index_t begin_index, index_t end_index,
|
|
DistributedCache* points_inout, DistributedCache* nodes_create,
|
|
TreeDecomposition* decomposition) {
|
|
param_ = param_in;
|
|
n_points_ = end_index - begin_index;
|
|
|
|
points_.Init(points_inout, BlockDevice::M_MODIFY);
|
|
nodes_.Init(nodes_create, BlockDevice::M_CREATE);
|
|
|
|
index_t dimension;
|
|
|
|
{
|
|
CacheRead<Point> first_point(&points_, points_.begin_index());
|
|
dimension = first_point->vec().length();
|
|
}
|
|
|
|
leaf_size_ = fx_param_int(module, "leaf_size", 32);
|
|
chunk_size_ = points_.n_block_elems();
|
|
if (chunk_size_ <= leaf_size_) {
|
|
NONFATAL("Decreasing leaf size from %d to %d due to block size!\n",
|
|
int(leaf_size_), int(chunk_size_));
|
|
leaf_size_ = chunk_size_;
|
|
}
|
|
|
|
fx_timer_start(module, "tree_build");
|
|
DecompNode* decomp_root;
|
|
Bound bound;
|
|
bound.Init(dimension);
|
|
FindBoundingBox_(begin_index, end_index, &bound);
|
|
Build_(begin_index, end_index, 0, rpc::n_peers(), bound, NULL, &decomp_root);
|
|
decomposition->Init(decomp_root);
|
|
fx_timer_stop(module, "tree_build");
|
|
}
|
|
|
|
template<typename TPoint, typename TNode, typename TParam>
|
|
void KdTreeHybridBuilder<TPoint, TNode, TParam>::FindBoundingBox_(
|
|
index_t begin_index, index_t end_index, Bound* bound) {
|
|
CacheReadIter<Point> point(&points_, begin_index);
|
|
for (index_t i = end_index - begin_index; i--; point.Next()) {
|
|
*bound |= point->vec();
|
|
}
|
|
}
|
|
|
|
template<typename TPoint, typename TNode, typename TParam>
|
|
index_t KdTreeHybridBuilder<TPoint, TNode, TParam>::Build_(
|
|
index_t begin_col, index_t end_col,
|
|
int begin_rank, int end_rank, const Bound& bound,
|
|
Node* parent, DecompNode** decomp_pp) {
|
|
index_t node_i = nodes_.AllocD(begin_rank, 1);
|
|
Node* node = nodes_.StartWrite(node_i);
|
|
DecompNode* left_decomp = NULL;
|
|
DecompNode* right_decomp = NULL;
|
|
|
|
node->set_range(begin_col, end_col - begin_col);
|
|
node->bound().Reset();
|
|
node->bound() |= bound;
|
|
|
|
if (node->count() > leaf_size_) {
|
|
index_t split_dim = BIG_BAD_NUMBER;
|
|
double max_width = -1;
|
|
|
|
// Short loop to find widest dimension
|
|
for (index_t d = 0; d < node->bound().dim(); d++) {
|
|
double w = node->bound().get(d).width();
|
|
|
|
if (unlikely(w > max_width)) {
|
|
max_width = w;
|
|
split_dim = d;
|
|
}
|
|
}
|
|
|
|
DEBUG_ASSERT_MSG(max_width >= 0, "max_width = %f, dim = %"LI"d, n = %"LI"d",
|
|
max_width, node->bound().dim(), node->count());
|
|
|
|
// even if the max width is zero, we still* must* split it!
|
|
Split_(node, begin_rank, end_rank, split_dim, parent,
|
|
&left_decomp, &right_decomp);
|
|
} else {
|
|
node->set_leaf();
|
|
// ensure leaves don't straddle block boundaries
|
|
DEBUG_SAME_INT(node->begin() / points_.n_block_elems(),
|
|
(node->end() - 1) / points_.n_block_elems());
|
|
for (index_t i = node->begin(); i < node->end(); i++) {
|
|
CacheRead<Point> point(&points_, i);
|
|
node->stat().Accumulate(*param_, *point);
|
|
}
|
|
}
|
|
|
|
if (parent != NULL) {
|
|
// accumulate self to parent's statistics
|
|
parent->stat().Accumulate(*param_,
|
|
node->stat(), node->bound(), node->count());
|
|
}
|
|
|
|
node->stat().Postprocess(*param_, node->bound(), node->count());
|
|
|
|
if (decomp_pp) {
|
|
*decomp_pp = new DecompNode(
|
|
typename TreeDecomposition::Info(begin_rank, end_rank),
|
|
&nodes_, node_i, nodes_.end_index());
|
|
DEBUG_ASSERT((left_decomp == NULL) == (right_decomp == NULL));
|
|
if (left_decomp != NULL) {
|
|
(*decomp_pp)->set_child(0, left_decomp);
|
|
(*decomp_pp)->set_child(1, right_decomp);
|
|
}
|
|
}
|
|
|
|
nodes_.StopWrite(node_i);
|
|
|
|
return node_i;
|
|
}
|
|
|
|
template<typename TPoint, typename TNode, typename TParam>
|
|
void KdTreeHybridBuilder<TPoint, TNode, TParam>::Split_(
|
|
Node* node, int begin_rank, int end_rank, int split_dim, Node *parent,
|
|
DecompNode** left_decomp_pp, DecompNode** right_decomp_pp) {
|
|
index_t split_col;
|
|
index_t begin_col = node->begin();
|
|
index_t end_col = node->end();
|
|
int split_rank = (begin_rank + end_rank) / 2;
|
|
double split_val;
|
|
DRange current_range = node->bound().get(split_dim);
|
|
typename Node::Bound final_left_bound;
|
|
typename Node::Bound final_right_bound;
|
|
|
|
final_left_bound.Init(node->bound().dim());
|
|
final_right_bound.Init(node->bound().dim());
|
|
|
|
if ((node->begin() & points_.n_block_elems_mask()) == 0
|
|
&& (!parent || parent->begin() != node->begin())) {
|
|
// We got one block of points! Let's give away ownership.
|
|
points_.cache()->GiveOwnership(
|
|
points_.Blockid(node->begin()), begin_rank);
|
|
// This is also a convenient time to display status.
|
|
percent_indicator("tree built", node->begin(), n_points_);
|
|
}
|
|
|
|
if (node->count() <= chunk_size_) {
|
|
split_val = current_range.mid();
|
|
if (current_range.width() == 0) {
|
|
// All points are equal. As a point of diligence, we still divide it,
|
|
// into two overlapping nodes.
|
|
split_col = (node->begin() + node->end()) / 2;
|
|
} else {
|
|
// perform a midpoint split
|
|
split_col = thor::Partition(
|
|
HrectPartitionCondition(split_dim, split_val),
|
|
begin_col, end_col - begin_col,
|
|
&points_, &final_left_bound, &final_right_bound);
|
|
}
|
|
} else {
|
|
index_t goal_col;
|
|
typename Node::Bound left_bound;
|
|
typename Node::Bound right_bound;
|
|
|
|
left_bound.Init(node->bound().dim());
|
|
right_bound.Init(node->bound().dim());
|
|
|
|
if (end_rank <= begin_rank + 1) {
|
|
// All points will go on the same machine, so do median split.
|
|
goal_col = (begin_col + end_col) / 2;
|
|
} else {
|
|
// We're distributing these between machines. Let's make sure
|
|
// we give roughly even work to the machines. What we do is
|
|
// pretend the points are distributed as equally as possible, by
|
|
// using the global number of machines and points, to avoid errors
|
|
// introduced by doing this split computation recursively.
|
|
goal_col = (uint64(split_rank) * 2 * n_points_ + rpc::n_peers())
|
|
/ rpc::n_peers() / 2;
|
|
}
|
|
|
|
// Round the goal to the nearest block.
|
|
goal_col = (goal_col + chunk_size_ / 2) / chunk_size_ * chunk_size_;
|
|
|
|
for (;;) {
|
|
// use linear interpolation to guess the value to split on.
|
|
// this typically leads to convergence rather quickly.
|
|
split_val = current_range.interpolate(
|
|
(goal_col - begin_col) / double(end_col - begin_col));
|
|
|
|
left_bound.Reset();
|
|
right_bound.Reset();
|
|
split_col = thor::Partition(
|
|
HrectPartitionCondition(split_dim, split_val),
|
|
begin_col, end_col - begin_col,
|
|
&points_, &left_bound, &right_bound);
|
|
|
|
if (split_col == goal_col) {
|
|
final_left_bound |= left_bound;
|
|
final_right_bound |= right_bound;
|
|
break;
|
|
} else if (split_col < goal_col) {
|
|
final_left_bound |= left_bound;
|
|
current_range = right_bound.get(split_dim);
|
|
if (current_range.width() == 0) {
|
|
break; // identical elements
|
|
}
|
|
begin_col = split_col;
|
|
} else if (split_col > goal_col) {
|
|
final_right_bound |= right_bound;
|
|
current_range = left_bound.get(split_dim);
|
|
if (current_range.width() == 0) {
|
|
break; // identical elements
|
|
}
|
|
end_col = split_col;
|
|
}
|
|
}
|
|
if (split_col != goal_col) {
|
|
// we got identical elements in that dimension, compute actual bound
|
|
FindBoundingBox_(begin_col, goal_col, &final_left_bound);
|
|
FindBoundingBox_(goal_col, end_col, &final_right_bound);
|
|
}
|
|
split_col = goal_col;
|
|
}
|
|
|
|
if (end_rank - begin_rank <= 1) {
|
|
// I'm only one machine, don't need to expand children.
|
|
left_decomp_pp = right_decomp_pp = NULL;
|
|
}
|
|
|
|
node->set_child(0, Build_(node->begin(), split_col,
|
|
begin_rank, split_rank, final_left_bound, node, left_decomp_pp));
|
|
node->set_child(1, Build_(split_col, node->end(),
|
|
split_rank, end_rank, final_right_bound, node, right_decomp_pp));
|
|
}
|
|
|
|
template<typename Point, typename Node, typename Param>
|
|
void thor::CreateKdTreeMaster(const Param& param,
|
|
int nodes_channel, int block_size_kb, double megs, datanode *module,
|
|
index_t n_points,
|
|
DistributedCache *points_cache, DistributedCache *nodes_cache,
|
|
ThorTreeDecomposition<Node> *decomposition) {
|
|
Node example_node;
|
|
|
|
example_node.stat().Init(param);
|
|
Point example_point;
|
|
CacheArray<Point>::GetDefaultElement(points_cache, &example_point);
|
|
example_node.bound().Init(example_point.vec().length());
|
|
|
|
CacheArray<Node>::CreateCacheMaster(nodes_channel,
|
|
CacheArray<Node>::ConvertBlockSize(example_node, block_size_kb),
|
|
example_node, megs, nodes_cache);
|
|
KdTreeHybridBuilder<Point, Node, Param> builder;
|
|
builder.Doit(module, ¶m, 0, n_points, points_cache, nodes_cache,
|
|
decomposition);
|
|
}
|
|
|
|
template<typename Point, typename Node, typename Param>
|
|
void thor::CreateKdTree(const Param& param,
|
|
int nodes_channel, int extra_channel,
|
|
datanode *module, index_t n_points,
|
|
DistributedCache *points_cache,
|
|
ThorTree<Param, Point, Node> *tree_out) {
|
|
double megs = fx_param_double(module, "megs", 1000);
|
|
DistributedCache *nodes_cache = new DistributedCache();
|
|
Broadcaster<ThorTreeDecomposition<Node> > broadcaster;
|
|
|
|
if (rpc::is_root()) {
|
|
ThorTreeDecomposition<Node> decomposition;
|
|
int block_size_kb = fx_param_int(module, "block_size_kb", 64);
|
|
CreateKdTreeMaster<Point, Node>(param,
|
|
nodes_channel, block_size_kb, megs, module, n_points,
|
|
points_cache, nodes_cache, &decomposition);
|
|
broadcaster.SetData(decomposition);
|
|
} else {
|
|
CacheArray<Node>::CreateCacheWorker(nodes_channel, megs, nodes_cache);
|
|
}
|
|
|
|
points_cache->Sync();
|
|
nodes_cache->Sync();
|
|
broadcaster.Doit(extra_channel); // broadcast the decomposition
|
|
|
|
tree_out->Init(param, broadcaster.get(), points_cache, nodes_cache);
|
|
}
|