masahi commented on code in PR #10876:
URL: https://github.com/apache/tvm/pull/10876#discussion_r842231131
##########
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 don't think so. The purpose of this workload commit is to match against
the unmodified mod during `ApplyHistoryBest`. So we want to commit the original
mod as is.
--
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]
