Browse Source

docs(api/lite): add global setting docstring for megenginelite

GitOrigin-RevId: 764a479cab
release-1.10
Megvii Engine Team 3 years ago
parent
commit
92ded7216b
1 changed files with 97 additions and 4 deletions
  1. +97
    -4
      lite/pylite/megenginelite/global_setting.py

+ 97
- 4
lite/pylite/megenginelite/global_setting.py View File

@@ -16,7 +16,7 @@ LiteDecryptionFunc = CFUNCTYPE(


class _GlobalAPI(_LiteCObjBase): class _GlobalAPI(_LiteCObjBase):
""" """
get the api from the lib
Get APIs from the lib
""" """


_api_ = [ _api_ = [
@@ -42,8 +42,12 @@ class _GlobalAPI(_LiteCObjBase):




def decryption_func(func): 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) @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): class LiteGlobal(object):
""" """
some global config in lite
Some global config in lite
""" """


_api = _GlobalAPI()._lib _api = _GlobalAPI()._lib


@staticmethod @staticmethod
def register_decryption_and_key(decryption_name, decryption_func, key): 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")) c_name = c_char_p(decryption_name.encode("utf-8"))
key_length = len(key) key_length = len(key)
c_key = (c_uint8 * key_length)(*key) c_key = (c_uint8 * key_length)(*key)
@@ -79,6 +93,14 @@ class LiteGlobal(object):


@staticmethod @staticmethod
def update_decryption_key(decryption_name, key): 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")) c_name = c_char_p(decryption_name.encode("utf-8"))
key_length = len(key) key_length = len(key)
c_key = (c_uint8 * key_length)(*key) c_key = (c_uint8 * key_length)(*key)
@@ -86,42 +108,92 @@ class LiteGlobal(object):


@staticmethod @staticmethod
def set_loader_lib_path(path): 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")) c_path = c_char_p(path.encode("utf-8"))
LiteGlobal._api.LITE_set_loader_lib_path(c_path) LiteGlobal._api.LITE_set_loader_lib_path(c_path)


@staticmethod @staticmethod
def set_persistent_cache(path, always_sync=False): 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")) c_path = c_char_p(path.encode("utf-8"))
LiteGlobal._api.LITE_set_persistent_cache(c_path, always_sync) LiteGlobal._api.LITE_set_persistent_cache(c_path, always_sync)


@staticmethod @staticmethod
def set_tensorrt_cache(path): 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")) c_path = c_char_p(path.encode("utf-8"))
LiteGlobal._api.LITE_set_tensorrt_cache(c_path) LiteGlobal._api.LITE_set_tensorrt_cache(c_path)


@staticmethod @staticmethod
def dump_persistent_cache(path): 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")) c_path = c_char_p(path.encode("utf-8"))
LiteGlobal._api.LITE_dump_persistent_cache(c_path) LiteGlobal._api.LITE_dump_persistent_cache(c_path)


@staticmethod @staticmethod
def dump_tensorrt_cache(): def dump_tensorrt_cache():
"""Dump the TensorRT cache to the file set in :py:func:`~.set_tensorrt_cache`."""
LiteGlobal._api.LITE_dump_tensorrt_cache() LiteGlobal._api.LITE_dump_tensorrt_cache()


@staticmethod @staticmethod
def get_device_count(device_type): 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() count = c_size_t()
LiteGlobal._api.LITE_get_device_count(device_type, byref(count)) LiteGlobal._api.LITE_get_device_count(device_type, byref(count))
return count.value return count.value


@staticmethod @staticmethod
def try_coalesce_all_free_memory(): 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() LiteGlobal._api.LITE_try_coalesce_all_free_memory()


@staticmethod @staticmethod
def register_memory_pair( def register_memory_pair(
vir_ptr, phy_ptr, length, device, backend=LiteBackend.LITE_DEFAULT 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( assert isinstance(vir_ptr, c_void_p) and isinstance(
phy_ptr, c_void_p phy_ptr, c_void_p
), "clear memory pair only accept c_void_p type." ), "clear memory pair only accept c_void_p type."
@@ -131,6 +203,17 @@ class LiteGlobal(object):


@staticmethod @staticmethod
def clear_memory_pair(vir_ptr, phy_ptr, device, backend=LiteBackend.LITE_DEFAULT): 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( assert isinstance(vir_ptr, c_void_p) and isinstance(
phy_ptr, c_void_p phy_ptr, c_void_p
), "clear memory pair only accept c_void_p type." ), "clear memory pair only accept c_void_p type."
@@ -138,6 +221,16 @@ class LiteGlobal(object):


@staticmethod @staticmethod
def lookup_physic_ptr(vir_ptr, device, backend=LiteBackend.LITE_DEFAULT): 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( assert isinstance(
vir_ptr, c_void_p vir_ptr, c_void_p
), "lookup physic ptr only accept c_void_p type." ), "lookup physic ptr only accept c_void_p type."


Loading…
Cancel
Save