#########################################################
# Copyright (C) 2022 SiMa Technologies, Inc.
#
# This material is SiMa proprietary and confidential.
#
# This material may not be copied or distributed without
# the express prior written permission of SiMa.
#
# All rights reserved.
#########################################################
# Code owner: Stefan Jovic
#########################################################
from dataclasses import dataclass, field
from typing import Any
from afe.apis.release_v1 import get_model_sdk_version
from sima_utils.logging import sima_logger
# The parent class is only used to provide documentation
@dataclass
[docs]
class AFEMPKData:
"""
Provides data that is needed for reconstructing the MPK JSON file.
The data is convertible to JSON format.
All changes to these data structures must adhere to changes in MPK
file specification.
"""
pass
@dataclass
[docs]
class TensorTypeMPKData(AFEMPKData):
"""
Provides the information about a tensor type.
Attributes:
scalar: The tensor's scalar element type, as a string that numpy can interpret as a type.
shape: The tensor's shape.
"""
@dataclass
[docs]
class InOutNodesMPKData(AFEMPKData):
"""
Provides the information about a node.
Attributes:
name: Node name which uniquely identifies the node over the scope of one model file.
type: Data type. From AFE's perspective, the type of all nodes is "buffer".
size: Node's data size in bytes.
"""
@dataclass
@dataclass
[docs]
class ConfigParamsMPKData(AFEMPKData):
"""
Provides the configuration parameters for the plugin.
Specification depends on the plugin's processor type.
"""
pass
@dataclass
[docs]
class EV74ConfigParamsMPKData(ConfigParamsMPKData):
"""
Configuration parameters for EV74 plugin.
Attributes:
desired_batch_size: Batch size requested by user.
actual_batch_size: Batch size used in code generation.
kernel: EV74 function.
params: Parameters of the EV74 function.
"""
[docs]
desired_batch_size: int
[docs]
params: dict[str, Any] = field(default_factory=dict)
@dataclass
[docs]
class APUConfigParamsMPKData(ConfigParamsMPKData):
"""
Configuration parameters for APU plugin.
Attributes:
input_names: Names of the parameters of the APU code's entry point function. These
names may be used when passing parameters at runtime. Must have the same length as input_types.
input_types: Shapes of input tensors, as they should be interpreted by the APU code.
output_types: Shapes of output tensors, as they should be interpreted by the APU code.
"""
[docs]
output_types: list[TensorTypeMPKData]
@dataclass
[docs]
class MLAConfigParamsMPKData(ConfigParamsMPKData):
"""
Configuration parameters for MLA plugin.
Attributes:
desired_batch_size: Batch size requested by user.
actual_batch_size: Batch size used in code generation.
number_of_quads_to_user: Number of quads used in MLA Production Compiler. Must be 4,
other values might be supported in the future.
"""
[docs]
desired_batch_size: int = 1
[docs]
actual_batch_size: int = 1
[docs]
number_of_quads_to_user: int = 4
@dataclass
@dataclass
[docs]
class PluginResourcesMPKData(AFEMPKData):
"""
Provides information about resources used by plugin.
Attributes:
executable: Name of the generated executable file. The file format depends on
which backend the resource belongs to.
"""
@dataclass
[docs]
class PluginMPKData(AFEMPKData):
"""
Provides data containing all information needed to generate MPK JSON data for a single plugin.
Attributes:
name: Plugin name.
sequence: Plugin's position in the model's execution sequence.
processor: Execution processor for the plugin. It can be either "MLA", "EV74" or "A65".
config_params: Configuration parameters for the plugin.
input_nodes: JSON objects providing information on the inputs for this plugin.
output_nodes: JSON objects providing information on the outputs for this plugin.
type: Plugin type. From the AFE's perspective, it always has the value "sgpProcess".
resources: JSON object providing information about resources used by this plugin.
"""
[docs]
config_params: ConfigParamsMPKData
[docs]
output_nodes: list[InOutNodesMPKData]
[docs]
type: str = "sgpProcess"
[docs]
resources: PluginResourcesMPKData | None = None
@dataclass
[docs]
class AwesomeNetMPKData(AFEMPKData):
"""
Provides data containing all information needed to generate the MPK JSON data for a single model.
Attributes:
name: Model name.
sequence: Model's position in an execution sequence.
model_path: Original model path.
model_checksum: Original model checksum.
input_nodes: An array of JSON objects providing information on the pipeline's input nodes.
plugins: An array of JSON objects, each of which define an invocation of a plugin either
on MLA, EV74 or A65.
"""
[docs]
model_sdk_version: str = get_model_sdk_version()
[docs]
plugins: list[PluginMPKData] = field(default_factory=list)
[docs]
def get_actual_batch_size(self) -> int | None:
"""
Returns actual batch size from plugins data.
Returns None if all plugins are assigned to A65.
Display warnings if plugins have different actual batch sizes.
"""
actual_batch_size = None
# Extract actual batch size from plugins data if it exists.
for plugin in self.plugins:
if isinstance(plugin.config_params, (EV74ConfigParamsMPKData, MLAConfigParamsMPKData)):
plugin_batch_size = plugin.config_params.actual_batch_size
if actual_batch_size is None:
actual_batch_size = plugin_batch_size
elif actual_batch_size != plugin_batch_size:
sima_logger.sima_log_warning('Some plugins does not match actual batch size')
break
return actual_batch_size
[docs]
def get_plugins_backend_distribution(self) -> dict:
"""
Returns number of plugins assigned to MLA, EV74 and A65.
"""
plugin_distribution = {key: 0 for key in ["MLA", "EV74", "A65"]}
for plugin in self.plugins:
plugin_distribution[plugin.processor] += 1
return plugin_distribution