perf(PULPOpen): let Gemm take a one-dimensional bias instead of widening it

Deeploy/Targets/PULPOpen/Layers.py

@@ -5,7 +5,7 @@
 from typing import List, Tuple
 
 from Deeploy.DeeployTypes import NodeMapper, Shape
-from Deeploy.Targets.Generic.Layers import RQGEMMLayer, RQSConvLayer
+from Deeploy.Targets.Generic.Layers import GEMMLayer, RQGEMMLayer, RQSConvLayer
 
 
 class PULPRQSConvLayer(RQSConvLayer):
@@ -41,3 +41,28 @@ def computeShapes(self, inputShapes: Shape, outputShapes: Shape, operatorReprese
         inputShapes[3] = [inputShapes[1][channelDim]]  # Channels out dimension of Kernel
 
         return (inputShapes, outputShapes)
+
+
+class PULPGEMMLayer(GEMMLayer):
+    """GEMM that keeps a one-dimensional bias one-dimensional.
+
+    The generic layer widens the C operand to [M, N] unconditionally. That is what
+    the kernel used to require, and for a transformer it is expensive: every Linear
+    bias is O values shared by all M rows, so the widened operand stores the same
+    vector M times. On CCT-2 at 64 tokens each 128-wide bias occupies 32 KB to carry
+    512 B, and the eight of them live at the peak account for 252 KB, around 13% of
+    the measured peak under LoRA and linear probing.
+
+    PULPFloatGEMMTemplate now passes the bias row stride to the kernel, so a
+    one-dimensional C is read with stride 0 and needs no widening. A C operand that
+    arrives with a real second dimension is left alone and behaves as before.
+    """
+
+    def __init__(self, maps: List[NodeMapper]):
+        super().__init__(maps)
+
+    def computeShapes(self, inputShapes: Shape, outputShapes: Shape, operatorRepresentation,
+                      channels_first) -> Tuple[Shape, Shape]:
+        if len(inputShapes) == 3 and len([d for d in inputShapes[2] if d != 1]) <= 1:
+            return (inputShapes, outputShapes)  # broadcast bias: leave as [O]
+        return super().computeShapes(inputShapes, outputShapes, operatorRepresentation, channels_first)

Deeploy/Targets/PULPOpen/Platform.py

@@ -15,7 +15,7 @@
     BasicRQIntegerDivBinding
 from Deeploy.Targets.Generic.Layers import AddLayer, AveragePoolGradLayer, AveragePoolLayer, \
     BatchNormalizationGradLayer, BatchNormInternalLayer, ConcatLayer, ConvGradBLayer, ConvGradWLayer, ConvGradXLayer, \
-    ConvLayer, GatherLayer, GELUGradLayer, GELULayer, GEMMLayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, \
+    ConvLayer, GatherLayer, GELUGradLayer, GELULayer, GlobalAveragePoolGradLayer, GlobalAveragePoolLayer, \
     InPlaceAccumulatorV2Layer, LayerNormGradLayer, LayerNormLayer, MatMulLayer, MaxPoolGradLayer, MaxPoolLayer, \
     MSELossGradLayer, MSELossLayer, MulLayer, PadLayer, QuantLayer, ReduceMeanLayer, ReduceSumLayer, ReluGradLayer, \
     ReluLayer, RequantShiftLayer, ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, RQSiHardswishLayer, SGDLayer, \
@@ -35,7 +35,7 @@
     MergeConstAddAndRequantPass, MergeTrueIntegerDivRequantShiftPass, QuantPatternPass, RQSSplitPass, \
     SkipEmptyConcatPass, SkipUnityRequantPass, iGELURequantMergePass, iHardswishRequantMergePass
 from Deeploy.Targets.PULPOpen.Bindings import BasicQuantBindings, PULPDMASliceBindings, PULPDWConv1DBinding
-from Deeploy.Targets.PULPOpen.Layers import PULPRQSConvLayer, PULPRQSGEMMLayer
+from Deeploy.Targets.PULPOpen.Layers import PULPGEMMLayer, PULPRQSConvLayer, PULPRQSGEMMLayer
 from Deeploy.Targets.PULPOpen.Parsers import PULPConv1DParser, PULPConv2DParser, PULPConvGradW2DParser, \
     PULPConvGradX2DParser, PULPDWConv1DParser, PULPDWConv2DParser, PULPDWConvGradW2DParser, PULPDWConvGradX2DParser, \
     PULPFPConv2DParser, PULPFPDWConv2DParser, PULPGEMMParser, PULPMatrixVecParser, PULPPWConvGradW2DParser, \
@@ -151,7 +151,7 @@
     'ConvGradB': ConvGradBLayer([ConvGradBMapper]),
     'RequantizedConv': PULPRQSConvLayer([Conv2DMapper, DWConv2DMapper, Conv1DMapper, DWConv1DMapper]),
     'RequantizedGemm': PULPRQSGEMMLayer([MatrixVecMapper, TallGEMMMapper, GEMMMapper]),
-    'Gemm': GEMMLayer([FloatGEMMMapper, GEMMDequantMapper]),
+    'Gemm': PULPGEMMLayer([FloatGEMMMapper, GEMMDequantMapper]),
     'Gelu': GELULayer([GELUMapper]),
     'GeluGrad': GELUGradLayer([GELUGradMapper]),
     'LayerNormalization': LayerNormLayer([LayerNormMapper]),

Deeploy/Targets/PULPOpen/Templates/FloatGemmTemplate.py

@@ -22,6 +22,18 @@ def alignToContext(self, ctxt: NetworkContext,
             operatorRepresentation['C'] = None
             operatorRepresentation['C_type'] = PointerClass(float32_t)  # Default to fp32 type
             operatorRepresentation['C_batched'] = False
+            operatorRepresentation['C_stride'] = 0
+            return ctxt, operatorRepresentation, []
+
+        # Row stride of the bias. A bias the parser left one-dimensional holds O
+        # values shared by every output row, so the kernel re-reads the same vector
+        # (stride 0) instead of the graph carrying M identical copies of it. A
+        # two-dimensional bias keeps the original [M,O] walk.
+        biasShape = ctxt.lookup(operatorRepresentation['C']).shape
+        broadcast = len([d for d in biasShape if d != 1]) <= 1
+        operatorRepresentation['C_stride'] = 0 if broadcast else operatorRepresentation['O']
+        if broadcast:
+            operatorRepresentation['C_batched'] = False
 
         return ctxt, operatorRepresentation, []
 
@@ -48,7 +60,8 @@ def alignToContext(self, ctxt: NetworkContext,
         ${N},
         ${O},
         ${transA},
-        ${transB}
+        ${transB},
+        ${C_stride}
     );
     % else:
     PULP_Gemm_fp${A_type.referencedType.typeWidth}_fp${B_type.referencedType.typeWidth}_fp${C_type.referencedType.typeWidth}_fp${data_out_type.referencedType.typeWidth}(
@@ -60,7 +73,8 @@ def alignToContext(self, ctxt: NetworkContext,
         ${N},
         ${O},
         ${transA},
-        ${transB}
+        ${transB},
+        ${C_stride}
     );
     % endif
     % if A_batched:

Deeploy/Targets/PULPOpen/TileConstraints/GEMMTileConstraint.py

@@ -233,11 +233,21 @@ def addGeometricalConstraint(tilerModel: TilerModel, parseDict: Dict, ctxt: Netw
 
         # Add bias constraints only if bias is present
         if has_bias:
-            dimOffsetC = len(bufferC.shape) - 2
-            addDimVar_1 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC)
-            addDimVar_2 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC + 1)
-            tilerModel.addConstraint(outputFirstDimVar == addDimVar_1)
-            tilerModel.addConstraint(outputSecondDimVar == addDimVar_2)
+            # A one-dimensional bias is broadcast over the output rows: it has no M
+            # extent to tie to the output's, and the kernel re-reads the same vector
+            # for every row (bias row stride 0). Constraining only the trailing
+            # dimension is what lets it stay [O] instead of being widened to [M,O].
+            biasDims = [d for d in bufferC.shape if d != 1]
+            if len(biasDims) <= 1:
+                lastDimIdx = len(bufferC.shape) - 1
+                addDimVar = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = lastDimIdx)
+                tilerModel.addConstraint(outputSecondDimVar == addDimVar)
+            else:
+                dimOffsetC = len(bufferC.shape) - 2
+                addDimVar_1 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC)
+                addDimVar_2 = tilerModel.getTensorDimVar(tensorName = bufferC.name, dimIdx = dimOffsetC + 1)
+                tilerModel.addConstraint(outputFirstDimVar == addDimVar_1)
+                tilerModel.addConstraint(outputSecondDimVar == addDimVar_2)
 
         return tilerModel
 
@@ -340,7 +350,18 @@ def serializeTilingSolution(
             inputBCubes.append(BCube)
 
             if has_bias:
-                CCube = HyperRectangle(tuple(cube.offset), tuple(cube.dims))
+                # A widened [M, O] bias is tiled exactly like the output. A broadcast
+                # bias is not: it only has the trailing O extent, so asking for the
+                # output's cube would have the DMA read M * O elements out of an
+                # O-element buffer. That is invisible when everything already sits in
+                # L2 and there is no transfer, and it is why this only ever showed up
+                # on the L3 and promote configurations.
+                biasShape = ctxt.lookup(operatorRepresentation['C']).shape
+                if len([d for d in biasShape if d != 1]) <= 1:
+                    leading = len(biasShape) - 1
+                    CCube = HyperRectangle(tuple([0] * leading + [OOffset]), tuple([1] * leading + [OSize]))
+                else:
+                    CCube = HyperRectangle(tuple(cube.offset), tuple(cube.dims))
                 inputAddCubes.append(CCube)
 
         inputLoadSchedule = []

TargetLibraries/GAP9/inc/kernel/GAP9Kernels.h

@@ -150,7 +150,7 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
                                    const float32_t *__restrict__ pDstC,
                                    float32_t *__restrict__ pDstY, uint32_t M,
                                    uint32_t N, uint32_t O, uint32_t transA,
-                                   uint32_t transB);
+                                   uint32_t transB, uint32_t biasStride);
 
 void PULP_GlobalAveragePoolGrad_fp32(const float32_t *dY, float32_t *dX,
                                      uint32_t N, uint32_t C, uint32_t H,

TargetLibraries/PULPOpen/inc/kernel/gemm.h

@@ -14,6 +14,6 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
                                    const float32_t *__restrict__ pDstC,
                                    float32_t *__restrict__ pDstY, uint32_t M,
                                    uint32_t N, uint32_t O, uint32_t transA,
-                                   uint32_t transB);
+                                   uint32_t transB, uint32_t biasStride);
 
 #endif // __DEEPLOY_MATH_GEMM_KERNEL_HEADER_

TargetLibraries/PULPOpen/src/Gemm.c

@@ -7,12 +7,19 @@
 #include "DeeployPULPMath.h"
 #include "pmsis.h"
 
+// biasStride is the row stride of pDstC: O for a full [M,O] bias, and 0 for a
+// broadcast [O] bias, where every output row adds the same vector. A
+// transformer stores its Linear bias as [O] but GEMMLayer.computeShapes used to
+// widen the C operand to [M,O], materialising the same O values M times -- 32
+// KB per 128-wide bias at 64 tokens, against 512 B of data. Passing a stride
+// keeps one kernel for both layouts, and biasStride == O reproduces the
+// previous behaviour exactly.
 void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
                                    const float32_t *__restrict__ pSrcB,
                                    const float32_t *__restrict__ pDstC,
                                    float32_t *__restrict__ pDstY, uint32_t M,
                                    uint32_t N, uint32_t O, uint32_t transA,
-                                   uint32_t transB) {
+                                   uint32_t transB, uint32_t biasStride) {
 
   int8_t core_id = pi_core_id();
   int8_t log2Core = LOG2(NUM_CORES);
@@ -35,7 +42,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
     for (uint32_t i = M_start; i < M_end; ++i) {
       const float32_t *__restrict__ a_row = &pSrcA[i * N];
       float32_t *__restrict__ y_row = &pDstY[i * O];
-      const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL;
+      const float32_t *__restrict__ c_row =
+          has_bias ? &pDstC[i * biasStride] : NULL;
 
       uint32_t j = 0;
 
@@ -85,7 +93,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
 
     for (uint32_t i = M_start; i < M_end; ++i) {
       float32_t *__restrict__ y_row = &pDstY[i * O];
-      const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL;
+      const float32_t *__restrict__ c_row =
+          has_bias ? &pDstC[i * biasStride] : NULL;
 
       uint32_t j = 0;
       for (; j < O_unroll; j += 6) {
@@ -181,7 +190,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
     for (uint32_t i = M_start; i < M_end; ++i) {
       const float32_t *__restrict__ a_row = &pSrcA[i * N];
       float32_t *__restrict__ y_row = &pDstY[i * O];
-      const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL;
+      const float32_t *__restrict__ c_row =
+          has_bias ? &pDstC[i * biasStride] : NULL;
 
       uint32_t j = 0;
       for (; j < O_unroll; j += 6) {
@@ -266,7 +276,8 @@ void PULP_Gemm_fp32_fp32_fp32_fp32(const float32_t *__restrict__ pSrcA,
 
     for (uint32_t i = M_start; i < M_end; ++i) {
       float32_t *__restrict__ y_row = &pDstY[i * O];
-      const float32_t *__restrict__ c_row = has_bias ? &pDstC[i * O] : NULL;
+      const float32_t *__restrict__ c_row =
+          has_bias ? &pDstC[i * biasStride] : NULL;
 
       uint32_t j = 0;
       for (; j < O_unroll; j += 6) {