[DRAFT] Spatz integration

CMakeLists.txt

@@ -20,8 +20,8 @@ if(TOOLCHAIN STREQUAL GCC)
   set(CMAKE_INTERPROCEDURAL_OPTIMIZATION TRUE)
 endif()
 
-set(platform MemPool CACHE STRING "Platform (MemPool, SoftHier, QEMU, Siracusa, Siracusa_w_neureka, PULP-Open, GAP9, Generic, Snitch)")
-set_property(CACHE platform PROPERTY STRINGS MemPool SoftHier QEMU Siracusa Siracusa_w_neureka PULP-Open GAP9 Generic Snitch)
+set(platform MemPool CACHE STRING "Platform (MemPool, SoftHier, QEMU, Siracusa, Siracusa_w_neureka, PULP-Open, GAP9, Generic, Snitch, Spatz)")
+set_property(CACHE platform PROPERTY STRINGS MemPool SoftHier QEMU Siracusa Siracusa_w_neureka PULP-Open GAP9 Generic Snitch Spatz)
 
 if(platform STREQUAL MemPool)
   message(STATUS "Building for platform 'MemPool'")
@@ -46,6 +46,8 @@ elseif(platform STREQUAL SoftHier)
   message(STATUS "Building for platform 'SoftHier'")
 elseif(platform STREQUAL Chimera)
   message(STATUS "Building for platform 'Chimera'")
+elseif(platform STREQUAL Spatz)
+  message(STATUS "Building for platform 'Spatz'")
 else()
   message(FATAL_ERROR "Invalid platform '${platform}' specified!")
 endif()
@@ -300,4 +302,34 @@ if(platform STREQUAL Chimera)
 endif()
 
 
+
+if(platform STREQUAL Spatz)
+
+  if(NOT DEFINED ENV{SPATZ_HOME})
+    message(FATAL_ERROR "Environment variable SPATZ_HOME not set.")
+  endif()
+
+  set(SPATZ_HOME $ENV{SPATZ_HOME})
+
+  set(CMAKE_TOOLCHAIN_FILE ${CMAKE_CURRENT_LIST_DIR}/cmake/spatz/toolchain_llvm.cmake)
+
+  include(${CMAKE_CURRENT_LIST_DIR}/cmake/spatz/spatz.cmake)
+
+  project(deeploy LANGUAGES C ASM)
+
+  message(STATUS "============================= ${platform} Configuration ============================")
+  message(STATUS "[cMake  ]   ISA                      = " ${ISA})
+  message(STATUS "================================================================================")
+  message(STATUS "")
+
+  add_subdirectory(TargetLibraries/Generic)
+  add_subdirectory(TargetLibraries/Spatz)
+  target_include_directories(deeployspatz PUBLIC TargetLibraries/Generic/inc)
+
+  add_subdirectory(DeeployTest)
+  target_link_libraries(deeploylib INTERFACE deeploybasic deeployspatz)
+
+endif()
+
+
 print_simulation_config()

Deeploy/Targets/Generic/Bindings.py

@@ -19,12 +19,12 @@
     GatherTemplate, GemmTemplate, IntegerDivTemplate, ITAMaxTemplate, ITAPartialMaxTemplate, MatMulTemplate, \
     MaxPoolTemplate, MulTemplate, PadTemplate, QuantTemplate, ReduceMeanTemplate, ReduceSumTemplate, \
     RequantShiftTemplate, ReshapeTemplate, RQIntegerDivTemplate, RQSiGELUTemplate, SliceTemplate, TransposeTemplate, \
-    iGELUTemplate, iLayernormTemplate, iRMSNormTemplate, iSoftmaxTemplate
+    iGELUTemplate, iLayernormTemplate, iRMSNormTemplate, iSoftmaxTemplate, TopKTemplate
 from Deeploy.Targets.Generic.TypeCheckers import AddChecker, BatchNormChecker, ConcatChecker, ConvChecker, \
     DebugPrintChecker, DequantChecker, DivChecker, DummyChecker, GatherChecker, GELUChecker, GEMMChecker, \
     LayerNormChecker, MatMulChecker, MaxPoolChecker, MulChecker, PadChecker, QuantChecker, ReduceMeanChecker, \
     ReduceSumChecker, ReluChecker, RequantShiftChecker, ReshapeChecker, RQIntegerDivChecker, SliceChecker, \
-    SoftmaxChecker, TransposeChecker
+    SoftmaxChecker, TransposeChecker, TopKChecker
 
 BasicTransformer = CodeTransformation([ArgumentStructGeneration(), MemoryManagementGeneration(), FutureGeneration()])
 
@@ -327,3 +327,14 @@
         ConvTransposeTemplate.referenceTemplate,
         BasicTransformer) for type in FloatDataTypes
 ]
+
+BasicTopKBindings = [
+    NodeBinding(
+        TopKChecker(
+            [PointerClass(float32_t), PointerClass(int8_t)], # inputs
+            [PointerClass(float32_t), PointerClass(int8_t)] # outputs
+        ),
+        TopKTemplate.referenceTemplate,
+        BasicTransformer,
+    )
+]

Deeploy/Targets/Generic/Layers.py

@@ -709,3 +709,15 @@ def computeOps(self):
             numPx = opRep['dim_im_out_x']
 
         return numPx * opsPerPx
+
+
+class TopKLayer(ONNXLayer):
+
+    def __init__(self, maps: List[NodeMapper]):
+        super().__init__(maps)
+
+#     def computeOps(self):
+#         ???
+# 
+#     def computeShapes(self):
+#         ???

Deeploy/Targets/Generic/Parsers.py

@@ -982,7 +982,7 @@ def parseNode(self, node: gs.Node) -> (bool):
             return False
 
         indices_shape = node.inputs[1].shape
-        assert np.prod(indices_shape) == 1, f"Only indices of size 1 supported. Got indices of shape {indices_shape}"
+        self.operatorRepresentation['num_indices'] = int(np.prod(indices_shape))
 
         self.operatorRepresentation['axis'] = node.attrs['axis'] if 'axis' in node.attrs else 0
         return True
@@ -1002,10 +1002,17 @@ def parseNodeCtxt(self,
 
         axis = self.operatorRepresentation['axis']
         shape = ctxt.lookup(node.inputs[0].name).shape
-        self.operatorRepresentation['batch'] = np.prod(shape[:axis])
-        self.operatorRepresentation['batch_length'] = np.prod(shape[axis:])
-        self.operatorRepresentation['axis_length'] = np.prod(shape[axis + 1:])
-        self.operatorRepresentation['index'] = int(node.inputs[1].values.item())
+        self.operatorRepresentation['batch'] = int(np.prod(shape[:axis])) if axis > 0 else 1
+        self.operatorRepresentation['batch_length'] = int(np.prod(shape[axis:]))
+        self.operatorRepresentation['axis_length'] = int(np.prod(shape[axis + 1:])) if axis + 1 < len(shape) else 1
+
+        if self.operatorRepresentation['num_indices'] == 1:
+            try:
+                self.operatorRepresentation['index'] = int(node.inputs[1].values.item())
+            except Exception:
+                self.operatorRepresentation['index'] = f"{self.operatorRepresentation['indices']}[0]"
+        else:
+            self.operatorRepresentation['index'] = 0 # in this case is not used but is needed for mako template
 
         return ctxt, True
 
@@ -2886,3 +2893,30 @@ def parseNodeCtxt(self,
         self.operatorRepresentation['size'] = int(np.prod(data_in.shape))
 
         return ctxt, True
+
+# TopKParser: selects the largest k elements from a vector
+class TopKParser(NodeParser):
+    def __init__(self):
+        super().__init__()
+
+    def parseNode(self, node: gs.Node) -> bool:
+        return len(node.inputs)==2 and len(node.outputs)==2 and node.op=='TopK'
+
+    def parseNodeCtxt(self,
+                      ctxt: NetworkContext,
+                      node: gs.Node,
+                      channels_first: bool = True) -> Tuple[NetworkContext, bool]:
+        data_in = ctxt.lookup(node.inputs[0].name)
+        k_in = ctxt.lookup(node.inputs[1].name)
+        values_out = ctxt.lookup(node.outputs[0].name)
+        indices_out = ctxt.lookup(node.outputs[1].name)
+
+        self.operatorRepresentation['data_in'] = data_in.name
+        self.operatorRepresentation['data_in_size'] = int(np.prod(data_in.shape))
+        self.operatorRepresentation['k_value'] = int(k_in.values[0])
+        self.operatorRepresentation['values_out'] = values_out.name
+        self.operatorRepresentation['indices_out'] = indices_out.name
+
+        print("parsenodectxt returning true")
+
+        return ctxt, True

Deeploy/Targets/Generic/Platform.py

@@ -14,19 +14,19 @@
     BasicPad1DBindings, BasicPad2DBindings, BasicPowBindings, BasicQuantBindings, BasicReduceMeanBindings, \
     BasicReduceSumBindings, BasicReluBinding, BasicReshapeBindings, BasicRQIntegerDivBinding, BasicRQSBindings, \
     BasicRQSGELUBinding, BasicSliceBindings, BasicSoftmaxBindings, BasicSqrtBindings, BasicTransposeBindings, \
-    DummyBinding
+    DummyBinding, BasicTopKBindings
 from Deeploy.Targets.Generic.Layers import AddLayer, BatchNormalizationLayer, ConcatLayer, ConvLayer, \
     ConvTransposeLayer, DebugPrintLayer, DequantLayer, DivLayer, GatherLayer, GELULayer, GEMMLayer, ITAMaxLayer, \
     LayerNormLayer, MatMulLayer, MaxPoolLayer, MulLayer, PadLayer, PowLayer, QuantLayer, ReduceMeanLayer, \
     ReduceSumLayer, ReluLayer, RequantShiftLayer, ReshapeLayer, RQIntegerDivLayer, RQSiGELULayer, SliceLayer, \
-    SoftmaxLayer, SqrtLayer, TransposeLayer
+    SoftmaxLayer, SqrtLayer, TransposeLayer, TopKLayer
 from Deeploy.Targets.Generic.Parsers import AddParser, BatchNormParser, ConcatParser, ConvTranspose1DParser, \
     DebugParser, DequantParser, DivParser, DummyParser, FlattenParser, GatherParser, GELUParser, GenericConv1DParser, \
     GenericConv2DParser, GenericDWConv1DParser, GenericDWConv2DParser, GenericGEMMParser, GenericMaxPool2DParser, \
     IntegerDivParser, ITAMaxParser, ITAPartialMaxParser, LayerNormParser, MatMulParser, MaxPool1DParser, MulParser, \
     Pad1DParser, Pad2DParser, PowParser, QuantParser, ReduceMeanParser, ReduceSumParser, ReluParser, \
     RequantShiftParser, ReshapeParser, RQIntegerDivParser, RQSiGELUParser, SliceParser, SoftmaxParser, SqrtParser, \
-    TransposeParser, UnsqueezeParser, iLayerNormParser, iSoftmaxParser
+    TransposeParser, TopKParser, UnsqueezeParser, iLayerNormParser, iSoftmaxParser
 from Deeploy.Targets.Generic.Templates import AllocateTemplate, FreeTemplate
 from Deeploy.Targets.Generic.TopologyOptimizationPasses.Passes import DequantPatternPass, ExtractPaddingFromConvPass, \
     ExtractPaddingFromPoolPass, MatMulAddMergePass, MergeConstAddAndRequantPass, QuantPatternPass, \
@@ -67,6 +67,7 @@
 SoftmaxMapper = NodeMapper(SoftmaxParser(), BasicSoftmaxBindings)
 iSoftmaxMapper = NodeMapper(iSoftmaxParser(), BasicSoftmaxBindings)
 TransposeMapper = NodeMapper(TransposeParser(), BasicTransposeBindings)
+TopKMapper = NodeMapper(TopKParser(), BasicTopKBindings)
 UnsqueezeMapper = NodeMapper(UnsqueezeParser(), BasicReshapeBindings)
 QuantMapper = NodeMapper(QuantParser(), BasicQuantBindings)
 DequantMapper = NodeMapper(DequantParser(), BasicDequantBindings)
@@ -113,6 +114,7 @@
     'RQIntegerDiv': RQIntegerDivLayer([RQIntegerDivMapper]),
     'Squeeze': ReshapeLayer([UnsqueezeMapper]),
     'Transpose': TransposeLayer([TransposeMapper]),
+    'TopK': TopKLayer([TopKMapper]),
     'Unsqueeze': ReshapeLayer([UnsqueezeMapper]),
     'Slice': SliceLayer([SliceMapper]),
     'Quant': QuantLayer([QuantMapper]),

Deeploy/Targets/Generic/Templates/GatherTemplate.py

@@ -10,8 +10,18 @@
 width = int(data_in_type.referencedType.typeWidth/8)
 %>
 BEGIN_SINGLE_CORE
+% if num_indices == 1:
 for (uint32_t i=0; i<${batch}; ++i) {
     memcpy(${data_out} + i * ${axis_length}, ${data_in} + i * ${batch_length} + ${index} * ${axis_length}, ${axis_length} * ${width});
 }
+% else:
+for (uint32_t i=0; i<${batch}; ++i) {
+    for (uint32_t j=0; j<${num_indices}; ++j) {
+        memcpy(${data_out} + i * (${num_indices} * ${axis_length}) + j * ${axis_length}, 
+               ${data_in} + i * ${batch_length} + ${indices}[j] * ${axis_length}, 
+               ${axis_length} * ${width});
+    }
+}
+% endif
 END_SINGLE_CORE
 """)

Deeploy/Targets/Generic/Templates/TopKTemplate.py

@@ -0,0 +1,40 @@
+from typing import Dict, List, Tuple
+
+from Deeploy.DeeployTypes import NetworkContext, NodeTemplate, OperatorRepresentation
+
+
+referenceTemplate = NodeTemplate("""
+// TopK (Name: ${nodeName}, Op: ${nodeOp})
+BEGIN_SINGLE_CORE
+// Find the top ${k_value} values and their indices
+// Assumes 1D input for simplicity
+typedef struct {
+	${data_in_type.referencedType.typeName} value;
+	uint32_t index;
+} topk_pair_t;
+
+topk_pair_t pairs[${data_in_size}];
+for (uint32_t i = 0; i < ${data_in_size}; ++i) {
+	pairs[i].value = ((${data_in_type.referencedType.typeName}*)${data_in})[i];
+	pairs[i].index = i;
+}
+// Simple selection sort for top-k
+for (uint32_t i = 0; i < ${k_value}; ++i) {
+	uint32_t max_idx = i;
+	for (uint32_t j = i + 1; j < ${data_in_size}; ++j) {
+		if (pairs[j].value > pairs[max_idx].value) {
+			max_idx = j;
+		}
+	}
+	// Swap
+	if (max_idx != i) {
+		topk_pair_t tmp = pairs[i];
+		pairs[i] = pairs[max_idx];
+		pairs[max_idx] = tmp;
+	}
+	// Write output
+	((${values_out_type.referencedType.typeName}*)${values_out})[i] = pairs[i].value;
+	((${indices_out_type.referencedType.typeName}*)${indices_out})[i] = pairs[i].index;
+}
+END_SINGLE_CORE
+""")

Deeploy/Targets/Generic/TypeCheckers.py

@@ -610,3 +610,17 @@ def _inferNumLevels(self, inputs: List[VariableBuffer],
     def _inferSignedness(self, inputs: List[VariableBuffer],
                          operatorRepresentation: OperatorRepresentation) -> List[bool]:
         return [True]
+
+# TopKChecker: infers types for both values and indices outputs of TopK operation
+class TopKChecker(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]:
+        # Output 0: values (same as input), Output 1: indices (integer, usually not quantized)
+        # We assume indices output is not quantized (set to 0 or 1)
+        return [inputs[0].nLevels, 1]
+
+    def _inferSignedness(self, inputs: List[VariableBuffer], operatorRepresentation: OperatorRepresentation) -> List[bool]:
+        # Output 0: values (same signedness as input), Output 1: indices (unsigned)
+        return [inputs[0]._signed, False]

Deeploy/Targets/Spatz/Bindings.py

@@ -0,0 +1,29 @@
+from Deeploy.DeeployTypes import CodeTransformation, NodeBinding
+from Deeploy.CommonExtensions.CodeTransformationPasses.MemoryAllocation import ArgumentStructGeneration, \
+    MemoryManagementGeneration
+from Deeploy.FutureExtension.CodeTransformationPasses.FutureCodeTransformation import FutureGeneration
+from Deeploy.AbstractDataTypes import PointerClass
+from Deeploy.CommonExtensions.DataTypes import FloatDataTypes, IntegerDataTypes, SignedIntegerDataTypes, float32_t, int8_t, int32_t, uint8_t
+from Deeploy.Targets.Generic.TypeCheckers import GatherChecker
+from Deeploy.Targets.Spatz.Templates import GatherTemplate
+
+BasicTransformer = CodeTransformation([ArgumentStructGeneration(), MemoryManagementGeneration(), FutureGeneration()])
+
+SpatzGatherBindings = [
+    NodeBinding(
+        GatherChecker(
+            [PointerClass(type), PointerClass(int32_t)],
+            [PointerClass(type)]
+        ),
+        GatherTemplate.referenceTemplate,
+        BasicTransformer
+    ) for type in SignedIntegerDataTypes] + [
+    NodeBinding(
+        GatherChecker(
+            [PointerClass(float32_t), PointerClass(type)],
+            [PointerClass(float32_t)]
+        ),
+        GatherTemplate.referenceTemplate,
+        BasicTransformer
+    ) for type in IntegerDataTypes
+]

Deeploy/Targets/Spatz/Deployer.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: 2023 ETH Zurich and University of Bologna
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable, Dict, Type
+
+import onnx_graphsurgeon as gs
+
+from Deeploy.AbstractDataTypes import Pointer
+from Deeploy.CommonExtensions.NetworkDeployers.SignPropDeployer import SignPropDeployer
+from Deeploy.CommonExtensions.OptimizationPasses.TopologyOptimizationPasses.DebugPasses import DebugPrintMergePass
+from Deeploy.CommonExtensions.OptimizationPasses.TopologyOptimizationPasses.LoweringOptimizationPasses import \
+    NCHWtoNHWCPass, TransposeMatmulInputsPass
+from Deeploy.DeeployTypes import DeploymentPlatform, TopologyOptimizer
+from Deeploy.Targets.Generic.TopologyOptimizationPasses.Passes import TransposeConstOptPass, TransposeMergePass
+
+
+class SpatzDeployer(SignPropDeployer):
+
+    def __init__(self,
+                 graph: gs.Graph,
+                 deploymentPlatform: DeploymentPlatform,
+                 inputTypes: Dict[str, Type[Pointer]],
+                 loweringOptimizer: TopologyOptimizer,
+                 scheduler: Callable = lambda x: x,
+                 name: str = 'DeeployNetwork',
+                 default_channels_first = False,
+                 deeployStateDir: str = "DeeployStateDir",
+                 inputOffsets: Dict[str, int] = {}):
+
+        super().__init__(graph,
+                         deploymentPlatform,
+                         inputTypes,
+                         loweringOptimizer,
+                         scheduler,
+                         name,
+                         default_channels_first = default_channels_first,
+                         deeployStateDir = deeployStateDir)
+
+#         self.inputOffsets = inputOffsets
+# 
+#         self.loweringOptimizer.passes += [
+#             TransposeMatmulInputsPass(),
+#             NCHWtoNHWCPass(self.default_channels_first),
+#             TransposeMergePass(),
+#             TransposeConstOptPass(),
+#             DebugPrintMergePass()
+#         ]