tqchen commented on a change in pull request #7918:
URL: https://github.com/apache/tvm/pull/7918#discussion_r620538434
##########
File path: src/target/codegen.cc
##########
@@ -110,32 +115,100 @@ class ModuleSerializer {
// invariance: root module is always at location 0.
// The module order is collected via DFS
+ // This function merges all the DSO exportable module into
+ // a single one as this is also what happens in the final hierachy
void CreateModuleIndex() {
std::unordered_set<const runtime::ModuleNode*> visited{mod_.operator->()};
std::vector<runtime::ModuleNode*> stack{mod_.operator->()};
uint64_t module_index = 0;
- while (!stack.empty()) {
- runtime::ModuleNode* n = stack.back();
- stack.pop_back();
- mod2index_[n] = module_index++;
- mod_vec_.emplace_back(n);
- for (runtime::Module m : n->imports()) {
+ auto fpush_imports_to_stack = [&](runtime::ModuleNode* node) {
+ for (runtime::Module m : node->imports()) {
runtime::ModuleNode* next = m.operator->();
if (visited.count(next) == 0) {
visited.insert(next);
stack.push_back(next);
}
}
+ };
+
+ std::vector<runtime::ModuleNode*> dso_exportable_boundary;
+
+ // Create module index that merges all dso module into a single group.
+ //
+ // Do a two phase visit, to ensure dso module's index
+ // is always bigger than a parent of any dso module
+ // and smaller than children of any dso module.
+ //
+ // Error will be raised in CreateImportTree
+ // if merging dso module causes a cycle in the import tree
+
+ // Phase 0: only expand non-dso-module and record the boundary.
+ while (!stack.empty()) {
+ runtime::ModuleNode* n = stack.back();
+ stack.pop_back();
+ if (DSOExportable(n)) {
+ // do not recursively expand dso modules
+ // we will expand in phase 1
+ dso_exportable_boundary.emplace_back(n);
+ } else {
+ // expand the non-dso modules
+ mod2index_[n] = module_index++;
+ mod_group_vec_.emplace_back(std::vector<runtime::ModuleNode*>({n}));
+ fpush_imports_to_stack(n);
+ }
+ }
+ if (dso_exportable_boundary.size() == 0) return;
+
+ // create the slot for dso exportable modules
+ uint64_t dso_module_index = module_index++;
+ mod_group_vec_.emplace_back(std::vector<runtime::ModuleNode*>());
+
+ stack = std::move(dso_exportable_boundary);
Review comment:
I still think this is better given we are still manipulating the
stack(while dso_exportable boundary is not the stack)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
