tvm_ffi.register_global_func

tvm_ffi.register_global_func#

tvm_ffi.register_global_func(func_name, f=None, override=False)[source]#

Register global function.

Parameters:
  • func_name (str | Callable[..., Any]) – The function name

  • f (Callable[..., Any] | None, default: None) – The function to be registered.

  • override (bool, default: False) – Whether override existing entry.

Return type:

Any

Returns:

fregister – Register function if f is not specified.

Examples

import tvm_ffi

# we can use decorator to register a function
@tvm_ffi.register_global_func("mytest.echo")
def echo(x):
    return x
# After registering, we can get the function by its name
f = tvm_ffi.get_global_func("mytest.echo")
assert f(1) == 1

# we can also directly register a function
tvm_ffi.register_global_func("mytest.add_one", lambda x: x + 1)
f = tvm_ffi.get_global_func("mytest.add_one")
assert f(1) == 2