tinymlc 0.1.0__py3-none-any.whl

TinyMLC/templates/model.c.tpl

@@ -0,0 +1,314 @@
+/* TinyMLC - Tiny Machine Learning Compiler
+ *
+ * Copyright (c) 2026 Jia Liu & TinyMLC Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * This file is part of TinyMLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at:
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Auto-generated code, do not modify manually
+// Generated by tinymlc
+
+#include "tinymlc.h"
+{{ includes }}
+#include "model.h"
+#include "debug_print.h"
+#include <stddef.h>
+
+// Input/output sizes
+#define INPUT_SIZE {{ input_size }}
+#define OUTPUT_SIZE {{ output_size }}
+
+// Intermediate tensor memory (static allocation, placed outside function)
+{% for tensor in tensors_to_define %}
+{% if target == "host" %}
+{{ tensor.type }} tensor_{{ tensor.index }}[{{ tensor.size }}];
+{% else %}
+{{ tensor.type }} tensor_{{ tensor.index }}[{{ tensor.size }}] \
+    __attribute__((section(".bss")));
+{% endif %}
+{% endfor %}
+
+{% if has_lstm %}
+// LSTM parameters
+#define LSTM_TIME_STEPS {{ lstm_time_steps }}
+#define LSTM_HIDDEN_SIZE {{ lstm_hidden_size }}
+#define TINYMLC_HAS_LSTM
+{% endif %}
+
+// Inference function
+{% if inputs_count == 1 %}
+void {{ inference_func }}(const int8_t* input, int8_t* output) {
+    // Input tensor mapping
+    int8_t* tensor_{{ input_tensor_indices[0] }} = (int8_t*)input;
+{% elif inputs_count == 2 %}
+void {{ inference_func }}(
+    const int8_t* input1,
+    const int8_t* input2,
+    int8_t* output
+) {
+    // Input tensor mapping
+    int8_t* tensor_{{ input_tensor_indices[0] }} = (int8_t*)input1;
+    int8_t* tensor_{{ input_tensor_indices[1] }} = (int8_t*)input2;
+{% endif %}
+    // Execute operators in order
+    {% for op in execution_order %}
+        {% if op.op_name == "UNIDIRECTIONAL_SEQUENCE_LSTM" %}
+        tmlc_unidirectional_sequence_lstm_s8(
+            tensor_{{ op.input_indices[0] }},
+            lstm_input_weights,
+            lstm_recurrent_weights,
+            lstm_bias,
+            tensor_{{ op.output_indices[0] }},
+            NULL, NULL,
+            {{ op.lstm_params.time_steps }},
+            {{ op.lstm_params.batch_size }},
+            {{ op.lstm_params.input_size }},
+            {{ op.lstm_params.hidden_size }}
+        );
+        {% elif op.op_name == "FULLY_CONNECTED" %}
+        tmlc_fully_connected_s8(
+            tensor_{{ op.input_indices[0] }},
+            fc_weights,
+            fc_bias,
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.input_indices[0]] }},
+            {{ tensor_sizes[op.output_indices[0]] }},
+            {{ fc_multiplier }},
+            {{ fc_shift }}
+        );
+        {% elif op.op_name == "SOFTMAX" %}
+        tmlc_softmax_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.input_indices[0]] }}
+        );
+        {% elif op.op_name == "RESHAPE" %}
+        {
+            static const int reshape_target[] = {
+                {% for s in op.reshape_target_shape %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            };
+            int reshape_input_size = {{ tensor_sizes[op.input_indices[0]] }};
+            tmlc_reshape_s8(
+                tensor_{{ op.input_indices[0] }},
+                tensor_{{ op.output_indices[0] }},
+                reshape_input_size,
+                reshape_target,
+                {{ op.reshape_target_shape | length }}
+            );
+        }
+        {% elif op.op_name == "ADD" %}
+        tmlc_add_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.input_indices[1] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.output_indices[0]] }}
+        );
+        {% elif op.op_name == "SVDF" %}
+        tmlc_svdf_s8(
+            tensor_{{ op.data_input_idx }},
+            svdf_weights,
+            svdf_bias,
+            tensor_{{ op.output_indices[0] }},
+            {{ op.svdf_params.time_steps }},
+            {{ op.svdf_params.input_size }},
+            {{ op.svdf_params.rank }},
+            {{ op.svdf_params.units }}
+        );
+        {% elif op.op_name == "CONV_2D" %}
+        tmlc_conv2d_s8(
+            tensor_{{ op.data_input_idx }},
+            conv_weights,
+            conv_bias,
+            tensor_{{ op.output_indices[0] }},
+            {{ op.conv_params.input_h }},
+            {{ op.conv_params.input_w }},
+            {{ op.conv_params.input_c }},
+            {{ op.conv_params.output_h }},
+            {{ op.conv_params.output_w }},
+            {{ op.conv_params.output_c }},
+            {{ op.conv_params.kernel_h }},
+            {{ op.conv_params.kernel_w }},
+            {{ op.conv_params.stride_h }},
+            {{ op.conv_params.stride_w }},
+            {{ op.conv_params.padding_h }},
+            {{ op.conv_params.padding_w }},
+            {{ conv_multiplier }},
+            {{ conv_shift }}
+        );
+        {% elif op.op_name == "MAX_POOL_2D" %}
+        tmlc_max_pool_2d_s8(
+            tensor_{{ op.data_input_idx }},
+            tensor_{{ op.output_indices[0] }},
+            {{ op.pool_params.input_h }},
+            {{ op.pool_params.input_w }},
+            {{ op.pool_params.input_c }},
+            {{ op.pool_params.output_h }},
+            {{ op.pool_params.output_w }},
+            {{ op.pool_params.output_c }},
+            {{ op.pool_params.pool_size_h }},
+            {{ op.pool_params.pool_size_w }},
+            {{ op.pool_params.stride_h }},
+            {{ op.pool_params.stride_w }},
+            0, 0
+        );
+        {% elif op.op_name == "DEPTHWISE_CONV_2D" %}
+        tmlc_depthwise_conv_2d_s8(
+            tensor_{{ op.data_input_idx }},
+            dw_weights,
+            dw_bias,
+            tensor_{{ op.output_indices[0] }},
+            {{ op.dw_params.input_h }},
+            {{ op.dw_params.input_w }},
+            {{ op.dw_params.input_c }},
+            {{ op.dw_params.output_h }},
+            {{ op.dw_params.output_w }},
+            {{ op.dw_params.output_c }},
+            {{ op.dw_params.kernel_h }},
+            {{ op.dw_params.kernel_w }},
+            {{ op.dw_params.stride_h }},
+            {{ op.dw_params.stride_w }},
+            {{ op.dw_params.depth_multiplier }},
+            {{ op.dw_params.padding_h }},
+            {{ op.dw_params.padding_w }},
+            {{ conv_multiplier }},
+            {{ conv_shift }}
+        );
+        {% elif op.op_name == "RELU" %}
+        tmlc_relu_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.input_indices[0]] }}
+        );
+        {% elif op.op_name == "AVERAGE_POOL_2D" %}
+        tmlc_avg_pool_2d_s8(
+            tensor_{{ op.data_input_idx }},
+            tensor_{{ op.output_indices[0] }},
+            {{ op.pool_params.input_h }},
+            {{ op.pool_params.input_w }},
+            {{ op.pool_params.input_c }},
+            {{ op.pool_params.output_h }},
+            {{ op.pool_params.output_w }},
+            {{ op.pool_params.output_c }},
+            {{ op.pool_params.pool_size_h }},
+            {{ op.pool_params.pool_size_w }},
+            {{ op.pool_params.stride_h }},
+            {{ op.pool_params.stride_w }},
+            0, 0
+        );
+        {% elif op.op_name == "TRANSPOSE" %}
+        tmlc_transpose_s8(
+            tensor_{{ op.data_input_idx }},
+            NULL,
+            tensor_{{ op.output_indices[0] }},
+            {{ op.transpose_params.input_dims }},
+            (const int[]){
+                {% for s in tensor_shapes[op.data_input_idx] %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            }
+        );
+        {% elif op.op_name == "PAD" %}
+        tmlc_pad_s8(
+            tensor_{{ op.data_input_idx }},
+            NULL,  // paddings, currently NULL
+            tensor_{{ op.output_indices[0] }},
+            4,
+            (const int[]){
+                {% for s in tensor_shapes[op.data_input_idx] %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            },
+            (const int[]){
+                {% for s in tensor_shapes[op.output_indices[0]] %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            }
+        );
+        {% elif op.op_name == "MEAN" %}
+        tmlc_mean_s8(
+            tensor_{{ op.data_input_idx }},
+            tensor_{{ op.output_indices[0] }},
+            {{ op.mean_params.input_dims }},
+            (const int[]){
+                {% for s in tensor_shapes[op.data_input_idx] %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            },
+            (const int[]){
+                {% for s in tensor_shapes[op.output_indices[0]] %}
+                    {{ s }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            },
+            NULL,
+            0,
+            0
+        );
+        {% elif op.op_name == "MULTIPLY" %}
+        tmlc_multiply_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.input_indices[1] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.output_indices[0]] }}
+        );
+        {% elif op.op_name == "SIGMOID" %}
+        tmlc_sigmoid_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.input_indices[0]] }}
+        );
+        {% elif op.op_name == "CONCAT" %}
+        {
+            static const int8_t* concat_inputs[] = {
+                {% for idx in op.input_indices %}
+                    tensor_{{ idx }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            };
+            static const int concat_sizes[] = {
+                {% for idx in op.input_indices %}
+                    {{ tensor_sizes[idx] }}{% if not loop.last %},{% endif %}
+                {% endfor %}
+            };
+            tmlc_concat_s8(
+                concat_inputs,
+                concat_sizes,
+                {{ op.input_indices | length }},
+                tensor_{{ op.output_indices[0] }}
+            );
+        }
+        {% elif op.op_name == "SUB" %}
+        tmlc_sub_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.input_indices[1] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.output_indices[0]] }}
+        );
+        {% elif op.op_name == "TANH" %}
+        tmlc_tanh_s8(
+            tensor_{{ op.input_indices[0] }},
+            tensor_{{ op.output_indices[0] }},
+            {{ tensor_sizes[op.input_indices[0]] }}
+        );
+        {% endif %}
+    {% endfor %}
+
+    // Output tensor mapping
+    if (output != NULL) {
+        for (int i = 0; i < OUTPUT_SIZE; i++) {
+            output[i] = tensor_{{ last_output_tensor }}[i];
+        }
+    }
+}

TinyMLC/templates/model.h.tpl

@@ -0,0 +1,66 @@
+/* TinyMLC - Tiny Machine Learning Compiler
+ *
+ * Copyright (c) 2026 Jia Liu & TinyMLC Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * This file is part of TinyMLC.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at:
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// Auto-generated code, do not modify manually
+// Generated by tinymlc
+
+#ifndef TINYMLC_MODEL_H
+#define TINYMLC_MODEL_H
+
+#include <stdint.h>
+
+#define INPUT_SIZE {{ input_size }}
+#define OUTPUT_SIZE {{ output_size }}
+
+{% if inputs_count >= 1 %}
+#define INPUT_SIZE_1 {{ INPUT_SIZE_1 }}
+{% endif %}
+{% if inputs_count >= 2 %}
+#define INPUT_SIZE_2 {{ INPUT_SIZE_2 }}
+{% endif %}
+
+{% if target == "arm" %}
+// FC quantization parameters
+#define FC_INPUT_OFFSET 0
+#define FC_OUTPUT_OFFSET 0
+#define FC_MULTIPLIER {{ fc_multiplier }}
+#define FC_SHIFT {{ fc_shift }}
+
+// Softmax quantization parameters
+#define SOFTMAX_MULTIPLIER {{ softmax_multiplier }}
+#define SOFTMAX_SHIFT {{ softmax_shift }}
+#define SOFTMAX_DIFF_MIN -128
+{% endif %}
+
+// LSTM right shift bits (calculated from model quantization parameters)
+#define LSTM_SHIFT_I {{ lstm_shifts[0] }}
+#define LSTM_SHIFT_F {{ lstm_shifts[1] }}
+#define LSTM_SHIFT_G {{ lstm_shifts[2] }}
+#define LSTM_SHIFT_O {{ lstm_shifts[3] }}
+
+// Inference function declaration
+{% if inputs_count == 1 %}
+    void {{ inference_func }}(const int8_t* input, int8_t* output);
+{% elif inputs_count == 2 %}
+    void {{ inference_func }}(const int8_t* input1,
+                              const int8_t* input2,
+                              int8_t* output);
+{% endif %}
+
+#endif // TINYMLC_MODEL_H

TinyMLC/transform/algebraic.py

@@ -0,0 +1,286 @@
+# -*- coding: utf-8 -*-
+# TinyMLC - Tiny Machine Learning Compiler
+#
+# Copyright (c) 2026 Jia Liu & TinyMLC Contributors
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file is part of TinyMLC.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at:
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Algebraic simplification: transforms arithmetic expressions
+# to simpler equivalent forms.
+#
+# Patterns handled:
+#   1. ADD(x, -x) -> 0 (or constant zero)
+#   2. SUB(x, x) -> 0
+#   3. ADD(x, constant) -> x + constant (evaluate if possible)
+#   4. SUB(x, constant) -> x - constant
+#   5. ADD(x, ADD(y, z)) -> ADD(ADD(x, y), z)  (associative)
+#   6. ADD(x, y) -> SUB(x, -y)  (if y is negative constant)
+import numpy as np
+
+from typing import Dict, Any, List, Optional
+from TinyMLC.transform.base import Pass
+
+
+class AlgebraicSimplify(Pass):
+    """
+    Algebraic simplification for arithmetic operations.
+
+    Runs after Simplify to catch patterns that require numeric analysis.
+    """
+
+    def __init__(self, name: str = "AlgebraicSimplify"):
+        super().__init__(name)
+        self._simplified_count = 0
+
+    def run(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        model_info = self._copy_model(model_info)
+        self._simplified_count = 0
+
+        changed = True
+        iteration = 0
+
+        while changed:
+            changed = False
+            iteration += 1
+
+            # 1. SUB(x, x) -> 0
+            changed |= self._simplify_sub_same(model_info)
+
+            # 2. ADD(x, -x) -> 0 (if -x is a constant)
+            changed |= self._simplify_add_neg(model_info)
+
+            # 3. ADD(x, constant) with constant that can be folded
+            changed |= self._simplify_add_constant(model_info)
+
+            # 4. SUB(x, constant) with constant that can be folded
+            changed |= self._simplify_sub_constant(model_info)
+
+            if changed:
+                self._log_change(f"Iteration {iteration}: applied")
+
+        if self._simplified_count > 0:
+            self._log_change(
+                f"Total: {self._simplified_count} algebraic simplifications"
+            )
+
+        return model_info
+
+    # ================================================================
+    # 1. SUB(x, x) -> 0 (forward constant zero)
+    # ================================================================
+
+    def _simplify_sub_same(self, model_info: Dict[str, Any]) -> bool:
+        """
+        SUB(x, x) -> 0
+        """
+        ops = model_info.get("ops", [])
+        changed = False
+        i = 0
+
+        while i < len(ops):
+            op = ops[i]
+            if op.get("op_name") == "SUB":
+                input_indices = op.get("input_indices", [])
+                if (len(input_indices) == 2
+                        and input_indices[0] == input_indices[1]):
+                    # Create a constant zero tensor
+                    zero_idx = self._create_constant_zero(model_info)
+                    self._remove_op_and_forward(ops, i, zero_idx, model_info)
+                    changed = True
+                    self._simplified_count += 1
+                    self._log_change("  SUB(x, x) -> 0")
+                    continue
+
+            i += 1
+
+        if changed:
+            model_info["ops"] = ops
+
+        return changed
+
+    # ================================================================
+    # 2. ADD(x, -x) -> 0 (if -x is a constant)
+    # ================================================================
+
+    def _simplify_add_neg(self, model_info: Dict[str, Any]) -> bool:
+        """
+        ADD(x, -x) -> 0 when -x is a constant with opposite sign.
+        """
+        ops = model_info.get("ops", [])
+        changed = False
+        i = 0
+
+        while i < len(ops):
+            op = ops[i]
+            if op.get("op_name") == "ADD":
+                input_indices = op.get("input_indices", [])
+                if len(input_indices) == 2:
+                    a = input_indices[0]
+                    b = input_indices[1]
+                    # Check if a is constant and b is the same constant
+                    # with opposite sign
+                    const_val = self._get_constant_value(model_info, a)
+                    if const_val is not None:
+                        const_b = self._get_constant_value(model_info, b)
+                        if const_b is not None and const_b == -const_val:
+                            zero_idx = self._create_constant_zero(model_info)
+                            self._remove_op_and_forward(
+                                ops, i, zero_idx, model_info
+                            )
+                            changed = True
+                            self._simplified_count += 1
+                            self._log_change(
+                                f"  ADD({const_val}, {-const_val}) -> 0"
+                            )
+                            continue
+
+            i += 1
+
+        if changed:
+            model_info["ops"] = ops
+
+        return changed
+
+    # ================================================================
+    # 3. ADD(x, constant) with constant that can be folded
+    # ================================================================
+
+    def _simplify_add_constant(self, model_info: Dict[str, Any]) -> bool:
+        """
+        ADD(x, constant) -> x + constant (if constant is a scalar).
+        """
+        ops = model_info.get("ops", [])
+        changed = False
+        i = 0
+
+        while i < len(ops):
+            op = ops[i]
+            if op.get("op_name") == "ADD":
+                input_indices = op.get("input_indices", [])
+                if len(input_indices) == 2:
+                    # Check if one input is a constant
+                    const_idx = None
+                    other_idx = None
+                    for idx in input_indices:
+                        if self._is_scalar_constant(model_info, idx):
+                            const_idx = idx
+                        else:
+                            other_idx = idx
+
+                    if const_idx is not None and other_idx is not None:
+                        # Fold constant into the op's params
+                        const_val = self._get_constant_value(
+                            model_info, const_idx
+                        )
+                        if const_val is not None:
+                            # Replace with a new op that has the constant
+                            # baked in
+                            # For now, we just keep the constant as a tensor and
+                            # let the constant folding pass handle it.
+                            # But we can mark it for later folding.
+                            pass
+            i += 1
+
+        if changed:
+            model_info["ops"] = ops
+
+        return changed
+
+    # ================================================================
+    # 4. SUB(x, constant) -> x - constant (if constant is a scalar)
+    # ================================================================
+
+    def _simplify_sub_constant(self, model_info: Dict[str, Any]) -> bool:
+        """
+        SUB(x, constant) -> x - constant (if constant is a scalar).
+        """
+        ops = model_info.get("ops", [])
+        changed = False
+        i = 0
+
+        while i < len(ops):
+            op = ops[i]
+            if op.get("op_name") == "SUB":
+                input_indices = op.get("input_indices", [])
+                if len(input_indices) == 2:
+                    # Check if the second input is a constant
+                    const_idx = input_indices[1]
+                    if self._is_scalar_constant(model_info, const_idx):
+                        # This can be folded by constant folding
+                        pass
+            i += 1
+
+        return changed
+
+    # ================================================================
+    # Helper functions
+    # ================================================================
+
+    def _create_constant_zero(self, model_info: Dict[str, Any]) -> int:
+        """Create a constant zero tensor."""
+        # Find the max tensor index
+        tensors = model_info.get("tensors", {})
+        max_idx = max(tensors.keys()) if tensors else 0
+        new_idx = max_idx + 1
+
+        # Create zero tensor (scalar)
+        tensors[new_idx] = {
+            "name": f"zero_{new_idx}",
+            "shape": [1],
+            "dtype": "int8",
+            "scale": 1.0,
+            "zero_point": 0,
+        }
+        model_info["weights"][new_idx] = np.array([0], dtype=np.int8)
+
+        return new_idx
+
+    def _get_constant_value(self, model_info: Dict, idx: int) -> Optional[int]:
+        """Get scalar constant value if tensor at idx is a scalar constant."""
+        weights = model_info.get("weights", {})
+        if idx in weights:
+            weight = weights[idx]
+            if hasattr(weight, "size") and weight.size == 1:
+                return int(weight.item())
+            if isinstance(weight, list) and len(weight) == 1:
+                return weight[0]
+        return None
+
+    def _is_scalar_constant(self, model_info: Dict, idx: int) -> bool:
+        """Check if tensor at idx is a scalar constant."""
+        return self._get_constant_value(model_info, idx) is not None
+
+    def _remove_op_and_forward(
+        self,
+        ops: List[Dict],
+        op_idx: int,
+        forward_idx: int,
+        model_info: Dict
+    ) -> None:
+        """Remove op at op_idx and forward its output from forward_idx."""
+        op = ops[op_idx]
+        output_idx = op.get("output_indices", [])[0]
+
+        for other_op in ops:
+            for i, idx in enumerate(other_op.get("input_indices", [])):
+                if idx == output_idx:
+                    other_op["input_indices"][i] = forward_idx
+
+        del ops[op_idx]
+
+        if output_idx in model_info.get("tensors", {}):
+            del model_info["tensors"][output_idx]
+        if output_idx in model_info.get("weights", {}):
+            del model_info["weights"][output_idx]