From 92ded7216b95da79740ba007889e669f646df681 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 21 Jun 2022 16:07:36 +0800 Subject: [PATCH] docs(api/lite): add global setting docstring for megenginelite GitOrigin-RevId: 764a479cab61b99dde15ab4f012cabbdd408f325 --- lite/pylite/megenginelite/global_setting.py | 101 +++++++++++++++++++- 1 file changed, 97 insertions(+), 4 deletions(-) diff --git a/lite/pylite/megenginelite/global_setting.py b/lite/pylite/megenginelite/global_setting.py index 139f3429d..e8a72eea7 100644 --- a/lite/pylite/megenginelite/global_setting.py +++ b/lite/pylite/megenginelite/global_setting.py @@ -16,7 +16,7 @@ LiteDecryptionFunc = CFUNCTYPE( class _GlobalAPI(_LiteCObjBase): """ - get the api from the lib + Get APIs from the lib """ _api_ = [ @@ -42,8 +42,12 @@ class _GlobalAPI(_LiteCObjBase): def decryption_func(func): - """the decryption function decorator - :type func: a function accept three array, in_arr, key_arr and out_arr, if out_arr is None, just query the out array lenght in byte + """the decryption function decorator. + + .. note:: + + The function accept three array: ``in_arr``, ``key_arr`` and ``out_arr``. + If ``out_arr`` is None, just query the out array length in byte. """ @CFUNCTYPE(c_size_t, c_void_p, c_size_t, POINTER(c_uint8), c_size_t, c_void_p) @@ -63,13 +67,23 @@ def decryption_func(func): class LiteGlobal(object): """ - some global config in lite + Some global config in lite """ _api = _GlobalAPI()._lib @staticmethod def register_decryption_and_key(decryption_name, decryption_func, key): + """Register a custom decryption method and key to lite + + Args: + decryption_name: the name of the decryption, which will act as the hash + key to find the decryption method. + decryption_func: the decryption function, which will decrypt the model with + the registered key, then return the decrypted model. + See :py:func:`~.decryption_func` for more details. + key: the decryption key of the method. + """ c_name = c_char_p(decryption_name.encode("utf-8")) key_length = len(key) c_key = (c_uint8 * key_length)(*key) @@ -79,6 +93,14 @@ class LiteGlobal(object): @staticmethod def update_decryption_key(decryption_name, key): + """Update decryption key of a custom decryption method. + + Args: + decrypt_name: the name of the decryption, + which will act as the hash key to find the decryption method. + key: the decryption key of the method, + if the length of key is zero, the key will not be updated. + """ c_name = c_char_p(decryption_name.encode("utf-8")) key_length = len(key) c_key = (c_uint8 * key_length)(*key) @@ -86,42 +108,92 @@ class LiteGlobal(object): @staticmethod def set_loader_lib_path(path): + """Set the loader path to be used in lite. + + Args: + path: the file path which store the loader library. + """ c_path = c_char_p(path.encode("utf-8")) LiteGlobal._api.LITE_set_loader_lib_path(c_path) @staticmethod def set_persistent_cache(path, always_sync=False): + """Set the algo policy cache file for CPU/CUDA, + the algo policy cache is produced by MegEngine fast-run. + + Args: + path: the file path which store the cache. + always_sync: always update the cache file when model runs. + """ c_path = c_char_p(path.encode("utf-8")) LiteGlobal._api.LITE_set_persistent_cache(c_path, always_sync) @staticmethod def set_tensorrt_cache(path): + """Set the TensorRT engine cache path for serialized prebuilt ICudaEngine. + + Args: + path: the cache file path to set + """ c_path = c_char_p(path.encode("utf-8")) LiteGlobal._api.LITE_set_tensorrt_cache(c_path) @staticmethod def dump_persistent_cache(path): + """Dump the PersistentCache policy cache to the specific file. + If the network is set to profile when forward, + though this the algo policy will dump to file. + + Args: + path: the cache file path to be dumped. + """ c_path = c_char_p(path.encode("utf-8")) LiteGlobal._api.LITE_dump_persistent_cache(c_path) @staticmethod def dump_tensorrt_cache(): + """Dump the TensorRT cache to the file set in :py:func:`~.set_tensorrt_cache`.""" LiteGlobal._api.LITE_dump_tensorrt_cache() @staticmethod def get_device_count(device_type): + """Get the number of device of the given device type in current context. + + Args: + device_type: the device type to be counted. + + Returns: + the number of device. + """ count = c_size_t() LiteGlobal._api.LITE_get_device_count(device_type, byref(count)) return count.value @staticmethod def try_coalesce_all_free_memory(): + """Try to coalesce all free memory in MegEngine. + When call it MegEnine Lite will try to free all the unused memory + thus decrease the runtime memory usage. + """ LiteGlobal._api.LITE_try_coalesce_all_free_memory() @staticmethod def register_memory_pair( vir_ptr, phy_ptr, length, device, backend=LiteBackend.LITE_DEFAULT ): + """Register the physical and virtual address pair to the MegEngine, + some device need the map from physical to virtual. + + Args: + vir_ptr: the virtual ptr to set to MegEngine. + phy_ptr: the physical ptr to set to MegEngine. + length: the length of bytes to set pair memory. + device: the the device to set the pair memory. + backend: the backend to set the pair memory + + Return: + Whether the register operation is successful. + """ assert isinstance(vir_ptr, c_void_p) and isinstance( phy_ptr, c_void_p ), "clear memory pair only accept c_void_p type." @@ -131,6 +203,17 @@ class LiteGlobal(object): @staticmethod def clear_memory_pair(vir_ptr, phy_ptr, device, backend=LiteBackend.LITE_DEFAULT): + """Clear the physical and virtual address pair in MegEngine. + + Args: + vir_ptr: the virtual ptr to set to MegEngine. + phy_ptr: the physical ptr to set to MegEngine. + device: the the device to set the pair memory. + backend: the backend to set the pair memory. + + Return: + Whether the clear is operation successful. + """ assert isinstance(vir_ptr, c_void_p) and isinstance( phy_ptr, c_void_p ), "clear memory pair only accept c_void_p type." @@ -138,6 +221,16 @@ class LiteGlobal(object): @staticmethod def lookup_physic_ptr(vir_ptr, device, backend=LiteBackend.LITE_DEFAULT): + """Get the physic address by the virtual address in MegEngine. + + Args: + vir_ptr: the virtual ptr to set to MegEngine. + device: the the device to set the pair memory. + backend: the backend to set the pair memory. + + Return: + The physic address to lookup. + """ assert isinstance( vir_ptr, c_void_p ), "lookup physic ptr only accept c_void_p type." -- GitLab