tinymlc 0.1.0__py3-none-any.whl

TinyMLC/ANG/model_info.py

@@ -0,0 +1,283 @@
+# -*- 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.
+
+# Unified Intermediate Representation (IR) for TinyMLC.
+# This structure is used as the contract between all frontends
+# (LiteRT, ONNX, ANG) and the backend code generator.
+
+from typing import Dict, List, Optional, Any
+import numpy as np
+
+
+class TensorSpec:
+    """
+    Specification of a tensor in the model.
+
+    Attributes:
+        name: Unique identifier for the tensor.
+        shape: List of dimensions (batch, height, width, channels, ...).
+        dtype: Data type of tensor values (e.g., "int8", "int32", "float32").
+        scale: Quantization scale factor (optional, None for float models).
+        zero_point: Quantization zero point (optional, None for float models).
+    """
+
+    def __init__(
+        self,
+        name: str,
+        shape: List[int],
+        dtype: str,
+        tensor_index: Optional[int] = None,
+        scale: Optional[float] = None,
+        zero_point: Optional[int] = None,
+    ):
+        self.name = name
+        self.shape = shape
+        self.dtype = dtype
+        self.tensor_index = tensor_index
+        self.scale = scale
+        self.zero_point = zero_point
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary representation."""
+        return {
+            "name": self.name,
+            "shape": self.shape,
+            "dtype": self.dtype,
+            "tensor_index": self.tensor_index,
+            "scale": self.scale,
+            "zero_point": self.zero_point,
+        }
+
+
+class Op:
+    """
+    A single operation/layer in the model.
+
+    Attributes:
+        index: Unique index for this operation.
+        op_name: Type of operation (e.g., "CONV_2D", "FULLY_CONNECTED").
+        input_indices: List of tensor indices that are inputs to this op.
+        output_indices: List of tensor indices that are outputs of this op.
+        params: Op-specific parameters (conv_params, fc_params, etc.).
+        state: Operator state ("created", "translated", "generated").
+        pass_flags: Dictionary of pass flags.
+    """
+
+    def __init__(
+        self,
+        op_name: str,
+        input_indices: List[int],
+        output_indices: List[int],
+        params: Optional[Dict[str, Any]] = None,
+        index: Optional[int] = None,
+        state: str = "created",
+    ):
+        self.index = index
+        self.op_name = op_name
+        self.input_indices = input_indices
+        self.output_indices = output_indices
+        self.params = params or {}
+        self.state = state
+        self.pass_flags = {}
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to a dictionary representation."""
+        result = {
+            "index": self.index,
+            "op_name": self.op_name,
+            "input_indices": self.input_indices,
+            "output_indices": self.output_indices,
+            "state": self.state,
+            "pass_flags": self.pass_flags,
+        }
+        result.update(self.params)
+        return result
+
+
+class ModelInfo:
+    """
+    Unified Intermediate Representation for neural network models.
+
+    This is the central data structure used throughout TinyMLC.
+    All frontends (LiteRT, ONNX, ANG) produce this structure.
+    The backend code generator consumes this structure.
+
+    Attributes:
+        inputs: List of input tensor specifications.
+        outputs: List of output tensor specifications.
+        ops: List of operations in execution order.
+        tensors: Dictionary mapping tensor index to TensorSpec.
+        weights: Dictionary mapping tensor index to numpy array data.
+        quant_scales: Global quantization scales (if not per-tensor).
+    """
+
+    def __init__(
+        self,
+        inputs: List[TensorSpec],
+        outputs: List[TensorSpec],
+        ops: List[Op],
+        tensors: Dict[int, TensorSpec],
+        weights: Dict[int, np.ndarray],
+        quant_scales: Optional[Dict[str, Any]] = None,
+    ):
+        self.inputs = inputs
+        self.outputs = outputs
+        self.ops = ops
+        self.tensors = tensors
+        self.weights = weights
+        self.quant_scales = quant_scales or {}
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the entire ModelInfo to a dictionary."""
+        return {
+            "input": [t.to_dict() for t in self.inputs],
+            "output": [t.to_dict() for t in self.outputs],
+            "ops": [op.to_dict() for op in self.ops],
+            "tensors": {
+                idx: spec.to_dict() for idx, spec in self.tensors.items()
+            },
+            "weights": {
+                str(idx): weight.tolist()
+                for idx, weight in self.weights.items()
+            },
+            "quant_scales": self.quant_scales,
+        }
+
+    def get_tensor(self, index: int) -> Optional[TensorSpec]:
+        """Get tensor specification by index."""
+        return self.tensors.get(index)
+
+    def get_weight(self, index: int) -> Optional[np.ndarray]:
+        """Get weight data by tensor index."""
+        return self.weights.get(index)
+
+    def add_tensor(
+        self,
+        index: int,
+        spec: TensorSpec,
+        weight: Optional[np.ndarray] = None,
+    ) -> None:
+        """
+        Add a tensor to the model.
+
+        Args:
+            index: Unique index for the tensor.
+            spec: TensorSpec describing the tensor.
+            weight: Optional numpy array with weight data.
+        """
+        self.tensors[index] = spec
+        if weight is not None:
+            self.weights[index] = weight
+
+    def add_op(self, op: Op) -> None:
+        """Add an operation to the model."""
+        self.ops.append(op)
+
+    def validate(self) -> bool:
+        """
+        Validate the model for consistency.
+
+        Checks:
+            - All indices in ops refer to valid tensors.
+            - All weights have matching shapes with their tensor specs.
+            - Input/output lists are not empty.
+
+        Returns:
+            True if the model is valid, False otherwise.
+        """
+        # Check that input and output are not empty
+        if not self.inputs or not self.outputs:
+            return False
+
+        # Check that all tensor indices in ops exist
+        all_tensor_indices = set(self.tensors.keys())
+        for op in self.ops:
+            for idx in op.input_indices + op.output_indices:
+                if idx not in all_tensor_indices:
+                    return False
+
+        # Check that weights match tensor shapes
+        for idx, weight in self.weights.items():
+            if idx not in self.tensors:
+                return False
+            expected_shape = self.tensors[idx].shape
+            if list(weight.shape) != expected_shape:
+                # Allow broadcasting for scalar weights (shape mismatch is OK)
+                if len(weight.shape) != 0:
+                    return False
+
+        return True
+
+
+def default_quant_scale() -> float:
+    """
+    Return the default quantization scale for int8 models.
+
+    This is the scale corresponding to 1/256, which is a common default
+    when no specific scale is provided.
+
+    Returns:
+        Default scale value.
+    """
+    # 1/256 = 0.00390625
+    return 1.0 / 256.0
+
+
+def default_zero_point() -> int:
+    """
+    Return the default zero point for int8 models.
+
+    Returns:
+        Default zero point (0 for symmetric quantization).
+    """
+    return 0
+
+
+def create_default_tensor_spec(
+    name: str,
+    shape: List[int],
+    dtype: str = "int8",
+) -> TensorSpec:
+    """
+    Create a tensor specification with default quantization parameters.
+
+    Args:
+        name: Tensor name.
+        shape: Tensor shape.
+        dtype: Data type (default: "int8").
+
+    Returns:
+        A TensorSpec with default scale and zero point if quantized.
+    """
+    if dtype.startswith("int"):
+        return TensorSpec(
+            name=name,
+            shape=shape,
+            dtype=dtype,
+            scale=default_quant_scale(),
+            zero_point=default_zero_point(),
+        )
+    else:
+        return TensorSpec(
+            name=name,
+            shape=shape,
+            dtype=dtype,
+            scale=None,
+            zero_point=None,
+        )

TinyMLC/ANG/utils.py

@@ -0,0 +1,420 @@
+# -*- 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.
+
+# Common utility functions used across ANG.
+
+import copy
+import json
+import hashlib
+from typing import Dict, Any, List, Optional
+import numpy as np
+
+
+def hash_structure(structure: Dict[str, Any]) -> str:
+    """
+    Compute a hash of a network structure (op sequence only, not weights).
+
+    This is used for caching and duplicate detection.
+
+    Args:
+        structure: Dictionary containing the network structure.
+
+    Returns:
+        SHA256 hash string.
+    """
+    # Extract only the structure-defining parts
+    ops = structure.get("ops", [])
+    # Normalize to a stable representation
+    normalized = []
+    for op in ops:
+        op_copy = {
+            "op_name": op.get("op_name"),
+            "params": op.get("params", {}),
+        }
+        normalized.append(op_copy)
+
+    # Sort to ensure stability
+    normalized.sort(key=lambda x: str(x))
+    json_str = json.dumps(normalized, sort_keys=True)
+    return hashlib.sha256(json_str.encode()).hexdigest()
+
+
+def calculate_macs(model_info: Dict[str, Any]) -> int:
+    """
+    Calculate the total number of multiply-accumulate operations.
+
+    Args:
+        model_info: ModelInfo dictionary.
+
+    Returns:
+        Total MACs count.
+    """
+    total = 0
+    ops = model_info.get("ops", [])
+    tensors = model_info.get("tensors", {})
+
+    for op in ops:
+        op_name = op.get("op_name")
+
+        if op_name == "CONV_2D":
+            conv_params = op.get("conv_params", {})
+            kernel_size = conv_params.get("kernel_size", 3)
+            stride = conv_params.get("stride", 1)
+
+            # Get input and output shapes
+            input_idx = op.get("input_indices", [0])[0]
+            output_idx = op.get("output_indices", [0])[0]
+
+            input_shape = tensors.get(input_idx, {}).get(
+                "shape", [1, 1, 1, 1]
+            )
+            output_shape = tensors.get(output_idx, {}).get(
+                "shape", [1, 1, 1, 1]
+            )
+
+            # NHWC format: [batch, height, width, channels]
+            if len(input_shape) >= 4 and len(output_shape) >= 4:
+                h = output_shape[1] if output_shape[1] else 1
+                w = output_shape[2] if output_shape[2] else 1
+                c_in = input_shape[3] if input_shape[3] else 1
+                c_out = output_shape[3] if output_shape[3] else 1
+                # MACs = H * W * C_in * C_out * K * K
+                # Bias is not counted as a MAC
+                total += h * w * c_in * c_out * kernel_size * kernel_size
+
+        elif op_name == "FULLY_CONNECTED":
+            fc_params = op.get("fc_params", {})
+            input_size = 1
+            output_size = fc_params.get("units", 64)
+
+            # Get input shape
+            input_idx = op.get("input_indices", [0])[0]
+            input_shape = tensors.get(input_idx, {}).get("shape", [])
+            for dim in input_shape:
+                input_size *= dim
+
+            # MACs = input_size * output_size
+            total += input_size * output_size
+
+        elif op_name == "DEPTHWISE_CONV_2D":
+            # Similar to CONV_2D but with groups = C_in = C_out
+            dw_params = op.get("dw_params", {})
+            kernel_size = dw_params.get("kernel_size", 3)
+
+            input_idx = op.get("input_indices", [0])[0]
+            output_idx = op.get("output_indices", [0])[0]
+
+            input_shape = tensors.get(input_idx, {}).get(
+                "shape", [1, 1, 1, 1]
+            )
+            output_shape = tensors.get(output_idx, {}).get(
+                "shape", [1, 1, 1, 1]
+            )
+
+            if len(input_shape) >= 4 and len(output_shape) >= 4:
+                h = output_shape[1] if output_shape[1] else 1
+                w = output_shape[2] if output_shape[2] else 1
+                c = input_shape[3] if input_shape[3] else 1
+                # MACs = H * W * C * K * K (no C_out multiplier)
+                total += h * w * c * kernel_size * kernel_size
+        elif op_name == "UPSAMPLE_2D":
+            # Upsample has no MACs, just interpolation.
+            # Computation is 0, but reserved for potential future weighting.
+            total += 0
+
+        elif op_name == "CONCAT":
+            # Concat has no MACs; it's just memory concatenation.
+            total += 0
+
+        elif op_name == "ADD":
+            # Add is element-wise and typically doesn't count as MACs
+            # (though some implementations count 1 op/element).
+            # We conservatively estimate it as 0 here.
+            total += 0
+
+        elif op_name == "DETECTION_HEAD":
+            # Detection heads usually have 1-2 conv layers.
+            # We simplify this by reading from params if available.
+            # Otherwise, we estimate it as 0 for generality,
+            # as the true MACs are covered by internal convolutions.
+            total += 0
+
+        elif op_name in ["CONV_TRANSPOSE", "TRANSPOSED_CONV"]:
+            # Transposed convolution: Same calculation as standard convolution,
+            # but reversed, enlarging H/W.
+            conv_params = op.get("conv_params", {})
+            kernel_size = conv_params.get("kernel_size", 3)
+            stride = conv_params.get("stride", 2)
+
+            input_idx = op.get("input_indices", [0])[0]
+            output_idx = op.get("output_indices", [0])[0]
+
+            input_shape = tensors.get(input_idx, {}).get("shape", [1, 1, 1, 1])
+            output_shape = tensors.get(output_idx, {}).get("shape",
+                                                           [1, 1, 1, 1])
+
+            if len(input_shape) >= 4 and len(output_shape) >= 4:
+                h = output_shape[1] if output_shape[1] else 1
+                w = output_shape[2] if output_shape[2] else 1
+                c_in = input_shape[3] if input_shape[3] else 1
+                c_out = output_shape[3] if output_shape[3] else 1
+                # MACs of transposed convolution: H * W * C_in * C_out * K * K
+                total += h * w * c_in * c_out * kernel_size * kernel_size
+
+    return total
+
+
+def calculate_params(model_info: Dict[str, Any]) -> int:
+    """
+    Calculate the total number of parameters.
+
+    Args:
+        model_info: ModelInfo dictionary.
+
+    Returns:
+        Total parameter count.
+    """
+    total = 0
+    weights = model_info.get("weights", {})
+
+    for weight in weights.values():
+        # Weight is a numpy array or a list
+        if hasattr(weight, "size"):
+            total += weight.size
+        elif isinstance(weight, list):
+            total += len(weight)
+
+    return total
+
+
+def calculate_peak_ram(model_info: Dict[str, Any]) -> int:
+    """
+    Calculate the peak RAM usage of the model.
+
+    This is the maximum size of all tensors that need to be held in
+    memory simultaneously at any point during inference.
+
+    Args:
+        model_info: ModelInfo dictionary.
+
+    Returns:
+        Peak RAM usage in bytes.
+    """
+    ops = model_info.get("ops", [])
+    tensors = model_info.get("tensors", {})
+
+    # Track which tensors are live at each op
+    tensor_sizes = {}
+    for idx, spec in tensors.items():
+        shape = spec.get("shape", [])
+        size = 1
+        for dim in shape:
+            size *= dim
+        # Assume int8 = 1 byte per element
+        tensor_sizes[idx] = size
+
+    # Simulate execution to find peak memory usage
+    live_tensors = set()
+    peak = 0
+
+    for op in ops:
+        # Input tensors become live before the op
+        for idx in op.get("input_indices", []):
+            live_tensors.add(idx)
+
+        # Output tensors become live after the op
+        for idx in op.get("output_indices", []):
+            live_tensors.add(idx)
+
+        # Calculate current memory usage
+        current = 0
+        for idx in live_tensors:
+            current += tensor_sizes.get(idx, 0)
+
+        peak = max(peak, current)
+
+        # Tensors that are no longer needed can be freed
+        # For simplicity, we assume only the last output is kept
+        # More precise analysis would require liveness tracking
+
+    return peak
+
+
+def calculate_flash(model_info):
+    # Estimated Flash usage = Params + Code size (approximate)
+    params = calculate_params(model_info)
+    # Code size is roughly 1KB, but can be calculated more precisely if needed
+    return params + 1024
+
+
+def flatten_weights(weights: Dict[int, np.ndarray]) -> Dict[int, List[int]]:
+    """
+    Flatten all weight tensors into lists.
+
+    Args:
+        weights: Dictionary of weight tensors (numpy arrays).
+
+    Returns:
+        Dictionary of flattened weight lists.
+    """
+    result = {}
+    for idx, weight in weights.items():
+        if hasattr(weight, "flatten"):
+            result[idx] = weight.flatten().tolist()
+        elif isinstance(weight, list):
+            result[idx] = weight
+        else:
+            result[idx] = [weight]
+    return result
+
+
+def generate_random_weights_from_structure(
+    structure: Dict[str, Any],
+    seed: Optional[int] = None,
+) -> Dict[int, np.ndarray]:
+    """
+    Generate random weights for a network structure.
+
+    Args:
+        structure: Structure dict with 'layers', 'input_shape', 'output_shape'
+        seed: Random seed for reproducibility
+
+    Returns:
+        Dict mapping tensor index to weight array
+    """
+    if seed is not None:
+        np.random.seed(seed)
+
+    weights = {}
+    layers = structure.get("layers", [])
+    input_shape = structure.get("input_shape", [1, 28, 28, 1])
+
+    next_idx = 1
+    current_shape = list(input_shape)
+    current_channels = input_shape[-1] if len(input_shape) >= 2 else 1
+
+    for layer in layers:
+        layer_type = layer.get("type")
+
+        if layer_type == "conv":
+            kernel = layer.get("kernel", 3)
+            channels_out = layer.get("channels", 16)
+
+            # Weight: [K, K, C_in, C_out]
+            weight_shape = [kernel, kernel, current_channels, channels_out]
+            weight = np.random.uniform(-0.5, 0.5, size=weight_shape)
+            weight = np.clip(np.round(weight * 256), -128, 127).astype(np.int8)
+            weights[next_idx] = weight
+
+            # Bias: [C_out]
+            bias_shape = [channels_out]
+            bias = np.random.uniform(-0.5, 0.5, size=bias_shape)
+            bias = np.clip(np.round(bias * 256), -128, 127).astype(np.int32)
+            weights[next_idx + 1] = bias
+
+            next_idx += 2
+            current_channels = channels_out
+
+        elif layer_type == "fc":
+            units = layer.get("units", 64)
+
+            # Flatten current shape to get input size
+            input_size = 1
+            for d in current_shape:
+                input_size *= d
+
+            # Weight: [input_size, units]
+            weight_shape = [input_size, units]
+            weight = np.random.uniform(-0.5, 0.5, size=weight_shape)
+            weight = np.clip(np.round(weight * 256), -128, 127).astype(np.int8)
+            weights[next_idx] = weight
+
+            # Bias: [units]
+            bias_shape = [units]
+            bias = np.random.uniform(-0.5, 0.5, size=bias_shape)
+            bias = np.clip(np.round(bias * 256), -128, 127).astype(np.int32)
+            weights[next_idx + 1] = bias
+
+            next_idx += 2
+
+        elif layer_type == "pool":
+            # Pool layers have no weights
+            pass
+
+        elif layer_type == "detection_head":
+            # Detection head is a placeholder, no weights for now
+            pass
+
+    return weights
+
+
+def _random_int8_weight(shape: List[int]) -> np.ndarray:
+    """
+    Generate random int8 weight tensor.
+
+    Values are uniformly distributed in [-128, 127] range.
+
+    Args:
+        shape: Shape of the weight tensor.
+
+    Returns:
+        int8 numpy array.
+    """
+    # Generate random values in [-128, 127] range
+    # Use uniform distribution centered at 0
+    weight = np.random.uniform(-0.5, 0.5, size=shape)
+    # Scale to int8 range: [-128, 127]
+    # 0.5 * 256 = 128, so this maps [-0.5, 0.5] to [-128, 127]
+    weight = np.round(weight * 256).astype(np.int8)
+    return weight
+
+
+def _random_int32_bias(shape: List[int]) -> np.ndarray:
+    """
+    Generate random int32 bias tensor.
+
+    Values are uniformly distributed in [-128, 127] range.
+
+    Args:
+        shape: Shape of the bias tensor.
+
+    Returns:
+        int32 numpy array.
+    """
+    # Use same range as weights for consistency
+    bias = np.random.uniform(-0.5, 0.5, size=shape)
+    bias = np.round(bias * 256).astype(np.int32)
+    return bias
+
+
+def fill_model_info_with_weights(
+    model_info: Dict[str, Any],
+    weights: Dict[int, np.ndarray],
+) -> Dict[str, Any]:
+    """
+    Fill a model_info dictionary with weight data.
+    """
+    result = copy.deepcopy(model_info)
+    result["weights"] = {}
+    for idx, weight in weights.items():
+        if isinstance(weight, np.ndarray):
+            result["weights"][str(idx)] = weight.tolist()
+        else:
+            result["weights"][str(idx)] = weight
+    return result