6 #include <dmlc/registry.h> 12 DMLC_REGISTRY_FILE_TAG(split);
14 int count_tu_nodes(ASTNode* node) {
15 int accum = (
dynamic_cast<TranslationUnitNode*
>(node)) ? 1 : 0;
16 for (ASTNode* child : node->children) {
17 accum += count_tu_nodes(child);
22 template <
typename ThresholdType,
typename LeafOutputType>
24 ASTBuilder<ThresholdType, LeafOutputType>::Split(
int parallel_comp) {
25 if (parallel_comp <= 0) {
26 LOG(INFO) <<
"Parallel compilation disabled; all member trees will be " 27 <<
"dumped to a single source file. This may increase " 28 <<
"compilation time and memory usage.";
31 LOG(INFO) <<
"Parallel compilation enabled; member trees will be " 32 <<
"divided into " << parallel_comp <<
" translation units.";
33 CHECK_EQ(this->main_node->children.size(), 1);
34 ASTNode* top_ac_node = this->main_node->children[0];
35 CHECK(dynamic_cast<AccumulatorContextNode*>(top_ac_node));
38 std::vector<ASTNode*> tree_head;
39 for (ASTNode* node : top_ac_node->children) {
40 CHECK(dynamic_cast<ConditionNode*>(node) ||
dynamic_cast<OutputNode<LeafOutputType>*
>(node)
41 || dynamic_cast<CodeFolderNode*>(node));
42 tree_head.push_back(node);
47 const int ntree =
static_cast<int>(tree_head.size());
48 const int nunit = parallel_comp;
49 const int unit_size = (ntree + nunit - 1) / nunit;
50 std::vector<ASTNode*> tu_list;
51 const int current_num_tu = count_tu_nodes(this->main_node);
52 for (
int unit_id = 0; unit_id < nunit; ++unit_id) {
53 const int tree_begin = unit_id * unit_size;
54 const int tree_end = std::min((unit_id + 1) * unit_size, ntree);
55 if (tree_begin < tree_end) {
56 TranslationUnitNode* tu
57 = AddNode<TranslationUnitNode>(top_ac_node, current_num_tu + unit_id);
58 tu_list.push_back(tu);
59 AccumulatorContextNode* ac = AddNode<AccumulatorContextNode>(tu);
60 tu->children.push_back(ac);
61 for (
int tree_id = tree_begin; tree_id < tree_end; ++tree_id) {
62 ASTNode* tree_head_node = tree_head[tree_id];
63 tree_head_node->parent = ac;
64 ac->children.push_back(tree_head_node);
68 top_ac_node->children = tu_list;
71 template void ASTBuilder<float, uint32_t>::Split(
int);
72 template void ASTBuilder<float, float>::Split(
int);
73 template void ASTBuilder<double, uint32_t>::Split(
int);
74 template void ASTBuilder<double, double>::Split(
int);