[Metaschedule] Add utility API to ease using manual schedules #10876
[Metaschedule] Add utility API to ease using manual schedules #10876junrushao merged 9 commits intoapache:mainfrom
Conversation
masahi
force-pushed
the
tir-manual-schedule-api
branch
from
819ca8c to
bdfd8e4
Compare
|
@junrushao1994 ready for review. |
zxybazh
left a comment
zxybazh
left a comment
There was a problem hiding this comment.
Hi I have some nitpicks over the use of schedule_fn other than that LGTM.
| sch = Schedule(mod) | ||
|
|
||
| if schedule_fn(task, sch): | ||
| workload = database.commit_workload(mod) |
There was a problem hiding this comment.
I think it should be sch.mod given it has gone through a schedule function.
There was a problem hiding this 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.
| 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] |
There was a problem hiding this 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.
There was a problem hiding this comment.
That's possible, but that will force users to write
mod = Parse._mod(task.dispatched[0])
sch = Schedule(mod)
in every schedule_fn callback. I think this boilerplate is non-trivial (users shouldn't care about dispatched or Parse stuff).
There was a problem hiding this comment.
That make sense, and I wonder if using the task name and schedule would suffice, if we don't want users to care about details inside of the extracted task.
There was a problem hiding this comment.
There are some cases where having the relay mod is required. For example, I want to be able to skip tasks based on the output dtype of the compute, which can be retrieved from the relay mod.
There was a problem hiding this comment.
Sounds good. Thanks for the explaination.
| print(sch.trace) | ||
|
|
||
|
|
||
| @pytest.mark.skip("Integeration test") |
There was a problem hiding this comment.
May I ask why this test is removed?
There was a problem hiding this comment.
This test depends on auto-tensorization for tensorcore, which is not in main
There was a problem hiding this comment.
I see. We may upstream it later then.
There was a problem hiding this comment.
Ooops...Thanks for spotting this!
…#10876) As discussed in apache#10856 (comment), add a utility under `meta_schedule/testing/utils.py` to clean up the database boilerplate. Also using `DummyDatabase` instead of `JsonDatabase` for further clean up, as suggested by @junrushao1994 .
…#10876) As discussed in apache#10856 (comment), add a utility under `meta_schedule/testing/utils.py` to clean up the database boilerplate. Also using `DummyDatabase` instead of `JsonDatabase` for further clean up, as suggested by @junrushao1994 .
…#10876) As discussed in apache#10856 (comment), add a utility under `meta_schedule/testing/utils.py` to clean up the database boilerplate. Also using `DummyDatabase` instead of `JsonDatabase` for further clean up, as suggested by @junrushao1994 .
3 participants
As discussed in #10856 (comment), add a utility under
meta_schedule/testing/utils.pyto clean up the database boilerplate. Also usingDummyDatabaseinstead ofJsonDatabasefor further clean up, as suggested by @junrushao1994 .@junrushao1994 @zxybazh