zxybazh commented on code in PR #10876:
URL: https://github.com/apache/tvm/pull/10876#discussion_r842200643
##########
python/tvm/meta_schedule/testing/utils.py:
##########
@@ -110,3 +116,46 @@ def initialize_with_tune_context(self, context:
"TuneContext") -> None:
def apply(self, trace: Trace, _) -> Optional[Trace]:
return Trace(trace.insts, {})
+
+
+def apply_fixed_schedules(
+ relay_mod: Union[RelayFunc, IRModule],
+ target: Union[str, Target],
+ params: Optional[Dict[str, NDArray]],
+ schedule_fn: Callable[[ExtractedTask, Schedule], bool],
+):
+ """Apply fixed schedules (manually written, without any tunable knobs) as
specified by
+ schedule_fn to extracted tasks, and return a database that can be passed
to ApplyHistoryBest.
+
+ Parameters
+ ----------
+ mod : Union[RelayFunc, IRModule]
+ The Relay module to apply fixed schedules.
+ target : Union[str, Target]
+ The target used to extract tasks.
+ params : Optional[Dict[str, tvm.runtime.NDArray]]
+ The associated parameters of the module.
+ schedule_fn : Callable[[ExtractedTask, Schedule], bool]
Review Comment:
What about we change it to `Callable[[ExtractedTask], Schedule]`, i.e., for
input the dispatched IRModule is inside of task, and we can return a schedule
if it matches our rule, otherwise return None. I think it might be better to
avoid confusion.
##########
python/tvm/meta_schedule/testing/utils.py:
##########
@@ -110,3 +116,46 @@ def initialize_with_tune_context(self, context:
"TuneContext") -> None:
def apply(self, trace: Trace, _) -> Optional[Trace]:
return Trace(trace.insts, {})
+
+
+def apply_fixed_schedules(
+ relay_mod: Union[RelayFunc, IRModule],
+ target: Union[str, Target],
+ params: Optional[Dict[str, NDArray]],
+ schedule_fn: Callable[[ExtractedTask, Schedule], bool],
+):
+ """Apply fixed schedules (manually written, without any tunable knobs) as
specified by
+ schedule_fn to extracted tasks, and return a database that can be passed
to ApplyHistoryBest.
+
+ Parameters
+ ----------
+ mod : Union[RelayFunc, IRModule]
+ The Relay module to apply fixed schedules.
+ target : Union[str, Target]
+ The target used to extract tasks.
+ params : Optional[Dict[str, tvm.runtime.NDArray]]
+ The associated parameters of the module.
+ schedule_fn : Callable[[ExtractedTask, Schedule], bool]
+ A callable that is applied for each extracted task and the
corresponding default schedule.
+ Returns True if the given schedule should be committed to the
database, False otherwise.
+
+ Returns
+ -------
+ database : Database
+ The database containing dummy tuning records for manually scheduled
traces.
+ """
+ target = Target(target) if isinstance(target, str) else target
+ extracted_tasks = extract_task_from_relay(relay_mod, target, params)
+
+ database = DummyDatabase()
+
+ for task in extracted_tasks:
+ mod = Parse._mod(task.dispatched[0])
+ sch = Schedule(mod)
+
+ if schedule_fn(task, sch):
+ workload = database.commit_workload(mod)
Review Comment:
I think it should be `sch.mod` given it has gone through a schedule function.
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -91,127 +91,6 @@ def test_tune_matmul_cuda():
print(sch.trace)
[email protected]("Integeration test")
Review Comment:
May I ask why this test is removed?
--
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]
