Skip to content

Commit

Permalink
[API] Added remove_global_func to the Python API (apache#6787)
Browse files Browse the repository at this point in the history
This is useful for unregistering functions after a test.

Change-Id: Ic39499aa8f36bfe5470bc1f058ad3b96cf52b49c
  • Loading branch information
mbaret authored and trevor-m committed Dec 4, 2020
1 parent 304f9ca commit 3488a0f
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 0 deletions.
6 changes: 6 additions & 0 deletions include/tvm/runtime/c_runtime_api.h
Original file line number Diff line number Diff line change
Expand Up @@ -366,6 +366,12 @@ TVM_DLL int TVMFuncGetGlobal(const char* name, TVMFunctionHandle* out);
*/
TVM_DLL int TVMFuncListGlobalNames(int* out_size, const char*** out_array);

/*!
* \brief Remove a global function.
* \param name The name of the function.
*/
TVM_DLL int TVMFuncRemoveGlobal(const char* name);

// Array related apis for quick proptyping
/*!
* \brief Allocate a nd-array's memory,
Expand Down
11 changes: 11 additions & 0 deletions python/tvm/_ffi/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -262,6 +262,17 @@ def _list(name, func):
return fdict


def remove_global_func(name):
"""Remove a global function by name
Parameters
----------
name : str
The name of the global function
"""
check_call(_LIB.TVMFuncRemoveGlobal(c_str(name)))


def _get_api(f):
flocal = f
flocal.is_global = True
Expand Down
6 changes: 6 additions & 0 deletions src/runtime/registry.cc
Original file line number Diff line number Diff line change
Expand Up @@ -146,3 +146,9 @@ int TVMFuncListGlobalNames(int* out_size, const char*** out_array) {
*out_size = static_cast<int>(ret->ret_vec_str.size());
API_END();
}

int TVMFuncRemoveGlobal(const char* name) {
API_BEGIN();
tvm::runtime::Registry::Remove(name);
API_END();
}

0 comments on commit 3488a0f

Please sign in to comment.