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(so that
is not too much of a difference as long as we confirm weakref works). 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 how ModulePass is exposed as an object, while the
decorator returns an object that is represented as a ModulePass
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(Note that because TaskSchedule corresponds to the C++ class, passing
self is the intended approach there and no outer is needed)
--
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]
