support of some fp32 operators

code_generator/CodeGenerator.py

@@ -18,7 +18,10 @@
 
 import os
 
+import numpy as np
+
 from .constant import FUSE_SGD_UPDATE_STR, FUSHION_CONFIG
+from .operators.basic_utils import tensor
 from .OpGenerator import OpGenerator
 
 Codegen_root = "./codegen/"
@@ -34,8 +37,6 @@ class CodeGenerator:
     """Provide utilities to generate C code for a given model and memory schdeule."""
 
     parse_count = 0
-    header_handle = None
-    source_handle = None
 
     def __init__(
         self,
@@ -585,12 +586,13 @@ def _parseTrainable(self):
                     )
                 else:
                     self._parseBias(self.parse_count, layer_info["bias"].flatten())
-                self._parseEffectivescales(self.parse_count, layer_info["effective_scale"].flatten())
-                self._parseRequantize(
-                    self.parse_count,
-                    layer_info["shift"].flatten(),
-                    layer_info["multiplier"].flatten(),
-                )
+                if layer_info["input_dtype"] == "int8":
+                    self._parseEffectivescales(self.parse_count, layer_info["effective_scale"].flatten())
+                    self._parseRequantize(
+                        self.parse_count,
+                        layer_info["shift"].flatten(),
+                        layer_info["multiplier"].flatten(),
+                    )
 
                 layer_info["parsed_trainable"] = self.parse_count
                 self.parse_count += 1
@@ -770,6 +772,37 @@ def _parseTrainable(self):
 
                     layer_info["parsed_trainable"] = self.parse_count
                     self.parse_count += 1
+            else:
+                # Parse constants of inputs
+                for t in op.input_tensors:
+                    if t.constant():
+                        # for TTE compatible
+                        if "constant" in layer_info and layer_info["constant"] is not None:
+                            continue
+                        if t.data is None:
+                            raise ValueError(f"constant tensor data not found for op:{layer_info['op']}")
+                        self._parseConstant(t)
+
+    def _parseConstant(self, t: tensor):
+        def type_to_c_type(type: str) -> str:
+            if type == "int8":
+                return "unsigned char"
+            elif type == "float32":
+                return "float"
+            elif type == "int32":
+                return "int32_t"
+            elif type == "bool":
+                return "char"  # Using bytes to store boolean
+            raise NotImplementedError
+
+        fp = self.header_handle
+        # 8bit implementation
+        string = f"const {type_to_c_type(t.dtype)} {t.graph_idx}" + "[" + str(np.prod(t.size)) + "] = {"
+        flat_data = t.data.flatten()
+        for d in flat_data:
+            string += f"{d}, "
+        string += "};\n"
+        fp.write(string)
 
     def _parseCWHWeight(self, Lindex, weight, height, width, channel):
         fp = self.header_handle
@@ -818,34 +851,43 @@ def _parseEffectivescales(self, Lindex, scales):
     def _parseWeight(self, Lindex, weight, weight_name=None, is_const=True):
         fp = self.header_handle
         const_str = "const " if is_const else ""
-        string = f"{const_str}unsigned char weight" + str(Lindex) + "[" + str(len(weight)) + "] = {"
-        fp.write(string)
-        for _, value in enumerate(weight):
-            value = int(value)
-            if value < 0:
-                value += 256
-            fp.write(str(format(value, "#04x")) + ", ")
-        fp.write("};\n")
-
-        if self.is_training:
-            string = f"{const_str}float weight_fp" + str(Lindex) + "[" + str(len(weight)) + "] = {"
+        if weight.dtype == "float32":
+            string = f"{const_str}unsigned char weight_fp" + str(Lindex) + "[" + str(len(weight)) + "] = {"
+            for _, value in enumerate(weight):
+                string += f"{value}, "
+            string += "};\n"
+            fp.write(string)
+        elif weight.dtype == "int8":
+            string = f"{const_str}unsigned char weight" + str(Lindex) + "[" + str(len(weight)) + "] = {"
             fp.write(string)
-            for _, w in enumerate(weight):
-                value = float(w)
-                fp.write(str(value) + ", ")
+            for _, value in enumerate(weight):
+                value = int(value)
+                if value < 0:
+                    value += 256
+                fp.write(str(format(value, "#04x")) + ", ")
             fp.write("};\n")
 
-        if weight_name is not None:
-            for r in self.trainSRAMTable:
-                if r.name == weight_name:
-                    return
-            self.trainSRAMTable.append(tensorRecorder(weight_name, len(weight), "unknown"))
-
-            if weight.dtype == "int8":
-                string = f"{const_str}unsigned char* {weight_name}=weight" + str(Lindex) + ";\n"
-            else:
-                raise NotImplementedError
-            fp.write(string)
+            if self.is_training:
+                string = f"{const_str}float weight_fp" + str(Lindex) + "[" + str(len(weight)) + "] = {"
+                fp.write(string)
+                for _, w in enumerate(weight):
+                    value = float(w)
+                    fp.write(str(value) + ", ")
+                fp.write("};\n")
+
+            if weight_name is not None:
+                for r in self.trainSRAMTable:
+                    if r.name == weight_name:
+                        return
+                self.trainSRAMTable.append(tensorRecorder(weight_name, len(weight), "unknown"))
+
+                if weight.dtype == "int8":
+                    string = f"{const_str}unsigned char* {weight_name}=weight" + str(Lindex) + ";\n"
+                else:
+                    raise NotImplementedError
+                fp.write(string)
+        else:
+            raise NotImplementedError
 
     def _parseWeightPartial(self, Lindex, weight, first_k_channel=None, weight_name=None, is_const=True):
         fp = self.header_handle

code_generator/TTEParser.py

@@ -1405,6 +1405,7 @@ def _convert_sub(self, op):
             "input2_idx": input2_info["name"],
             "output_idx": output_info["name"],
             "input_size": input0_h * input0_w * input0_c,
+            "input2_size": input0_h * input0_w * input0_c,
             "input_dtype": input_dtype,
             "input2_dtype": input2_dtype,
             "output_dtype": output_dtype,

code_generator/TfliteConvertor.py

@@ -17,14 +17,38 @@
 # ----------------------------------------------------------------------
 
 import logging
+from typing import List
 
 import code_generator.converters.tflite_parser as TF_Parser
 from code_generator.converters.tflite_parser.mean1dto2d import MEAN2D
+from code_generator.converters.tflite_parser.reshape import tensor_mapping
 from code_generator.converters.tflite_parser.utils import get_input_tensors, get_output_tensors, getOpCodeStr
 
 from .constant import SKIP_OPs
 from .tflite import Model
 
+regular_opconverter = {
+    "ADD": TF_Parser.parse_add,
+    "AVERAGE_POOL_2D": TF_Parser.parse_avgpool,
+    "RESIZE_NEAREST_NEIGHBOR": TF_Parser.parse_upsample,
+    "MAX_POOL_2D": TF_Parser.parse_maxpool,
+    "FULLY_CONNECTED": TF_Parser.parse_fc,
+    "DIV": TF_Parser.parse_div,
+    "BATCH_MATMUL": TF_Parser.parse_batchmatmul,
+    "NOT_EQUAL": TF_Parser.parse_notequal,
+    "EQUAL": TF_Parser.parse_equal,
+    "CONCATENATION": TF_Parser.parse_concat,
+    "CAST": TF_Parser.parse_cast,
+    "SUB": TF_Parser.parse_sub,
+    "MUL": TF_Parser.parse_mul,
+    "SOFTMAX": TF_Parser.parse_softmax,
+    "SQUARED_DIFFERENCE": TF_Parser.parse_squarddiff,
+    "RSQRT": TF_Parser.parse_rsqrt,
+    "SLICE": TF_Parser.parse_slice,
+    "MEAN": TF_Parser.parse_mean1d,
+    "TRANSPOSE": TF_Parser.parse_transpose,
+}
+
 
 # Parse tflite model into TinyEngine IR format
 class TfliteConvertor(object):
@@ -37,6 +61,7 @@ def __init__(self, filepath):
         self.tmpPADIndice = None
         self.skip_transpose = None
         self.average_1D_to_2D_holder = MEAN2D()  # For merging 1D to 2D
+        self.inplace_reshape_table: List[tensor_mapping] = []  # A list of tensor_mapping
 
     # public functions
     def loadTFmodel(self, filepath):
@@ -104,41 +129,62 @@ def parseOperatorInfo(self):
 
                     self.layer.append(SEelementmult_op)
                     continue
+            if i + 2 < operators_len - 2:
+                next_op = self.subgraph.Operators(i + 1)
+                next_next_op = self.subgraph.Operators(i + 2)
+                three_op_sequence = [op, next_op, next_next_op]
+
+                if self.checkIfMergeTransposeTwoMean1d(three_op_sequence):
+                    logging.info("found target to merge transpose and two mean1d")
+                    self._convert_TRANSPOSE(op)
+                    skip_next_ops = 2
+                    ret_op = TF_Parser.parse_mean1dto2d(next_op, self.model, self.average_1D_to_2D_holder)
+                    ret_op = TF_Parser.parse_mean1dto2d(next_next_op, self.model, self.average_1D_to_2D_holder)
+                    if ret_op is not None:
+                        if self.skip_transpose is not None:
+                            ret_op.params["input_idx"] = self.skip_transpose.input_idx
+                            ret_op.input_tensors[0].graph_idx = self.skip_transpose.input_idx
+                        self.layer.append(ret_op)
+                    continue
+            if i + 1 < operators_len - 1:
+                next_op = self.subgraph.Operators(i + 1)
+                two_op_sequence = [op, next_op]
+
+                if self.checkIfMergeTwoMean1d(two_op_sequence):
+                    logging.info("found target to merge two mean1d")
+                    skip_next_ops = 1
+                    ret_op = TF_Parser.parse_mean1dto2d(op, self.model, self.average_1D_to_2D_holder)
+                    ret_op = TF_Parser.parse_mean1dto2d(next_op, self.model, self.average_1D_to_2D_holder)
+                    if ret_op is not None:
+                        if self.skip_transpose is not None:
+                            ret_op.params["input_idx"] = self.skip_transpose.input_idx
+                            ret_op.input_tensors[0].graph_idx = self.skip_transpose.input_idx
+                        self.layer.append(ret_op)
+                    continue
 
             # parse the op
             self._handleOperator(op)
 
+        # Handle inplace_reshape_table here
+        logging.error("Please handle inplace_reshape_table here for fused tensors.")
+
     # handle one op and parse it into layers[] for supported operators
     def _handleOperator(self, op):
         op_code_str = getOpCodeStr(op, self.model)
         if op_code_str == "CONV_2D":
             self.layer.append(TF_Parser.parse_conv2d(op, self.model, self.tmpPADIndice))
             self.tmpPADIndice = None
-        elif op_code_str == "ADD":
-            self.layer.append(TF_Parser.parse_add(op, self.model))
-        elif op_code_str == "AVERAGE_POOL_2D":
-            self.layer.append(TF_Parser.parse_avgpool(op, self.model))
         elif op_code_str == "DEPTHWISE_CONV_2D":
             self.layer.append(TF_Parser.parse_conv2d(op, self.model, self.tmpPADIndice))
             self.tmpPADIndice = None
         elif op_code_str == "PAD":
             self._convert_PAD(op)
-        elif op_code_str == "RESIZE_NEAREST_NEIGHBOR":
-            self.layer.append(TF_Parser.parse_upsample(op, self.model))
-        elif op_code_str == "MAX_POOL_2D":
-            self.layer.append(TF_Parser.parse_maxpool(op, self.model))
-        elif op_code_str in "MEAN":
-            ret_op = TF_Parser.parse_mead1dto2d(op, self.model, self.average_1D_to_2D_holder)
-            if ret_op is not None:
-                # TODO: This only handle a specific graph: TRANSPOSE -> MEAN -> MEANS
-                if self.skip_transpose is not None:
-                    ret_op.params["input_idx"] = self.skip_transpose.input_idx
-                    ret_op.input_tensors[0].graph_idx = self.skip_transpose.input_idx
-                self.layer.append(ret_op)
-        elif op_code_str == "TRANSPOSE":
-            self._convert_TRANSPOSE(op)
-        elif op_code_str in "FULLY_CONNECTED":
-            self.layer.append(TF_Parser.parse_fc(op, self.model))
+        # elif op_code_str == "TRANSPOSE":
+        # self._convert_TRANSPOSE(op)
+        elif op_code_str == "RESHAPE":
+            self.inplace_reshape_table.append(TF_Parser.parse_reshape_fuse_tensor_tuple(op, self.model))
+        elif op_code_str in regular_opconverter:
+            self.layer.append(regular_opconverter[op_code_str](op, self.model))
         elif op_code_str in SKIP_OPs:
             pass
         else:
@@ -156,6 +202,29 @@ def checkIfRequireSEelementmult(self, three_op_sequence):
             return True
         return False
 
+    #  | TRANSPOSE -> MEAN -> MEAN ->             |
+    #  |                           -> AVG_POOL_2D |
+    #  |                Fuse Target               |
+    def checkIfMergeTwoMean1d(self, two_op_sequence):
+        if (
+            getOpCodeStr(two_op_sequence[0], self.model) == "MEAN"
+            and getOpCodeStr(two_op_sequence[1], self.model) == "MEAN"
+        ):
+            return True
+        return False
+
+    #  | MEAN -> MEAN ->              |
+    #  |               -> AVG_POOL_2D |
+    #  |       Fuse Target            |
+    def checkIfMergeTransposeTwoMean1d(self, three_op_sequence):
+        if (
+            getOpCodeStr(three_op_sequence[1], self.model) == "TRANSPOSE"
+            and getOpCodeStr(three_op_sequence[1], self.model) == "MEAN"
+            and getOpCodeStr(three_op_sequence[2], self.model) == "MEAN"
+        ):
+            return True
+        return False
+
     def _convert_PAD(self, op):
         # get input, weight, and output tensors
         input_tensors = get_input_tensors(op, self.model)

code_generator/converters/tflite_parser/__init__.py

@@ -1,8 +1,23 @@
 from .add import parse_add
 from .avgpool import parse_avgpool
+from .batchmatmul import parse_batchmatmul
+from .cast import parse_cast
+from .concat import parse_concat
 from .conv2d import parse_conv2d
+from .div import parse_div
+from .equal import parse_equal
 from .fc import parse_fc
 from .maxpool import parse_maxpool
-from .mean1dto2d import parse_mead1dto2d
+from .mean1d import parse_mean1d
+from .mean1dto2d import parse_mean1dto2d
+from .mul import parse_mul
+from .notequal import parse_notequal
+from .reshape import parse_reshape_fuse_tensor_tuple
+from .rsqrt import parse_rsqrt
 from .SEelement import parse_SEelement
+from .slice import parse_slice
+from .softmax import parse_softmax
+from .squarddiff import parse_squarddiff
+from .sub import parse_sub
+from .transpose import parse_transpose
 from .upsample import parse_upsample