[PYTHON] Make decorator optional for runtime (#1350)

python/tvm/_ffi/base.py

@@ -98,3 +98,18 @@ def c_array(ctype, values):
         Created ctypes array
     """
     return (ctype * len(values))(*values)
+
+
+def decorate(func, fwrapped):
+    """A wrapper call of decorator package, differs to call time
+
+    Parameters
+    ----------
+    func : function
+        The original function
+
+    fwrapped : function
+        The wrapped function
+    """
+    import decorator
+    return decorator.decorate(func, fwrapped)

python/tvm/hybrid/api.py

@@ -2,21 +2,31 @@
 from __future__ import absolute_import as _abs
 
 import types
-import decorator
+from .._ffi.base import decorate
 from .parser import parse_python
 
-@decorator.decorator
-def script(func, *args):
-    """If the arguments are tvm types, compile it to HalideIR.
-    O.W. return the python emulated result"""
-    from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types
-    if _is_tvm_arg_types(args):
-        return parse(func, args)
-    else:
+
+def script(pyfunc):
+    """Decorate a python function function as  hybrid script.
+
+    The hybrid function support emulation mode and parsing to
+    the internal language IR.
+
+    Returns
+    -------
+    hybrid_func : function
+        A decorated hybrid script function.
+    """
+    def wrapped_func(func, *args, **kwargs):
+        from .util import _enter_hybrid_runtime, _restore_runtime, _is_tvm_arg_types
+        if _is_tvm_arg_types(args):
+            return parse(func, args)
+
         intersect = _enter_hybrid_runtime(func)
-        func(*args)
+        value = func(*args, **kwargs)
         _restore_runtime(func, intersect)
-    return func
+        return value
+    return decorate(pyfunc, wrapped_func)
 
 
 def parse(func, args):

python/tvm/tag.py

@@ -1,11 +1,5 @@
 """Tag class for TVM operators."""
-from ._ffi.base import _LIB_NAME
-try:
-    from decorator import decorate
-except ImportError as err_msg:
-    # Allow decorator to be missing in runtime
-    if _LIB_NAME != "libtvm_runtime.so":
-        raise err_msg
+from ._ffi.base import decorate
 
 class TagScope(object):
     """Tag scope object to set tag for operators, working as context