FrozenGene commented on a change in pull request #5753:
URL: https://github.com/apache/incubator-tvm/pull/5753#discussion_r450051536
##########
File path: python/tvm/runtime/module.py
##########
@@ -222,29 +223,31 @@ def evaluator(*args):
except NameError:
raise NameError("time_evaluate is only supported when RPC is
enabled")
- def _collect_dso_modules(self):
- """Helper function to collect dso modules, then return it."""
- visited, stack, dso_modules = set(), [], []
+ def _collect_modules(self, module_type_keys):
+ """Helper function to collect specifit modules, then return it."""
+ visited, stack, modules = set(), [], []
+ type_keys = module_type_keys if isinstance(module_type_keys, (list,
tuple)) else [module_type_keys]
# append root module
visited.add(self)
stack.append(self)
while stack:
module = stack.pop()
- if module._dso_exportable():
- dso_modules.append(module)
+ if module.type_key in type_keys:
+ modules.append(module)
for m in module.imported_modules:
if m not in visited:
visited.add(m)
stack.append(m)
- return dso_modules
+ return modules
def _dso_exportable(self):
Review comment:
Right. We could name it like `_dso_exportable_types`. What about this?
----------------------------------------------------------------
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:
us...@infra.apache.org
