Skip to content

Commit 1fe6638

Browse files
committed
Improve code style, python wrapper and test cases (apache#11)
* Update c++ code style and unit test * Update python State wrapper and test cases
1 parent e52135f commit 1fe6638

43 files changed

Lines changed: 2461 additions & 2595 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

python/tvm/ansor/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,14 @@
1717
# pylint: disable=unused-import, redefined-builtin
1818
"""Namespace for Ansor autoSchedule"""
1919

20+
from . import compute_dag
21+
from . import measure
22+
from . import serialization
23+
from . import loop_state
24+
from . import task
25+
from . import utils
26+
27+
# Shortcut
2028
from .compute_dag import ComputeDAG
2129
from .task import SearchTask, MetaTileRewritePolicy, TuneOption
2230
from .task import auto_schedule

python/tvm/ansor/_ffi_api.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,8 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17-
"""FFI APIs for tvm.ansor"""
17+
18+
"""Register FFI APIs from C++ for the namespace tvm.ansor"""
1819
import tvm._ffi
1920

2021

python/tvm/ansor/compute_dag.py

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,12 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17-
# pylint: disable=unused-import
18-
""" ... """
17+
18+
""" Computational graph and its analysis tools """
1919

2020
import tvm._ffi
2121
from tvm.runtime import Object
22-
23-
from .state import State
24-
22+
from .loop_state import State
2523
from . import _ffi_api
2624

2725

@@ -50,13 +48,13 @@ def get_init_state(self):
5048
-------
5149
state : State
5250
"""
53-
return _ffi_api.ComputeDAGGetInitState(self)
51+
return State(_ffi_api.ComputeDAGGetInitState(self))
5452

5553
def apply_steps_from_state(self, state, layout_rewrite_level=None):
5654
"""
5755
Parameters
5856
----------
59-
state : State
57+
state : StateObject
6058
layout_rewrite_level : LayoutRewriteLevel(***)
6159
6260
Returns
@@ -71,7 +69,7 @@ def print_python_code_from_state(self, state):
7169
"""
7270
Parameters
7371
----------
74-
state : State
72+
state : StateObject
7573
7674
Returns
7775
-------
@@ -83,10 +81,10 @@ def infer_bound_from_state(self, state):
8381
"""
8482
Parameters
8583
----------
86-
state : State
84+
state : StateObject
8785
8886
Returns
8987
-------
90-
state : State
88+
state : StateObject
9189
"""
9290
return _ffi_api.ComputeDAGInferBoundFromState(self, state)

python/tvm/ansor/cost_model/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,6 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717
# pylint: disable=unused-import, redefined-builtin
18-
""" ... """
18+
""" Cost model that estimates the performance of programs """
1919

2020
from .cost_model import RandomModel

python/tvm/ansor/cost_model/cost_model.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,13 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17-
# pylint: disable=unused-import
18-
""" ... """
17+
18+
""" Cost model that estimates the performance of programs """
1919
import ctypes
2020
import numpy as np
2121

2222
import tvm._ffi
2323
from tvm.runtime import Object
24-
2524
from .. import _ffi_api
2625

2726

@@ -32,9 +31,6 @@ class CostModel(Object):
3231

3332
@tvm._ffi.register_object("ansor.RandomModel")
3433
class RandomModel(Object):
35-
"""
36-
"""
37-
3834
def __init__(self):
3935
self.__init_handle_by_constructor__(_ffi_api.RandomModel)
4036

0 commit comments

Comments
 (0)