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,13 +229,14 @@ 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
         else:
             raise ValueError(f"Cannot determine spatialDims for node {node.name} with operator {node.op}")
 
+        # Insert Transpose nodes around the op's activation input/output to convert the data layout.
         permuteIn = _transformLayoutPermutation(len(tensorIn.shape), spatialDims, default_channels_first)
         graph.nodes.append(_appendTranspose(tensorIn, node, permuteIn))
 
@@ -245,7 +246,14 @@ 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):
+                    # Inference graph: weight is a fixed constant — permute its data in-place.
+                    _transformLayoutConst(tensor, spatialDims, default_channels_first)
+                elif isinstance(tensor, gs.Variable) and tensor.shape is not None and len(tensor.shape) >= 2:
+                    # Training graph: weight is a Variable (updated by the optimizer) — cannot permute
+                    # in-place, so insert an explicit Transpose node that will run at inference/forward time.
+                    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 +269,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 +389,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 +404,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 +563,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

@@ -336,14 +336,14 @@ def has_live_aliases(self, ctxt: NetworkContext) -> bool:
             True if this VariableBuffer has any live aliases, False otherwise
         """
         # Do a breadth-first search across the aliasing double-linked list
-        live = self._live
+        live = self._live or self.is_input or self.is_output
         queue = set(self.aliases)
         visited = set(self.name)
         while len(queue) > 0:
             next = queue.pop()
             buffNext = ctxt.lookup(next)
             assert isinstance(buffNext, VariableBuffer)
-            live |= buffNext._live
+            live |= buffNext._live or buffNext.is_input or buffNext.is_output
             visited.add(next)
             queue |= buffNext.aliases - visited
         return live

Deeploy/Targets/GAP9/Bindings.py

@@ -23,25 +23,33 @@
 # Import templates from PULPOpen and Generic
 from Deeploy.Targets.Generic.Templates import AddTemplate, ConcatTemplate, DequantTemplate, FloatReduceMeanTemplate, \
     FloatReduceSumTemplate, GatherTemplate, QuantTemplate, RQSiGELUTemplate, SliceTemplate, iHardswishTemplate
-from Deeploy.Targets.Generic.TypeCheckers import AddChecker, ConcatChecker, ConvChecker, DequantChecker, \
-    GatherChecker, GELUChecker, GEMMChecker, HardswishChecker, LayerNormChecker, MatMulChecker, MulChecker, \
-    QuantChecker, ReduceMeanChecker, ReluChecker, ReshapeChecker, RQAddChecker, RQHardswishChecker, SGDChecker, \
-    SliceChecker, SoftmaxChecker, SoftmaxCrossEntropyLossChecker, TransposeChecker
+from Deeploy.Targets.Generic.TypeCheckers import AddChecker, BatchNormInternalChecker, \
+    BatchNormalizationGradChecker, BNGradNormalizeChecker, BNGradReduceChecker, \
+    ChannelNormalizeChecker, ConcatChecker, ConvChecker, DequantChecker, \
+    WelfordReduceChecker, \
+    GatherChecker, GELUChecker, GEMMChecker, GlobalAveragePoolChecker, GlobalAveragePoolGradChecker, \
+    HardswishChecker, InPlaceAccumulatorV2Checker, LayerNormChecker, MatMulChecker, MaxPoolGradChecker, MulChecker, \
+    MSELossChecker, QuantChecker, ReduceMeanChecker, ReluChecker, ReshapeChecker, RQAddChecker, RQHardswishChecker, \
+    SGDChecker, SliceChecker, SoftmaxChecker, SoftmaxCrossEntropyLossChecker, TransposeChecker, \
+    PULPConvGradBChecker
 from Deeploy.Targets.PULPOpen.Bindings import ForkClosure, L3MemoryAwareFunctionCallClosure, \
     MemoryAwareForkTransformer, MemoryAwareFunctionCallClosure, TilingCallClosure
 from Deeploy.Targets.PULPOpen.CodeTransformationPasses.PULPClusterSynch import PULPSynchCoresPass
 from Deeploy.Targets.PULPOpen.CodeTransformationPasses.PULPClusterTiling import PULPClusterTiling
 from Deeploy.Targets.PULPOpen.CodeTransformationPasses.PULPL3Tiling import PULPL3Tiling
 from Deeploy.Targets.PULPOpen.CodeTransformationPasses.PULPProfileUntiled import PULPProfileUntiled
 from Deeploy.Targets.PULPOpen.DataTypes import PULPDMAFuture
-from Deeploy.Targets.PULPOpen.Templates import ConvTemplate, DMASliceTemplate, FloatAddTemplate, FloatConvTemplate, \
-    FloatGELUTemplate, FloatGemmTemplate, FloatLayernormTemplate, FloatMatMulTemplate, FloatMaxPoolTemplate, \
-    FloatMulTemplate, FloatReluTemplate, FloatSoftmaxTemplate, GEMMTemplate, MatrixVectorTemplate, MaxPoolTemplate, \
-    MulTemplate, ReduceMeanTemplate, RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, \
+from Deeploy.Targets.PULPOpen.Templates import ConvTemplate, DMASliceTemplate, FloatAddTemplate, \
+    FloatAveragePoolTemplate, FloatBatchNormTemplate, FloatConvGradTemplate, FloatConvTemplate, \
+    FloatGELUTemplate, FloatGemmTemplate, FloatGlobalAveragePoolTemplate, \
+    FloatInPlaceAccumulatorV2Template, FloatLayernormTemplate, FloatMatMulTemplate, \
+    FloatMaxPoolTemplate, FloatMulTemplate, FloatReluTemplate, FloatSoftmaxTemplate, GEMMTemplate, \
+    MatrixVectorTemplate, MaxPoolTemplate, MSELossTemplate, MulTemplate, ReduceMeanTemplate, \
+    RequantShiftTemplate, ReshapeTemplate, RQAddTemplate, RQSiHardswishTemplate, \
     SGDTemplate, SoftmaxCrossEntropyLossTemplate, TallGEMMTemplate, TransposeTemplate, UniformRequantShiftTemplate, \
     iRMSNormTemplate, iSoftmaxTemplate
-from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPLinearChecker, PULPMaxPoolChecker, \
-    PULPRequantShiftChecker
+from Deeploy.Targets.PULPOpen.TypeCheckers import PULPConvChecker, PULPLinearChecker, \
+    PULPMaxPoolChecker, PULPRequantShiftChecker
 from Deeploy.TilingExtension.CodeTransformationPasses.TilingVariableReplacement import TilingVariableReplacement, \
     TilingVariableReplacementUpdate
 
@@ -306,6 +314,14 @@
         SoftmaxCrossEntropyLossTemplate.referenceTemplate, GAP9Transformer) for type in IntegerDataTypes
 ]
 
+# Dual-output binding: outputs[0]=loss (scalar), outputs[1]=log_prob
+GAP9SoftmaxCrossEntropyLossDualOutputBindings = [
+    NodeBinding(
+        SoftmaxCrossEntropyLossChecker([PointerClass(float32_t), PointerClass(type)],
+                                       [PointerClass(float32_t), PointerClass(float32_t)]),
+        SoftmaxCrossEntropyLossTemplate.referenceDualOutputTemplate, GAP9Transformer) for type in IntegerDataTypes
+]
+
 GAP9SoftmaxCrossEntropyLossGradBindings = [
     NodeBinding(
         SoftmaxCrossEntropyLossChecker([PointerClass(float32_t), PointerClass(type)], [PointerClass(float32_t)]),
@@ -317,6 +333,179 @@
                 SGDTemplate.referenceTemplate, GAP9Transformer)
 ]
 
+# ── Training / Gradient bindings ─────────────────────────────────────────
+
+GAP9ReluGradBinding = NodeBinding(
+    ReluChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+    FloatReluTemplate.referenceGradTemplate, GAP9Transformer)
+
+GAP9FloatGELUGradBinding = NodeBinding(
+    GELUChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+    FloatGELUTemplate.referenceGradTemplate, GAP9Transformer)
+
+GAP9LayernormGradBinding = NodeBinding(
+    LayerNormChecker(
+        [PointerClass(float32_t),
+         PointerClass(float32_t),
+         PointerClass(float32_t),
+         PointerClass(float32_t),
+         PointerClass(float32_t)],
+        [PointerClass(float32_t),
+         PointerClass(float32_t),
+         PointerClass(float32_t)]), FloatLayernormTemplate.referenceGradTemplate,
+    GAP9Transformer)
+
+GAP9FloatConvGradW2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referenceConvGradW2DIm2ColTemplate,
+        GAP9ClusterTransformer)
+]
+
+GAP9FloatConvGradX2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referenceConvGradX2DIm2ColTiledTemplate,
+        GAP9Transformer)
+]
+
+GAP9FloatDWConvGradX2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referenceDWConvGradX2DTiledTemplate,
+        GAP9Transformer)
+]
+
+GAP9FloatDWConvGradW2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referenceDWConvGradW2DTemplate,
+        GAP9ClusterTransformer)
+]
+
+GAP9FloatPWConvGradW2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referencePWConvGradW2DTemplate,
+        GAP9ClusterTransformer)
+]
+
+GAP9FloatPWConvGradX2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvGradTemplate.referencePWConvGradX2DTemplate,
+        GAP9ClusterTransformer)
+]
+
+GAP9FloatConvGradBBindings = [
+    NodeBinding(
+        PULPConvGradBChecker([PointerClass(float32_t)],
+                             [PointerClass(float32_t)]), FloatConvGradTemplate.referenceConvGradB2DTemplate,
+        GAP9ClusterTransformer)
+]
+
+GAP9MaxPoolGrad2DBindings = [
+    NodeBinding(
+        MaxPoolGradChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+        FloatMaxPoolTemplate.referenceGradTemplate, GAP9Transformer)
+]
+
+GAP9AveragePool2DBindings = [
+    NodeBinding(PULPMaxPoolChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatAveragePoolTemplate.referenceTemplate, GAP9Transformer)
+]
+
+GAP9AveragePoolGrad2DBindings = [
+    NodeBinding(PULPMaxPoolChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatAveragePoolTemplate.referenceGradTemplate, GAP9Transformer)
+]
+
+GAP9GlobalAveragePool2DBindings = [
+    NodeBinding(
+        GlobalAveragePoolChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+        FloatGlobalAveragePoolTemplate.globalAveragePoolTemplate,
+        GAP9Transformer)
+]
+
+GAP9GlobalAveragePoolGrad2DBindings = [
+    NodeBinding(
+        GlobalAveragePoolGradChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
+        FloatGlobalAveragePoolTemplate.globalAveragePoolGradTemplate,
+        GAP9Transformer)
+]
+
+GAP9MSELossBindings = [
+    NodeBinding(MSELossChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+                MSELossTemplate.referenceTemplate, GAP9Transformer)
+]
+
+GAP9MSELossGradBindings = [
+    NodeBinding(MSELossChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+                MSELossTemplate.referenceGradientTemplate, GAP9Transformer)
+]
+
+GAP9InPlaceAccumulatorV2Bindings = [
+    NodeBinding(
+        InPlaceAccumulatorV2Checker(
+            [PointerClass(float32_t), PointerClass(float32_t), PointerClass(uint8_t)], [PointerClass(float32_t)]),
+        FloatInPlaceAccumulatorV2Template.referenceTemplate, GAP9Transformer)
+]
+
+GAP9InPlaceAccumulatorV2TiledBindings = [
+    NodeBinding(
+        InPlaceAccumulatorV2Checker(
+            [PointerClass(float32_t), PointerClass(float32_t), PointerClass(uint8_t)], [PointerClass(float32_t)]),
+        FloatInPlaceAccumulatorV2Template.tiledReferenceTemplate, GAP9Transformer)
+]
+
+GAP9BatchNormInternalBindings = [
+    NodeBinding(
+        BatchNormInternalChecker(
+            [PointerClass(float32_t)] * 5,
+            [PointerClass(float32_t)] * 5), FloatBatchNormTemplate.batchNormInternalTemplate,
+        GAP9Transformer)
+]
+
+GAP9BatchNormalizationGradBindings = [
+    NodeBinding(
+        BatchNormalizationGradChecker(
+            [PointerClass(float32_t)] * 5,
+            [PointerClass(float32_t)] * 3), FloatBatchNormTemplate.batchNormGradTemplate,
+        GAP9Transformer)
+]
+
+GAP9WelfordReduceBindings = [
+    NodeBinding(
+        WelfordReduceChecker(
+            [PointerClass(float32_t)] * 1,
+            [PointerClass(float32_t)] * 2), FloatBatchNormTemplate.welfordReduceTemplate,
+        GAP9Transformer)
+]
+
+GAP9ChannelNormalizeBindings = [
+    NodeBinding(
+        ChannelNormalizeChecker(
+            [PointerClass(float32_t)] * 5,
+            [PointerClass(float32_t)] * 1), FloatBatchNormTemplate.channelNormalizeTemplate,
+        GAP9Transformer)
+]
+
+GAP9BNGradReduceBindings = [
+    NodeBinding(
+        BNGradReduceChecker(
+            [PointerClass(float32_t)] * 4,
+            [PointerClass(float32_t)] * 2), FloatBatchNormTemplate.bnGradReduceTemplate,
+        GAP9Transformer)
+]
+
+GAP9BNGradNormalizeBindings = [
+    NodeBinding(
+        BNGradNormalizeChecker(
+            [PointerClass(float32_t)] * 7,
+            [PointerClass(float32_t)] * 1), FloatBatchNormTemplate.bnGradNormalizeTemplate,
+        GAP9Transformer)
+]
+
 GAP9TransposeBindings = [
     NodeBinding(TransposeChecker([PointerClass(type)], [PointerClass(type)]), TransposeTemplate.referenceTemplate,
                 GAP9Transformer) for type in IntegerDataTypes
@@ -328,6 +517,9 @@
 GAP9ConcatBindings = [
     NodeBinding(ConcatChecker([PointerClass(type), PointerClass(type)], [PointerClass(type)]),
                 ConcatTemplate.referenceTemplate, GAP9ClusterTransformer) for type in IntegerDataTypes
+] + [
+    NodeBinding(ConcatChecker([PointerClass(float_type), PointerClass(float_type)], [PointerClass(float_type)]),
+                ConcatTemplate.referenceTemplate, GAP9ClusterTransformer) for float_type in FloatDataTypes
 ]
 
 GAP9iRMSNormBindings = [
@@ -370,7 +562,18 @@
 GAP9ReluBinding = NodeBinding(ReluChecker([PointerClass(float32_t)], [PointerClass(float32_t)]),
                               FloatReluTemplate.referenceTemplate, GAP9Transformer)
 
+# Forward LayerNorm: 3 inputs (data, weight, bias), 3 outputs (Y, mean_stash, inv_std_stash)
+# The 3-output version is needed for training (backward pass needs mean and inv_std stashes).
 GAP9LayernormBinding = NodeBinding(
+    LayerNormChecker(
+        [PointerClass(float32_t), PointerClass(float32_t),
+         PointerClass(float32_t)],
+        [PointerClass(float32_t), PointerClass(float32_t),
+         PointerClass(float32_t)]), FloatLayernormTemplate.referenceTemplate,
+    GAP9Transformer)
+
+# Inference-only LayerNorm: 3 inputs, 1 output (Y only, no stashes)
+GAP9LayernormInferenceBinding = NodeBinding(
     LayerNormChecker(
         [PointerClass(float32_t), PointerClass(float32_t),
          PointerClass(float32_t)], [PointerClass(float32_t)]), FloatLayernormTemplate.referenceTemplate,

Deeploy/Targets/GAP9/DMA/L3Dma.py

@@ -29,7 +29,7 @@ class GAP9L3Dma(AsyncDma):
     _transferTemplates = {
         2:
             NodeTemplate(
-                "pi_cl_ram_copy_2d(get_ram_ptr(), ${ext}, ${loc}, ${transfer_size}, ${stride}, ${length}, ${ext2loc}, &${future});"
+                "pi_cl_ram_copy_2d(get_ram_ptr(), (uint32_t)${ext}, ${loc}, ${transfer_size}, ${stride}, ${length}, ${ext2loc}, &${future});"
             )
     }
     _waitingStrategy = PerTensorWaitingStrategy(GAP9L3DmaFuture)