tinymlc 0.1.0__py3-none-any.whl

TinyMLC/transform/base.py

@@ -0,0 +1,58 @@
+# -*- 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.
+
+# Base class for all optimization passes.
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any
+import copy
+
+
+class Pass(ABC):
+    """
+    Base class for all optimization passes.
+
+    Each pass takes a model_info dict, transforms it, and returns
+    the transformed model_info.
+    """
+
+    def __init__(self, name: str = None):
+        self.name = name or self.__class__.__name__
+        self._stats = {
+            "before": {},
+            "after": {},
+            "changes": [],
+        }
+
+    @abstractmethod
+    def run(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        """Run the pass on model_info and return transformed model_info."""
+        pass
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Return statistics about the pass execution."""
+        return self._stats
+
+    def _log_change(self, msg: str) -> None:
+        """Record a change made by this pass."""
+        self._stats["changes"].append(msg)
+
+    def _copy_model(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        """Deep copy model_info to avoid mutating the original."""
+        return copy.deepcopy(model_info)

TinyMLC/transform/constant_folding.py

@@ -0,0 +1,260 @@
+# -*- 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.
+
+# Constant folding optimization pass.
+
+from typing import Dict, Any
+import numpy as np
+
+from TinyMLC.transform.base import Pass
+
+
+class ConstantFolding(Pass):
+    """
+    Constant folding optimization pass.
+
+    This pass evaluates operations at compile time when all inputs
+    are constants (known at compile time).
+
+    Currently supports:
+        - Reshape with constant shape
+        - Transpose with constant permutation
+        - Concat with constant axis
+        - Add, Multiply, Subtract with constants
+
+    Future extensions:
+        - Softmax with constant input
+        - Mean with constant axis
+    """
+
+    def __init__(self, name: str = "ConstantFolding"):
+        super().__init__(name)
+        self._const_tensors: Dict[int, np.ndarray] = {}
+
+    def run(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        """Run constant folding on model_info."""
+        model_info = self._copy_model(model_info)
+
+        # 1. Find all constant tensors (weights, bias, etc.)
+        self._collect_constants(model_info)
+
+        # 2. Scan ops and fold constants
+        self._fold_ops(model_info)
+
+        # 3. Prune unused tensors
+        self._prune_unused_tensors(model_info)
+
+        return model_info
+
+    def _collect_constants(self, model_info: Dict[str, Any]) -> None:
+        """Collect all tensors that are constant (weights, biases, etc.)."""
+        weights = model_info.get("weights", {})
+        self._const_tensors = {}
+
+        for idx, weight in weights.items():
+            if isinstance(weight, np.ndarray):
+                self._const_tensors[idx] = weight
+            elif isinstance(weight, list):
+                self._const_tensors[idx] = np.array(weight)
+            else:
+                # Scalar or other type
+                self._const_tensors[idx] = np.array([weight])
+
+        self._log_change(f"Found {len(self._const_tensors)} constant tensors")
+
+    def _fold_ops(self, model_info: Dict[str, Any]) -> None:
+        """Fold operations where all inputs are constants."""
+        ops = model_info.get("ops", [])
+        new_ops = []
+        folded_count = 0
+
+        for op in ops:
+            op_name = op.get("op_name")
+            folded = False
+
+            # Only fold if we can evaluate it now
+            if op_name == "RESHAPE":
+                folded = self._fold_reshape(model_info, op)
+            elif op_name == "TRANSPOSE":
+                folded = self._fold_transpose(model_info, op)
+            elif op_name in ("ADD", "MULTIPLY", "SUB"):
+                folded = self._fold_binary_op(model_info, op)
+            elif op_name == "MEAN":
+                folded = self._fold_mean(model_info, op)
+
+            if folded:
+                folded_count += 1
+                self._log_change(f"Folded {op_name}")
+            else:
+                new_ops.append(op)
+
+        if folded_count > 0:
+            model_info["ops"] = new_ops
+            self._log_change(f"Folded {folded_count} ops")
+
+    def _fold_reshape(
+        self, model_info: Dict[str, Any], op: Dict[str, Any]
+    ) -> bool:
+        """Fold reshape if input is constant."""
+        input_idx = op.get("input_indices", [])[0]
+        output_idx = op.get("output_indices", [])[0]
+
+        if input_idx in self._const_tensors:
+            try:
+                # Get the target shape from params
+                params = op.get("reshape_params", {})
+                target_shape = params.get("shape", [])
+                if not target_shape:
+                    target_shape = params.get("target_shape", [])
+
+                data = self._const_tensors[input_idx]
+                folded = data.reshape(target_shape)
+
+                # Store as constant tensor
+                self._const_tensors[output_idx] = folded
+
+                # Add to weights so it gets written out
+                model_info["weights"][output_idx] = folded
+
+                self._log_change(
+                    f"  Reshape constant: {data.shape} -> {folded.shape}"
+                )
+                return True
+            except Exception as e:
+                print(f"  Warning: failed to fold reshape: {e}")
+                return False
+        return False
+
+    def _fold_transpose(
+        self, model_info: Dict[str, Any], op: Dict[str, Any]
+    ) -> bool:
+        """Fold transpose if input is constant."""
+        input_idx = op.get("input_indices", [])[0]
+        output_idx = op.get("output_indices", [])[0]
+
+        if input_idx in self._const_tensors:
+            try:
+                params = op.get("transpose_params", {})
+                perm = params.get("perm", [])
+
+                data = self._const_tensors[input_idx]
+                folded = np.transpose(data, axes=perm or None)
+
+                self._const_tensors[output_idx] = folded
+                model_info["weights"][output_idx] = folded
+
+                self._log_change(
+                    f"  Transpose constant: {data.shape} -> {folded.shape}"
+                )
+                return True
+            except Exception as e:
+                print(f"  Warning: failed to fold transpose: {e}")
+                return False
+        return False
+
+    def _fold_binary_op(
+        self, model_info: Dict[str, Any], op: Dict[str, Any]
+    ) -> bool:
+        """Fold binary ops (ADD, MULTIPLY, SUB) if all inputs constant."""
+        op_name = op.get("op_name")
+        input_indices = op.get("input_indices", [])
+        output_idx = op.get("output_indices", [0])[0]
+
+        if all(idx in self._const_tensors for idx in input_indices):
+            try:
+                a = self._const_tensors[input_indices[0]]
+                b = self._const_tensors[input_indices[1]]
+
+                if op_name == "ADD":
+                    folded = a + b
+                elif op_name == "MULTIPLY":
+                    folded = a * b
+                elif op_name == "SUB":
+                    folded = a - b
+                else:
+                    return False
+
+                self._const_tensors[output_idx] = folded
+                model_info["weights"][output_idx] = folded
+
+                self._log_change(
+                    f"  {op_name} constant: {a.shape} + {b.shape} "
+                    f"-> {folded.shape}"
+                )
+                return True
+            except Exception as e:
+                print(f"  Warning: failed to fold {op_name}: {e}")
+                return False
+        return False
+
+    def _fold_mean(
+        self, model_info: Dict[str, Any], op: Dict[str, Any]
+    ) -> bool:
+        """Fold MEAN if input is constant."""
+        input_idx = op.get("input_indices", [0])[0]
+        output_idx = op.get("output_indices", [0])[0]
+
+        if input_idx in self._const_tensors:
+            try:
+                params = op.get("mean_params", {})
+                axis = params.get("axis", None)
+                keepdims = params.get("keepdims", False)
+
+                data = self._const_tensors[input_idx]
+                folded = np.mean(data, axis=axis, keepdims=keepdims)
+
+                self._const_tensors[output_idx] = folded
+                model_info["weights"][output_idx] = folded
+
+                self._log_change(
+                    f"  Mean constant: {data.shape} -> {folded.shape}"
+                )
+                return True
+            except Exception as e:
+                print(f"  Warning: failed to fold mean: {e}")
+                return False
+        return False
+
+    def _prune_unused_tensors(self, model_info: Dict[str, Any]) -> None:
+        """Remove tensors that are no longer used."""
+        # Get all used tensor indices from ops
+        used_indices = set()
+        for op in model_info.get("ops", []):
+            for idx in op.get("input_indices", []):
+                used_indices.add(idx)
+            for idx in op.get("output_indices", []):
+                used_indices.add(idx)
+
+        # Get input/output indices
+        for inp in model_info.get("input", []):
+            # Inputs don't have indices in this representation
+            pass
+
+        # Remove unused tensors
+        all_indices = set(model_info["tensors"].keys())
+        unused = all_indices - used_indices
+
+        for idx in unused:
+            if idx in model_info["tensors"]:
+                del model_info["tensors"][idx]
+            if idx in model_info["weights"]:
+                del model_info["weights"][idx]
+
+        if unused:
+            self._log_change(f"Removed {len(unused)} unused tensors")

TinyMLC/transform/cse.py

@@ -0,0 +1,192 @@
+# -*- 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 Subexpression Elimination.
+
+from typing import Dict, Any
+import hashlib
+import json
+
+from TinyMLC.transform.base import Pass
+
+
+class CommonSubexpressionElimination(Pass):
+    """
+    Common Subexpression Elimination.
+
+    Finds and eliminates duplicate computations:
+        - Same op with same inputs and same params
+        - Same constant tensor being computed multiple times
+
+    Strategy:
+        1. Compute a signature for each op (op_name + input_indices + params)
+        2. If two ops have the same signature, keep the first one
+        3. Replace all uses of the later op's outputs with
+           the first op's outputs
+    """
+
+    def __init__(self, name: str = "CommonSubexpressionElimination"):
+        super().__init__(name)
+        self._signature_map: Dict[str, int] = {}  # signature -> op_index
+        self._replace_map: Dict[int, int] = {}    # old_tensor -> new_tensor
+
+    def run(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        """Run CSE on model_info."""
+        model_info = self._copy_model(model_info)
+
+        self._signature_map.clear()
+        self._replace_map.clear()
+
+        ops = model_info.get("ops", [])
+        new_ops = []
+        eliminated_count = 0
+
+        for op in ops:
+            signature = self._compute_signature(op)
+
+            if signature in self._signature_map:
+                # Duplicate found: replace outputs with the original
+                orig_op_idx = self._signature_map[signature]
+                orig_op = new_ops[orig_op_idx]
+                self._replace_outputs(op, orig_op)
+                eliminated_count += 1
+                self._log_change(
+                    f"  Eliminated duplicate {op.get('op_name')} "
+                    f"(outputs: {op.get('output_indices')} "
+                    f"-> {orig_op.get('output_indices')})"
+                )
+                # Don't add this op to new_ops
+            else:
+                # New op: record its signature and keep it
+                self._signature_map[signature] = len(new_ops)
+                new_ops.append(op)
+
+        if eliminated_count > 0:
+            model_info["ops"] = new_ops
+            # Update all tensor references in remaining ops
+            self._update_tensor_refs(model_info)
+            self._log_change(
+                f"Eliminated {eliminated_count} duplicate expressions"
+            )
+
+        return model_info
+
+    def _compute_signature(self, op: Dict[str, Any]) -> str:
+        """
+        Compute a unique signature for an op.
+
+        The signature includes:
+            - op_name
+            - input_indices (sorted for commutativity)
+            - output_indices (for ops with multiple outputs)
+            - params (sorted, excluding irrelevant fields)
+        """
+        op_name = op.get("op_name", "UNKNOWN")
+
+        # Input indices: sorted for commutative ops? Not always safe.
+        # For now, keep the order as-is, as ops like SUB are not commutative.
+        input_indices = op.get("input_indices", [])
+
+        # Params: filter out fields that don't affect computation
+        params = op.get("params", {})
+        # Remove fields that are just metadata
+        skip_keys = {"name", "index", "state", "pass_flags"}
+        filtered_params = {
+            k: v for k, v in params.items()
+            if k not in skip_keys and not k.startswith("_")
+        }
+
+        # Build signature dict
+        sig = {
+            "op_name": op_name,
+            "input_indices": input_indices,
+            "params": filtered_params,
+        }
+
+        # Hash to a string
+        sig_str = json.dumps(sig, sort_keys=True)
+        return hashlib.sha256(sig_str.encode()).hexdigest()[:16]
+
+    def _replace_outputs(
+        self, dup_op: Dict[str, Any], orig_op: Dict[str, Any]
+    ) -> None:
+        """
+        Map outputs of dup_op to outputs of orig_op.
+
+        Assumes the output indices are in the same order.
+        """
+        dup_outputs = dup_op.get("output_indices", [])
+        orig_outputs = orig_op.get("output_indices", [])
+
+        if len(dup_outputs) != len(orig_outputs):
+            # Different number of outputs, can't replace
+            self._log_change(
+                f"  Warning: output count mismatch "
+                f"({len(dup_outputs)} vs {len(orig_outputs)})"
+            )
+            return
+
+        for dup_idx, orig_idx in zip(dup_outputs, orig_outputs):
+            self._replace_map[dup_idx] = orig_idx
+
+    def _update_tensor_refs(self, model_info: Dict[str, Any]) -> None:
+        """
+        Update all tensor references in ops:
+            - Replace old tensor indices with new ones
+            - Remove any ops that now have duplicate inputs/outputs
+        """
+        if not self._replace_map:
+            return
+
+        ops = model_info.get("ops", [])
+
+        for op in ops:
+            # Update input_indices
+            input_indices = op.get("input_indices", [])
+            new_inputs = [
+                self._replace_map.get(idx, idx) for idx in input_indices
+            ]
+            op["input_indices"] = new_inputs
+
+            # Update output_indices
+            output_indices = op.get("output_indices", [])
+            new_outputs = [
+                self._replace_map.get(idx, idx) for idx in output_indices
+            ]
+            op["output_indices"] = new_outputs
+
+        # Update tensors dict: remove replaced tensors
+        tensors = model_info.get("tensors", {})
+        for old_idx in self._replace_map.keys():
+            if old_idx in tensors:
+                del tensors[old_idx]
+
+        # Update weights dict
+        weights = model_info.get("weights", {})
+        for old_idx in self._replace_map.keys():
+            if old_idx in weights:
+                del weights[old_idx]
+
+        # Update tensor references in input/output specs
+        # (tensor_index is metadata, we don't need to update it for CSE)
+        # But if we want to keep consistency, we could update it.
+
+        self._log_change(
+            f"  Replaced {len(self._replace_map)} tensor references"
+        )

TinyMLC/transform/dce.py

@@ -0,0 +1,182 @@
+# -*- 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.
+
+# Dead Code Elimination pass.
+
+from typing import Dict, Any, Set, List
+from TinyMLC.transform.base import Pass
+
+
+class DeadCodeElimination(Pass):
+    """
+    Dead Code Elimination.
+
+    Removes:
+        - Tensors that are never used as inputs to any op
+        - Ops whose outputs are never used
+        - Unreachable ops (in case of control flow, not implemented yet)
+
+    This pass should be run after each pass that may create dead code.
+    """
+
+    def __init__(self, name: str = "DeadCodeElimination"):
+        super().__init__(name)
+
+    def run(self, model_info: Dict[str, Any]) -> Dict[str, Any]:
+        """Run dead code elimination on model_info."""
+        model_info = self._copy_model(model_info)
+
+        changed = True
+        iteration = 0
+
+        while changed:
+            changed = False
+            iteration += 1
+
+            # 1. Find all used tensor indices
+            used_indices = self._collect_used_indices(model_info)
+
+            # 2. Remove dead tensors
+            dead_tensors = self._remove_dead_tensors(model_info, used_indices)
+            if dead_tensors:
+                changed = True
+                self._log_change(
+                    f"Iteration {iteration}: removed {len(dead_tensors)} "
+                    f"dead tensors"
+                )
+
+            # 3. Remove dead ops
+            dead_ops = self._remove_dead_ops(model_info, used_indices)
+            if dead_ops:
+                changed = True
+                self._log_change(
+                    f"Iteration {iteration}: removed {len(dead_ops)} dead ops"
+                )
+
+        return model_info
+
+    def _collect_used_indices(self, model_info: Dict[str, Any]) -> Set[int]:
+        """
+        Collect all tensor indices that are used as inputs to any op.
+        """
+        used = set()
+
+        # 1. All input tensors are used
+        for inp in model_info.get("input", []):
+            idx = inp.get("tensor_index")
+            if idx is not None:
+                used.add(idx)
+
+        # 2. All output tensors are used
+        for out in model_info.get("output", []):
+            idx = out.get("tensor_index")
+            if idx is not None:
+                used.add(idx)
+
+        # 3. All tensor indices referenced by ops
+        for op in model_info.get("ops", []):
+            for idx in op.get("input_indices", []):
+                used.add(idx)
+            for idx in op.get("output_indices", []):
+                used.add(idx)
+
+        return used
+
+    def _remove_dead_tensors(
+        self,
+        model_info: Dict[str, Any],
+        used_indices: Set[int]
+    ) -> Set[int]:
+        """
+        Remove tensors that are not used as inputs to any op,
+        except outputs (they must be preserved).
+        """
+        # Output tensors must be preserved (they are the final result)
+        output_indices = set()
+        for out in model_info.get("output", []):
+            # Outputs are identified by name, not index
+            # We need to find which tensor index corresponds to each output
+            # For now, assume outputs are in tensors dict with some mapping
+            pass
+
+        # For simplicity: find tensors that are never used as inputs
+        all_indices = set(model_info.get("tensors", {}).keys())
+        dead = all_indices - used_indices
+
+        # Don't delete tensors that are explicitly marked as outputs
+        # This requires knowing which tensors are outputs.
+        # In our model_info, outputs are separate from tensors.
+        # For now, we keep all tensors that are outputs.
+
+        # Actually delete them
+        for idx in dead:
+            if idx in model_info.get("tensors", {}):
+                del model_info["tensors"][idx]
+            if idx in model_info.get("weights", {}):
+                del model_info["weights"][idx]
+
+        return dead
+
+    def _remove_dead_ops(
+        self,
+        model_info: Dict[str, Any],
+        used_indices: Set[int]
+    ) -> List[Dict[str, Any]]:
+        """
+        Remove ops whose output indices are never used as inputs.
+        """
+        # For each op, check if any of its outputs are used
+        ops = model_info.get("ops", [])
+        tensors = model_info.get("tensors", {})
+        dead_ops = []
+        alive_ops = []
+
+        for op in ops:
+            outputs = op.get("output_indices", [])
+            # Check if all output indices are in tensors.
+            all_outputs_valid = all(idx in tensors for idx in outputs)
+            # An op is alive if any of its outputs is used
+            is_alive = (
+                any(idx in used_indices for idx in outputs)
+                and all_outputs_valid
+            )
+
+            # Also: if this op produces an output tensor that is
+            # the final output
+            # For now, keep it if it's the last op in the graph
+            # (we'll use a more sophisticated analysis later)
+
+            if is_alive:
+                alive_ops.append(op)
+            else:
+                dead_ops.append(op)
+
+        if dead_ops:
+            model_info["ops"] = alive_ops
+
+            # Remove any tensors that were only produced by dead ops
+            # (they'll be caught by the tensor removal in the next iteration)
+            for op in dead_ops:
+                for idx in op.get("output_indices", []):
+                    if idx in model_info.get("tensors", {}):
+                        # Don't delete right away, let the tensor removal
+                        # handle it
+                        pass
+
+        return dead_ops