tqchen commented on a change in pull request #10368:
URL: https://github.com/apache/tvm/pull/10368#discussion_r817802314
##########
File path: python/tvm/meta_schedule/task_scheduler/task_scheduler.py
##########
@@ -187,10 +160,114 @@ def f_join_running_task(task_id: int) -> None:
database,
cost_model,
measure_callbacks,
- f_tune,
- f_initialize_task,
- f_set_task_stopped,
- f_is_task_running,
- f_join_running_task,
- f_next_task_id,
+ *methods,
)
+
+
+class PyTaskScheduler:
+ """
+ An abstract task scheduler with customized methods on the python-side.
+ This is the user facing class for function overloading inheritance.
+
+ Note: @derived_object is required for proper usage of any inherited class.
+ """
+
+ _tvm_metadata = {
+ "cls": _PyTaskScheduler,
+ "members": [
+ "tasks",
+ "builder",
+ "runner",
+ "database",
+ "cost_model",
+ "measure_callbacks",
+ ],
+ "methods": [
+ "tune",
+ "initialize_task",
+ "set_task_stopped",
+ "is_task_running",
+ "join_running_task",
+ "next_task_id",
+ ],
+ }
+
+ def __init__(
+ self,
+ tasks: List[TuneContext],
+ builder: Builder,
+ runner: Runner,
+ database: Database,
+ cost_model: Optional[CostModel] = None,
+ measure_callbacks: Optional[List[MeasureCallback]] = None,
+ ):
+ self.tasks = tasks
+ self.builder = builder
+ self.runner = runner
+ self.database = database
+ self.cost_model = cost_model
+ self.measure_callbacks = measure_callbacks
+
+ def tune(self) -> None:
+ """Auto-tuning."""
+ # Using self._outer to replace the self pointer
+ _ffi_api.TaskSchedulerTune(self._outer()) # type: ignore # pylint:
disable=no-member
Review comment:
Assigning handle to inst is kinda of similar to weakref already. The
main reason was in your inst class implementation there is a need to pass in
the object itself.
One possible alternative is to let tune to take a second parameter that
records as outer. Then have the wrapper class pass the function in.
```python
def tune(self, outer):
pass
```
The weakref might be an OK approach here if we property document it.
Alternatively, do not implement methods like tune in PyTaskScheduler, and
expose a TaskScheduler object(that corresponds to the C++ object) which exposes
the FFI function. Just like
https://github.com/apache/tvm/blob/main/python/tvm/ir/transform.py#L160
Then the derived class only need to specify some of the related methods like
next_task_id, get wrapped as a TaskScheduler which comes with the tune method
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
