docs(api/lite): add lite global.h and pylite utils.py doc

GitOrigin-RevId: a7dfcf8914819e8dcea5585be61178cd56eea80f

lite/include/lite/global.h

@@ -10,65 +10,97 @@
 namespace lite {
 
 /**
- * \brief Model decryption function
+ * @brief Model decryption function
  *
- * \param[in] const void* is the decrypted model memory pointer
- * \param[in] size_t the size the decrypted model memory in byte
- * \param[in] const std::vector<uint8_t>& the decryption key vector
+ * @param arg1 the to be decrypted model memory pointer
+ * @param arg2 the byte size of the decrypted model memory
+ * @param arg3 the decryption key in vector
+ * @return the decrypted model in vector format, it's length and content can get by
+ * the operators of vector
  */
 using DecryptionFunc = std::function<std::vector<uint8_t>(
         const void*, size_t, const std::vector<uint8_t>&)>;
 
 /**
- * \brief register a custom decryption method and key to lite.
+ * @brief register a custom decryption method and key to lite
  *
- * \param[in] decrypt_name the name of the decryption, which will act as the
- * hash key to find the decryption method.
+ * @param decrypt_name the name of the decryption, which will act as the
+ * hash key to find the decryption method
  *
- * \param[in] func the decryption function, which will decrypt the model with
- * the registered key, return a vector that contain the decrypted model.
+ * @param func the decryption function, which will decrypt the model with
+ * the registered key, return a vector that contain the decrypted model
  *
- * \param[in] key the decryption key of the method
+ * @param key the decryption key of the method
+ *
+ * @return Whether or not the decryption method register successful
  */
 LITE_API bool register_decryption_and_key(
         std::string decrypt_name, const DecryptionFunc& func,
         const std::vector<uint8_t>& key);
 
 /**
- * \brief update decryption function or key of a custom decryption method.
+ * @brief update decryption function or key of a custom decryption method, in
+ * lite the decryption function and the key store in pair, user can change one of which
+ * by this function
  *
- * \param[in] decrypt_name the name of the decryption, which will act as the
- * hash key to find the decryption method.
+ * @param decrypt_name the name of the decryption, which will act as the
+ * hash key to find the decryption method
  *
- * \param[in] func the decryption function, which will decrypt the model with
+ * @param func the decryption function, which will decrypt the model with
  * the registered key, return a vector that contain the decrypted model. if
- * function is nullptr, it will not be updated.
+ * the function is nullptr, it will not be updated
+ *
+ * @param key the decryption key of the method, if the size of key is zero,
+ * the key will not be updated
  *
- * \param[in] key the decryption key of the method, if the size of key is zero,
- * it will not be updated
+ * @return Whether or not the decryption method update successful
  */
 LITE_API bool update_decryption_or_key(
         std::string decrypt_name, const DecryptionFunc& func,
         const std::vector<uint8_t>& key);
 
 /**
- * \brief Model information parse function
- *
- * \param[in] const void* is the information memory
- * \param[in] size_t the size the information memory
- * \param[in] const std::string the model name used for check whether the
- * infomation match the model
- * \param[in] Config the model config, ParseInfoFunc can fill it with the
- * information in json, the config will influence Network loading later
- * \param[in] NetworkIO the model IO, ParseInfoFunc can fill it with the
- * information in json, the networkio will influence Network forwarding later
- * \param[in] std::unordered_map<std::string, LiteAny>& isolated_config_map, the
- * other config not inclue in config and networkIO, ParseInfoFunc can fill it
- * with the information in json, now support:
- * "device_id" : int, default 0
- * "number_threads" : uint32_t, default 1
- * "is_inplace_model" : bool, default false
- * "use_tensorrt" : bool, default false
+ * @brief Model information parse function, MegEngine Lite model may pack some
+ * information with the model to configure the model inference processing conveniently,
+ * this function is used to parse the information packed with model, and store
+ * the parsed result into the params
+ *
+ * @param arg1 the information memory pointer
+ * @param arg2 the size the information memory
+ * @param arg3 the model name used for check whether the name in the information
+ * @param arg4 the model configuration, ParseInfoFunc fill it with the
+ * parsed information, the configuration will influence Network inference later
+ *
+ * @param arg5 the model IO information, ParseInfoFunc fill it with the parsed
+ * information, the networkio will influence Network inference later
+ *
+ * @param arg6 the other configurations do not include in configuration and networkIO,
+ * ParseInfoFunc fill it with the parsed information pair, now support:
+ *
+ * \verbatim
+ * embed:rst:leading-asterisk
+ *  .. list-table::
+ *      :widths: 20 10 30
+ *      :header-rows: 1
+ *
+ *      * - name
+ *        - type
+ *        - default
+ *      * - "device_id"
+ *        - int
+ *        - 0
+ *      * - "number_threads"
+ *        - uint32_t
+ *        - 1
+ *      * - "is_inplace_model"
+ *        - bool
+ *        - false
+ *      * - "use_tensorrt"
+ *        - bool
+ *        - false
+ * \endverbatim
+ *
+ * @return Whether or not the parse function parse successfully
  */
 using ParseInfoFunc = std::function<bool(
         const void*, size_t, const std::string model_name, Config& config,
@@ -77,89 +109,122 @@ using ParseInfoFunc = std::function<bool(
         std::string& extra_info)>;
 
 /**
- * \brief register a custom parser function to lite.
+ * @brief register a custom parser function to lite
  *
- * \param[in] info_type the name of the parser function, which will act as the
+ * @param info_type  the name of the parser function, which will act as the
  * hash key to find the parser method.
  *
- * \param[in] parse_func the parser function, which will parse the given
- * information and modify the Network Config and IO.
+ * @param parse_func  the parser function, which will parse the given
+ * information and modify the Network configuration and IO information.
  *
+ * @return Whether or not the parse function register successful
  */
 LITE_API bool register_parse_info_func(
         std::string info_type, const ParseInfoFunc& parse_func);
 
-/*! \brief Get version
+/** @brief get megengint lite version
+ *
+ * @param major  the major version of megengine lite
+ * @param minor  the minor version of megengine lite
+ * @param patch  the patch version of megengine lite
  */
 LITE_API void get_version(int& major, int& minor, int& patch);
 
-/*! \brief Set the current log level.
- * \param[in] level The new log level
+/**
+ * @brief set the current log level
+ * @param level the new log level to be set
  */
 LITE_API void set_log_level(LiteLogLevel level);
 
-/*! \brief Get the current log level.
- * \return The current log level
+/** @brief get the current log level
+ * @return the current log level
  */
 LITE_API LiteLogLevel get_log_level();
 
-/*! \brief Get device count
- * \param[in] device_type device type
- * \return the device count
+/** @brief get the number of device of the given device type in current context
+ * @param device_type the to be count device type
+ * @return the number of device
  */
 LITE_API size_t get_device_count(LiteDeviceType device_type);
 
-/*! \brief try to coalesce all free memory in megenine
+/** @brief try to coalesce all free memory in megenine, when call it MegEnine Lite
+ * will try to free all the unused memory thus decrease the runtime memory usage
  */
 LITE_API void try_coalesce_all_free_memory();
 
-/*!
- * \brief Set the loader to the lite
- * \param loader_path is the file path which store the cache
+/**
+ * @brief set the loader path to be used in lite
+ * @param loader_path the file path which store the loader library
  */
 LITE_API void set_loader_lib_path(const std::string& loader_path);
 
-/*!
- * \brief Set the algo policy cache file for CPU/CUDA ...
- * \param cache_path is the file path which store the cache
- * \param always_sync sync the cache when model run
+/**
+ * @brief Set the algo policy cache file for CPU/CUDA, the algo policy cache is
+ * produced by megengine fast-run
+ *
+ * @param cache_path  the file path which store the cache
+ * @param always_sync  always update the cache file when model run
  */
 LITE_API void set_persistent_cache(
         const std::string& cache_path, bool always_sync = false);
 
-/*!
- * \brief dump the PersistentCache policy cache to file, if the network is set
- * to profile when forward, though this the algo policy will dump to file
+/**
+ * @brief 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
+ *
+ * @param cache_path  the cache file path to be dump
  */
 LITE_API void dump_persistent_cache(const std::string& cache_path);
 
-/*!
- * \brief Set the TensorRT engine cache path for serialized prebuilt ICudaEngine
+/**
+ * @brief set the TensorRT engine cache path for serialized prebuilt ICudaEngine
+ *
+ * @param cache_path  the cache file path to set
  */
 LITE_API void set_tensor_rt_cache(std::string tensorrt_cache_path);
 
-/*!
- * \brief dump the TensorRT cache to the file set in set_tensor_rt_cache
+/**
+ * @brief dump the TensorRT cache to the file set in set_tensor_rt_cache
  */
 LITE_API void dump_tensor_rt_cache();
 
 /**
- * register the physical and virtual address pair to the mge, some device
- * need the map from physical to virtual.
+ * @brief register the physical and virtual address pair to the mge, some device
+ * need the map from physical to virtual
+ *
+ * @param vir_ptr - the virtual ptr to set to megenine
+ * @param phy_ptr - the physical ptr to set to megenine
+ * @param device - the device to set the pair memory
+ * @param backend - the backend to set the pair memory
+ *
+ * @return Whether the register is successful
  */
 LITE_API bool register_memory_pair(
         void* vir_ptr, void* phy_ptr, size_t length, LiteDeviceType device,
         LiteBackend backend = LiteBackend::LITE_DEFAULT);
 
 /**
- * clear the physical and virtual address pair in mge.
+ * @brief clear the physical and virtual address pair in mge
+ *
+ * @param vir_ptr - the virtual ptr to set to megenine
+ * @param phy_ptr - the physical ptr to set to megenine
+ * @param device - the device to set the pair memory
+ * @param backend - the backend to set the pair memory
+ *
+ * @return Whether the clear is successful
  */
 LITE_API bool clear_memory_pair(
         void* vir_ptr, void* phy_ptr, LiteDeviceType device,
         LiteBackend backend = LiteBackend::LITE_DEFAULT);
 
 /**
- * get the physic address by the virtual address in mge.
+ * @brief get the physic address by the virtual address in mge.
+ *
+ * @param vir_ptr - the virtual ptr to set to megenine
+ * @param device - the device to set the pair memory
+ * @param backend - the backend to set the pair memory
+ *
+ * @return The physic address to lookup
  */
 void* lookup_physic_ptr(void* vir_ptr, LiteDeviceType device, LiteBackend backend);
 

lite/pylite/megenginelite/utils.py

@@ -11,7 +11,53 @@ from .tensor import *
 
 class TensorBatchCollector:
     """
-    this is a tensor utils to collect subtensor in batch continuous
+    A tensor utils is used to collect many single batch tensor to a multi batch
+    size tensor, when the multi batch size tensor collect finish, the result
+    tensor can be get and send to the model input for forwarding.
+
+    when collect single batch tensor, the single batch tensor is no need in the
+    same device_type and device_id with the result tensor, however the dtype must
+    match and the shape must match except the highest dimension.
+
+    Args:
+        shape: the multi batch size tensor shape, After collection, the result
+            tensor shape.
+        dtype(LiteDataType): the datatype of the single batch tensor and the
+            result tensor, default value is LiteDataType.LITE_INT8.
+        device_type(LiteDeviceType): the target device type the result tensor
+            will allocate, default value is LiteDeviceType.LITE_CUDA.
+        device_id: the device id the result tensor will allocate, default 0.
+        is_pinned_host: Whether the memory is pinned memory, refer to CUDA
+            pinned memory, default False.
+        tensor(LiteTensor): the result tensor, user can also create the multi
+            batch size tensor and then create the TensorBatchColletor, if tensor is
+            not None, all the member, such as shape, dtype, device_type,
+            device_id, is_pinned_host will get from the tensor, if the tensor is
+            None and the result tensor will create by the TensorBatchCollector,
+            default is None.
+
+    Note:
+        when collect tensor, the single batch tensor or array shape must match the 
+        result tensor shape except the batch size dimension (the highest dimension)
+
+    Examples:
+
+        .. code-block:: python
+
+            import numpy as np
+            batch_tensor = TensorBatchCollector([4, 8, 8])
+            arr = np.ones([8, 8], "int8")
+            for i in range(4):
+                batch_tensor.collect(arr)
+                arr += 1
+            data = batch_tensor.to_numpy()
+            assert data.shape[0] == 4
+            assert data.shape[1] == 8
+            assert data.shape[2] == 8
+            for i in range(4):
+                for j in range(64):
+                    assert data[i][j // 8][j % 8] == i + 1
+
     """
 
     def __init__(
@@ -45,6 +91,17 @@ class TensorBatchCollector:
             )
 
     def collect_id(self, array, batch_id):
+        """
+        Collect a single batch through an array and store the array data to the
+        specific batch_id.
+
+        Args:
+            array: an array maybe LiteTensor or numpy ndarray, the shape of
+                array must match the result tensor shape except the highest
+                dimension.
+            batch_id: the batch id to store the array data to the result tensor,
+                if the batch_id has already collected, a warning will generate.
+        """
         # get the batch index
         with self._mutex:
             if batch_id in self._free_list:
@@ -87,6 +144,14 @@ class TensorBatchCollector:
         return batch_id
 
     def collect(self, array):
+        """
+        Collect a single batch through an array and store the array data to an
+        empty batch, the empty batch is the front batch id in free list.
+
+        Args:
+            array: an array maybe LiteTensor or numpy ndarray, the shape must
+                match the result tensor shape except the highest dimension
+        """
         with self._mutex:
             if len(self._free_list) == 0:
                 warnings.warn(
@@ -98,7 +163,13 @@ class TensorBatchCollector:
 
     def collect_by_ctypes(self, data, length):
         """
-        collect with ctypes data input
+        Collect a single batch through an ctypes memory buffer and store the
+        ctypes memory data to an empty batch, the empty batch is the front
+        batch id in free list.
+
+        Args:
+            array: an array maybe LiteTensor or numpy ndarray, the shape must
+                match the result tensor shape except the highest dimension
         """
         with self._mutex:
             if len(self._free_list) == 0:
@@ -116,6 +187,13 @@ class TensorBatchCollector:
             subtensor.copy_from(pinned_tensor)
 
     def free(self, indexes):
+        """
+        free the batch ids in the indexes, after the batch id is freed, it can
+        be collected again without warning.
+
+        Args:
+            indexes: a list of to be freed batch id
+        """
         with self._mutex:
             for i in indexes:
                 if i in self._free_list:
@@ -126,7 +204,13 @@ class TensorBatchCollector:
             self._free_list.extend(indexes)
 
     def get(self):
+        """
+        After finish collection, get the result tensor
+        """
         return self._tensor
 
     def to_numpy(self):
+        """
+        Convert the result tensor to a numpy ndarray
+        """
         return self._tensor.to_numpy()