[API/Refactor] Unified PackedFunc for API and Generated Functions (#26)
This commit is contained in:
Tianqi Chen
@@ -1,7 +1,7 @@
|
||||
# pylint: disable=redefined-builtin, wildcard-import
|
||||
"""C++ backend related python scripts"""
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import register_node
|
||||
from ._ctypes._node import register_node
|
||||
|
||||
from . import tensor
|
||||
from . import expr
|
||||
|
||||
@@ -91,45 +91,3 @@ def c_array(ctype, values):
|
||||
Created ctypes array
|
||||
"""
|
||||
return (ctype * len(values))(*values)
|
||||
|
||||
|
||||
def ctypes2docstring(num_args, arg_names, arg_types, arg_descs, remove_dup=True):
|
||||
"""Convert ctypes returned doc string information into parameters docstring.
|
||||
|
||||
num_args : nn_uint
|
||||
Number of arguments.
|
||||
|
||||
arg_names : ctypes.POINTER(ctypes.c_char_p)
|
||||
Argument names.
|
||||
|
||||
arg_types : ctypes.POINTER(ctypes.c_char_p)
|
||||
Argument type information.
|
||||
|
||||
arg_descs : ctypes.POINTER(ctypes.c_char_p)
|
||||
Argument description information.
|
||||
|
||||
remove_dup : boolean, optional
|
||||
Whether remove duplication or not.
|
||||
|
||||
Returns
|
||||
-------
|
||||
docstr : str
|
||||
Python docstring of parameter sections.
|
||||
"""
|
||||
param_keys = set()
|
||||
param_str = []
|
||||
for i in range(num_args.value):
|
||||
key = py_str(arg_names[i])
|
||||
if key in param_keys and remove_dup:
|
||||
continue
|
||||
param_keys.add(key)
|
||||
type_info = py_str(arg_types[i])
|
||||
ret = '%s : %s' % (key, type_info)
|
||||
if len(arg_descs[i]) != 0:
|
||||
ret += '\n ' + py_str(arg_descs[i])
|
||||
param_str.append(ret)
|
||||
doc_str = ('Parameters\n' +
|
||||
'----------\n' +
|
||||
'%s\n')
|
||||
doc_str = doc_str % ('\n'.join(param_str))
|
||||
return doc_str
|
||||
|
||||
@@ -1,416 +0,0 @@
|
||||
# coding: utf-8
|
||||
# pylint: disable=invalid-name, protected-access, too-many-arguments, too-many-lines
|
||||
# pylint: disable=attribute-defined-outside-init, no-member, missing-docstring, too-many-return-statements
|
||||
"""Symbolic configuration API."""
|
||||
from __future__ import absolute_import as _abs
|
||||
|
||||
import ctypes
|
||||
import sys
|
||||
from numbers import Number, Integral
|
||||
|
||||
from .._base import _LIB
|
||||
from .._base import c_str, py_str, string_types
|
||||
from .._base import check_call, ctypes2docstring
|
||||
from .. import _api_internal
|
||||
from . import _runtime_api
|
||||
from ._types import TVMValue, TypeCode, TVMPackedCFunc, TVMCFuncFinalizer
|
||||
|
||||
# type definitions
|
||||
APIFuncHandle = ctypes.c_void_p
|
||||
NodeHandle = ctypes.c_void_p
|
||||
FunctionHandle = ctypes.c_void_p
|
||||
|
||||
class APIType(object):
|
||||
"""TVMType used in API calls"""
|
||||
INT = ctypes.c_int(TypeCode.INT)
|
||||
UINT = ctypes.c_int(TypeCode.UINT)
|
||||
FLOAT = ctypes.c_int(TypeCode.FLOAT)
|
||||
HANDLE = ctypes.c_int(TypeCode.HANDLE)
|
||||
NULL = ctypes.c_int(TypeCode.NULL)
|
||||
NODE_HANDLE = ctypes.c_int(TypeCode.NODE_HANDLE)
|
||||
STR = ctypes.c_int(TypeCode.STR)
|
||||
FUNC_HANDLE = ctypes.c_int(TypeCode.FUNC_HANDLE)
|
||||
|
||||
|
||||
NODE_TYPE = {
|
||||
}
|
||||
|
||||
def _return_node(x):
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, NodeHandle):
|
||||
handle = NodeHandle(handle)
|
||||
ret_val = TVMValue()
|
||||
ret_type_code = ctypes.c_int()
|
||||
ret_success = ctypes.c_int()
|
||||
check_call(_LIB.TVMNodeGetAttr(
|
||||
handle, c_str("type_key"),
|
||||
ctypes.byref(ret_val),
|
||||
ctypes.byref(ret_type_code),
|
||||
ctypes.byref(ret_success)))
|
||||
return NODE_TYPE.get(py_str(ret_val.v_str), NodeBase)(handle)
|
||||
|
||||
|
||||
def _return_func(x):
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, FunctionHandle):
|
||||
handle = FunctionHandle(handle)
|
||||
return _runtime_api._function_cls(handle)
|
||||
|
||||
|
||||
def _return_handle(x):
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, ctypes.c_void_p):
|
||||
handle = ctypes.c_void_p(handle)
|
||||
return handle
|
||||
|
||||
|
||||
RET_SWITCH = {
|
||||
TypeCode.NULL: lambda x: None,
|
||||
TypeCode.INT: lambda x: x.v_int64,
|
||||
TypeCode.FLOAT: lambda x: x.v_float64,
|
||||
TypeCode.STR: lambda x: py_str(x.v_str),
|
||||
TypeCode.NODE_HANDLE: _return_node,
|
||||
TypeCode.FUNC_HANDLE: _return_func
|
||||
}
|
||||
|
||||
PACK_ARG_SWITCH = {
|
||||
TypeCode.NULL: lambda x: None,
|
||||
TypeCode.INT: lambda x: x.v_int64,
|
||||
TypeCode.FLOAT: lambda x: x.v_float64,
|
||||
TypeCode.STR: lambda x: py_str(x.v_str),
|
||||
TypeCode.HANDLE: lambda x: _return_handle,
|
||||
}
|
||||
|
||||
|
||||
class SliceBase(object):
|
||||
"""base class of slice object"""
|
||||
pass
|
||||
|
||||
class NodeBase(object):
|
||||
"""Symbol is symbolic graph."""
|
||||
__slots__ = ["handle"]
|
||||
# pylint: disable=no-member
|
||||
def __init__(self, handle):
|
||||
"""Initialize the function with handle
|
||||
|
||||
Parameters
|
||||
----------
|
||||
handle : SymbolHandle
|
||||
the handle to the underlying C++ Symbol
|
||||
"""
|
||||
self.handle = handle
|
||||
|
||||
def __repr__(self):
|
||||
return _api_internal._format_str(self)
|
||||
|
||||
def __del__(self):
|
||||
check_call(_LIB.TVMNodeFree(self.handle))
|
||||
|
||||
def __getattr__(self, name):
|
||||
ret_val = TVMValue()
|
||||
ret_type_code = ctypes.c_int()
|
||||
ret_success = ctypes.c_int()
|
||||
check_call(_LIB.TVMNodeGetAttr(
|
||||
self.handle, c_str(name),
|
||||
ctypes.byref(ret_val),
|
||||
ctypes.byref(ret_type_code),
|
||||
ctypes.byref(ret_success)))
|
||||
value = RET_SWITCH[ret_type_code.value](ret_val)
|
||||
if not ret_success.value:
|
||||
raise AttributeError(
|
||||
"'%s' object has no attribute '%s'" % (str(type(self)), name))
|
||||
return value
|
||||
|
||||
def __hash__(self):
|
||||
return _api_internal._raw_ptr(self)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, NodeBase):
|
||||
return False
|
||||
return self.__hash__() == other.__hash__()
|
||||
|
||||
def __ne__(self, other):
|
||||
return not self.__eq__(other)
|
||||
|
||||
def __dir__(self):
|
||||
plist = ctypes.POINTER(ctypes.c_char_p)()
|
||||
size = ctypes.c_uint()
|
||||
check_call(_LIB.TVMNodeListAttrNames(
|
||||
self.handle, ctypes.byref(size), ctypes.byref(plist)))
|
||||
names = []
|
||||
for i in range(size.value):
|
||||
names.append(py_str(plist[i]))
|
||||
return names
|
||||
|
||||
def __reduce__(self):
|
||||
return (type(self), (None,), self.__getstate__())
|
||||
|
||||
def __getstate__(self):
|
||||
handle = self.handle
|
||||
if handle is not None:
|
||||
return {'handle': _api_internal._save_json(self)}
|
||||
else:
|
||||
return {'handle': None}
|
||||
|
||||
def __setstate__(self, state):
|
||||
# pylint: disable=assigning-non-slot
|
||||
handle = state['handle']
|
||||
if handle is not None:
|
||||
json_str = handle
|
||||
_push_arg(json_str)
|
||||
other = _api_internal._load_json(json_str)
|
||||
self.handle = other.handle
|
||||
other.handle = None
|
||||
else:
|
||||
self.handle = None
|
||||
|
||||
|
||||
def const(value, dtype=None):
|
||||
"""construct a constant"""
|
||||
if dtype is None:
|
||||
if isinstance(value, Integral):
|
||||
dtype = 'int32'
|
||||
else:
|
||||
dtype = 'float32'
|
||||
return _api_internal._const(value, dtype)
|
||||
|
||||
|
||||
def _ctypes_free_resource(rhandle):
|
||||
"""callback to free resources when it it not needed."""
|
||||
pyobj = ctypes.cast(rhandle, ctypes.py_object)
|
||||
ctypes.pythonapi.Py_DecRef(pyobj)
|
||||
|
||||
# Global callback that is always alive
|
||||
TVM_FREE_PYOBJ = TVMCFuncFinalizer(_ctypes_free_resource)
|
||||
ctypes.pythonapi.Py_IncRef(ctypes.py_object(TVM_FREE_PYOBJ))
|
||||
|
||||
def convert_to_tvm_func(pyfunc):
|
||||
"""Convert a python function to TVM function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pyfunc : python function
|
||||
The python function to be converted.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tvmfunc: tvm.nd.Function
|
||||
The converted tvm function.
|
||||
"""
|
||||
local_pyfunc = pyfunc
|
||||
def cfun(args, type_codes, num_args, _):
|
||||
""" ctypes function """
|
||||
num_args = num_args.value if isinstance(num_args, ctypes.c_int) else num_args
|
||||
pyargs = [PACK_ARG_SWITCH[type_codes[i]](args[i]) for i in range(num_args)]
|
||||
local_pyfunc(*pyargs)
|
||||
handle = FunctionHandle()
|
||||
f = TVMPackedCFunc(cfun)
|
||||
# NOTE: We will need to use python-api to increase ref count of the f
|
||||
# TVM_FREE_PYOBJ will be called after it is no longer needed.
|
||||
pyobj = ctypes.py_object(f)
|
||||
ctypes.pythonapi.Py_IncRef(pyobj)
|
||||
check_call(_LIB.TVMFuncCreateFromCFunc(
|
||||
f, pyobj, TVM_FREE_PYOBJ, ctypes.byref(handle)))
|
||||
return _runtime_api._function_cls(handle)
|
||||
|
||||
|
||||
def convert(value):
|
||||
"""Convert a value to expression."""
|
||||
if isinstance(value, (NodeBase, _runtime_api.FunctionBase)):
|
||||
return value
|
||||
elif isinstance(value, Number):
|
||||
return const(value)
|
||||
elif isinstance(value, string_types):
|
||||
return _api_internal._str(value)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
value = [convert(x) for x in value]
|
||||
return _api_internal._Array(*value)
|
||||
elif isinstance(value, dict):
|
||||
vlist = []
|
||||
for it in value.items():
|
||||
if not isinstance(it[0], NodeBase):
|
||||
raise ValueError("key of map must already been a container type")
|
||||
vlist.append(it[0])
|
||||
vlist.append(convert(it[1]))
|
||||
return _api_internal._Map(*vlist)
|
||||
elif isinstance(value, SliceBase):
|
||||
return value.tensor(*value.indices)
|
||||
elif callable(value):
|
||||
return convert_to_tvm_func(value)
|
||||
else:
|
||||
raise ValueError("don't know how to handle type %s" % type(value))
|
||||
return value
|
||||
|
||||
|
||||
def _push_arg(arg):
|
||||
a = TVMValue()
|
||||
if arg is None:
|
||||
_LIB.TVMAPIPushStack(a, APIType.NULL)
|
||||
elif isinstance(arg, NodeBase):
|
||||
a.v_handle = arg.handle
|
||||
_LIB.TVMAPIPushStack(a, APIType.NODE_HANDLE)
|
||||
elif isinstance(arg, Integral):
|
||||
a.v_int64 = ctypes.c_int64(arg)
|
||||
_LIB.TVMAPIPushStack(a, APIType.INT)
|
||||
elif isinstance(arg, Number):
|
||||
a.v_double = ctypes.c_double(arg)
|
||||
_LIB.TVMAPIPushStack(a, APIType.FLOAT)
|
||||
elif isinstance(arg, string_types):
|
||||
a.v_str = c_str(arg)
|
||||
_LIB.TVMAPIPushStack(a, APIType.STR)
|
||||
else:
|
||||
raise TypeError("Don't know how to handle type %s" % type(arg))
|
||||
|
||||
|
||||
def _make_function(handle, name):
|
||||
"""Create an atomic symbol function by handle and funciton name."""
|
||||
real_name = ctypes.c_char_p()
|
||||
desc = ctypes.c_char_p()
|
||||
num_args = ctypes.c_int()
|
||||
arg_names = ctypes.POINTER(ctypes.c_char_p)()
|
||||
arg_types = ctypes.POINTER(ctypes.c_char_p)()
|
||||
arg_descs = ctypes.POINTER(ctypes.c_char_p)()
|
||||
ret_type = ctypes.c_char_p()
|
||||
|
||||
check_call(_LIB.TVMGetAPIFuncInfo(
|
||||
handle, ctypes.byref(real_name), ctypes.byref(desc),
|
||||
ctypes.byref(num_args),
|
||||
ctypes.byref(arg_names),
|
||||
ctypes.byref(arg_types),
|
||||
ctypes.byref(arg_descs),
|
||||
ctypes.byref(ret_type)))
|
||||
|
||||
param_str = ctypes2docstring(num_args, arg_names, arg_types, arg_descs)
|
||||
func_name = name
|
||||
desc = py_str(desc.value)
|
||||
|
||||
doc_str = ('%s\n\n' +
|
||||
'%s\n')
|
||||
doc_str = doc_str % (desc, param_str)
|
||||
arg_names = [py_str(arg_names[i]) for i in range(num_args.value)]
|
||||
|
||||
def func(*args):
|
||||
"""TVM function"""
|
||||
cargs = []
|
||||
for x in args:
|
||||
if isinstance(x, (list, tuple, dict, SliceBase)):
|
||||
cargs.append(convert(x))
|
||||
else:
|
||||
cargs.append(x)
|
||||
|
||||
for arg in cargs:
|
||||
_push_arg(arg)
|
||||
ret_val = TVMValue()
|
||||
ret_type_code = ctypes.c_int()
|
||||
check_call(_LIB.TVMAPIFuncCall(
|
||||
handle, ctypes.byref(ret_val), ctypes.byref(ret_type_code)))
|
||||
return RET_SWITCH[ret_type_code.value](ret_val)
|
||||
|
||||
func.__name__ = func_name
|
||||
func.__doc__ = doc_str
|
||||
return func
|
||||
|
||||
|
||||
def register_node(type_key=None):
|
||||
"""register node type
|
||||
|
||||
Parameters
|
||||
----------
|
||||
type_key : str or cls
|
||||
The type key of the node
|
||||
"""
|
||||
if isinstance(type_key, str):
|
||||
def register(cls):
|
||||
"""internal register function"""
|
||||
NODE_TYPE[type_key] = cls
|
||||
return cls
|
||||
return register
|
||||
else:
|
||||
cls = type_key
|
||||
NODE_TYPE[cls.__name__] = cls
|
||||
return cls
|
||||
|
||||
|
||||
def register_func(func_name, f=None):
|
||||
"""Register global function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_name : str or function
|
||||
The function name
|
||||
|
||||
f : function
|
||||
The function to be registered.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fregister : function
|
||||
Register function if f is not specified.
|
||||
"""
|
||||
if callable(func_name):
|
||||
f = func_name
|
||||
func_name = f.__name__
|
||||
|
||||
if not isinstance(func_name, str):
|
||||
raise ValueError("expect string function name")
|
||||
def register(myf):
|
||||
"""internal register function"""
|
||||
if not isinstance(myf, _runtime_api.FunctionBase):
|
||||
myf = convert_to_tvm_func(myf)
|
||||
check_call(_LIB.TVMFuncRegisterGlobal(
|
||||
c_str(func_name), myf.handle))
|
||||
if f:
|
||||
register(f)
|
||||
else:
|
||||
return register
|
||||
|
||||
|
||||
def get_global_func(name):
|
||||
"""Get a global function by name
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name of the global function
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : tvm.nd.Function
|
||||
The function to be returned.
|
||||
"""
|
||||
handle = FunctionHandle()
|
||||
check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))
|
||||
return _runtime_api._function_cls(handle)
|
||||
|
||||
|
||||
def _init_api_module(root_namespace):
|
||||
"""List and add all the functions to current module."""
|
||||
plist = ctypes.POINTER(ctypes.c_char_p)()
|
||||
size = ctypes.c_uint()
|
||||
|
||||
check_call(_LIB.TVMListAPIFuncNames(ctypes.byref(size),
|
||||
ctypes.byref(plist)))
|
||||
op_names = []
|
||||
for i in range(size.value):
|
||||
op_names.append(py_str(plist[i]))
|
||||
|
||||
module_obj = sys.modules["%s.api" % root_namespace]
|
||||
module_internal = sys.modules["%s._api_internal" % root_namespace]
|
||||
namespace_match = {
|
||||
"_make_": sys.modules["%s.make" % root_namespace],
|
||||
"_pass_": sys.modules["%s.ir_pass" % root_namespace],
|
||||
"_codegen_": sys.modules["%s.codegen" % root_namespace],
|
||||
"_schedule_": sys.modules["%s.schedule" % root_namespace]
|
||||
}
|
||||
|
||||
for name in op_names:
|
||||
hdl = APIFuncHandle()
|
||||
check_call(_LIB.TVMGetAPIFuncHandle(c_str(name), ctypes.byref(hdl)))
|
||||
fname = name
|
||||
target_module = module_internal if name.startswith('_') else module_obj
|
||||
for k, v in namespace_match.items():
|
||||
if name.startswith(k):
|
||||
fname = name[len(k):]
|
||||
target_module = v
|
||||
function = _make_function(hdl, fname)
|
||||
setattr(target_module, function.__name__, function)
|
||||
@@ -0,0 +1,250 @@
|
||||
# coding: utf-8
|
||||
# pylint: disable=invalid-name, protected-access
|
||||
"""Symbolic configuration API."""
|
||||
from __future__ import absolute_import
|
||||
|
||||
import ctypes
|
||||
import sys
|
||||
from numbers import Number, Integral
|
||||
|
||||
from .._base import _LIB, check_call
|
||||
from .._base import c_str, py_str, string_types
|
||||
from ._types import TVMValue, TypeCode, TVMType
|
||||
from ._types import TVMPackedCFunc, TVMCFuncFinalizer
|
||||
from ._types import RETURN_SWITCH, C_TO_PY_ARG_SWITCH
|
||||
from ._node import NodeBase, SliceBase, convert_to_node
|
||||
from ._ndarray import NDArrayBase
|
||||
|
||||
FunctionHandle = ctypes.c_void_p
|
||||
TVMRetValueHandle = ctypes.c_void_p
|
||||
|
||||
def _ctypes_free_resource(rhandle):
|
||||
"""callback to free resources when it it not needed."""
|
||||
pyobj = ctypes.cast(rhandle, ctypes.py_object)
|
||||
ctypes.pythonapi.Py_DecRef(pyobj)
|
||||
|
||||
# Global callback that is always alive
|
||||
TVM_FREE_PYOBJ = TVMCFuncFinalizer(_ctypes_free_resource)
|
||||
ctypes.pythonapi.Py_IncRef(ctypes.py_object(TVM_FREE_PYOBJ))
|
||||
|
||||
def convert_to_tvm_func(pyfunc):
|
||||
"""Convert a python function to TVM function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pyfunc : python function
|
||||
The python function to be converted.
|
||||
|
||||
Returns
|
||||
-------
|
||||
tvmfunc: tvm.nd.Function
|
||||
The converted tvm function.
|
||||
"""
|
||||
local_pyfunc = pyfunc
|
||||
def cfun(args, type_codes, num_args, ret, _):
|
||||
""" ctypes function """
|
||||
num_args = num_args.value if isinstance(num_args, ctypes.c_int) else num_args
|
||||
pyargs = [C_TO_PY_ARG_SWITCH[type_codes[i]](args[i]) for i in range(num_args)]
|
||||
rv = local_pyfunc(*pyargs)
|
||||
if rv is not None:
|
||||
if isinstance(rv, tuple):
|
||||
raise ValueError("PackedFunction can only support one reurn value")
|
||||
temp_args = []
|
||||
values, tcodes, _ = _make_tvm_args((rv,), temp_args)
|
||||
if not isinstance(ret, TVMRetValueHandle):
|
||||
ret = TVMRetValueHandle(ret)
|
||||
check_call(_LIB.TVMCFuncSetReturn(ret, values[0], ctypes.c_int(tcodes[0])))
|
||||
_ = temp_args
|
||||
_ = rv
|
||||
|
||||
handle = FunctionHandle()
|
||||
f = TVMPackedCFunc(cfun)
|
||||
# NOTE: We will need to use python-api to increase ref count of the f
|
||||
# TVM_FREE_PYOBJ will be called after it is no longer needed.
|
||||
pyobj = ctypes.py_object(f)
|
||||
ctypes.pythonapi.Py_IncRef(pyobj)
|
||||
check_call(_LIB.TVMFuncCreateFromCFunc(
|
||||
f, pyobj, TVM_FREE_PYOBJ, ctypes.byref(handle)))
|
||||
return Function(handle)
|
||||
|
||||
|
||||
def _make_tvm_args(args, temp_args):
|
||||
"""Pack arguments into c args tvm call accept"""
|
||||
num_args = len(args)
|
||||
values = (TVMValue * num_args)()
|
||||
type_codes = (ctypes.c_int * num_args)()
|
||||
for i, arg in enumerate(args):
|
||||
if arg is None:
|
||||
values[i].v_handle = None
|
||||
type_codes[i] = TypeCode.NULL
|
||||
elif isinstance(arg, NDArrayBase):
|
||||
values[i].v_handle = ctypes.cast(arg.handle, ctypes.c_void_p)
|
||||
type_codes[i] = TypeCode.ARRAY_HANDLE
|
||||
elif isinstance(arg, NodeBase):
|
||||
values[i].v_handle = arg.handle
|
||||
type_codes[i] = TypeCode.NODE_HANDLE
|
||||
elif isinstance(arg, Integral):
|
||||
values[i].v_int64 = arg
|
||||
type_codes[i] = TypeCode.INT
|
||||
elif isinstance(arg, Number):
|
||||
values[i].v_float64 = arg
|
||||
type_codes[i] = TypeCode.FLOAT
|
||||
elif isinstance(arg, TVMType):
|
||||
values[i].v_type = arg
|
||||
type_codes[i] = TypeCode.TVM_TYPE
|
||||
elif isinstance(arg, string_types):
|
||||
values[i].v_str = c_str(arg)
|
||||
type_codes[i] = TypeCode.STR
|
||||
elif isinstance(arg, (list, tuple, dict, SliceBase)):
|
||||
arg = convert_to_node(arg)
|
||||
values[i].v_handle = arg.handle
|
||||
type_codes[i] = TypeCode.NODE_HANDLE
|
||||
temp_args.append(arg)
|
||||
elif isinstance(arg, Function):
|
||||
values[i].v_handle = arg.handle
|
||||
type_codes[i] = TypeCode.FUNC_HANDLE
|
||||
elif callable(arg):
|
||||
arg = convert_to_tvm_func(arg)
|
||||
values[i].v_handle = arg.handle
|
||||
type_codes[i] = TypeCode.FUNC_HANDLE
|
||||
temp_args.append(arg)
|
||||
else:
|
||||
raise TypeError("Don't know how to handle type %s" % type(arg))
|
||||
return values, type_codes, num_args
|
||||
|
||||
|
||||
class Function(object):
|
||||
"""A function object at runtime."""
|
||||
__slots__ = ["handle", "is_global"]
|
||||
# pylint: disable=no-member
|
||||
def __init__(self, handle, is_global=False):
|
||||
"""Initialize the function with handle
|
||||
|
||||
Parameters
|
||||
----------
|
||||
handle : FunctionHandle
|
||||
the handle to the underlying function.
|
||||
|
||||
is_global : bool, optional
|
||||
Whether it is global function
|
||||
"""
|
||||
self.handle = handle
|
||||
self.is_global = is_global
|
||||
|
||||
def __del__(self):
|
||||
if not self.is_global:
|
||||
check_call(_LIB.TVMFuncFree(self.handle))
|
||||
|
||||
def __call__(self, *args):
|
||||
temp_args = []
|
||||
values, tcodes, num_args = _make_tvm_args(args, temp_args)
|
||||
ret_val = TVMValue()
|
||||
ret_tcode = ctypes.c_int()
|
||||
check_call(_LIB.TVMFuncCall(
|
||||
self.handle, values, tcodes, ctypes.c_int(num_args),
|
||||
ctypes.byref(ret_val), ctypes.byref(ret_tcode)))
|
||||
_ = temp_args
|
||||
_ = args
|
||||
return RETURN_SWITCH[ret_tcode.value](ret_val)
|
||||
|
||||
|
||||
def _handle_return_func(x):
|
||||
"""Return function"""
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, FunctionHandle):
|
||||
handle = FunctionHandle(handle)
|
||||
return Function(handle, False)
|
||||
|
||||
# setup return handle for function type
|
||||
RETURN_SWITCH[TypeCode.FUNC_HANDLE] = _handle_return_func
|
||||
|
||||
def register_func(func_name, f=None):
|
||||
"""Register global function
|
||||
|
||||
Parameters
|
||||
----------
|
||||
func_name : str or function
|
||||
The function name
|
||||
|
||||
f : function
|
||||
The function to be registered.
|
||||
|
||||
Returns
|
||||
-------
|
||||
fregister : function
|
||||
Register function if f is not specified.
|
||||
"""
|
||||
if callable(func_name):
|
||||
f = func_name
|
||||
func_name = f.__name__
|
||||
|
||||
if not isinstance(func_name, str):
|
||||
raise ValueError("expect string function name")
|
||||
def register(myf):
|
||||
"""internal register function"""
|
||||
if not isinstance(myf, Function):
|
||||
myf = convert_to_tvm_func(myf)
|
||||
check_call(_LIB.TVMFuncRegisterGlobal(
|
||||
c_str(func_name), myf.handle))
|
||||
if f:
|
||||
register(f)
|
||||
else:
|
||||
return register
|
||||
|
||||
|
||||
def get_global_func(name):
|
||||
"""Get a global function by name
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
The name of the global function
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : tvm.nd.Function
|
||||
The function to be returned.
|
||||
"""
|
||||
handle = FunctionHandle()
|
||||
check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))
|
||||
return Function(handle, True)
|
||||
|
||||
|
||||
def list_global_func_names():
|
||||
"""Get list of global functions registered.
|
||||
|
||||
Returns
|
||||
-------
|
||||
names : list
|
||||
List of global functions names.
|
||||
"""
|
||||
plist = ctypes.POINTER(ctypes.c_char_p)()
|
||||
size = ctypes.c_uint()
|
||||
|
||||
check_call(_LIB.TVMFuncListGlobalNames(ctypes.byref(size),
|
||||
ctypes.byref(plist)))
|
||||
fnames = []
|
||||
for i in range(size.value):
|
||||
fnames.append(py_str(plist[i]))
|
||||
return fnames
|
||||
|
||||
|
||||
def _init_api_functions(root_namespace):
|
||||
"""List and add all the functions to current module."""
|
||||
module_obj = sys.modules["%s.api" % root_namespace]
|
||||
module_internal = sys.modules["%s._api_internal" % root_namespace]
|
||||
namespace_match = {
|
||||
"_make_": sys.modules["%s.make" % root_namespace],
|
||||
"_pass_": sys.modules["%s.ir_pass" % root_namespace],
|
||||
"_codegen_": sys.modules["%s.codegen" % root_namespace],
|
||||
"_schedule_": sys.modules["%s.schedule" % root_namespace]
|
||||
}
|
||||
for name in list_global_func_names():
|
||||
fname = name
|
||||
target_module = module_internal if name.startswith('_') else module_obj
|
||||
for k, v in namespace_match.items():
|
||||
if name.startswith(k):
|
||||
fname = name[len(k):]
|
||||
target_module = v
|
||||
f = get_global_func(name)
|
||||
setattr(target_module, fname, f)
|
||||
@@ -1,18 +1,15 @@
|
||||
# pylint: disable=invalid-name, protected-access, too-many-arguments, global-statement
|
||||
# pylint: disable=attribute-defined-outside-init, no-member, missing-docstring
|
||||
"""Symbolic configuration API."""
|
||||
from __future__ import absolute_import as _abs
|
||||
from __future__ import absolute_import
|
||||
|
||||
import ctypes
|
||||
from numbers import Number, Integral
|
||||
import numpy as np
|
||||
|
||||
from .._base import _LIB
|
||||
from .._base import c_array, c_str, string_types
|
||||
from .._base import check_call
|
||||
from ._types import TVMValue, TypeCode, TVMType
|
||||
from .._base import _LIB, check_call
|
||||
from .._base import c_array, c_str
|
||||
from ._types import TVMType, tvm_index_t
|
||||
|
||||
tvm_index_t = ctypes.c_uint32
|
||||
|
||||
class TVMContext(ctypes.Structure):
|
||||
"""TVM context strucure."""
|
||||
@@ -39,6 +36,19 @@ class TVMContext(ctypes.Structure):
|
||||
return ret.value != 0
|
||||
|
||||
|
||||
class TVMArray(ctypes.Structure):
|
||||
"""TVMValue in C API"""
|
||||
_fields_ = [("data", ctypes.c_void_p),
|
||||
("shape", ctypes.POINTER(tvm_index_t)),
|
||||
("strides", ctypes.POINTER(tvm_index_t)),
|
||||
("ndim", tvm_index_t),
|
||||
("dtype", TVMType),
|
||||
("ctx", TVMContext)]
|
||||
|
||||
|
||||
TVMArrayHandle = ctypes.POINTER(TVMArray)
|
||||
|
||||
|
||||
def cpu(dev_id=0):
|
||||
"""Construct a CPU device
|
||||
|
||||
@@ -72,18 +82,6 @@ def opencl(dev_id=0):
|
||||
return TVMContext(4, dev_id)
|
||||
|
||||
|
||||
class TVMArray(ctypes.Structure):
|
||||
"""TVMValue in C API"""
|
||||
_fields_ = [("data", ctypes.c_void_p),
|
||||
("shape", ctypes.POINTER(tvm_index_t)),
|
||||
("strides", ctypes.POINTER(tvm_index_t)),
|
||||
("ndim", tvm_index_t),
|
||||
("dtype", TVMType),
|
||||
("ctx", TVMContext)]
|
||||
|
||||
TVMArrayHandle = ctypes.POINTER(TVMArray)
|
||||
|
||||
|
||||
def numpyasarray(np_data):
|
||||
"""Return a TVMArray representation of a numpy array.
|
||||
"""
|
||||
@@ -102,7 +100,6 @@ def numpyasarray(np_data):
|
||||
|
||||
|
||||
_ndarray_cls = None
|
||||
_function_cls = None
|
||||
|
||||
|
||||
def empty(shape, dtype="float32", ctx=cpu(0)):
|
||||
@@ -275,51 +272,6 @@ class NDArrayBase(object):
|
||||
return target
|
||||
|
||||
|
||||
class FunctionBase(object):
|
||||
"""A function object at runtim."""
|
||||
__slots__ = ["handle"]
|
||||
# pylint: disable=no-member
|
||||
def __init__(self, handle):
|
||||
"""Initialize the function with handle
|
||||
|
||||
Parameters
|
||||
----------
|
||||
handle : FunctionHandle
|
||||
the handle to the underlying function.
|
||||
"""
|
||||
self.handle = handle
|
||||
|
||||
def __del__(self):
|
||||
check_call(_LIB.TVMFuncFree(self.handle))
|
||||
|
||||
def __call__(self, *args):
|
||||
num_args = len(args)
|
||||
tvm_args = (TVMValue * num_args)()
|
||||
tvm_type_code = (ctypes.c_int * num_args)()
|
||||
for i, arg in enumerate(args):
|
||||
if arg is None:
|
||||
tvm_args[i].v_handle = None
|
||||
tvm_type_code[i] = TypeCode.NULL
|
||||
elif isinstance(arg, NDArrayBase):
|
||||
tvm_args[i].v_handle = ctypes.cast(arg.handle, ctypes.c_void_p)
|
||||
tvm_type_code[i] = TypeCode.HANDLE
|
||||
elif isinstance(arg, Integral):
|
||||
tvm_args[i].v_int64 = arg
|
||||
tvm_type_code[i] = TypeCode.INT
|
||||
elif isinstance(arg, Number):
|
||||
tvm_args[i].v_float64 = arg
|
||||
tvm_type_code[i] = TypeCode.FLOAT
|
||||
elif isinstance(arg, string_types):
|
||||
tvm_args[i].v_str = c_str(arg)
|
||||
tvm_type_code[i] = TypeCode.STR
|
||||
else:
|
||||
raise TypeError("Don't know how to handle type %s" % type(arg))
|
||||
check_call(_LIB.TVMFuncCall(
|
||||
self.handle, tvm_args, tvm_type_code, ctypes.c_int(num_args)))
|
||||
|
||||
|
||||
def _init_runtime_module(ndarray_class, function_class):
|
||||
def _init_ndarray_module(ndarray_class):
|
||||
global _ndarray_cls
|
||||
global _function_cls
|
||||
_ndarray_cls = ndarray_class
|
||||
_function_cls = function_class
|
||||
@@ -0,0 +1,184 @@
|
||||
# coding: utf-8
|
||||
# pylint: disable=invalid-name, protected-access
|
||||
# pylint: disable=no-member, missing-docstring
|
||||
"""Symbolic configuration API."""
|
||||
from __future__ import absolute_import
|
||||
|
||||
import ctypes
|
||||
from numbers import Number, Integral
|
||||
|
||||
from .._base import _LIB, check_call
|
||||
from .._base import c_str, py_str, string_types
|
||||
from .. import _api_internal
|
||||
from ._types import TVMValue, TypeCode, RETURN_SWITCH
|
||||
|
||||
NodeHandle = ctypes.c_void_p
|
||||
|
||||
"""Maps node type to its constructor"""
|
||||
NODE_TYPE = {
|
||||
}
|
||||
|
||||
def _return_node(x):
|
||||
"""Return function"""
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, NodeHandle):
|
||||
handle = NodeHandle(handle)
|
||||
ret_val = TVMValue()
|
||||
ret_type_code = ctypes.c_int()
|
||||
ret_success = ctypes.c_int()
|
||||
check_call(_LIB.TVMNodeGetAttr(
|
||||
handle, c_str("type_key"),
|
||||
ctypes.byref(ret_val),
|
||||
ctypes.byref(ret_type_code),
|
||||
ctypes.byref(ret_success)))
|
||||
return NODE_TYPE.get(py_str(ret_val.v_str), NodeBase)(handle)
|
||||
|
||||
|
||||
RETURN_SWITCH[TypeCode.NODE_HANDLE] = _return_node
|
||||
|
||||
|
||||
class SliceBase(object):
|
||||
"""base class of slice object"""
|
||||
pass
|
||||
|
||||
class NodeBase(object):
|
||||
"""Symbol is symbolic graph."""
|
||||
__slots__ = ["handle"]
|
||||
# pylint: disable=no-member
|
||||
def __init__(self, handle):
|
||||
"""Initialize the function with handle
|
||||
|
||||
Parameters
|
||||
----------
|
||||
handle : SymbolHandle
|
||||
the handle to the underlying C++ Symbol
|
||||
"""
|
||||
self.handle = handle
|
||||
|
||||
def __repr__(self):
|
||||
return _api_internal._format_str(self)
|
||||
|
||||
def __del__(self):
|
||||
check_call(_LIB.TVMNodeFree(self.handle))
|
||||
|
||||
def __getattr__(self, name):
|
||||
ret_val = TVMValue()
|
||||
ret_type_code = ctypes.c_int()
|
||||
ret_success = ctypes.c_int()
|
||||
check_call(_LIB.TVMNodeGetAttr(
|
||||
self.handle, c_str(name),
|
||||
ctypes.byref(ret_val),
|
||||
ctypes.byref(ret_type_code),
|
||||
ctypes.byref(ret_success)))
|
||||
if not ret_success.value:
|
||||
raise AttributeError(
|
||||
"'%s' object has no attribute '%s'" % (str(type(self)), name))
|
||||
return RETURN_SWITCH[ret_type_code.value](ret_val)
|
||||
|
||||
def __hash__(self):
|
||||
return _api_internal._raw_ptr(self)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, NodeBase):
|
||||
return False
|
||||
return self.__hash__() == other.__hash__()
|
||||
|
||||
def __ne__(self, other):
|
||||
return not self.__eq__(other)
|
||||
|
||||
def __dir__(self):
|
||||
plist = ctypes.POINTER(ctypes.c_char_p)()
|
||||
size = ctypes.c_uint()
|
||||
check_call(_LIB.TVMNodeListAttrNames(
|
||||
self.handle, ctypes.byref(size), ctypes.byref(plist)))
|
||||
names = []
|
||||
for i in range(size.value):
|
||||
names.append(py_str(plist[i]))
|
||||
return names
|
||||
|
||||
def __reduce__(self):
|
||||
return (type(self), (None,), self.__getstate__())
|
||||
|
||||
def __getstate__(self):
|
||||
handle = self.handle
|
||||
if handle is not None:
|
||||
return {'handle': _api_internal._save_json(self)}
|
||||
else:
|
||||
return {'handle': None}
|
||||
|
||||
def __setstate__(self, state):
|
||||
# pylint: disable=assigning-non-slot
|
||||
handle = state['handle']
|
||||
if handle is not None:
|
||||
json_str = handle
|
||||
other = _api_internal._load_json(json_str)
|
||||
self.handle = other.handle
|
||||
other.handle = None
|
||||
else:
|
||||
self.handle = None
|
||||
|
||||
|
||||
def const(value, dtype=None):
|
||||
"""construct a constant"""
|
||||
if dtype is None:
|
||||
if isinstance(value, Integral):
|
||||
dtype = 'int32'
|
||||
else:
|
||||
dtype = 'float32'
|
||||
return _api_internal._const(value, dtype)
|
||||
|
||||
|
||||
def convert_to_node(value):
|
||||
"""Convert a python value to corresponding node type.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : str
|
||||
The value to be inspected.
|
||||
|
||||
Returns
|
||||
-------
|
||||
node : Node
|
||||
The corresponding node value.
|
||||
"""
|
||||
if isinstance(value, NodeBase):
|
||||
return value
|
||||
elif isinstance(value, Number):
|
||||
return const(value)
|
||||
elif isinstance(value, string_types):
|
||||
return _api_internal._str(value)
|
||||
elif isinstance(value, (list, tuple)):
|
||||
value = [convert_to_node(x) for x in value]
|
||||
return _api_internal._Array(*value)
|
||||
elif isinstance(value, dict):
|
||||
vlist = []
|
||||
for it in value.items():
|
||||
if not isinstance(it[0], NodeBase):
|
||||
raise ValueError("key of map must already been a container type")
|
||||
vlist.append(it[0])
|
||||
vlist.append(convert_to_node(it[1]))
|
||||
return _api_internal._Map(*vlist)
|
||||
elif isinstance(value, SliceBase):
|
||||
return value.tensor(*value.indices)
|
||||
else:
|
||||
raise ValueError("don't know how to convert type %s to node" % type(value))
|
||||
|
||||
|
||||
def register_node(type_key=None):
|
||||
"""register node type
|
||||
|
||||
Parameters
|
||||
----------
|
||||
type_key : str or cls
|
||||
The type key of the node
|
||||
"""
|
||||
if isinstance(type_key, str):
|
||||
def register(cls):
|
||||
"""internal register function"""
|
||||
NODE_TYPE[type_key] = cls
|
||||
return cls
|
||||
return register
|
||||
else:
|
||||
cls = type_key
|
||||
NODE_TYPE[cls.__name__] = cls
|
||||
return cls
|
||||
@@ -4,13 +4,9 @@ from __future__ import absolute_import as _abs
|
||||
|
||||
import ctypes
|
||||
import numpy as np
|
||||
from .._base import py_str
|
||||
|
||||
class TVMValue(ctypes.Union):
|
||||
"""TVMValue in C API"""
|
||||
_fields_ = [("v_int64", ctypes.c_int64),
|
||||
("v_float64", ctypes.c_double),
|
||||
("v_handle", ctypes.c_void_p),
|
||||
("v_str", ctypes.c_char_p)]
|
||||
tvm_index_t = ctypes.c_uint32
|
||||
|
||||
class TypeCode(object):
|
||||
"""Type code used in API calls"""
|
||||
@@ -19,9 +15,11 @@ class TypeCode(object):
|
||||
FLOAT = 2
|
||||
HANDLE = 3
|
||||
NULL = 4
|
||||
NODE_HANDLE = 5
|
||||
STR = 6
|
||||
FUNC_HANDLE = 7
|
||||
ARRAY_HANDLE = 5
|
||||
TVM_TYPE = 6
|
||||
NODE_HANDLE = 7
|
||||
STR = 8
|
||||
FUNC_HANDLE = 9
|
||||
|
||||
def _api_type(code):
|
||||
"""create a type accepted by API"""
|
||||
@@ -40,13 +38,13 @@ class TVMType(ctypes.Structure):
|
||||
CODE2STR = {
|
||||
0 : 'int',
|
||||
1 : 'uint',
|
||||
2 : 'float'
|
||||
2 : 'float',
|
||||
4 : 'handle'
|
||||
}
|
||||
def __init__(self, type_str, lanes=1):
|
||||
super(TVMType, self).__init__()
|
||||
if isinstance(type_str, np.dtype):
|
||||
type_str = str(type_str)
|
||||
|
||||
if type_str.startswith("int"):
|
||||
self.type_code = 0
|
||||
bits = int(type_str[3:])
|
||||
@@ -56,6 +54,9 @@ class TVMType(ctypes.Structure):
|
||||
elif type_str.startswith("float"):
|
||||
self.type_code = 2
|
||||
bits = int(type_str[5:])
|
||||
elif type_str.startswith("handle"):
|
||||
self.type_code = 4
|
||||
bits = 64
|
||||
else:
|
||||
raise ValueError("Donot know how to handle type %s" % type_str)
|
||||
|
||||
@@ -71,15 +72,61 @@ class TVMType(ctypes.Structure):
|
||||
x += "x%d" % self.lanes
|
||||
return x
|
||||
|
||||
def __eq__(self, other):
|
||||
return (self.bits == other.bits and
|
||||
self.type_code == other.type_code and
|
||||
self.lanes == other.lanes)
|
||||
|
||||
def __ne__(self, other):
|
||||
return not self.__eq__(other)
|
||||
|
||||
|
||||
class TVMValue(ctypes.Union):
|
||||
"""TVMValue in C API"""
|
||||
_fields_ = [("v_int64", ctypes.c_int64),
|
||||
("v_float64", ctypes.c_double),
|
||||
("v_handle", ctypes.c_void_p),
|
||||
("v_str", ctypes.c_char_p),
|
||||
("v_type", TVMType)]
|
||||
|
||||
|
||||
TVMPackedCFunc = ctypes.CFUNCTYPE(
|
||||
None,
|
||||
ctypes.POINTER(TVMValue),
|
||||
ctypes.POINTER(ctypes.c_int),
|
||||
ctypes.c_int,
|
||||
ctypes.c_void_p,
|
||||
ctypes.c_void_p)
|
||||
|
||||
|
||||
TVMCFuncFinalizer = ctypes.CFUNCTYPE(
|
||||
None,
|
||||
ctypes.c_void_p)
|
||||
|
||||
|
||||
def _return_handle(x):
|
||||
"""return handle"""
|
||||
handle = x.v_handle
|
||||
if not isinstance(handle, ctypes.c_void_p):
|
||||
handle = ctypes.c_void_p(handle)
|
||||
return handle
|
||||
|
||||
|
||||
RETURN_SWITCH = {
|
||||
TypeCode.INT: lambda x: x.v_int64,
|
||||
TypeCode.FLOAT: lambda x: x.v_float64,
|
||||
TypeCode.HANDLE: _return_handle,
|
||||
TypeCode.NULL: lambda x: None,
|
||||
TypeCode.TVM_TYPE: lambda x: x.v_type,
|
||||
TypeCode.STR: lambda x: py_str(x.v_str)
|
||||
}
|
||||
|
||||
|
||||
C_TO_PY_ARG_SWITCH = {
|
||||
TypeCode.INT: lambda x: x.v_int64,
|
||||
TypeCode.FLOAT: lambda x: x.v_float64,
|
||||
TypeCode.HANDLE: _return_handle,
|
||||
TypeCode.NULL: lambda x: None,
|
||||
TypeCode.TVM_TYPE: lambda x: x.v_type,
|
||||
TypeCode.STR: lambda x: py_str(x.v_str)
|
||||
}
|
||||
|
||||
+33
-5
@@ -2,16 +2,23 @@
|
||||
# pylint: disable=redefined-builtin, undefined-variable, unused-import
|
||||
"""Functions defined in TVM."""
|
||||
from __future__ import absolute_import as _abs
|
||||
|
||||
from numbers import Integral as _Integral
|
||||
from ._ctypes._api import _init_api_module, convert, register_func, get_global_func
|
||||
|
||||
from ._ctypes._types import TVMType
|
||||
from ._ctypes._node import register_node, NodeBase
|
||||
from ._ctypes._node import convert_to_node as _convert_to_node
|
||||
from ._ctypes._function import Function
|
||||
from ._ctypes._function import _init_api_functions, register_func, get_global_func
|
||||
from ._ctypes._function import convert_to_tvm_func as _convert_tvm_func
|
||||
from . import _api_internal
|
||||
from . import make as _make
|
||||
from . import expr as _expr
|
||||
from . import collections as _collections
|
||||
|
||||
int32 = "int32"
|
||||
float32 = "float32"
|
||||
handle = "handle"
|
||||
int32 = TVMType("int32")
|
||||
float32 = TVMType("float32")
|
||||
handle = TVMType("handle")
|
||||
|
||||
def const(value, dtype=None):
|
||||
"""construct a constant"""
|
||||
@@ -266,4 +273,25 @@ def Schedule(ops):
|
||||
return _api_internal._Schedule(ops)
|
||||
|
||||
|
||||
_init_api_module("tvm")
|
||||
def convert(value):
|
||||
"""Convert value to TVM node or function.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
value : python value
|
||||
|
||||
Returns
|
||||
-------
|
||||
tvm_val : Node or function
|
||||
Converted value in TVM
|
||||
"""
|
||||
if isinstance(value, (Function, NodeBase)):
|
||||
return value
|
||||
|
||||
if callable(value):
|
||||
return _convert_tvm_func(value)
|
||||
else:
|
||||
return _convert_to_node(value)
|
||||
|
||||
|
||||
_init_api_functions("tvm")
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# pylint: disable=protected-access, no-member
|
||||
"""Collection structure in the high level DSL."""
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import NodeBase, register_node
|
||||
from ._ctypes._node import NodeBase, register_node
|
||||
from . import _api_internal
|
||||
from . import expr as _expr
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
# pylint: disable=protected-access, no-member, missing-docstring
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import NodeBase, register_node
|
||||
from ._ctypes._node import NodeBase, register_node
|
||||
from . import make as _make
|
||||
|
||||
class ExprOp(object):
|
||||
|
||||
+6
-11
@@ -6,11 +6,11 @@ This is a simplified runtime API for quick testing and proptyping.
|
||||
from __future__ import absolute_import as _abs
|
||||
import numpy as _np
|
||||
|
||||
from ._ctypes._runtime_api import TVMContext, TVMType, NDArrayBase, FunctionBase
|
||||
from ._ctypes._runtime_api import cpu, gpu, opencl, empty, sync
|
||||
from ._ctypes._runtime_api import _init_runtime_module
|
||||
from ._ctypes._runtime_api import init_opencl
|
||||
|
||||
from ._ctypes._ndarray import TVMContext, TVMType, NDArrayBase
|
||||
from ._ctypes._ndarray import cpu, gpu, opencl, empty, sync
|
||||
from ._ctypes._ndarray import _init_ndarray_module
|
||||
from ._ctypes._ndarray import init_opencl
|
||||
from ._ctypes._function import Function
|
||||
|
||||
class NDArray(NDArrayBase):
|
||||
"""Lightweight NDArray class of TVM runtime.
|
||||
@@ -26,11 +26,6 @@ class NDArray(NDArrayBase):
|
||||
pass
|
||||
|
||||
|
||||
class Function(FunctionBase):
|
||||
"""Function class that can executed a generated code."""
|
||||
pass
|
||||
|
||||
|
||||
def array(arr, ctx=cpu(0)):
|
||||
"""Create an array from source arr.
|
||||
|
||||
@@ -54,4 +49,4 @@ def array(arr, ctx=cpu(0)):
|
||||
return ret
|
||||
|
||||
|
||||
_init_runtime_module(NDArray, Function)
|
||||
_init_ndarray_module(NDArray)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# pylint: disable=protected-access, no-member
|
||||
"""Collection structure in the high level DSL."""
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import NodeBase, register_node
|
||||
from ._ctypes._node import NodeBase, register_node
|
||||
from . import _api_internal
|
||||
from . import tensor as _tensor
|
||||
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
# pylint: disable=protected-access, no-member, missing-docstring
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import NodeBase, register_node
|
||||
from ._ctypes._node import NodeBase, register_node
|
||||
|
||||
class Stmt(NodeBase):
|
||||
pass
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# pylint: disable=protected-access, no-member, invalid-name
|
||||
"""Tensor related abstractions"""
|
||||
from __future__ import absolute_import as _abs
|
||||
from ._ctypes._api import NodeBase, SliceBase, register_node, convert
|
||||
from ._ctypes._node import NodeBase, SliceBase, register_node, convert_to_node
|
||||
from . import collections as _collections
|
||||
from . import _api_internal
|
||||
from . import make as _make
|
||||
@@ -26,7 +26,7 @@ class Tensor(NodeBase):
|
||||
ndim = self.ndim
|
||||
if len(indices) != ndim:
|
||||
raise ValueError("Need to provide %d index in tensor slice" % ndim)
|
||||
indices = convert(indices)
|
||||
indices = convert_to_node(indices)
|
||||
args = []
|
||||
for x in indices:
|
||||
if isinstance(x, _collections.IterVar):
|
||||
|
||||
Reference in New Issue
Block a user