explicit CollapseNode(const ParentDIA& parent)
: Super(parent.ctx(), "Collapse", { parent.id() }, { parent.node() }),
parent_stack_empty_(ParentDIA::stack_empty) {
auto propagate_fn = [this](const ValueType& input) {
this->PushItem(input);
};
auto lop_chain = parent.stack().push(propagate_fn).fold();
parent.node()->AddChild(this, lop_chain);
}
explicit RebalanceNode(const ParentDIA& parent)
: Super(parent.ctx(), "Rebalance", { parent.id() }, { parent.node() }),
parent_stack_empty_(ParentDIA::stack_empty) {
auto save_fn = [this](const ValueType& input) {
writer_.Put(input);
};
auto lop_chain = parent.stack().push(save_fn).fold();
parent.node()->AddChild(this, lop_chain);
}
PrefixSumNode(const ParentDIA& parent,
const char* label,
const SumFunction& sum_function,
const ValueType& initial_element)
: Super(parent.ctx(), label, { parent.id() }, { parent.node() }),
sum_function_(sum_function),
local_sum_(),
initial_element_(initial_element),
parent_stack_empty_(ParentDIA::stack_empty) {
// Hook PreOp(s)
auto pre_op_fn = [this](const ValueType& input) {
PreOp(input);
};
auto lop_chain = parent.stack().push(pre_op_fn).fold();
parent.node()->AddChild(this, lop_chain);
}
GatherNode(const ParentDIA& parent, const char* label,
size_t target_id,
std::vector<ValueType>* out_vector)
: Super(parent.ctx(), label,
{ parent.id() }, { parent.node() }),
target_id_(target_id),
out_vector_(out_vector)
{
assert(target_id_ < context_.num_workers());
auto pre_op_fn = [this](const ValueType& input) {
emitters_[target_id_].Put(input);
};
// close the function stack with our pre op and register it at parent
// node for output
auto lop_chain = parent.stack().push(pre_op_fn).fold();
parent.node()->AddChild(this, lop_chain);
}