Improve code style, python wrapper and test cases (apache#11)

* Update c++ code style and unit test

* Update python State wrapper and test cases

python/tvm/ansor/__init__.py

@@ -17,6 +17,14 @@
 # pylint: disable=unused-import, redefined-builtin
 """Namespace for Ansor autoSchedule"""
 
+from . import compute_dag
+from . import measure
+from . import serialization
+from . import loop_state
+from . import task
+from . import utils
+
+# Shortcut
 from .compute_dag import ComputeDAG
 from .task import SearchTask, MetaTileRewritePolicy, TuneOption
 from .task import auto_schedule

python/tvm/ansor/_ffi_api.py

@@ -14,7 +14,8 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-"""FFI APIs for tvm.ansor"""
+
+"""Register FFI APIs from C++ for the namespace tvm.ansor"""
 import tvm._ffi
 
 

python/tvm/ansor/compute_dag.py

@@ -14,14 +14,12 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=unused-import
-""" ... """
+
+""" Computational graph and its analysis tools """
 
 import tvm._ffi
 from tvm.runtime import Object
-
-from .state import State
-
+from .loop_state import State
 from . import _ffi_api
 
 
@@ -50,13 +48,13 @@ def get_init_state(self):
         -------
         state : State
         """
-        return _ffi_api.ComputeDAGGetInitState(self)
+        return State(_ffi_api.ComputeDAGGetInitState(self))
 
     def apply_steps_from_state(self, state, layout_rewrite_level=None):
         """
         Parameters
         ----------
-        state : State
+        state : StateObject
         layout_rewrite_level : LayoutRewriteLevel(***)
 
         Returns
@@ -71,7 +69,7 @@ def print_python_code_from_state(self, state):
         """
         Parameters
         ----------
-        state : State
+        state : StateObject
 
         Returns
         -------
@@ -83,10 +81,10 @@ def infer_bound_from_state(self, state):
         """
         Parameters
         ----------
-        state : State
+        state : StateObject
 
         Returns
         -------
-        state : State
+        state : StateObject
         """
         return _ffi_api.ComputeDAGInferBoundFromState(self, state)

python/tvm/ansor/cost_model/__init__.py

@@ -15,6 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=unused-import, redefined-builtin
-""" ... """
+""" Cost model that estimates the performance of programs """
 
 from .cost_model import RandomModel

python/tvm/ansor/cost_model/cost_model.py

@@ -14,14 +14,13 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# pylint: disable=unused-import
-""" ... """
+
+""" Cost model that estimates the performance of programs """
 import ctypes
 import numpy as np
 
 import tvm._ffi
 from tvm.runtime import Object
-
 from .. import _ffi_api
 
 
@@ -32,9 +31,6 @@ class CostModel(Object):
 
 @tvm._ffi.register_object("ansor.RandomModel")
 class RandomModel(Object):
-    """
-    """
-
     def __init__(self):
         self.__init_handle_by_constructor__(_ffi_api.RandomModel)