AttentionTraining Support

Author: runwangdl

Deeploy/Targets/Generic/Layers.py

@@ -84,6 +84,13 @@ class GELUGradLayer(ONNXLayer):
 
     def __init__(self, maps: List[NodeMapper]):
         super().__init__(maps)
+
+    def computeOps(self):
+        size = self.mapper.parser.operatorRepresentation['size']
+        ops_per_element = 9  
+        gelu_grad_ops = size * ops_per_element
+        return gelu_grad_ops
+        
 
 class iHardswishLayer(ONNXLayer):
 
@@ -490,6 +497,11 @@ class SGDLayer(ONNXLayer):
     def __init__(self, maps: List[NodeMapper]):
         super().__init__(maps)
 
+    def computeOps(self):
+        
+        size = self.mapper.parser.operatorRepresentation['size']
+        return size*2
+
 
 class LinearAttentionLayer(ONNXLayer):
 

Deeploy/Targets/Generic/Parsers.py

@@ -991,7 +991,7 @@ def __init__(self):
 
     def parseNode(self, node: gs.Node) -> (bool):
 
-        ret = all([len(node.inputs) == 1, len(node.outputs) == 1])
+        ret = all([len(node.inputs) >= 1, len(node.outputs) == 1])
 
         return ret
 

Deeploy/Targets/PULPOpen/Bindings.py

@@ -180,7 +180,10 @@
 ] + [
     NodeBinding(AddChecker([PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
                 FloatAddTemplate.referenceTemplate, ForkTransformer)
-]
+] + [
+    NodeBinding(AddChecker([PointerClass(float32_t), PointerClass(float32_t), PointerClass(float32_t)], [PointerClass(float32_t)]),
+                FloatAddTemplate.referenceTemplate, ForkTransformer)
+] 
 
 PULPRQSConv2DBindings = [
     NodeBinding(

Deeploy/Targets/PULPOpen/Templates/SGDTemplate.py

@@ -26,16 +26,45 @@
 from Deeploy.DeeployTypes import NodeTemplate
 
 referenceTemplate = NodeTemplate("""
-// SGD Weight Update (Name: ${nodeName}, Op: ${nodeOp})
-BEGIN_SINGLE_CORE
-    ${weight_type.typeName} ref_${weight} = ${weight};
-    ${grad_type.typeName} ref_${grad} = ${grad};
-    ${weight_type.typeName} ref_${weight_updated} = ${weight_updated};
+// SGD Weight Update with Separated Multiplication and Subtraction Unrolling
+// (Name: ${nodeName}, Op: ${nodeOp})
+int8_t ${nodeName}_core_id = pi_core_id();
+int8_t ${nodeName}_log2Core = log2(NUM_CORES);
+int16_t ${nodeName}_chunk = (${size} >> ${nodeName}_log2Core) + ((${size} & (NUM_CORES-1))!=0);
+int16_t ${nodeName}_chunk_start = MIN(${nodeName}_chunk*${nodeName}_core_id, ${size});
+int16_t ${nodeName}_chunk_stop = MIN(${nodeName}_chunk_start + ${nodeName}_chunk, ${size});
+
+${weight_type.typeName} ref_${weight} = ${weight};
+${grad_type.typeName} ref_${grad} = ${grad};
+${weight_type.typeName} ref_${weight_updated} = ${weight_updated};
+
+float32_t learning_rate = ${lr};
+
+// Temporary buffer for multiplication results
+float32_t temp_mul[6];
+
+uint32_t i = ${nodeName}_chunk_start;
+for (; i+5 < ${nodeName}_chunk_stop; i+=6) {
+    // Unrolled multiplication operations
+    temp_mul[0] = learning_rate * ref_${grad}[i];
+    temp_mul[1] = learning_rate * ref_${grad}[i+1];
+    temp_mul[2] = learning_rate * ref_${grad}[i+2];
+    temp_mul[3] = learning_rate * ref_${grad}[i+3];
+    temp_mul[4] = learning_rate * ref_${grad}[i+4];
+    temp_mul[5] = learning_rate * ref_${grad}[i+5];
     
-    float32_t learning_rate = ${lr}; 
+    // Unrolled subtraction operations
+    ref_${weight_updated}[i] = ref_${weight}[i] - temp_mul[0];
+    ref_${weight_updated}[i+1] = ref_${weight}[i+1] - temp_mul[1];
+    ref_${weight_updated}[i+2] = ref_${weight}[i+2] - temp_mul[2];
+    ref_${weight_updated}[i+3] = ref_${weight}[i+3] - temp_mul[3];
+    ref_${weight_updated}[i+4] = ref_${weight}[i+4] - temp_mul[4];
+    ref_${weight_updated}[i+5] = ref_${weight}[i+5] - temp_mul[5];
+}
 
-    for (uint32_t i=0; i<${size}; ++i) {
-        ref_${weight_updated}[i] = ref_${weight}[i] - learning_rate * ref_${grad}[i];
-    }
-END_SINGLE_CORE
-""")
+// Handle remaining elements
+for (; i < ${nodeName}_chunk_stop; i++) {
+    float32_t temp_grad = learning_rate * ref_${grad}[i];
+    ref_${weight_updated}[i] = ref_${weight}[i] - temp_grad;
+}
+""")