Add FP32 operators and tiling support for MicroLlama on Snitch

.github/workflows/ci-platform-snitch-tiled.yml

@@ -36,3 +36,11 @@ jobs:
       runner: ${{ needs.select-env.outputs.runner }}
       docker-image: ${{ needs.select-env.outputs.image }}
       pytest-marker: "kernels and singlebuffer and l2"
+
+  snitch-models-tiled-singlebuffer-L2:
+    needs: select-env
+    uses: ./.github/workflows/_runner-snitch-tiled-sequential.yml
+    with:
+      runner: ${{ needs.select-env.outputs.runner }}
+      docker-image: ${{ needs.select-env.outputs.image }}
+      pytest-marker: "models and singlebuffer and l2"

.github/workflows/ci-platform-snitch.yml

@@ -36,3 +36,11 @@ jobs:
       runner: ${{ needs.select-env.outputs.runner }}
       docker-image: ${{ needs.select-env.outputs.image }}
       pytest-marker: "kernels"
+
+  snitch-models:
+    needs: select-env
+    uses: ./.github/workflows/_runner-snitch.yml
+    with:
+      runner: ${{ needs.select-env.outputs.runner }}
+      docker-image: ${{ needs.select-env.outputs.image }}
+      pytest-marker: "models"

.yamllint

@@ -31,3 +31,5 @@ ignore:
   - "**/toolchain/"
   # Ignore all files in .git
   - "**/.git/**"
+  # Ignore all files in .venv
+  - "**/.venv/"

Deeploy/DeeployTypes.py

@@ -3121,7 +3121,7 @@ def _exportGraph(self, folderPath, fileName):
         # VJUNG: ONNX-Graphsurgeon needs tensors to be in their export types
         constTensors = [tensor for tensor in self.graph.tensors().values() if isinstance(tensor, gs.Constant)]
         for tensor in constTensors:
-            if tensor.dtype != tensor.export_dtype:
+            if hasattr(tensor, 'export_dtype') and tensor.dtype != tensor.export_dtype:
                 tensor.values = tensor.values.astype(tensor.export_dtype)
 
         model = gs.export_onnx(self.graph)

Deeploy/Targets/Generic/Bindings.py

@@ -300,6 +300,9 @@
 BasicConcatBindings = [
     NodeBinding(ConcatChecker([PointerClass(type), PointerClass(type)], [PointerClass(type)]),
                 ConcatTemplate.referenceTemplate, BasicTransformer) for type in IntegerDataTypes
+] + [
+    NodeBinding(ConcatChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+                ConcatTemplate.referenceTemplate, BasicTransformer)
 ]
 
 BasicQuantBindings = [

Deeploy/Targets/Generic/Layers.py

@@ -703,6 +703,43 @@ def computeOps(self):
         return numPx * opsPerPx
 
 
+class RMSNormLayer(ONNXLayer):
+    """Layer support for the ONNX RMSNormalization operator.
+
+    Supported opset: 23
+
+    It is computed as follows:
+        - XSquared = Mul(X, X)
+        - XSquaredMean = ReduceMean<axes=normalized_axes>(XSquared)
+        - MeanSquareEpsilon = Add(XSquaredMean, epsilon)
+        - RMS = Sqrt(MeanSquareEpsilon)
+        - Normalized = Div(X, RMS)
+        - Y = Mul(Normalized, Scale)
+
+    For more details, this is the official ONNX documentation:
+    https://onnx.ai/onnx/operators/onnx__RMSNormalization.html#rmsnormalization-23
+    """
+
+    def __init__(self, maps: List[NodeMapper]):
+        super().__init__(maps)
+
+    def computeOps(self):
+        inputSize = self.mapper.parser.operatorRepresentation['inputSize']
+        NormalizedAxesSize = self.mapper.parser.operatorRepresentation['NormalizedAxesSize']
+        scale = self.mapper.parser.operatorRepresentation['scale']
+
+        # a. XSquared = Mul(X, X) => inputSize ops
+        # b. XSquaredMean = ReduceMean<axes=normalized_axes>(XSquared)
+        #    => inputSize ops (additions) + (inputSize - NormalizedAxesSize) ops (divisions)
+        # c. MeanSquareEpsilon = Add(XSquaredMean, epsilon) => (inputSize - NormalizedAxesSize) ops
+        # d. RMS = Sqrt(MeanSquareEpsilon) => (inputSize - NormalizedAxesSize) ops
+        # e. Normalized = Div(X, RMS) => inputSize ops
+        # f. Y = Mul(Normalized, Scale) => 0 if all(Scale == 1.0), else inputSize ops
+        scale_ops = 0 if (scale == 1.0).all() else inputSize
+        ops = 6 * inputSize - 3 * NormalizedAxesSize + scale_ops
+        return ops
+
+
 class CeilLayer(SingleOperationPerElementLayer):
     pass
 

Deeploy/Targets/Generic/Parsers.py

@@ -72,6 +72,10 @@ def parseNode(self, node: gs.Node) -> (bool):
             self.operatorRepresentation['n_levels'] = int(node.attrs['n_levels'])
             self.operatorRepresentation['log2D'] = int(math.log2(node.attrs['D']))
 
+            stash_type = node.attrs.get('stash_type', 1)
+            if stash_type != 1:
+                raise ValueError(f"iRMSNorm: only stash_type=1 (FP32) is supported, got {stash_type}")
+
         return ret
 
     def parseNodeCtxt(self,
@@ -87,8 +91,15 @@ def parseNodeCtxt(self,
         for idx, outputNode in enumerate(node.outputs):
             self.operatorRepresentation[outputs[idx]] = ctxt.lookup(outputNode.name).name
 
-        self.operatorRepresentation['size'] = np.prod(ctxt.lookup(node.inputs[0].name).shape)
-        self.operatorRepresentation['lastDimLength'] = ctxt.lookup(node.inputs[0].name).shape[-1]
+        input_shape = list(ctxt.lookup(node.inputs[0].name).shape)
+
+        axis = node.attrs.get('axis', -1)
+        if axis < 0:
+            axis = len(input_shape) + axis
+
+        self.operatorRepresentation['inputSize'] = int(np.prod(input_shape))
+        self.operatorRepresentation['NormalizedAxesSize'] = int(np.prod(input_shape[axis:]))
+        self.operatorRepresentation['scale'] = node.inputs[1].values
 
         return ctxt, True
 
@@ -199,6 +210,32 @@ def parseNodeCtxt(self,
         self.operatorRepresentation['data_in_size'] = np.prod(data_in.shape)
         self.operatorRepresentation['data_out_size'] = np.prod(data_out.shape)
 
+        # Transpose layout adaptation, derived purely from perm and the shapes.
+        # Kept here (in the parser) rather than in the per-platform templates:
+        # the multi-dim index strings and the outer-dim parallelization choice.
+        # The tiled Snitch/PULPOpen templates consume these; simple templates
+        # (CortexM/Generic/MemPool) just ignore the extra keys.
+        perm = self.operatorRepresentation['perm']
+        in_shape = list(data_in.shape)
+        out_shape = list(data_out.shape)
+
+        self.operatorRepresentation['shapeStr'] = "".join(f"[dimLen_{idx + 1}]" for idx in range(len(perm) - 1))
+        self.operatorRepresentation['outShapeStr'] = "".join(
+            f"[dimLen_{perm[idx + 1]}]" for idx in range(len(perm) - 1))
+        self.operatorRepresentation['dimStr'] = "".join(f"[{dim}]" for dim in in_shape)
+        self.operatorRepresentation['accessStr'] = "".join(f"[i_{idx}]" for idx in range(len(perm)))
+        self.operatorRepresentation['outAccessStr'] = "".join(f"[i_{i}]" for i in perm)
+
+        # Parallelize over the outermost dim with at least one element per core
+        # (>= 8): each core then owns a contiguous slab, minimizing per-core
+        # loop-control overhead. Fall back to the largest dim so few cores idle.
+        parallelDims = [idx for idx, dim in enumerate(out_shape) if dim >= 8]
+        self.operatorRepresentation['parallelDim'] = parallelDims[0] if parallelDims else out_shape.index(
+            max(out_shape))
+
+        for idx in range(len(perm)):
+            self.operatorRepresentation[f"dimLen_{idx}"] = in_shape[idx]
+
         return ctxt, True
 
 
@@ -488,23 +525,21 @@ def __init__(self):
         super().__init__()
 
     def parseNode(self, node: gs.Node) -> bool:
-
         ret = all([len(node.inputs) == 2, len(node.outputs) == 1])
-
         return ret
 
     def parseNodeCtxt(self,
                       ctxt: NetworkContext,
                       node: gs.Node,
                       channels_first: bool = True) -> Tuple[NetworkContext, bool]:
-
         data_in_1 = ctxt.lookup(node.inputs[0].name)
         data_in_2 = ctxt.lookup(node.inputs[1].name)
         data_out = ctxt.lookup(node.outputs[0].name)
+
         self.operatorRepresentation['data_in_1'] = data_in_1.name
         self.operatorRepresentation['data_in_2'] = data_in_2.name
         self.operatorRepresentation['data_out'] = data_out.name
-        self.operatorRepresentation['size'] = np.prod(data_in_1.shape)
+        self.operatorRepresentation['size'] = np.prod(data_out.shape)
 
         return ctxt, True
 
@@ -2097,15 +2132,15 @@ def parseNodeCtxt(self,
                       node: gs.Node,
                       channels_first: bool = True) -> Tuple[NetworkContext, bool]:
 
-        inputs = ["input1", "input2"]
-        outputs = ["output"]
+        inputs = ["A", "B"]
+        outputs = ["C"]
         for idx, inputNode in enumerate(node.inputs):
             if idx < len(inputs):
                 self.operatorRepresentation[inputs[idx]] = ctxt.lookup(inputNode.name).name
         for idx, outputNode in enumerate(node.outputs):
             self.operatorRepresentation[outputs[idx]] = ctxt.lookup(outputNode.name).name
 
-        self.operatorRepresentation['size'] = np.prod(ctxt.lookup(self.operatorRepresentation['input1']).shape)
+        self.operatorRepresentation['size'] = np.prod(ctxt.lookup(self.operatorRepresentation['A']).shape)
 
         return ctxt, True
 

Deeploy/Targets/Generic/Templates/FloatDivTemplate.py

@@ -6,5 +6,5 @@
 
 referenceTemplate = NodeTemplate("""
 // Division (Name: ${nodeName}, Op: ${nodeOp})
-SINGLE_CORE Div_fp${input1_type.referencedType.typeWidth}_fp${input2_type.referencedType.typeWidth}_fp${output_type.referencedType.typeWidth}(${input1}, ${input2}, ${output}, ${size});
+SINGLE_CORE Div_fp${A_type.referencedType.typeWidth}_fp${B_type.referencedType.typeWidth}_fp${C_type.referencedType.typeWidth}(${A}, ${B}, ${C}, ${size});
 """)

Deeploy/Targets/Generic/Templates/ReshapeTemplate.py

@@ -34,6 +34,12 @@ def alignToContext(self, ctxt: NetworkContext,
         bufferIn.aliases.add(bufferOut.name)
         bufferOut.aliases.add(bufferIn.name)
 
+        # Tiling still reads the legacy single-valued `_alias` attribute
+        # (TilerExtension / MemoryScheduler). Set it here so platforms that
+        # rely on Reshape pointer-passthrough during tiling don't each need
+        # to carry the same workaround in a subclass.
+        bufferOut._alias = bufferIn.name
+
         return ctxt, operatorRepresentation, []
 
 

Deeploy/Targets/Generic/Templates/iRMSNormTemplate.py

@@ -23,5 +23,5 @@ def alignToContext(self, ctxt: NetworkContext,
 
 referenceTemplate = _iRMSNormTemplate("""
 // iRMSnorm (Name: ${nodeName}, Op: ${nodeOp})
-SINGLE_CORE iRMSnorm_s${data_in_type.referencedType.typeWidth}_s${data_out_type.referencedType.typeWidth}(${data_in}, ${data_out}, ${weight}, ${input_offset}, ${size}, ${lastDimLength}, ${log2D});
+SINGLE_CORE iRMSnorm_s${data_in_type.referencedType.typeWidth}_s${data_out_type.referencedType.typeWidth}(${data_in}, ${data_out}, ${weight}, ${input_offset}, ${inputSize}, ${NormalizedAxesSize}, ${log2D});
 """)

Deeploy/Targets/Generic/TileConstraints/BOPTileConstraint.py

@@ -12,11 +12,26 @@
 from Deeploy.TilingExtension.MemoryConstraints import NodeMemoryConstraint
 from Deeploy.TilingExtension.TileConstraint import TileConstraint
 from Deeploy.TilingExtension.TilerModel import TilerModel
-from Deeploy.TilingExtension.TilingCodegen import AbsoluteHyperRectangle, TilingSchedule, VariableReplacementScheme
+from Deeploy.TilingExtension.TilingCodegen import AbsoluteHyperRectangle, HyperRectangle, TilingSchedule, \
+    VariableReplacementScheme
 
 
 class BOPTileConstraint(TileConstraint):
-    """Tile constraint class for binary operators, i.e. operators that use two input tensors of equal dimensions
+    """Tile constraint class for binary operators, i.e. operators that have exactly 2 inputs and 1 output.
+
+    When the second input is a scalar (total size 1), it is kept full-size and only
+    the first input and the output are tiled together. This supports ONNX
+    broadcasting in operators that have a corresponding scalar kernel.
+
+    Warning:
+        Broadcasting support is partial -- only the case of a fully-scalar
+        second input (np.prod(input2.shape) == 1) is handled. Other ONNX
+        broadcasting patterns -- input1 scalar, partial broadcasting such
+        as (N, 1) + (1, M), single-dim broadcasting such as (N, M, K) +
+        (N, 1, K), or rank-mismatched shapes such as (N, M) + (M,) --
+        fall through to the non-scalar branch, where the dim-equality
+        constraints will fail to satisfy. Operators that need full ONNX
+        broadcasting must use a different tile constraint.
     """
 
     dataIn1Name = 'data_in_1'  #: str: Name of the first input tensor as defined by the operator's parser
@@ -34,14 +49,27 @@ def addGeometricalConstraint(cls, tilerModel: TilerModel, parseDict: Dict, ctxt:
             tilerModel.addTensorDimToModel(ctxt, bufferName)
 
         input1Shape = ctxt.lookup(inputBuffer1Name).shape
-
-        for dim in range(len(input1Shape)):
-            inputDim1Var = tilerModel.getTensorDimVar(tensorName = inputBuffer1Name, dimIdx = dim)
-            inputDim2Var = tilerModel.getTensorDimVar(tensorName = inputBuffer2Name, dimIdx = dim)
-            outputDimVar = tilerModel.getTensorDimVar(tensorName = outputBufferName, dimIdx = dim)
-
-            tilerModel.addConstraint(inputDim1Var == inputDim2Var)
-            tilerModel.addConstraint(inputDim1Var == outputDimVar)
+        input2Shape = list(ctxt.lookup(inputBuffer2Name).shape)
+        input2_is_scalar = (np.prod(input2Shape) == 1)
+
+        if input2_is_scalar:
+            # Scalar broadcasting: tile input1 and output together; input2 stays full-size.
+            for dim in range(len(input1Shape)):
+                inputDim1Var = tilerModel.getTensorDimVar(tensorName = inputBuffer1Name, dimIdx = dim)
+                outputDimVar = tilerModel.getTensorDimVar(tensorName = outputBufferName, dimIdx = dim)
+                tilerModel.addConstraint(inputDim1Var == outputDimVar)
+            for dim in range(len(input2Shape)):
+                inputDim2Var = tilerModel.getTensorDimVar(tensorName = inputBuffer2Name, dimIdx = dim)
+                tilerModel.addConstraint(inputDim2Var == input2Shape[dim])
+        else:
+            # Element-wise: all three tensors tiled identically.
+            for dim in range(len(input1Shape)):
+                inputDim1Var = tilerModel.getTensorDimVar(tensorName = inputBuffer1Name, dimIdx = dim)
+                inputDim2Var = tilerModel.getTensorDimVar(tensorName = inputBuffer2Name, dimIdx = dim)
+                outputDimVar = tilerModel.getTensorDimVar(tensorName = outputBufferName, dimIdx = dim)
+
+                tilerModel.addConstraint(inputDim1Var == inputDim2Var)
+                tilerModel.addConstraint(inputDim1Var == outputDimVar)
 
         return tilerModel
 
@@ -64,11 +92,18 @@ def serializeTilingSolution(
             newSize = np.prod(cube.dims)
             replacements["size"].append(newSize)
 
+        input2Shape = list(ctxt.lookup(operatorRepresentation[cls.dataIn2Name]).shape)
+        input2_is_scalar = (np.prod(input2Shape) == 1)
+
         inputLoadSchedule = []
         outputLoadSchedule = []
 
         for cube in outputCubes:
-            inputLoadSchedule.append({cls.dataIn1Name: cube, cls.dataIn2Name: cube})
+            if input2_is_scalar:
+                in2Cube = HyperRectangle(tuple([0] * len(input2Shape)), tuple(input2Shape))
+                inputLoadSchedule.append({cls.dataIn1Name: cube, cls.dataIn2Name: in2Cube})
+            else:
+                inputLoadSchedule.append({cls.dataIn1Name: cube, cls.dataIn2Name: cube})
 
         for out in outputCubes:
             outputLoadSchedule.append({cls.dataOutName: out})

Deeploy/Targets/Generic/TileConstraints/iRMSNormTileConstraint.py

@@ -54,12 +54,12 @@ def serializeTilingSolution(
         addrNames = ['data_in', 'weight', 'data_out']
         inputBaseOffsets, outputBaseOffsets = cls.extractBaseAddr(tilingSolution, targetMemLevel,
                                                                   operatorRepresentation, addrNames)
-        replacements = {"size": []}
-        replacementTypes = {"size": PointerClass(uint16_t)}
+        replacements = {"inputSize": []}
+        replacementTypes = {"inputSize": PointerClass(uint16_t)}
 
         for cube in outputCubes:
             newSize = np.prod(cube.dims)
-            replacements["size"].append(newSize)
+            replacements["inputSize"].append(newSize)
 
         inputLoadSchedule = []
         outputLoadSchedule = []

Deeploy/Targets/Generic/TypeCheckers.py

@@ -6,7 +6,7 @@
 
 import numpy as np
 
-from Deeploy.AbstractDataTypes import Pointer
+from Deeploy.AbstractDataTypes import FloatImmediate, Pointer
 from Deeploy.CommonExtensions.TypeCheckers.SignPropTypeChecker import SignPropTypeChecker
 from Deeploy.DeeployTypes import ConstantBuffer, OperatorRepresentation, VariableBuffer
 
@@ -409,7 +409,10 @@ def __init__(self, input_types: Sequence[Type[Pointer]], output_types: Sequence[
 
     def _inferNumLevels(self, inputs: List[VariableBuffer],
                         operatorRepresentation: OperatorRepresentation) -> List[int]:
-        return [2**(4 * self.input_types[0].referencedType.typeWidth)]
+        input_type = self.input_types[0].referencedType
+        if issubclass(input_type, FloatImmediate):
+            return [2**(input_type.typeWidth)]
+        return [2**(4 * input_type.typeWidth)]
 
     def _inferSignedness(self, inputs: List[VariableBuffer],
                          operatorRepresentation: OperatorRepresentation) -> List[bool]:
@@ -610,3 +613,23 @@ def _inferNumLevels(self, inputs: List[VariableBuffer],
     def _inferSignedness(self, inputs: List[VariableBuffer],
                          operatorRepresentation: OperatorRepresentation) -> List[bool]:
         return [True]
+
+
+class RMSNormChecker(SignPropTypeChecker):
+
+    def __init__(self, input_types: Sequence[Type[Pointer]], output_types: Sequence[Type[Pointer]]):
+        super().__init__(input_types, output_types)
+
+    def _inferNumLevels(self, inputs: List[VariableBuffer],
+                        operatorRepresentation: OperatorRepresentation) -> List[int]:
+        # RMSNorm: square, mean, sqrt, reciprocal, multiply
+        # Output precision similar to input
+        return [2**(self.input_types[0].referencedType.typeWidth)]
+
+    def _inferSignedness(self, inputs: List[VariableBuffer],
+                         operatorRepresentation: OperatorRepresentation) -> List[bool]:
+        # RMSNorm output can be signed (depending on input signedness)
+        if inputs[0]._signed:
+            return [True]
+        else:
+            return [False]