PULPOpen Training Hardness: tiling-ready backward and optimizer deployment

.gitmodules

@@ -11,3 +11,6 @@
 [submodule "CMSIS-NN"]
 	path = TargetLibraries/CMSIS/third_party/CMSIS-NN
 	url = https://github.com/ARM-software/CMSIS-NN.git
+[submodule "pulp-trainlib"]
+	path = TargetLibraries/PULPOpen/third_party/pulp-trainlib
+	url = https://github.com/runwangdl/pulp-trainlib.git

Deeploy/CommonExtensions/OptimizationPasses/TopologyOptimizationPasses/LoweringOptimizationPasses.py

@@ -229,7 +229,7 @@ def _NCHWtoNHWC_fun(graph: gs.Graph, match: Match, name: str, default_channels_f
 
         if node.op in ["RequantizedConv", "Conv"]:
             spatialDims = len(node.inputs[1].shape) - 2
-        elif node.op == "MaxPool":
+        elif node.op in ["MaxPool", "AveragePool", "AveragePoolGrad"]:
             spatialDims = len(node.attrs["kernel_shape"])
         elif node.op == "Pad":
             spatialDims = 2  # Hack based on current status
@@ -245,7 +245,12 @@ def _NCHWtoNHWC_fun(graph: gs.Graph, match: Match, name: str, default_channels_f
         if node.op in ["Conv", "RequantizedConv"]:
             # In the case of Conv: [weights, opt. bias], RequantizedConv: [weights, mul, add, opt. shift]
             for tensor in node.inputs[1:]:
-                _transformLayoutConst(tensor, spatialDims, default_channels_first)
+                if isinstance(tensor, gs.Constant):
+                    _transformLayoutConst(tensor, spatialDims, default_channels_first)
+                elif isinstance(tensor, gs.Variable) and tensor.shape is not None and len(tensor.shape) >= 2:
+                    # Trainable weight (Variable input in training graph) — insert Transpose node
+                    perm = _transformLayoutPermutation(len(tensor.shape), spatialDims, default_channels_first)
+                    graph.nodes.append(_appendTranspose(tensor, node, perm))
 
         node.attrs["channels_first"] = default_channels_first
 
@@ -261,6 +266,24 @@ def __init__(self, default_channels_first: bool = True):
         super().__init__(graph, partial(_NCHWtoNHWC_fun, default_channels_first = default_channels_first), name)
 
 
+@contextagnostic
+class NCHWtoNHWCAveragePoolPass(ReplaceSequentialPatternPass):
+
+    def __init__(self, default_channels_first: bool = True):
+        graph = _singleNodePattern(op = "AveragePool")
+        name = "_NCHW_TO_NHWC_AVERAGEPOOL_PASS"
+        super().__init__(graph, partial(_NCHWtoNHWC_fun, default_channels_first = default_channels_first), name)
+
+
+@contextagnostic
+class NCHWtoNHWCAveragePoolGradPass(ReplaceSequentialPatternPass):
+
+    def __init__(self, default_channels_first: bool = True):
+        graph = _singleNodePattern(op = "AveragePoolGrad")
+        name = "_NCHW_TO_NHWC_AVERAGEPOOLGRAD_PASS"
+        super().__init__(graph, partial(_NCHWtoNHWC_fun, default_channels_first = default_channels_first), name)
+
+
 @contextagnostic
 class NCHWtoNHWCConvPass(ReplaceSequentialPatternPass):
 
@@ -363,6 +386,8 @@ def __init__(self, default_channels_first: bool = True):
         passes = [
             NCHWtoNHWCPadPass(default_channels_first),
             NCHWtoNHWCMaxPoolPass(default_channels_first),
+            NCHWtoNHWCAveragePoolPass(default_channels_first),
+            NCHWtoNHWCAveragePoolGradPass(default_channels_first),
             NCHWtoNHWCDwConvPass(default_channels_first),
             NCHWtoNHWCConvPass(default_channels_first),
         ]
@@ -376,6 +401,8 @@ def __init__(self, default_channels_first: bool = True):
         passes = [
             NCHWtoNHWCPadPass(default_channels_first),
             NCHWtoNHWCMaxPoolPass(default_channels_first),
+            NCHWtoNHWCAveragePoolPass(default_channels_first),
+            NCHWtoNHWCAveragePoolGradPass(default_channels_first),
             PULPNCHWtoNHWCDwConvPass(default_channels_first),
             NCHWtoNHWCConvPass(default_channels_first),
         ]
@@ -533,8 +560,10 @@ def _remove_only_singleton_reduce_mean(graph: gs.Graph, match: Match, name: str)
     # Delete node if only reduction over singleton dimensions
     if 'axis' in node.attrs:
         axis = node.attrs['axis']
-    else:
+    elif len(node.inputs) > 1 and node.inputs[1] is not None and hasattr(node.inputs[1], 'values') and node.inputs[1].values is not None:
         axis = node.inputs[1].values
+    else:
+        return graph  # axis unknown, skip
 
     # Check if shape information is available
     if node.inputs[0].shape is not None and all(node.inputs[0].shape[ax] == 1 for ax in axis):

Deeploy/DeeployTypes.py

@@ -343,7 +343,7 @@ def has_live_aliases(self, ctxt: NetworkContext) -> bool:
             next = queue.pop()
             buffNext = ctxt.lookup(next)
             assert isinstance(buffNext, VariableBuffer)
-            live |= buffNext._live
+            live |= buffNext._live or (next in ctxt.globalObjects)
             visited.add(next)
             queue |= buffNext.aliases - visited
         return live
@@ -1308,6 +1308,8 @@ def typeCheckNodeInputs(self, ctxt: NetworkContext, node: gs.Node) -> bool:
                         reference._type = _type
                         reference._instance = _type(inputNode.name, ctxt)
                 else:
+                    if not hasattr(reference, '_type'):
+                        return False
                     retCheck &= reference._type.referencedType == _type.referencedType
         return retCheck
 
@@ -1339,7 +1341,9 @@ def annotateDict(self, ctxt: NetworkContext, node: gs.Node, operatorRepresentati
         for key, value in operatorRepresentation.items():
             # check if the referenced buffer is in the environment
             if isinstance(value, str) and value in env:
-                self.typeDict[key + '_type'] = ctxt.lookup(value)._type
+                buf = ctxt.lookup(value)
+                if hasattr(buf, '_type'):
+                    self.typeDict[key + '_type'] = buf._type
 
     def typeCheck(self, ctxt: NetworkContext, node: gs.Node,
                   operatorRepresentation: OperatorRepresentation) -> Tuple[NetworkContext, bool]:
@@ -2086,11 +2090,15 @@ def bind(self, ctxt: NetworkContext) -> Tuple[NetworkContext, bool]:
             # Update shapes and types of tensors in onnx graph based on type inference after binding
             for node in (self.node.inputs + self.node.outputs):
                 if ctxt.is_local(node.name):
+                    if not hasattr(ctxt.localObjects[node.name], '_type'):
+                        continue  # skip untyped buffers (e.g. ReduceSum axes, MaxPool mask)
                     node.shape = ctxt.localObjects[node.name].shape
                     npType = self._broadcastToNpType(ctxt.localObjects[node.name]._type)
                     if npType is not None:
                         node.dtype = npType
                 elif ctxt.is_global(node.name):
+                    if not hasattr(ctxt.globalObjects[node.name], '_type'):
+                        continue  # skip untyped global buffers
                     npType = self._broadcastToNpType(ctxt.globalObjects[node.name]._type)
                     if isinstance(ctxt.globalObjects[node.name], ConstantBuffer):
                         if isinstance(node, gs.Constant):
@@ -2840,6 +2848,12 @@ def generateInferenceInitializationCode(self) -> str:
             if isinstance(node, StructBuffer):
                 continue
 
+            # Skip local buffers that were registered but never typed (e.g. optional ONNX
+            # outputs like the MaxPool indices/mask tensor).  These are not referenced by any
+            # template and must not be emitted as C declarations.
+            if not hasattr(node, '_type'):
+                continue
+
             name = node.name
             node.name = self.ctxt._mangle(node.name)
             callStack += node.init()
@@ -2940,6 +2954,8 @@ def generateBufferInitializationCode(self) -> str:
         callStack = ''
         for node in ctxt.globalObjects.values():
             if isinstance(node, VariableBuffer) and not isinstance(node, StructBuffer):
+                if not hasattr(node, '_type'):
+                    continue  # skip untyped buffers (e.g. ReduceSum axes constants)
                 assert issubclass(node._type, Pointer), f"Global VariableBuffer {node.name} is not a Pointer!"
                 if node._deploy:
                     name = node.name
@@ -2985,6 +3001,8 @@ def generateBufferAllocationCode(self) -> str:
 
         for node in ctxt.globalObjects.values():
             if isinstance(node, VariableBuffer) and not isinstance(node, StructBuffer):
+                if not hasattr(node, '_type'):
+                    continue  # skip untyped buffers (e.g. ReduceSum axes constants)
                 assert issubclass(node._type, Pointer), f"Global VariableBuffer {node.name} is not a Pointer!"
                 if node._deploy:
                     name = node.name
@@ -3522,6 +3540,8 @@ def _printMemorySummary(self):
                 if isinstance(_buffer, ConstantBuffer) or (isinstance(_buffer, VariableBuffer) and _buffer._deploy):
                     # SCHEREMO: We only
                     if (hasattr(_buffer, "_memoryLevel") and _buffer._memoryLevel == level) or level == "None":
+                        if not hasattr(_buffer, '_type'):
+                            continue  # skip untyped buffers (e.g. ReduceSum axes constants)
                         staticSize += int((np.prod(_buffer.shape) * _buffer._type.referencedType.typeWidth // 8))
                     else:
                         log.warning(f"Buffer {_buffer.name} does not have a valid memory level")