`_ffi.registry`#

FFI registry to register function and objects.

tvm._ffi.registry._init_api(namespace, target_module_name=None)[源代码]#

Initialize api for a given module name

namespacestr: The namespace of the source registry
target_module_namestr: The target module name if different from namespace

tvm._ffi.registry.extract_ext_funcs(finit)[源代码]#

Extract the extension PackedFuncs from a C module.

finitctypes function: a ctypes that takes signature of TVMExtensionDeclarer

fdictdict of str to Function: The extracted functions

tvm._ffi.registry.get_global_func(name, allow_missing=False)[源代码]#

Get a global function by name

namestr: The name of the global function
allow_missingbool: Whether allow missing function or raise an error.

funcPackedFunc: The function to be returned, None if function is missing.

tvm._ffi.registry.get_object_type_index(cls)[源代码]#

Get type index of object type

clstype: The object type to get type index for.

type_indexOptional[int]: The type index, or None if type not found in the registry.

tvm._ffi.registry.list_global_func_names()[源代码]#

Get list of global functions registered.

nameslist: List of global functions names.

tvm._ffi.registry.register_extension(cls, fcreate=None)[源代码]#

After the class is registered, the class will be able to directly pass as Function argument generated by TVM.

clsclass: The class object to be registered as extension.
fcreatefunction, optional: The creation function to create a class object given handle value.

The registered class is requires one property: _tvm_handle.

If the registered class is a subclass of NDArray, it is required to have a class attribute _array_type_code. Otherwise, it is required to have a class attribute _tvm_tcode.

`_tvm_handle` returns integer represents the address of the handle.
`_tvm_tcode` or `_array_type_code` gives integer represents type code of the class.

clsclass: The class being registered.

The following code registers user defined class MyTensor to be DLTensor compatible.

@tvm.register_extension
class MyTensor(object):
    _tvm_tcode = tvm.ArgTypeCode.ARRAY_HANDLE

    def __init__(self):
        self.handle = _LIB.NewDLTensor()

    @property
    def _tvm_handle(self):
        return self.handle.value

tvm._ffi.registry.register_func(func_name, f=None, override=False)[源代码]#

func_namestr or function: The function name
ffunction, optional: The function to be registered.
override: boolean optional: Whether override existing entry.

fregisterfunction: Register function if f is not specified.

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.

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.PackedFunc)
y = f(*targs)
assert y == 10

tvm._ffi.registry.register_object(type_key=None)[源代码]#

type_keystr or cls: The type key of the node

The following code registers MyObject using type key “test.MyObject”

@tvm.register_object("test.MyObject")
class MyObject(Object):
    pass

tvm._ffi.registry.remove_global_func(name)[源代码]#

Remove a global function by name

namestr: The name of the global function

TVM 文档

_ffi.registry

_ffi.registry#

`_ffi.registry`#