function.py 8.95 KB
Newer Older
1 2
# pylint: disable=invalid-name, unused-import
"""Function namespace."""
3 4 5
from __future__ import absolute_import

import sys
6
import ctypes
7
from .base import _LIB, check_call, py_str, c_str, string_types, _FFI_MODE
8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27

IMPORT_EXCEPT = RuntimeError if _FFI_MODE == "cython" else ImportError

try:
    # pylint: disable=wrong-import-position
    if _FFI_MODE == "ctypes":
        raise ImportError()
    if sys.version_info >= (3, 0):
        from ._cy3.core import _set_class_function, _set_class_module
        from ._cy3.core import FunctionBase as _FunctionBase
        from ._cy3.core import convert_to_tvm_func
    else:
        from ._cy2.core import _set_class_function, _set_class_module
        from ._cy2.core import FunctionBase as _FunctionBase
        from ._cy2.core import convert_to_tvm_func
except IMPORT_EXCEPT:
    # pylint: disable=wrong-import-position
    from ._ctypes.function import _set_class_function, _set_class_module
    from ._ctypes.function import FunctionBase as _FunctionBase
    from ._ctypes.function import convert_to_tvm_func
28 29 30

FunctionHandle = ctypes.c_void_p

31
class Function(_FunctionBase):
32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
    """The PackedFunc object used in TVM.

    Function plays an key role to bridge front and backend in TVM.
    Function provide a type-erased interface, you can call function with positional arguments.

    The compiled module returns Function.
    TVM backend also registers and exposes its API as Functions.
    For example, the developer function exposed in tvm.ir_pass are actually
    C++ functions that are registered as PackedFunc

    The following are list of common usage scenario of tvm.Function.

    - Automatic exposure of C++ API into python
    - To call PackedFunc from python side
    - To call python callbacks to inspect results in generated code
    - Bring python hook into C++ backend

    See Also
    --------
    tvm.register_func: How to register global function.
    tvm.get_global_func: How to get global function.
53
    """
54
    pass
55 56


57 58
class ModuleBase(object):
    """Base class for module"""
59 60
    __slots__ = ["handle", "_entry", "entry_name"]

61 62 63
    def __init__(self, handle):
        self.handle = handle
        self._entry = None
64
        self.entry_name = "__tvm_main__"
65

66 67 68
    def __del__(self):
        check_call(_LIB.TVMModFree(self.handle))

69 70 71 72 73 74 75 76 77 78 79
    @property
    def entry_func(self):
        """Get the entry function

        Returns
        -------
        f : Function
            The entry function if exist
        """
        if self._entry:
            return self._entry
80
        self._entry = self.get_function(self.entry_name)
81
        return self._entry
82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106

    def get_function(self, name, query_imports=False):
        """Get function from the module.

        Parameters
        ----------
        name : str
            The name of the function

        query_imports : bool
            Whether also query modules imported by this module.

        Returns
        -------
        f : Function
            The result function.
        """
        ret_handle = FunctionHandle()
        check_call(_LIB.TVMModGetFunction(
            self.handle, c_str(name),
            ctypes.c_int(query_imports),
            ctypes.byref(ret_handle)))
        if not ret_handle.value:
            raise AttributeError(
                "Module has no function '%s'" %  name)
107
        return Function(ret_handle, False)
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126

    def import_module(self, module):
        """Add module to the import list of current one.

        Parameters
        ----------
        module : Module
            The other module.
        """
        check_call(_LIB.TVMModImport(self.handle, module.handle))

    def __getitem__(self, name):
        if not isinstance(name, string_types):
            raise ValueError("Can only take string as function name")
        return self.get_function(name)

    def __call__(self, *args):
        if self._entry:
            return self._entry(*args)
127 128
        f = self.entry_func
        return f(*args)
129

130

131
def register_func(func_name, f=None, override=False):
132 133 134 135 136 137 138
    """Register global function

    Parameters
    ----------
    func_name : str or function
        The function name

139
    f : function, optional
140 141
        The function to be registered.

142 143 144
    override: boolean optional
        Whether override existing entry.

145 146 147 148
    Returns
    -------
    fregister : function
        Register function if f is not specified.
149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168

    Examples
    --------
    The following code registers my_packed_func as global function.
    Note that we simply get it back from global function table to invoke
    it from python side. However, we can also invoke the same function
    from C++ backend, or in the compiled TVM code.

    .. code-block:: python

      targs = (10, 10.0, "hello")
      @tvm.register_func
      def my_packed_func(*args):
          assert(tuple(args) == targs)
          return 10
      # Get it out from global function table
      f = tvm.get_global_func("my_packed_func")
      assert isinstance(f, tvm.nd.Function)
      y = f(*targs)
      assert y == 10
169 170 171 172 173 174 175
    """
    if callable(func_name):
        f = func_name
        func_name = f.__name__

    if not isinstance(func_name, str):
        raise ValueError("expect string function name")
176 177

    ioverride = ctypes.c_int(override)
178 179 180 181 182
    def register(myf):
        """internal register function"""
        if not isinstance(myf, Function):
            myf = convert_to_tvm_func(myf)
        check_call(_LIB.TVMFuncRegisterGlobal(
183
            c_str(func_name), myf.handle, ioverride))
184
        return myf
185
    if f:
186 187
        return register(f)
    return register
188 189


190
def get_global_func(name, allow_missing=False):
191 192 193 194 195 196 197
    """Get a global function by name

    Parameters
    ----------
    name : str
        The name of the global function

198 199 200
    allow_missing : bool
        Whether allow missing function or raise an error.

201 202
    Returns
    -------
203
    func : tvm.Function
204
        The function to be returned, None if function is missing.
205 206 207
    """
    handle = FunctionHandle()
    check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))
208 209 210 211 212 213 214 215
    if handle.value:
        return Function(handle, False)
    else:
        if allow_missing:
            return None
        else:
            raise ValueError("Cannot find global function %s" % name)

216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236


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


237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261
def extract_ext_funcs(finit):
    """
    Extract the extension PackedFuncs from a C module.

    Parameters
    ----------
    finit : ctypes function
        a ctypes that takes signature of TVMExtensionDeclarer

    Returns
    -------
    fdict : dict of str to Function
        The extracted functions
    """
    fdict = {}
    def _list(name, func):
        fdict[name] = func
    myf = convert_to_tvm_func(_list)
    ret = finit(myf.handle)
    _ = myf
    if ret != 0:
        raise RuntimeError("cannot initialize with %s" % finit)
    return fdict


262 263
def _get_api(f):
    flocal = f
264
    flocal.is_global = True
265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
    def my_api_func(*args):
        """

        This is a type erased API that calls into Global PackedFunc.
        These APIs corresponds to functions registered from C++ backend
        and can be used as developer functions.

        args : list
          The positional arguments to the function call.

        Returns
        -------
        value : int, float, None, Node or Function
        The result of the API function call.
        """
        return flocal(*args)
    return my_api_func

283
def _init_api(namespace, target_module_name=None):
284 285
    """Initialize api for a given module name

286 287 288 289 290
    namespace : str
       The namespace of the source registry

    target_module_name : str
       The target module name if different from namespace
291
    """
292 293 294 295 296 297 298
    target_module_name = (
        target_module_name if target_module_name else namespace)
    if namespace.startswith("tvm."):
        _init_api_prefix(target_module_name, namespace[4:])
    else:
        _init_api_prefix(target_module_name, namespace)

299 300 301

def _init_api_prefix(module_name, prefix):
    module = sys.modules[module_name]
302

303
    for name in list_global_func_names():
304 305 306 307 308 309
        if prefix == "api":
            fname = name
            if name.startswith("_"):
                target_module = sys.modules["tvm._api_internal"]
            else:
                target_module = module
310
        else:
311 312 313
            if not name.startswith(prefix):
                continue
            fname = name[len(prefix)+1:]
314 315
            target_module = module

316 317
        if fname.find(".") != -1:
            continue
318
        f = get_global_func(name)
319 320 321 322
        ff = _get_api(f)
        ff.__name__ = fname
        ff.__doc__ = ("TVM PackedFunc %s. " % fname)
        setattr(target_module, ff.__name__, ff)
323

324
_set_class_function(Function)