[GitHub] [tvm] zxybazh commented on a diff in pull request #12342: [MetaSchedule] Extend tune_tir to support tuning of specific blocks.

[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