368 lines
16 KiB
Python
368 lines
16 KiB
Python
"""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))
|