comaniac commented on a change in pull request #6184:
URL: https://github.com/apache/incubator-tvm/pull/6184#discussion_r468178808
##########
File path: python/tvm/auto_scheduler/auto_schedule.py
##########
@@ -88,10 +91,76 @@ class SearchPolicy(Object):
class EmptyPolicy(SearchPolicy):
""" This is an example empty search policy which will always generate
the init state of ComputeDAG.
+
+ Parameters
+ ----------
+ task : SearchTask
+ The SearchTask for the computation declaration.
+ init_search_callbacks : Optional[List[SearchCallback]]
+ Callback functions called before the search process.
+ """
+ def __init__(self, task, init_search_callbacks=None):
+ self.__init_handle_by_constructor__(_ffi_api.EmptyPolicy, task,
init_search_callbacks)
+
+
+@tvm._ffi.register_object("auto_scheduler.SketchSearchPolicy")
+class SketchSearchPolicy(SearchPolicy):
+ """ The search policy that searches in a hierarchical search space
defined by sketches.
+ The policy randomly samples programs from the space defined by sketches
+ and use evolutionary search to fine-tune them.
+
+ Parameters
+ ----------
+ task : SearchTask
+ The SearchTask for the computation declaration.
+ schedule_cost_model : CostModel = RandomModel()
+ The cost model to estimate the complete schedules.
+ params : Optional[Dict[str, value]]
Review comment:
s/value/Any/
##########
File path: python/tvm/auto_scheduler/auto_schedule.py
##########
@@ -175,17 +236,6 @@ def auto_schedule(task, search_policy='default',
tuning_options=None):
raise ValueError("Invalid task: " + task +
" . `auto_scheduler.auto_schedule` expects a
SearchTask.")
- if isinstance(search_policy, str):
- if search_policy == 'default':
- # TODO(jcf94): This is an example policy for minimum system, will
be upgrated to
- # formal search policy later.
- search_policy = EmptyPolicy()
- else:
- raise ValueError("Invalid search policy: " + search_policy)
- elif not isinstance(search_policy, SearchPolicy):
- raise ValueError("Invalid search policy: " + search_policy +
- " . `auto_scheduler.auto_schedule` expects a
SearchPolicy or a string.")
-
- sch, tensors = _ffi_api.AutoSchedule(task, search_policy,
- tuning_options if tuning_options else
TuningOptions())
+ # TODO(jcf94): Remove EmptyPolicy after finish the porting of
SketchSearchPolicy
Review comment:
Maybe we can rename it to `NoPolicy` or something like that to output
one record with no transform steps, so that we can use it to be a baseline.
##########
File path: python/tvm/auto_scheduler/auto_schedule.py
##########
@@ -175,17 +236,6 @@ def auto_schedule(task, search_policy='default',
tuning_options=None):
raise ValueError("Invalid task: " + task +
" . `auto_scheduler.auto_schedule` expects a
SearchTask.")
- if isinstance(search_policy, str):
- if search_policy == 'default':
- # TODO(jcf94): This is an example policy for minimum system, will
be upgrated to
- # formal search policy later.
- search_policy = EmptyPolicy()
- else:
- raise ValueError("Invalid search policy: " + search_policy)
- elif not isinstance(search_policy, SearchPolicy):
- raise ValueError("Invalid search policy: " + search_policy +
- " . `auto_scheduler.auto_schedule` expects a
SearchPolicy or a string.")
-
- sch, tensors = _ffi_api.AutoSchedule(task, search_policy,
- tuning_options if tuning_options else
TuningOptions())
+ # TODO(jcf94): Remove EmptyPolicy after finish the porting of
SketchSearchPolicy
+ sch, tensors = _ffi_api.AutoSchedule(search_policy or EmptyPolicy(task),
tuning_options)
Review comment:
Should we make sketch search policy as the default one? Otherwise I can
imagine lots of people will ask in the discuss forum about why auto scheduler
doesn't do any search...
##########
File path: python/tvm/auto_scheduler/auto_schedule.py
##########
@@ -88,10 +91,76 @@ class SearchPolicy(Object):
class EmptyPolicy(SearchPolicy):
""" This is an example empty search policy which will always generate
the init state of ComputeDAG.
+
+ Parameters
+ ----------
+ task : SearchTask
+ The SearchTask for the computation declaration.
+ init_search_callbacks : Optional[List[SearchCallback]]
+ Callback functions called before the search process.
+ """
+ def __init__(self, task, init_search_callbacks=None):
+ self.__init_handle_by_constructor__(_ffi_api.EmptyPolicy, task,
init_search_callbacks)
+
+
+@tvm._ffi.register_object("auto_scheduler.SketchSearchPolicy")
+class SketchSearchPolicy(SearchPolicy):
+ """ The search policy that searches in a hierarchical search space
defined by sketches.
+ The policy randomly samples programs from the space defined by sketches
+ and use evolutionary search to fine-tune them.
+
+ Parameters
+ ----------
+ task : SearchTask
+ The SearchTask for the computation declaration.
+ schedule_cost_model : CostModel = RandomModel()
+ The cost model to estimate the complete schedules.
+ params : Optional[Dict[str, value]]
+ Parameters of the search policy.
+ See `src/auto_scheduler/search_policy/sketch_search_policy.h` for the
definitions.
+ See `DEFAULT_PARAMS` below to find the default values.
+ seed : Optional[int]
+ Random seed.
+ verbose : int = 1
+ Verbosity level. 0 for silent, 1 to output information during schedule
search.
+ init_search_callbacks : Optional[List[SearchCallback]]
+ Callback functions called before the search process.
+ Candidates:
+ - auto_scheduler.PreloadMeasuredStates
+ - auto_scheduler.PreloadCustomSketchRule
+ TODO(jcf94): Add these implementation in later PRs.
Review comment:
```suggestion
Customized sketch rules called before the search process.
Possible callbacks:
- auto_scheduler.PreloadMeasuredStates
- auto_scheduler.PreloadCustomSketchRule
TODO(jcf94): Add search callback implementations.
```
##########
File path: tests/python/unittest/test_auto_scheduler_sketch_generation.py
##########
@@ -0,0 +1,107 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+""" Test sketch generation. """
+
+import tvm
+from tvm import te, auto_scheduler
+
+from test_auto_scheduler_common import (matmul_auto_scheduler_test,
conv2d_nchw_bn_relu_auto_scheduler_test,
+ max_pool2d_auto_scheduler_test,
min_nm_auto_scheduler_test,
+ softmax_nm_auto_scheduler_test,
softmax_abcd_auto_scheduler_test,
+
conv2d_winograd_nhwc_auto_scheduler_test)
+
+def print_sketches(sketches):
Review comment:
I guess this is for debugging purpose? Should we move it to sketch since
it might be useful for developers as well?
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
