[Metaschedule] Add utility API to ease using manual schedules #10876
Conversation
masahi
force-pushed
the
tir-manual-schedule-api
branch
from
819ca8c
to
bdfd8e4
Compare
|
@junrushao1994 ready for review. |
| 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.
| @@ -91,127 +91,6 @@ def test_tune_matmul_cuda(): | |||
| 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 .
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