#########################################################
# Copyright (C) 2021 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: Joey Chou
#########################################################
from enum import Enum
from typing import Optional, Any
from dataclasses import dataclass
from afe.ir.tensor_type import NodeType, set_node_type_batch_size
[docs]
class Backend(str, Enum):
"""
The Backend class contains all SiMa supported backends that can be partitioned via AFE's
graph partitioning pass. During graph partitioning AFE will assign each IR in the graph
to at least one of the backend type in this Backend class. The partitioned IRs will then
be sent to its corresponding backend and runtime later.
Currently AFE support 4 different backends graph annotation:
* MLA - SiMa.ai 1st Machine Learning Accerator
* APU - x86 or ARM A65
* EV - Synopsys' DSP IP
* CPU - compile on x86 for simulation purpose, used for unsupported operators
* NONE - For operations outside of partitions like tuple and tuple_getitem
In the TVM, strings 'llvm' and 'arm' have special meaning, while others do not.
"""
[docs]
APU = 'arm' # Set to ARM A65.
@classmethod
[docs]
def get_backend_from_name(cls, name: str) -> Optional["Backend"]:
assert isinstance(name, str)
for backend_name, backend in cls._value2member_map_.items():
if backend_name in name:
return backend
return None
@dataclass
[docs]
class BackendIR:
"""
IR of a computation that has been lowered to an external backend representation.
The computation has an input/output interface like a node.
:param graph: The lowered IR. The format of this IR depends on the backend.
For Backend.MLA, it is a ModelGraph.
For Backend.APU, it is a CompiledTVMObjectFile.
:param type: AwesomeNet type of the lowered IR
:param backend: Backend that the lowered IR is for.
:param tessellate_parameters: Tessellation parameters obtained while compiling the
BackendIR for MLA backend. Currently unused.
:param detessellate_parameters: Detessellation parameters obtained while compiling
the BackendIR for MLA backend. Used for creating a check file from untessellated
data while executing the compiled BackendIR using N2ACompiledBackendRunner.
:param stage: Stage number of the graph. Every graph have a unique stage number
that represents their order in AwesomeNet.
"""
[docs]
pack_parameters: Any | None = None
[docs]
unpack_parameters: Any | None = None
def __deepcopy__(self, memodict={}):
from copy import deepcopy
if self.backend == Backend.MLA:
# Cannot create a deep copy of ModelGraph object, so return its reference,
# as it shall not be mutated.
copy_graph = self.graph
else:
copy_graph = deepcopy(self.graph, memodict)
return BackendIR(
copy_graph, deepcopy(self.type), self.backend, deepcopy(self.pack_parameters),
deepcopy(self.unpack_parameters), self.stage
)
[docs]
def get_type(self) -> Optional[NodeType]:
"""
Get the type that this code has when it is used as a model graph node.
:return: The node's type
"""
return self.type
[docs]
def set_batch_size(self, batch_size: int):
"""
Modifies BackendIR's internal parameters to accommodate for a given batch size.
:param batch_size: Integer value representing the batch size of the inputs to the AwesomeNet.
"""
self.type = set_node_type_batch_size(self.type, batch_size)