Add OperatorDescriptor

CHANGELOG.md

@@ -177,9 +177,9 @@ This release containing major architectural changes, new platform support, enhan
 
 
 ### Added
-- BatchNorm kernel 
-- ConvTranspose kernel 
-- MaxPool1D kernel 
+- BatchNorm kernel
+- ConvTranspose kernel
+- MaxPool1D kernel
 - Template for 1D Convolution
 - Support for float32 data type in the previous kernels
 - Float binding for Pad1D kernel
@@ -318,7 +318,7 @@ This release containing major architectural changes, new platform support, enhan
 
 ### Changed
 - FloatConvTemplate file
-- Platform.py file  
+- Platform.py file
 - Bump the CMake version to 3.24 as required for the chimera-sdk
 - Bump GVSoC's version and add chimera simulation target
 - Rename the generic source util to utils to avoid name collision with chimera-sdk

Deeploy/CommonExtensions/DataTypes.py

@@ -87,11 +87,11 @@ class float64_t(FloatImmediate):
 
 SignedIntegerDataTypes: Tuple[Type[IntegerImmediate], ...] = (int8_t, int16_t, int32_t, int64_t)
 UnsignedIntegerDataTypes: Tuple[Type[IntegerImmediate], ...] = (uint8_t, uint16_t, uint32_t, uint64_t)
-IntegerDataTypes: Tuple[Type[IntegerImmediate], ...] = (sorted((
-    *SignedIntegerDataTypes,
-    *UnsignedIntegerDataTypes,
-),
-                                                               key = lambda _type: _type.typeWidth))
+IntegerDataTypes: Tuple[Type[IntegerImmediate], ...] = tuple(
+    sorted((
+        *SignedIntegerDataTypes,
+        *UnsignedIntegerDataTypes,
+    ), key = lambda _type: _type.typeWidth))
 FloatDataTypes: Tuple[Type[FloatImmediate], ...] = (bfloat16_t, float16_t, float32_t, float64_t)
 
 

Deeploy/CommonExtensions/NetworkDeployers/SignPropDeployer.py

@@ -7,7 +7,7 @@
 import onnx_graphsurgeon as gs
 
 from Deeploy.AbstractDataTypes import Pointer
-from Deeploy.DeeployTypes import DeploymentPlatform, NetworkDeployer, TopologyOptimizer
+from Deeploy.DeeployTypes import DeploymentPlatform, NetworkDeployer, OperatorDescriptor, TopologyOptimizer
 from Deeploy.Logging import DEFAULT_LOGGER as log
 
 
@@ -18,12 +18,13 @@ def __init__(self,
                  deploymentPlatform: DeploymentPlatform,
                  inputTypes: Dict[str, Type[Pointer]],
                  loweringOptimizer: TopologyOptimizer,
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable = lambda x: x,
                  name: str = 'DeeployNetwork',
                  default_channels_first: bool = True,
                  deeployStateDir: str = "DeeployState",
                  inputOffsets: Dict[str, int] = {}):
-        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, scheduler, name,
+        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, operatorDescriptors, scheduler, name,
                          default_channels_first, deeployStateDir)
 
         if inputOffsets == {}:

Deeploy/DeeployTypes.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import copy
+import itertools
 import math
 import os
 import pickle
@@ -1020,6 +1021,153 @@ def copy(self) -> NetworkContext:
         return copy.copy(self)
 
 
+class IoDesc:
+
+    def __init__(self, required: Union[str, List[str]], optional: Optional[Union[str, List[str]]] = None) -> None:
+        if isinstance(required, str):
+            required = [required]
+        self.required = required
+        optional = optional if optional is not None else []
+        if isinstance(optional, str):
+            optional = [optional]
+        self.optional = optional
+
+    def symbolicName(self, idx: int) -> str:
+        return (self.required + self.optional)[idx]
+
+    def checkTensors(self, tensors: Sequence[gs.Tensor]) -> bool:
+        return len(tensors) >= len(self.required) and \
+               len(tensors) <= len(self.required) + len(self.optional)
+
+
+class VariadicIoDesc(IoDesc):
+
+    def __init__(self, baseName: str, minNumTensors: int = 0) -> None:
+        self.baseName = baseName
+        self.minNumTensors = minNumTensors
+
+    def symbolicName(self, idx: int) -> str:
+        return f"{self.baseName}_{idx}"
+
+    def checkTensors(self, tensors: Sequence[gs.Tensor]) -> bool:
+        return len(tensors) >= self.minNumTensors
+
+
+@dataclass
+class AttrDesc:
+    name: str
+    unpacker: Callable[[Any], Any]
+    default: Optional[Union[Any, Callable[[gs.Node], Any]]] = None
+
+    @staticmethod
+    def _constUnpack(value: Any) -> Any:
+        if isinstance(value, gs.Constant):
+            return value.values.tolist()
+        elif isinstance(value, np.ndarray):
+            return value.tolist()
+        # LMACAN: hacky way to detect a 0-dim numpy array
+        elif hasattr(value, "ndim") and value.ndim == 0 and hasattr(value, "item"):
+            return value.item()
+        else:
+            return value
+
+    def unpack(self, value: Any) -> Union[int, float, List[int], List[float]]:
+        return self.unpacker(self._constUnpack(value))
+
+    def getDefault(self, node: gs.Node) -> Any:
+        if callable(self.default):
+            return self.default(node)
+        else:
+            return self.default
+
+
+@dataclass
+class OperatorDescriptor:
+    inputDescriptor: IoDesc
+    outputDescriptor: IoDesc
+    attrDescriptors: List[AttrDesc]
+
+    def check(self, node: gs.Node) -> bool:
+        """This method checks whether the node is valid.
+
+        Parameters
+        ----------
+        node : gs.Node
+            Graphsurgeon node to be validated
+
+        Returns
+        -------
+        bool : node validity
+
+        """
+        valid = True
+
+        if not self.inputDescriptor.checkTensors(node.inputs):
+            log.error(f"[OP {node.op}] Invalid input tensors: {[t.name for t in node.inputs]}")
+            valid = False
+
+        if not self.outputDescriptor.checkTensors(node.outputs):
+            log.error(f"[OP {node.op}] Invalid output tensors: {[t.name for t in node.outputs]}")
+            valid = False
+
+        for attrDesc in self.attrDescriptors:
+            if attrDesc.default is None and not attrDesc.name in node.attrs:
+                log.error(f"[OP {node.op}] Missing attribute {attrDesc.name}")
+                valid = False
+
+        return valid
+
+    def canonicalize(self, node: gs.Node, opset: int) -> bool:
+        _ = opset
+        for desc in self.attrDescriptors:
+            if desc.default is None:
+                value = node.attrs[desc.name]
+            else:
+                value = node.attrs.get(desc.name, desc.getDefault(node))
+            try:
+                node.attrs[desc.name] = desc.unpack(value)
+            except Exception as e:
+                raise ValueError(f"[OP {node.op}] Error unpacking the attribute {desc.name}. {e}") from e
+        return True
+
+    def parseTensors(self, ctxt: NetworkContext, tensors: Sequence[gs.Tensor],
+                     ioDesc: IoDesc) -> OperatorRepresentation:
+        opRepr = {}
+        for i, tensor in enumerate(tensors):
+            symName = ioDesc.symbolicName(i)
+            buffer = ctxt.lookup(tensor.name)
+            assert isinstance(buffer, VariableBuffer)
+            opRepr[symName] = buffer.name
+            opRepr[f"{symName}_shape"] = buffer.shape
+            opRepr[f"{symName}_size"] = math.prod(buffer.shape)
+            opRepr[f"{symName}_type"] = buffer._type
+        return opRepr
+
+    def parseAttrs(self, node: gs.Node) -> OperatorRepresentation:
+        return node.attrs.copy()
+
+    def parse(self, ctxt: NetworkContext, node: gs.Node) -> OperatorRepresentation:
+        opReprs = {
+            "input tensors": self.parseTensors(ctxt, node.inputs, self.inputDescriptor),
+            "output tensors": self.parseTensors(ctxt, node.outputs, self.outputDescriptor),
+            "attributes": self.parseAttrs(node),
+        }
+
+        for (firstName, firstOpRepr), (secondName, secondOpRepr) in itertools.combinations(opReprs.items(), 2):
+            firstKeySet = set(firstOpRepr.keys())
+            secondKeySet = set(secondOpRepr.keys())
+            assert firstKeySet.isdisjoint(secondKeySet), \
+                f"[OP {node.op}] Encourntered error while parsing node {node.name}. " \
+                f"Keys from parsing {firstName} clash with the keys from parsing {secondName}. "\
+                f"Overlapping keys: {firstKeySet ^ secondKeySet}"
+
+        resultOpRepr = {}
+        for opRepr in opReprs.values():
+            resultOpRepr.update(opRepr)
+
+        return resultOpRepr
+
+
 class NodeParser():
     """Deeploy's core Parser class. Analyzes network nodes and evaluates whether they can be mapped by it.
 
@@ -1143,7 +1291,9 @@ def _unpack_const(attr) -> Union[int, float]:
         The attributes can either be a numpy scalar value or a Constant tensor.
         This expects the numpy value to be of size 1.
         """
-        if isinstance(attr, gs.Constant):
+        if isinstance(attr, (int, float, bool, str)):
+            return attr
+        elif isinstance(attr, gs.Constant):
             value = attr.values
         elif isinstance(attr, np.ndarray):
             value = attr
@@ -2429,6 +2579,7 @@ def __init__(self,
                  graph: gs.Graph,
                  platform: DeploymentPlatform,
                  inputTypes: Dict[str, Type[Pointer]],
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable[[gs.Graph], Schedule] = lambda graph: list(graph.nodes),
                  name: str = 'DeeployNetwork',
                  deeployStateDir: str = "DeeployState"):
@@ -2453,6 +2604,7 @@ def __init__(self,
 
         """
         self.graph = graph
+        self.operatorDescriptors = operatorDescriptors
         self.scheduler = scheduler
         self.layerBinding: 'OrderedDict[str, ONNXLayer]' = OrderedDict()
         self.parsed = False
@@ -2582,6 +2734,16 @@ def _bindLayers(self):
                 flatSchedule += subGraph
 
         for node in flatSchedule:
+            assert node.op in self.operatorDescriptors, \
+                f"[ERROR] Error parsing node {node.name}. There is no descriptor for operator {node.op}."
+            desc = self.operatorDescriptors[node.op]
+            try:
+                desc.canonicalize(node, self.graph.opset)
+            except BaseException as e:
+                raise ValueError(f"[ERROR] Node {node.name} of op {node.op} could not be canonicalized.") from e
+            assert desc.check(node), \
+                f"[ERROR] Node {node.name} is not a valid instance of {node.op} operator"
+
             layer = self._mapNode(node)
             if isinstance(layer, ONNXLayer):
                 log.debug(f"   {SUCCESS_MARK} Bind {node.name} to layer {layer.__class__.__name__}")
@@ -3181,6 +3343,7 @@ def __init__(self,
                  deploymentPlatform: DeploymentPlatform,
                  inputTypes: Dict[str, Type[Pointer]],
                  loweringOptimizer: TopologyOptimizer,
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable[[gs.Graph], Schedule] = lambda graph: list(graph.nodes),
                  name: str = 'DeeployNetwork',
                  default_channels_first: bool = True,
@@ -3213,7 +3376,13 @@ def __init__(self,
 
 
         """
-        super().__init__(graph, deploymentPlatform, inputTypes, scheduler, name, deeployStateDir = deeployStateDir)
+        super().__init__(graph,
+                         deploymentPlatform,
+                         inputTypes,
+                         operatorDescriptors,
+                         scheduler,
+                         name,
+                         deeployStateDir = deeployStateDir)
 
         self.loweringOptimizer = loweringOptimizer
         self.default_channels_first = default_channels_first

Deeploy/EngineExtension/NetworkDeployers/EngineColoringDeployer.py

@@ -8,7 +8,8 @@
 
 from Deeploy.AbstractDataTypes import Pointer
 from Deeploy.CommonExtensions.NetworkDeployers.NetworkDeployerWrapper import NetworkDeployerWrapper
-from Deeploy.DeeployTypes import DeploymentPlatform, NetworkDeployer, ONNXLayer, Schedule, TopologyOptimizer
+from Deeploy.DeeployTypes import DeploymentPlatform, NetworkDeployer, ONNXLayer, OperatorDescriptor, Schedule, \
+    TopologyOptimizer
 from Deeploy.EngineExtension.OptimizationPasses.TopologyOptimizationPasses.EngineColoringPasses import \
     EngineColoringPass, EngineMapper
 
@@ -20,12 +21,13 @@ def __init__(self,
                  deploymentPlatform: DeploymentPlatform,
                  inputTypes: Dict[str, Type[Pointer]],
                  loweringOptimizer: TopologyOptimizer,
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable[[gs.Graph], Schedule] = lambda graph: list(graph.nodes),
                  name: str = 'DeeployNetwork',
                  default_channels_first: bool = True,
                  deeployStateDir: str = "DeeployState",
                  engineMapperCls: Type[EngineMapper] = EngineMapper):
-        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, scheduler, name,
+        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, operatorDescriptors, scheduler, name,
                          default_channels_first, deeployStateDir)
         self._initEngineColoringDeployer(engineMapperCls)
 

Deeploy/MemoryLevelExtension/NetworkDeployers/MemoryLevelDeployer.py

@@ -11,8 +11,8 @@
 from Deeploy.CommonExtensions.NetworkDeployers.NetworkDeployerWrapper import NetworkDeployerWrapper
 from Deeploy.CommonExtensions.NetworkDeployers.SignPropDeployer import SignPropDeployer
 from Deeploy.DeeployTypes import CodeGenVerbosity, ConstantBuffer, DeploymentEngine, DeploymentPlatform, \
-    NetworkContext, NetworkDeployer, NetworkOptimizationPass, NetworkOptimizer, Schedule, StructBuffer, \
-    TopologyOptimizer, TransientBuffer, VariableBuffer, _NoVerbosity
+    NetworkContext, NetworkDeployer, NetworkOptimizationPass, NetworkOptimizer, OperatorDescriptor, Schedule, \
+    StructBuffer, TopologyOptimizer, TransientBuffer, VariableBuffer, _NoVerbosity
 from Deeploy.Logging import DEFAULT_LOGGER as log
 from Deeploy.MemoryLevelExtension.MemoryLevels import MemoryHierarchy, MemoryLevel
 from Deeploy.MemoryLevelExtension.OptimizationPasses.MemoryLevelAnnotationPasses import AnnotateDefaultMemoryLevel
@@ -112,12 +112,13 @@ def __init__(self,
                  deploymentPlatform: Union[MemoryPlatform, MemoryPlatformWrapper],
                  inputTypes: Dict[str, Type[Pointer]],
                  loweringOptimizer: TopologyOptimizer,
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable[[gs.Graph], Schedule] = lambda graph: list(graph.nodes),
                  name: str = 'DeeployNetwork',
                  default_channels_first: bool = True,
                  deeployStateDir: str = "DeeployState",
                  memoryLevelAnnotationPasses: List[NetworkOptimizationPass] = []):
-        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, scheduler, name,
+        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, operatorDescriptors, scheduler, name,
                          default_channels_first, deeployStateDir)
         if len(memoryLevelAnnotationPasses) == 0:
             memoryLevelAnnotationPasses.append(AnnotateDefaultMemoryLevel(self.Platform.memoryHierarchy))
@@ -155,13 +156,14 @@ def __init__(self,
                  deploymentPlatform: Union[MemoryPlatform, MemoryPlatformWrapper],
                  inputTypes: Dict[str, Type[Pointer]],
                  loweringOptimizer: TopologyOptimizer,
+                 operatorDescriptors: Dict[str, OperatorDescriptor],
                  scheduler: Callable = lambda x: x,
                  name: str = 'DeeployNetwork',
                  default_channels_first: bool = True,
                  deeployStateDir: str = "DeeployState",
                  inputOffsets: Dict[str, int] = {},
                  memoryLevelAnnotationPasses: List[NetworkOptimizationPass] = []):
-        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, scheduler, name,
+        super().__init__(graph, deploymentPlatform, inputTypes, loweringOptimizer, operatorDescriptors, scheduler, name,
                          default_channels_first, deeployStateDir, inputOffsets)
         if len(memoryLevelAnnotationPasses) == 0:
             memoryLevelAnnotationPasses.append(AnnotateDefaultMemoryLevel(self.Platform.memoryHierarchy))