Commit 15d5801f by Tianqi Chen Committed by GitHub

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

parent e00dacb1
......@@ -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)
......@@ -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"""
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)
else:
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):
......
"""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
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment