function.py 9.23 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

18 19
# pylint: disable=invalid-name, unused-import
"""Function namespace."""
20 21 22
from __future__ import absolute_import

import sys
23
import ctypes
24
from .base import _LIB, check_call, py_str, c_str, string_types, _FFI_MODE
25
from . import vmobj as _vmobj
26 27 28 29 30 31 32 33

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):
34
        from ._cy3.core import _set_class_function, _set_class_module
35 36 37
        from ._cy3.core import FunctionBase as _FunctionBase
        from ._cy3.core import convert_to_tvm_func
    else:
38
        from ._cy2.core import _set_class_function, _set_class_module
39 40 41 42
        from ._cy2.core import FunctionBase as _FunctionBase
        from ._cy2.core import convert_to_tvm_func
except IMPORT_EXCEPT:
    # pylint: disable=wrong-import-position
43
    from ._ctypes.function import _set_class_function, _set_class_module
44 45
    from ._ctypes.function import FunctionBase as _FunctionBase
    from ._ctypes.function import convert_to_tvm_func
46 47 48

FunctionHandle = ctypes.c_void_p

49
class Function(_FunctionBase):
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
    """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.
71
    """
72 73


74 75
class ModuleBase(object):
    """Base class for module"""
76 77
    __slots__ = ["handle", "_entry", "entry_name"]

78 79 80
    def __init__(self, handle):
        self.handle = handle
        self._entry = None
81
        self.entry_name = "__tvm_main__"
82

83 84 85
    def __del__(self):
        check_call(_LIB.TVMModFree(self.handle))

86 87 88 89 90 91 92 93 94 95 96
    @property
    def entry_func(self):
        """Get the entry function

        Returns
        -------
        f : Function
            The entry function if exist
        """
        if self._entry:
            return self._entry
97
        self._entry = self.get_function(self.entry_name)
98
        return self._entry
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123

    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)
124
        return Function(ret_handle, False)
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143

    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)
144 145
        f = self.entry_func
        return f(*args)
146

147

148
def register_func(func_name, f=None, override=False):
149 150 151 152 153 154 155
    """Register global function

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

156
    f : function, optional
157 158
        The function to be registered.

159 160 161
    override: boolean optional
        Whether override existing entry.

162 163 164 165
    Returns
    -------
    fregister : function
        Register function if f is not specified.
166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185

    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
186 187 188 189 190 191 192
    """
    if callable(func_name):
        f = func_name
        func_name = f.__name__

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

    ioverride = ctypes.c_int(override)
195 196 197 198 199
    def register(myf):
        """internal register function"""
        if not isinstance(myf, Function):
            myf = convert_to_tvm_func(myf)
        check_call(_LIB.TVMFuncRegisterGlobal(
200
            c_str(func_name), myf.handle, ioverride))
201
        return myf
202
    if f:
203 204
        return register(f)
    return register
205 206


207
def get_global_func(name, allow_missing=False):
208 209 210 211 212 213 214
    """Get a global function by name

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

215 216 217
    allow_missing : bool
        Whether allow missing function or raise an error.

218 219
    Returns
    -------
220
    func : tvm.Function
221
        The function to be returned, None if function is missing.
222 223 224
    """
    handle = FunctionHandle()
    check_call(_LIB.TVMFuncGetGlobal(c_str(name), ctypes.byref(handle)))
225 226
    if handle.value:
        return Function(handle, False)
227 228 229 230 231

    if allow_missing:
        return None

    raise ValueError("Cannot find global function %s" % name)
232

233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253


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


254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278
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


279 280
def _get_api(f):
    flocal = f
281
    flocal.is_global = True
282
    return flocal
283

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

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

    target_module_name : str
       The target module name if different from namespace
292
    """
293 294 295 296 297 298 299
    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)

300 301 302

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

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

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

325
_set_class_function(Function)