You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2022/08/08 23:58:34 UTC
[GitHub] [tvm] jwfromm opened a new pull request, #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
jwfromm opened a new pull request, #12342:
URL: https://github.com/apache/tvm/pull/12342
This PR extends `PostOrderApply` to allow passing in a function that filters which blocks have schedules generated. Although this could be used in many ways, I implement one example where blocks are filtered by name. This can be useful for hand optimizing part of a primfunc, and using metaschedule to tune other parts. I've added a new `Blocks` argument to `tune_tir` that allows users to optionally specify the names of blocks that should be tuned. If not provided, the behavior of `tune_tir` is exactly the same as before this PR.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940878137
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -153,7 +153,7 @@ class SpaceGenerator : public runtime::ObjectRef {
* to blocks in post-DFS order.
* \return The design space generator created.
*/
- TVM_DLL static SpaceGenerator PostOrderApply();
+ TVM_DLL static SpaceGenerator PostOrderApply(runtime::PackedFunc f_block_filter = nullptr);
Review Comment:
You're right. It's suggested to leave default values in header files. Thanks for the correction.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940766464
##########
python/tvm/meta_schedule/space_generator/post_order_apply.py:
##########
@@ -29,8 +29,8 @@ class PostOrderApply(SpaceGenerator):
rules to blocks in post-DFS order.
"""
- def __init__(self):
+ def __init__(self, filter_fn=None):
Review Comment:
Nit, shall we unify the function name as `f_block_filter`?
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -153,7 +153,7 @@ class SpaceGenerator : public runtime::ObjectRef {
* to blocks in post-DFS order.
* \return The design space generator created.
*/
- TVM_DLL static SpaceGenerator PostOrderApply();
+ TVM_DLL static SpaceGenerator PostOrderApply(runtime::PackedFunc f_block_filter = nullptr);
Review Comment:
We don't need default value in header file IIRC.
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -66,6 +66,7 @@ def test_tune_matmul_cpu():
max_trials_global=32,
),
work_dir=work_dir,
+ blocks=["update"],
Review Comment:
Can you please add another `tune_block` test case instead to check if `tune_tir` can work for the case where block is specified? Specifically you can create a workload with 2 blocks and apply `ComputeInline` to one of the blocks. It should work if the api works as expected.
##########
python/tvm/meta_schedule/tune.py:
##########
@@ -392,6 +394,21 @@ def tune_tir(
The cost model to use.
measure_callbacks : Optional[List[MeasureCallback]]
The callbacks used during tuning.
+ space : Optional[FnSpaceGenerator]
+ The space generator to use.
+ blocks : Optional[List[str]]
+ A list of block names to tune. If provided, other blocks
Review Comment:
Nit, `A list of block names specifying blocks to be tuned. Note that if the list is not None, blocks outside this list will not be tuned and space can not be given.`
##########
python/tvm/meta_schedule/tune.py:
##########
@@ -407,6 +424,15 @@ def tune_tir(
params=[{"log_dir": log_dir, "logger_name": __name__ + f".task_{task_name}"}],
)
+ if blocks is not None:
+ assert space is None, "Only one of blocks and space can be specified."
+ # Create a filter function to identify named blocks.
+ def _filter_fn(block, target_names) -> bool:
Review Comment:
Nit, we may want to unify the function name.
##########
python/tvm/meta_schedule/tune.py:
##########
@@ -407,6 +424,15 @@ def tune_tir(
params=[{"log_dir": log_dir, "logger_name": __name__ + f".task_{task_name}"}],
)
+ if blocks is not None:
+ assert space is None, "Only one of blocks and space can be specified."
Review Comment:
Nit, we can be more specific with the warning here `Can not specify blocks to tune when search space is given.`
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940875918
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -153,7 +153,7 @@ class SpaceGenerator : public runtime::ObjectRef {
* to blocks in post-DFS order.
* \return The design space generator created.
*/
- TVM_DLL static SpaceGenerator PostOrderApply();
+ TVM_DLL static SpaceGenerator PostOrderApply(runtime::PackedFunc f_block_filter = nullptr);
Review Comment:
I think its a good idea to provide a default value for this function since using a filter is an unusual case. Since this will mostly be called from python it probably wont matter but just in case it's used from C++ this default will be helpful. So we need to either set the default here or in `post_order_apply.cc`. My understanding of good c++ hygiene is that defaults should be set in headers when possible.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940903984
##########
python/tvm/meta_schedule/space_generator/post_order_apply.py:
##########
@@ -29,8 +29,8 @@ class PostOrderApply(SpaceGenerator):
rules to blocks in post-DFS order.
"""
- def __init__(self):
+ def __init__(self, filter_fn=None):
Review Comment:
Ok that seems reasonable. I've updated all names to `f_block_filter` :).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940909594
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -66,6 +66,7 @@ def test_tune_matmul_cpu():
max_trials_global=32,
),
work_dir=work_dir,
+ blocks=["update"],
Review Comment:
Thanks for adding the test and it works in this case. Furthermore, we are assuming there's an AutoInline schedule rule which would inline block 'A' if available so that we can tell if `tune_block` api works. We could simplify and avoid this assumption by directly supplying it with a schedule rule to inline the block it gets so that it will trigger an error if blocks api is not correctly called (because we cannot inline both blocks right).
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940875918
##########
include/tvm/meta_schedule/space_generator.h:
##########
@@ -153,7 +153,7 @@ class SpaceGenerator : public runtime::ObjectRef {
* to blocks in post-DFS order.
* \return The design space generator created.
*/
- TVM_DLL static SpaceGenerator PostOrderApply();
+ TVM_DLL static SpaceGenerator PostOrderApply(runtime::PackedFunc f_block_filter = nullptr);
Review Comment:
I think its a good idea to provide a default value for this function since using a filter is an unusual case. So we need to either set the default here or in `post_order_apply.cc`. My understanding of good c++ hygiene is that defaults should be set in headers when possible.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940874467
##########
python/tvm/meta_schedule/space_generator/post_order_apply.py:
##########
@@ -29,8 +29,8 @@ class PostOrderApply(SpaceGenerator):
rules to blocks in post-DFS order.
"""
- def __init__(self):
+ def __init__(self, filter_fn=None):
Review Comment:
I actually think `filter_fn` is more clear than `f_block_filter` (using f for functions in python is unusual). I could change the c++ names if you want them to be consistent.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940881978
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -66,6 +66,7 @@ def test_tune_matmul_cpu():
max_trials_global=32,
),
work_dir=work_dir,
+ blocks=["update"],
Review Comment:
I attempted to add a new test that covers this. Can you confirm it's what you were thinking?
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940876627
##########
python/tvm/meta_schedule/space_generator/post_order_apply.py:
##########
@@ -29,8 +29,8 @@ class PostOrderApply(SpaceGenerator):
rules to blocks in post-DFS order.
"""
- def __init__(self):
+ def __init__(self, filter_fn=None):
Review Comment:
Just to clarify, in previous cases we've been using f_func_name as PackedFunc argument name, for example in ApplyHistoryBest dispatch function. But I would be okay with using filter_fn here as long as it's consistent with c++ side.
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r940910170
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -66,6 +66,7 @@ def test_tune_matmul_cpu():
max_trials_global=32,
),
work_dir=work_dir,
+ blocks=["update"],
Review Comment:
Here is an example that you may follow:
```
@pytest.mark.skip("Integration test")
def test_tune_block_cpu():
@derived_object
class RemoveBlock(PyScheduleRule):
def _initialize_with_tune_context(self, context: TuneContext) -> None:
pass
def apply(self, sch: Schedule, block: BlockRV) -> List[Schedule]:
if sch.get(block).name_hint == "root":
return [sch]
sch = sch.copy()
sch.compute_inline(block)
return [sch]
with tempfile.TemporaryDirectory() as work_dir:
sch: Schedule = tune_tir(
mod=two_step,
target=Target("llvm --num-cores=16"),
config=TuneConfig(
strategy="replay_trace",
num_trials_per_iter=32,
max_trials_per_task=32,
max_trials_global=32,
),
work_dir=work_dir,
blocks=["A"],
sch_rules=lambda *args: [RemoveBlock()],
)
assert sch is not None
```
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] zxybazh merged pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
zxybazh merged PR #12342:
URL: https://github.com/apache/tvm/pull/12342
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jwfromm commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.
Posted by GitBox <gi...@apache.org>.
jwfromm commented on code in PR #12342:
URL: https://github.com/apache/tvm/pull/12342#discussion_r941545325
##########
tests/python/unittest/test_meta_schedule_tune_tir.py:
##########
@@ -66,6 +66,7 @@ def test_tune_matmul_cpu():
max_trials_global=32,
),
work_dir=work_dir,
+ blocks=["update"],
Review Comment:
That makes sense. Thanks Xiyou!
--
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: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org