DRAFT: Parse refactor

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
 
 
 class SignPropDeployer(NetworkDeployer):
@@ -17,12 +17,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

@@ -1054,6 +1054,155 @@ def copy(self) -> NetworkContext:
         return copy.copy(self)
 
 
+class IoDesc:
+
+    def __init__(self, required: Union[str, List[str]], optional: Union[str, List[str]] = []) -> None:
+        if isinstance(required, str):
+            required = [required]
+        self.required = required
+        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):
+            # TODO: Change to logging
+            print(f"[ERROR OP {node.op}] Invalid input tensors: {[t.name for t in node.inputs]}")
+            valid = False
+
+        if not self.outputDescriptor.checkTensors(node.outputs):
+            # TODO: Change to logging
+            print(f"[ERROR 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:
+                # TODO: Change to logging
+                print(f"[ERROR 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"[ERROR 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 tesnors": 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"[PARSE ERROR] (Node: {node.name}, Op: {node.op}) " \
+                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.
 
@@ -1177,7 +1326,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
@@ -1898,44 +2049,38 @@ def broadcast(self, ctxt: NetworkContext, default_channels_first: bool = True) -
         inputShapes = [ctxt.lookup(node.name).shape for node in self.node.inputs]
         outputShapes = [ctxt.lookup(node.name).shape for node in self.node.outputs]
 
-        if not "channels_first" in self.mapper.parser.operatorRepresentation:
-            channels_first = default_channels_first
-        else:
-            channels_first = self.mapper.parser.operatorRepresentation['channels_first']
+        opRepr = self.mapper.parser.operatorRepresentation
+        channels_first = opRepr.get("channels_first", default_channels_first)
+        newInputShapes, newOutputShapes = self.computeShapes(inputShapes, outputShapes, opRepr, channels_first)
 
-        newInputShapes, newOutputShapes = self.computeShapes(inputShapes, outputShapes,
-                                                             self.mapper.parser.operatorRepresentation, channels_first)
+        for tensor, shape in zip(self.node.inputs + self.node.outputs, newInputShapes + newOutputShapes):
+            buffer = ctxt.lookup(tensor.name)
+            assert isinstance(buffer, VariableBuffer)
 
-        for node, newShape in zip(self.node.inputs + self.node.outputs, newInputShapes + newOutputShapes):
-            if ctxt.is_local(node.name):
-                ctxt.localObjects[node.name].shape = newShape
+            if ctxt.is_local(tensor.name):
+                buffer.shape = shape
                 # Update shape of tensors in onnx graph
-                node.shape = newShape
+                tensor.shape = shape
 
                 # WIESEP: It is possible that the type was not yet set, so we assume some default type
                 # At this state, we assume that all local buffers are float32 type inference is not yet done.
-                if node.dtype is None:
-                    node.dtype = np.float32
+                if tensor.dtype is None:
+                    tensor.dtype = np.float32
 
-            elif ctxt.is_global(node.name):
-                ctxt.globalObjects[node.name].shape = newShape
-                if isinstance(ctxt.globalObjects[node.name], ConstantBuffer):
+            elif ctxt.is_global(tensor.name):
+                buffer.shape = shape
+                if isinstance(buffer, ConstantBuffer):
 
                     # If the number of elements is equal, reshape
-                    if np.prod(ctxt.globalObjects[node.name].values.shape) == np.prod(newShape):
-                        ctxt.globalObjects[node.name].values.reshape(newShape)
+                    if np.prod(buffer.values.shape) == np.prod(shape):
+                        buffer.values.reshape(shape)
                     # The number of elements SHOULD be lower, and we broadcast
                     else:
                         try:
-                            ctxt.globalObjects[node.name].values = np.broadcast_to(ctxt.globalObjects[node.name].values,
-                                                                                   newShape)
-                        except:
-                            raise RuntimeError(
-                                f"Could not broadcast {node.name} from {ctxt.globalObjects[node.name].values.shape} to {newShape}!"
-                            )
-
-            else:
-                raise KeyError(f'Expected node {node.name} to be in context!')
+                            buffer.values = np.broadcast_to(buffer.values, shape)
+                        except ValueError as e:
+                            raise ValueError(
+                                f"Could not broadcast tensor {tensor.name} of node {self.node.name}.") from e
 
         return ctxt
 
@@ -2409,6 +2554,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"):
@@ -2433,6 +2579,7 @@ def __init__(self,
 
         """
         self.graph = graph
+        self.operatorDescriptors = operatorDescriptors
         self.scheduler = scheduler
         self.layerBinding: 'OrderedDict[str, ONNXLayer]' = OrderedDict()
         self.parsed = False
@@ -2604,19 +2751,34 @@ def parse(self, default_channels_first: bool = True) -> bool:
                                    constantBuffer = self.Platform.ConstantBuffer,
                                    structBuffer = self.Platform.StructBuffer,
                                    transientBuffer = self.Platform.TransientBuffer)
+        schedule = self.scheduler(self.graph)
+        flatSchedule = []
+        for subGraph in schedule:
+            if isinstance(subGraph, gs.Node):
+                flatSchedule.append(subGraph)
+            else:
+                flatSchedule += subGraph
 
-        self.ctxt = self._createIOBindings(self.ctxt, self.graph)
-
-        self._bindLayers()
+        self.layerBinding: 'OrderedDict[str, ONNXLayer]' = OrderedDict()
+        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]
+            desc.canonicalize(node, self.graph.opset)
+            assert desc.check(node), \
+            f"[ERROR] Node {node.name} is not a valid instance of {node.op} operator"
 
-        ctxt = self.ctxt.copy()
+            layer = self._mapNode(node)
+            if isinstance(layer, ONNXLayer):
+                self.layerBinding[layer.node.name] = layer
 
-        ctxtStack = deque()
         scheduledLayerList = list(self.layerBinding.values())
-        idx: int = 0
 
-        deepestIdx = 0
+        self.ctxt = self._createIOBindings(self.ctxt, self.graph)
 
+        ctxt = self.ctxt.copy()
+        ctxtStack = deque()
+        idx, deepestIdx = 0, 0
         while (idx < len(scheduledLayerList)):
             currentLayer = scheduledLayerList[idx]
 
@@ -3164,6 +3326,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,
@@ -3196,7 +3359,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)