"""Functions to help transfer data and communicate between Python E2E and C postprocessing. This module will initialize and prepare the memory for C postprocess function calls using the KDPImage classes as inputs. It will also grab those postprocess results from memory and return a class for access within Python code. Some of the memory initialized in these functions may not automatically be freed, so the user must call 'free_result' once the data is no longer needed. Typical usage exampe: kdp_image = init_kdp_image(520) run_c_function(kdp_image, results, emu_config, 520, 'c_func', result_class) free_result(kdp_image, 520) """ import ctypes import math import struct from typing import List, Mapping, Tuple, Union import numpy as np import numpy.typing as npt from c_interface import constants import c_interface.kdp_image as kdp from c_interface.wrappers import ResultClass from python_flow.common import exceptions from python_flow.utils import csim from python_flow.utils import utils def _prep_inference_results(np_results: List[npt.ArrayLike], platform: int, is_channel_first: bool, has_batch: bool, reordering: List[Union[int, str]], data_type: str, output_folder: str, platform_version: int) -> Tuple[List[npt.NDArray[np.int8]], List[int]]: """Converts the inference results into data that will be directly loaded into memory. First, it reorders the output if the results dumped out from calling CSIM on the model is in a different order than specified in the setup file. Next, it will convert the data to an integer type if the inference results are given as floats. Finally, it will convert the inference results into appropriate formats as follows: 520 - (height, channel, width_aligned) all other platforms - (channel, height, width_aligned) Width_aligned will be the closest multiple of 16 to the original width. Arguments: np_results: List of NumPy arraylikes indicating inference results. platform: An integer indicating the version of CSIM used. is_channel_first: Flag indicating if inference results are in channel first format. has_batch: Flag indicating if inference results have a batch dimension. reordering: List of indices that were used to previously reorder the output. data_type: String indicating data type of the inference results ('float' or 'fixed'). output_folder: String path to folder containing CSIM outputs, only needed if data_type is 'float'. platform_version: Integer version of the specific CSIM platform that was used. Returns: A tuple of Lists with the same size. The first list holds the prepared integer inference results in the correct order set to the appropriate channel orderings. The second list holds the original widths of the inference results before alignment. Raises: ValueError: If any of the inference results is not between 2 and 4 dimensions. """ if platform not in constants.SUPPORTED_PLATFORMS: raise ValueError(f"prep_inference_results: Platform must be one of " f"{constants.SUPPORTED_PLATFORMS}, not {platform}") # reorder to match output order in CSIM setup file if reordering: reordered = [] for index in range(len(reordering)): new_index = reordering.index(index) reordered.append(np.array(np_results[new_index])) else: reordered = [np.array(result) for result in np_results] if data_type == "float": reordered = csim.convert_csim_data_type( reordered, output_folder, platform, "fl2fx", platform_version) new_results = [] old_widths = [] # assume 2 dim - 4 dim, bchw or bhwc for inf_result in reordered: shape = inf_result.shape if not has_batch: shape = (1, *shape) # remove batch, set h/w if necessary if len(shape) == 2: if is_channel_first: reshape_dims = (shape[1], 1, 1) else: reshape_dims = (1, 1, shape[1]) elif len(shape) == 3: # TODO may need testing with chw and hwc reshape_dims = (shape[1], shape[2], 1) elif len(shape) == 4: reshape_dims = shape[1:] else: raise ValueError("Output is not between 2 and 4 dimensions.") new_result = np.reshape(inf_result, reshape_dims) if platform == 520: # set to hcw if is_channel_first: new_result = np.transpose(new_result, (1, 0, 2)) else: new_result = np.transpose(new_result, (0, 2, 1)) # set to 16 byte align height, channel, width = new_result.shape width_aligned = 16 * math.ceil(width / 16.0) new_aligned = np.zeros((height, channel, width_aligned)) else: # set to chw if not is_channel_first: new_result = np.transpose(new_result, (2, 0, 1)) # set to 16 byte align channel, height, width = new_result.shape width_aligned = 16 * math.ceil(width / 16.0) new_aligned = np.zeros((channel, height, width_aligned)) new_aligned[:, :, :width] = new_result new_results.append(new_aligned.astype(np.int8)) old_widths.append(width) return new_results, old_widths def _load_np_to_memory(kdp_image: kdp.KDPImageType, np_results: List[npt.NDArray[np.int8]], old_widths: List[int], platform: int, output_folder: str, setup_file: str) -> None: """Loads data into the KDPImage class needed for postprocessing. Arguments: kdp_image: A kdp.KDPImageType instance. np_results: List of NumPy arrays indicating inference results. Assume already prepared for the corresponding platform. old_widths: List of integers indicating each inference result width before alignment. Only needed for 520 and 720. platform: An integer indicating the version of CSIM used. output_folder: String path to folder containing CSIM outputs. setup_file: String path to input setup binary. """ if platform == 520: _csim_520_to_memory(kdp_image, np_results, old_widths, output_folder, setup_file) elif platform == 720: _csim_720_to_memory(kdp_image, np_results, old_widths, setup_file) elif platform in constants.SUPPORTED_PLATFORMS: _csim_to_memory(kdp_image, np_results, setup_file, platform) def _csim_520_to_memory(kdp_image: kdp.KDPImage520, np_results: List[npt.NDArray[np.int8]], old_widths: List[int], output_folder: str, setup_file: str) -> None: """Loads CSIM 520 output data into memory defined by the setup file. Assumes np_results is set to (height, channel, width_aligned) ordering. Args: kdp_image: A kdp.KDPImage520 instance. np_results: List of NumPy arrays indicating inference results. Assume 16 byte width aligned and in hcw format. old_widths: List of integers indicating each inference result width before alignment. output_folder: String path to folder containing CSIM outputs. setup_file: String path to input setup binary. """ _setup_csim(kdp_image, 520, setup_file) with open("".join([output_folder, "/radix_scale_info.txt"])) as params: values = [line.strip().split(" ") for line in params.readlines()] radices = [struct.unpack("i", struct.pack("I", int(radix)))[0] for _, radix, _ in values] # convert "int" scale value into float value using its bytes scales = [int(scale) for _, _, scale in values] for index, (np_result, old_width, radix, scale) in enumerate(zip(np_results, old_widths, radices, scales)): # Set output node params out_node = kdp.OutNode520(row_length=np_result.shape[0], col_length=old_width, ch_length=np_result.shape[1], output_radix=radix, output_scale=scale) kdp_image.postproc.node_p[index] = out_node _load_csim_data(kdp_image, 520, np_result.flatten(), index) def _csim_to_memory(kdp_image: kdp.KDPImage, np_results: List[npt.NDArray[np.int8]], setup_file: str, platform: int) -> None: """Loads CSIM non-520/720 output data into memory defined by the setup file. Assumes np_results is set to (channel, height, width_aligned) ordering. Args: kdp_image: A kdp.KDPImage instance. np_results: List of NumPy arrays indicating inference results. Assume 16 byte width aligned and in chw format. setup_file: String path to input setup binary. platform: An integer indicating the version of CSIM used. """ _setup_csim(kdp_image, platform, setup_file) for index, np_result in enumerate(np_results): _load_csim_data(kdp_image, platform, np_result.flatten(), index) def _csim_720_to_memory(kdp_image: kdp.KDPImage720, np_results: List[npt.NDArray[np.int8]], old_widths: List[int], setup_file: str) -> None: """Loads CSIM 720 output data into memory defined by the setup file. Assumes np_results is set to (channel, height, width_aligned) ordering. Arguments: kdp_image: A kdp.KDPImage720 instance. np_results: List of NumPy arrays indicating inference results. Assume 16 byte width aligned and in chw format. old_widths: List of integers indicating each inference result width before alignment. setup_file: String path to input setup binary. """ _setup_csim(kdp_image, 720, setup_file) nodes = [] for index, (np_result, old_width) in enumerate(zip(np_results, old_widths)): # Set output node params out_node = kdp.OutNode720(row_length=np_result.shape[1], col_length=old_width, ch_length=np_result.shape[0]) nodes.append(out_node) _load_csim_data(kdp_image, 720, np_result.flatten(), index) kdp_image.postproc.set_node_p(nodes) def _c_processing(kdp_image: kdp.KDPImageType, function_name: str, platform: int) -> None: """Calls the provided C function with the specified platform. Args: kdp_image: A kdp.KDPImageType instance. function_name: String of the exact name of the C function to call. platform: An integer indicating the version of CSIM used. """ library = constants.PROCESSMAP[platform] kdp_version = kdp.KDPIMAGEMAP[platform] c_function = getattr(library, function_name) c_function.argtypes = [ctypes.POINTER(kdp_version)] c_function.restype = None c_function(kdp_image) def run_c_function(kdp_image: kdp.KDPImageType, np_results: List[npt.ArrayLike], emu_config: Mapping[str, Union[int, str, List[Union[int, str]]]], platform: int, function_name: str, result_class: ResultClass, platform_version: int = 0) -> ResultClass: """Runs the specified C function and gets the result back from memory. First, the results will be prepared for postprocessing and loaded into memory. Then, the specified C function will be called. Finally, the result will be extracted from memory and casted into the specified class. Args: kdp_image: A kdp.KDPImageType instance. np_results: List of NumPy arraylikes indicating inference results. emu_config: Dictionary of emulator configurations passed in from the input JSON. platform: An integer indicating the version of CSIM used. function_name: String of the exact name of the C function to call. result_class: Class that the result data will be cast to. platform_version: Integer version of the specific CSIM platform that was used. Returns: A ResultClass instance containing data in memory saved by the called C function. """ if platform not in constants.SUPPORTED_PLATFORMS: raise ValueError(f"run_c_function: Platform must be one of " f"{constants.SUPPORTED_PLATFORMS}, not {platform}") results, widths = _prep_inference_results( np_results, platform, emu_config["channel_first"], True, emu_config["reordering"], emu_config["data_type"], emu_config["csim_output"], platform_version) _load_np_to_memory(kdp_image, results, widths, platform, emu_config["csim_output"], emu_config["setup_file"]) _c_processing(kdp_image, function_name, platform) return _get_result(kdp_image, result_class) def free_result(kdp_image: kdp.KDPImageType, platform: int) -> None: """Frees memory allocated needed to call the C postprocess. Args: kdp_image: A kdp.KDPImageType instance platform: An integer indicating the version of CSIM used. """ if platform not in constants.SUPPORTED_PLATFORMS: raise ValueError(f"free_result: Platform must be one of " f"{constants.SUPPORTED_PLATFORMS}, not {platform}") _free_kdp_image_data(kdp_image, platform) _free_csim_data(kdp_image, platform) def _get_result(kdp_image: kdp.KDPImageType, result_class: ResultClass) -> ResultClass: """Returns a result_class instance using data stored in the result address. Args: kdp_image: A kdp.KDPImageType instance. result_class: Class that the result data will be cast to. """ return ctypes.cast(kdp_image.postproc.result_mem_addr, ctypes.POINTER(result_class)).contents # C memory function wrappers def _load_csim_data(kdp_image: kdp.KDPImageType, platform: int, data: npt.NDArray[np.int8], index: int) -> None: """Wrapper to load CSIM output data into memory for postprocessing. Args: kdp_image: A kdp.KDPImageType instance platform: An integer indicating the version of CSIM used. data: A np.int8 NumPy array prepared for postprocessing. index: An integer indicating the output number the data should be loaded into. """ library = constants.LOADMAP[platform] kdp_version = kdp.KDPIMAGEMAP[platform] c_function = library.load_csim_data c_function.argtypes = [ctypes.POINTER(kdp_version), ctypes.c_int, ctypes.POINTER(ctypes.c_int8), ctypes.c_int] c_function.restype = None c_function(ctypes.byref(kdp_image), index, data.ctypes.data_as(ctypes.POINTER(ctypes.c_int8)), data.size) def _setup_csim(kdp_image: kdp.KDPImageType, platform: int, setup_file: str) -> None: """Wrapper to load the model setup file into memory for postprocessing. Args: kdp_image: A kdp.KDPImageType instance platform: An integer indicating the version of CSIM used. setup_file: String path to input setup binary Raises: exceptions.LibError: If C function fails """ library = constants.LOADMAP[platform] kdp_version = kdp.KDPIMAGEMAP[platform] c_function = library.parse_model c_function.argtypes = [ctypes.POINTER(kdp_version), ctypes.c_char_p] c_function.restype = ctypes.c_int if c_function(ctypes.byref(kdp_image), setup_file.encode()): raise exceptions.LibError(f"Loading {platform} setup file failed in C") def _free_csim_data(kdp_image: kdp.KDPImageType, platform: int) -> None: """Wrapper to free memory initialized for postprocessing. Args: kdp_image: A kdp.KDPImageType instance platform: An integer indicating the version of CSIM used. """ library = constants.LOADMAP[platform] kdp_version = kdp.KDPIMAGEMAP[platform] c_function = library.free_csim_data c_function.argtypes = [ctypes.POINTER(kdp_version)] c_function.restype = None c_function(ctypes.byref(kdp_image)) def _free_kdp_image_data(kdp_image: kdp.KDPImageType, platform: int) -> None: """Wrapper to free KDPImage data initialized for postprocessing. Args: kdp_image: A kdp.KDPImageType instance platform: An integer indicating the version of CSIM used. """ library = constants.LOADMAP[platform] kdp_version = kdp.KDPIMAGEMAP[platform] c_function = library.free_kdp_image c_function.argtypes = [ctypes.POINTER(kdp_version)] c_function.restype = None c_function(ctypes.byref(kdp_image))