merrymercy commented on a change in pull request #6671:
URL: https://github.com/apache/incubator-tvm/pull/6671#discussion_r511940229
##########
File path: python/tvm/auto_scheduler/workload_registry.py
##########
@@ -175,6 +175,41 @@ def workload_key_to_tensors(workload_key):
return lookup(*args)
+def get_workload(task):
+ """Get the workload function for a given task
+
+ Parameters
+ ----------
+ task : SearchTask
+ Task to get workload of.
+
+ Returns
+ -------
+ workload : callable
+ The registered workload function.
+ """
+ name = workload_name(task.workload_key)
+ lookup = WORKLOAD_FUNC_REGISTRY[name]
+ assert callable(lookup)
+ return lookup
+
+
+def workload_name(workload_key):
Review comment:
```suggestion
def workload_func_name(workload_key):
```
##########
File path: python/tvm/auto_scheduler/workload_registry.py
##########
@@ -175,6 +175,41 @@ def workload_key_to_tensors(workload_key):
return lookup(*args)
+def get_workload(task):
Review comment:
```suggestion
def get_workload_func(task):
```
----------------------------------------------------------------
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]
