diff --git a/tensorflow/lite/micro/compression/BUILD b/tensorflow/lite/micro/compression/BUILD index 375a42d7a49..4d3bc1ce45c 100644 --- a/tensorflow/lite/micro/compression/BUILD +++ b/tensorflow/lite/micro/compression/BUILD @@ -326,6 +326,55 @@ tflm_py_library( ], ) +tflm_py_library( + name = "decode_insert", + srcs = ["decode_insert.py"], + deps = [ + ":compressor", + ":model_editor", + "//tensorflow/lite/python:schema_py", + ], +) + +tflm_py_test( + name = "decode_insert_test", + size = "small", + srcs = ["decode_insert_test.py"], + tags = [ + "noasan", + "nomsan", + "noubsan", + ], + deps = [ + ":compressor", + ":decode", + ":decode_insert", + ":model_editor", + "//tensorflow/lite/python:schema_py", + requirement("numpy"), + ], +) + +tflm_py_test( + name = "decode_insert_runtime_test", + size = "small", + srcs = ["decode_insert_runtime_test.py"], + tags = [ + "noasan", + "nomsan", + "noubsan", + ], + deps = [ + ":decode_insert", + ":lut", + ":model_editor", + ":spec", + "//python/tflite_micro:runtime", + "//tensorflow/lite/python:schema_py", + requirement("numpy"), + ], +) + tflm_py_binary( name = "view", srcs = [ diff --git a/tensorflow/lite/micro/compression/decode_insert.py b/tensorflow/lite/micro/compression/decode_insert.py new file mode 100644 index 00000000000..7dcf4edf3a6 --- /dev/null +++ b/tensorflow/lite/micro/compression/decode_insert.py @@ -0,0 +1,345 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# 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. +"""DECODE operator insertion into TFLite model graphs. + +This module inserts DECODE operators into a compressed model. DECODE operators +transform encoded tensors (with their paired ancillary data tensors) into +tensors ready for use by downstream operators. + +The DECODE operator is registered as a custom operator named "TFLM_DECODE". +Each DECODE output requires two inputs: the encoded tensor and the ancillary +data tensor (containing the DCM header and decode-type-specific data). +""" + +import warnings +from collections import defaultdict +from dataclasses import dataclass + +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + +# Custom operator name for DECODE +DECODE_CUSTOM_OP_NAME = "TFLM_DECODE" + + +@dataclass +class _CompressedTensorInfo: + """Information about a compressed tensor for DECODE insertion.""" + subgraph_idx: int + tensor_idx: int + tensor: model_editor.Tensor + encoded_data: bytes + ancillary_data: bytes + consumers: list[model_editor.Operator] + is_output: bool + + +def _create_ancillary_tensor( + ancillary_data: bytes, + original_tensor: model_editor.Tensor, +) -> model_editor.Tensor: + """Create an ancillary data tensor for a compressed tensor. + + Args: + ancillary_data: The complete ancillary data (DCM + type-specific data). + original_tensor: The original tensor being decoded, for naming. + + Returns: + A new Tensor containing the ancillary data. + """ + name = None + if original_tensor.name: + name = f"{original_tensor.name}_ancillary" + + return model_editor.Tensor( + shape=(len(ancillary_data), ), + dtype=tflite.TensorType.UINT8, + data=ancillary_data, + name=name, + ) + + +def _create_output_tensor( + original_tensor: model_editor.Tensor, ) -> model_editor.Tensor: + """Create the output tensor for a DECODE operator. + + The output tensor is a copy of the original tensor, differing only in + name and in having no data: the DECODE operator produces the values + at runtime. + + Args: + original_tensor: The original tensor being decoded. + + Returns: + A new Tensor for the DECODE output. + """ + name = None + if original_tensor.name: + name = f"{original_tensor.name}_decoded" + + tensor = original_tensor.copy(name=name) + tensor.buffer = None + return tensor + + +def _rewire_consumers( + consumers: list[model_editor.Operator], + old_tensor: model_editor.Tensor, + new_tensor: model_editor.Tensor, +) -> None: + """Replace old_tensor with new_tensor in all consumer inputs.""" + for consumer in consumers: + consumer.inputs = [ + new_tensor if t is old_tensor else t for t in consumer.inputs + ] + + +def _rewrite_encoded_tensor( + tensor: model_editor.Tensor, + encoded_data: bytes, +) -> None: + """Rewrite a compressed tensor to hold encoded data. + + The original tensor contained uncompressed values with quantization. After + compression, it holds packed indices (or other encoded form) as raw bytes. + The tensor receives a fresh Buffer, leaving the original buffer and any + tensors aliasing it untouched; identical encodings converge again in the + final deduplication pass. + + Args: + tensor: The tensor to rewrite. + encoded_data: The compressed/encoded data bytes. + """ + tensor.shape = (len(encoded_data), ) + tensor.dtype = tflite.TensorType.UINT8 + tensor.quantization = None + tensor.buffer = model_editor.Buffer(data=encoded_data) + + +def _drop_partially_covered_buffers( + model: model_editor.Model, + compression_results: dict[tuple[int, int], compressor.CompressionResult], +) -> dict[tuple[int, int], compressor.CompressionResult]: + """Drop compressed tensors whose buffer an uncompressed tensor shares. + + The uncompressed tensor keeps the original data in the model, so + compressing any alias of its buffer adds encoded data, ancillary + data, and DECODE latency without reducing model size. Warn and + return the results without the dropped entries. + + Args: + model: The model the results apply to. + compression_results: Map from (subgraph_idx, tensor_idx) to + CompressionResult. + + Returns: + compression_results, minus entries for partially covered buffers. + """ + coordinates = { + id(model.subgraphs[s].tensors[t]): (s, t) + for (s, t) in compression_results + } + by_buffer: dict[int, list[model_editor.Tensor]] = defaultdict(list) + for tensor in model_editor.iter_tensors(model): + if tensor.buffer is not None: + by_buffer[id(tensor.buffer)].append(tensor) + + results = dict(compression_results) + for aliases in by_buffer.values(): + covered = [t for t in aliases if id(t) in coordinates] + if covered and len(covered) < len(aliases): + uncovered = [t for t in aliases if id(t) not in coordinates] + warnings.warn( + f"Not compressing tensor(s) " + f"{[t.name for t in covered]}: sharing a buffer with " + f"uncompressed tensor(s) {[t.name for t in uncovered]}, whose " + "data stays in the model, so compression cannot reduce model " + "size.", + stacklevel=3) + for tensor in covered: + del results[coordinates[id(tensor)]] + return results + + +def insert_decode_operators( + model: model_editor.Model, + compression_results: dict[tuple[int, int], compressor.CompressionResult], +) -> None: + """Insert DECODE operators for all compressed tensors. + + This function modifies the model in-place, inserting a DECODE operator + before any operator that uses a compressed tensor as input, and appending + a DECODE for compressed tensors listed as subgraph outputs. + + A separate DECODE is inserted before each consumer, and one DECODE + decodes all the compressed tensors its consumer reads. DECODE outputs + are tensors with a lifetime limited to the very next operator in the + subgraph, so sharing one DECODE among multiple consumers would + violate the lifetime rule. The DECODE operator trades increased + latency for decreased memory usage. + + For each consumer of compressed tensors: + 1. Create an ancillary data tensor (DCM + type-specific data) for each + compressed tensor the consumer reads + 2. Create an output tensor as a copy of each original tensor + 3. Insert one DECODE operator immediately before the consumer + 4. Rewire the consumer to use the DECODE outputs + + A subgraph's output list is treated as one more consumer, one which + reads its tensors only after the last operator runs: a calling + operator (IF, WHILE) copies subgraph outputs when the subgraph + returns. Compressed tensors in the output list are therefore decoded + by a single DECODE appended after the last operator, and their output + list entries are rewired to the decoded values. + + Distinct tensors can share one buffer, in the same or different + subgraphs, where the converter deduplicated identical constants. + Compressed tensors sharing a buffer with uncompressed tensors are + skipped with a warning: the uncompressed data must stay in the model, + so compressing an alias cannot reduce model size. Otherwise each + rewritten tensor and ancillary tensor receives its own buffer, and a + final deduplication pass merges byte-identical buffers and prunes + unreferenced ones, preserving the converter's sharing wherever + compression results allow and dissolving it where they diverge. + + Args: + model: The model to modify in-place. + compression_results: Map from (subgraph_idx, tensor_idx) to the + CompressionResult containing ancillary_data. + """ + compression_results = _drop_partially_covered_buffers( + model, compression_results) + + # Group compressed tensors by subgraph + by_subgraph: dict[int, list[_CompressedTensorInfo]] = defaultdict(list) + + for (sg_idx, tensor_idx), result in compression_results.items(): + subgraph = model.subgraphs[sg_idx] + tensor = subgraph.tensors[tensor_idx] + consumers = subgraph.consumers_of(tensor) + is_output = tensor in subgraph.outputs + + if not consumers and not is_output: + warnings.warn( + f"Compressed tensor {tensor.name!r} (subgraph {sg_idx}, " + f"tensor {tensor_idx}) has no consumers and is not a subgraph " + "output. No DECODE operator will be inserted.", + stacklevel=2) + continue + + info = _CompressedTensorInfo( + subgraph_idx=sg_idx, + tensor_idx=tensor_idx, + tensor=tensor, + encoded_data=result.encoded_data, + ancillary_data=result.ancillary_data, + consumers=consumers, + is_output=is_output, + ) + by_subgraph[sg_idx].append(info) + + # Process each subgraph + for sg_idx, tensor_infos in by_subgraph.items(): + subgraph = model.subgraphs[sg_idx] + + # Cache ancillary tensors by content to avoid duplicates within + # this subgraph. Each DECODE needs its own output tensor, but + # DECODEs whose ancillary data coincides can read one tensor. + ancillary_cache: dict[bytes, model_editor.Tensor] = {} + + # Track tensors to rewrite after all output tensors are created, since + # _create_output_tensor reads the original tensor's shape/dtype/quantization. + tensors_to_rewrite: dict[model_editor.Tensor, bytes] = {} + + def ancillary_for(info: _CompressedTensorInfo) -> model_editor.Tensor: + """Reuse or create the ancillary tensor for info's ancillary data.""" + ancillary = ancillary_cache.get(info.ancillary_data) + if ancillary is None: + ancillary = _create_ancillary_tensor(info.ancillary_data, info.tensor) + subgraph.tensors.append(ancillary) + ancillary_cache[info.ancillary_data] = ancillary + return ancillary + + def build_decode( + infos: list[_CompressedTensorInfo] + ) -> tuple[model_editor.Operator, list[model_editor.Tensor]]: + """Build one DECODE operator decoding all of infos' tensors. + + Returns the operator and its decoded output tensors, parallel to + infos. + """ + inputs = [] + outputs = [] + for info in infos: + ancillary_tensor = ancillary_for(info) + tensors_to_rewrite[info.tensor] = info.encoded_data + decoded = _create_output_tensor(info.tensor) + subgraph.tensors.append(decoded) + inputs.extend([info.tensor, ancillary_tensor]) + outputs.append(decoded) + op = model_editor.Operator( + opcode=tflite.BuiltinOperator.CUSTOM, + custom_code=DECODE_CUSTOM_OP_NAME, + inputs=inputs, + outputs=outputs, + ) + return op, outputs + + # Positions of the original operators, computed once so the sort and + # insertions below avoid a linear scan per lookup. + op_position = {op: i for i, op in enumerate(subgraph.operators)} + + # Group compressed tensors by consumer, then handle consumers in + # reverse position order so insertions don't invalidate positions: + # each insertion falls after every consumer still to be processed, + # leaving the recorded positions valid. + by_consumer: dict[model_editor.Operator, list[_CompressedTensorInfo]] = {} + for info in tensor_infos: + for consumer in info.consumers: + by_consumer.setdefault(consumer, []).append(info) + + for consumer in sorted(by_consumer, + key=lambda op: op_position[op], + reverse=True): + infos = by_consumer[consumer] + decode_op, decoded_tensors = build_decode(infos) + + # Insert DECODE immediately before this consumer + subgraph.operators.insert(op_position[consumer], decode_op) + + # Rewire only this consumer to use the decoded outputs + for info, decoded in zip(infos, decoded_tensors): + _rewire_consumers([consumer], info.tensor, decoded) + + # Decode compressed tensors read from the subgraph's output list, all + # with one DECODE appended after the last operator (see docstring). + output_infos = [info for info in tensor_infos if info.is_output] + if output_infos: + decode_op, decoded_tensors = build_decode(output_infos) + subgraph.operators.append(decode_op) + for info, decoded in zip(output_infos, decoded_tensors): + subgraph.outputs = [ + decoded if t is info.tensor else t for t in subgraph.outputs + ] + + # Rewrite encoded tensors after all output tensors are created + for tensor, encoded_data in tensors_to_rewrite.items(): + _rewrite_encoded_tensor(tensor, encoded_data) + + # Every rewrite and ancillary tensor made a fresh buffer; converge + # byte-identical ones and drop those left unreferenced, preserving + # the sharing the converter created wherever results allow. + model_editor.dedupe_buffers(model) + model_editor.prune_buffers(model) diff --git a/tensorflow/lite/micro/compression/decode_insert_runtime_test.py b/tensorflow/lite/micro/compression/decode_insert_runtime_test.py new file mode 100644 index 00000000000..0d3bafa9e62 --- /dev/null +++ b/tensorflow/lite/micro/compression/decode_insert_runtime_test.py @@ -0,0 +1,272 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# 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. +"""Runtime tests for DECODE outputs crossing subgraph boundaries. + +These tests build multi-subgraph models containing WHILE, compress selected +constant tensors with the LUT compressor, insert DECODE operators with +decode_insert, and run the result on the TFLM interpreter. They exercise the +two directions in which a DECODE output can cross a subgraph boundary: + +1. Out of a callee subgraph: a compressed constant listed in a callee + subgraph's output list, decoded by an appended DECODE, and copied to the + calling operator's outputs by the runtime. + +2. Into a callee subgraph: a compressed constant consumed by a WHILE + operator, whose decoded value the runtime copies into the cond and body + subgraph inputs. WHILE reads its inputs again after invoking the cond + subgraph, so a DECODE inside cond that shares alternate decompression + memory with the caller's DECODE can overwrite the value between reads. + +Each case runs both with the arena memory planner and with alternate +decompression memory, since the two allocate DECODE outputs differently: +the planner gives each output a distinct, lifetime-managed buffer, while +alternate memory restarts at the same base address for every DECODE. +""" + +import numpy as np +import unittest + +from tflite_micro.python.tflite_micro import runtime +from tflite_micro.tensorflow.lite.micro.compression import decode_insert +from tflite_micro.tensorflow.lite.micro.compression import lut +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.micro.compression import spec +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + +_SHAPE = (4, ) +_ARENA_SIZE = 65536 +_ALT_MEMORY_SIZE = 1024 + + +def _while_operator(cond_subgraph_idx, body_subgraph_idx, inputs, outputs): + """Create a WHILE operator with its subgraph indices. + + model_editor has no public API for builtin options, so set them on the + backing OperatorT directly. + """ + op = model_editor.Operator( + opcode=tflite.BuiltinOperator.WHILE, + inputs=inputs, + outputs=outputs, + ) + options = tflite.WhileOptionsT() + options.condSubgraphIndex = cond_subgraph_idx + options.bodySubgraphIndex = body_subgraph_idx + op._fb.builtinOptionsType = tflite.BuiltinOptions.WhileOptions + op._fb.builtinOptions = options + return op + + +def _float_tensor(name, data=None): + return model_editor.Tensor( + shape=_SHAPE, + dtype=tflite.TensorType.FLOAT32, + data=data, + name=name, + ) + + +def _cond_subgraph(threshold_values): + """Build a cond subgraph computing LESS(input, threshold_constant).""" + c_in = _float_tensor("cond_in") + threshold = _float_tensor("threshold", + np.array(threshold_values, dtype=np.float32)) + cond_out = model_editor.Tensor( + shape=_SHAPE, + dtype=tflite.TensorType.BOOL, + name="cond_out", + ) + return model_editor.Subgraph( + tensors=[threshold], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.LESS, + inputs=[c_in, threshold], + outputs=[cond_out], + ) + ], + inputs=[c_in], + outputs=[cond_out], + ) + + +def _build_body_output_model(): + """Model whose WHILE body subgraph outputs a constant. + + Subgraph 0 feeds input x into WHILE. The cond subgraph tests x < 5. The + body subgraph has no operators; its sole output is the constant K = + [7, 8, 7, 8]. With x = 0: cond is true, the body replaces x with K, cond + is then false (7 < 5), and the model outputs K. + + The tensor at coordinates (2, 0) is K, the body output constant. The + tensor at (1, 0) is the cond threshold. + """ + x0 = _float_tensor("x0") + y0 = _float_tensor("y0") + sg0 = model_editor.Subgraph( + operators=[_while_operator(1, 2, [x0], [y0])], + inputs=[x0], + outputs=[y0], + ) + + sg1 = _cond_subgraph([5.0, 6.0, 5.0, 6.0]) + + b_in = _float_tensor("body_in") + k = _float_tensor("k", np.array([7.0, 8.0, 7.0, 8.0], dtype=np.float32)) + sg2 = model_editor.Subgraph( + tensors=[k], + operators=[], + inputs=[b_in], + outputs=[k], + ) + + model = model_editor.Model(subgraphs=[sg0, sg1, sg2]) + model._fb.version = 3 + return model + + +def _build_while_input_model(): + """Model whose WHILE operator consumes a constant as its input. + + Subgraph 0 feeds the constant INIT = [10, 20, 10, 20] into WHILE as the + initial loop value. The cond subgraph tests x < 3, false from the start, + so the loop body (ADD 1) never runs and the model outputs INIT unchanged. + + The tensor at coordinates (0, 0) is INIT, the WHILE input constant. The + tensor at (1, 0) is the cond threshold. + """ + init = _float_tensor("init", + np.array([10.0, 20.0, 10.0, 20.0], dtype=np.float32)) + y0 = _float_tensor("y0") + sg0 = model_editor.Subgraph( + tensors=[init], + operators=[_while_operator(1, 2, [init], [y0])], + inputs=[], + outputs=[y0], + ) + + sg1 = _cond_subgraph([3.0, 4.0, 3.0, 4.0]) + + b_in = _float_tensor("body_in") + one = _float_tensor("one", np.array([1.0, 1.0, 1.0, 1.0], dtype=np.float32)) + b_out = _float_tensor("body_out") + sg2 = model_editor.Subgraph( + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.ADD, + inputs=[b_in, one], + outputs=[b_out], + ) + ], + inputs=[b_in], + outputs=[b_out], + ) + + model = model_editor.Model(subgraphs=[sg0, sg1, sg2]) + model._fb.version = 3 + return model + + +def _compress_and_insert(model, coordinates): + """LUT-compress the tensors at (subgraph, tensor) coordinates and insert + DECODE operators for them.""" + compressor_plugin = lut.LutCompressor() + method = spec.LookUpTableCompression(index_bitwidth=1) + results = {} + for sg_idx, tensor_idx in coordinates: + tensor = model.subgraphs[sg_idx].tensors[tensor_idx] + results[(sg_idx, tensor_idx)] = compressor_plugin.compress(tensor, method) + decode_insert.insert_decode_operators(model, results) + + +def _run(model, x0=None, alt_memory_size=0): + """Build the model and run one inference on the TFLM interpreter.""" + flatbuffer = bytes(model.build()) + interpreter = runtime.Interpreter.from_bytes( + flatbuffer, + custom_op_registerers=[], + arena_size=_ARENA_SIZE, + alt_decompression_memory_size=alt_memory_size, + ) + if x0 is not None: + interpreter.set_input(x0, 0) + interpreter.invoke() + return interpreter.get_output(0) + + +class TestDecodeOutOfSubgraph(unittest.TestCase): + """A DECODE output listed as a callee subgraph output. + + The calling operator copies callee outputs to its own outputs immediately + after the callee returns, before any other subgraph (and thus any other + DECODE) runs, so the decoded values must survive in both memory modes. + """ + + X0 = np.zeros(_SHAPE, dtype=np.float32) + EXPECTED = np.array([7.0, 8.0, 7.0, 8.0], dtype=np.float32) + + def test_arena(self): + model = _build_body_output_model() + _compress_and_insert(model, [(2, 0)]) + output = _run(model, x0=self.X0) + np.testing.assert_array_equal(output, self.EXPECTED) + + def test_alt_memory(self): + model = _build_body_output_model() + _compress_and_insert(model, [(2, 0)]) + output = _run(model, x0=self.X0, alt_memory_size=_ALT_MEMORY_SIZE) + np.testing.assert_array_equal(output, self.EXPECTED) + + def test_alt_memory_with_decode_in_cond(self): + # The cond threshold is also compressed, so a DECODE inside cond + # overwrites the shared alternate memory after the body's decoded + # output was copied out. The copy must have preserved the values. + model = _build_body_output_model() + _compress_and_insert(model, [(2, 0), (1, 0)]) + output = _run(model, x0=self.X0, alt_memory_size=_ALT_MEMORY_SIZE) + np.testing.assert_array_equal(output, self.EXPECTED) + + +class TestDecodeIntoSubgraph(unittest.TestCase): + """A DECODE output consumed as a WHILE operator input. + + WHILE reads its inputs once to seed the cond subgraph, then again after + cond returns, to seed the body subgraph and initialize its outputs. A + DECODE inside cond that shares alternate decompression memory with the + DECODE feeding WHILE overwrites the value between those reads. + """ + + EXPECTED = np.array([10.0, 20.0, 10.0, 20.0], dtype=np.float32) + + def test_arena(self): + model = _build_while_input_model() + _compress_and_insert(model, [(0, 0), (1, 0)]) + output = _run(model) + np.testing.assert_array_equal(output, self.EXPECTED) + + def test_alt_memory(self): + model = _build_while_input_model() + _compress_and_insert(model, [(0, 0)]) + output = _run(model, alt_memory_size=_ALT_MEMORY_SIZE) + np.testing.assert_array_equal(output, self.EXPECTED) + + def test_alt_memory_with_decode_in_cond(self): + model = _build_while_input_model() + _compress_and_insert(model, [(0, 0), (1, 0)]) + output = _run(model, alt_memory_size=_ALT_MEMORY_SIZE) + np.testing.assert_array_equal(output, self.EXPECTED) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorflow/lite/micro/compression/decode_insert_test.py b/tensorflow/lite/micro/compression/decode_insert_test.py new file mode 100644 index 00000000000..fe4067be3f8 --- /dev/null +++ b/tensorflow/lite/micro/compression/decode_insert_test.py @@ -0,0 +1,818 @@ +# Copyright 2026 The TensorFlow Authors. All Rights Reserved. +# +# 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. +"""Unit tests for DECODE operator insertion. + +The models these tests build are structural fixtures for exercising +insertion, not fully valid, runnable models. Likewise the compression +payloads are dummies: sized and structured plausibly, but not data +that decodes to the models' tensor contents. Insertion treats both as +opaque structure, so the tests do not depend on their validity. +""" + +import warnings + +import numpy as np +import unittest + +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import decode +from tflite_micro.tensorflow.lite.micro.compression import decode_insert +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + + +def _build_simple_fc_model(): + """Build a simple model with one FC operator and compressible weights.""" + # yapf: disable + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.array([[1, 2, 1, 2], + [3, 4, 3, 4], + [1, 2, 1, 2], + [3, 4, 3, 4]], dtype=np.int8), + name="weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + # yapf: enable + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, weights], + outputs=[output_t], + ) + ], + inputs=[input_t], + outputs=[output_t], + ) + ]) + return model + + +def _build_shared_weights_model(): + """Build model where one tensor is used by multiple operators.""" + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="shared_weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input1", + ) + input2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input2", + ) + output1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output1", + ) + output2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output2", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input1, weights], + outputs=[output1], + ), + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input2, weights], + outputs=[output2], + ), + ], + inputs=[input1, input2], + outputs=[output1, output2], + ) + ]) + return model + + +def _build_output_constant_model(): + """Build a model where a compressed constant is a subgraph output.""" + table = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="table", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[table], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, weights], + outputs=[output_t], + ) + ], + inputs=[input_t], + outputs=[output_t, table], + ) + ]) + return model + + +def _build_shared_buffer_model(subgraph_count): + """Build subgraphs whose weights tensors all share one Buffer. + + Models the TfLite converter's deduplication of identical constants: + distinct tensors, in the same or different subgraphs, backed by a + single Buffer. + """ + shared = model_editor.Buffer(data=np.ones((4, 4), dtype=np.int8).tobytes()) + subgraphs = [] + for i in range(subgraph_count): + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + buffer=shared, + name=f"weights{i}", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name=f"input{i}", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name=f"output{i}", + ) + subgraphs.append( + model_editor.Subgraph( + tensors=[weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, weights], + outputs=[output_t], + ) + ], + inputs=[input_t], + outputs=[output_t], + )) + return model_editor.Model(subgraphs=subgraphs) + + +def _make_dummy_compression_result( + element_count=16, + bitwidth=1, + value_table=b'\x01', +) -> compressor.CompressionResult: + """Create a CompressionResult with plausible dummy payloads. + + The ancillary data carries a LUT DCM describing the given index + bitwidth and value table, and the encoded data is sized for + element_count indices at that bitwidth. Tests pass values consistent + with the tensor they compress, but the payloads remain dummies (all + indices zero) that do not decode to the tensor's contents. + """ + dcm = decode.DecodeCommonMetadata( + decode_type=decode.DecodeType.LUT, + # lut_version, index bitwidth, value table stride in elements + user_data=bytes([1, bitwidth, len(value_table)]) + b'\x00' * 9, + ) + encoded_data = bytes((element_count * bitwidth + 7) // 8) + return compressor.CompressionResult( + encoded_data=encoded_data, + ancillary_data=dcm.to_bytes() + value_table, + ) + + +class TestDecodeInsertion(unittest.TestCase): + """Tests for insert_decode_operators function.""" + + def test_insert_single_decode_operator(self): + """DECODE operator inserted before FC that uses compressed weights.""" + model = _build_simple_fc_model() + + # Create compression result + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + # Insert DECODE operators + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # Should have 2 operators: DECODE then FC + self.assertEqual(len(sg.operators), 2) + self.assertEqual(sg.operators[0].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[0].custom_code, + decode_insert.DECODE_CUSTOM_OP_NAME) + self.assertEqual(sg.operators[1].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + + def test_decode_inputs_structure(self): + """DECODE operator has correct inputs: encoded tensor + ancillary.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + + # DECODE has 2 inputs + self.assertEqual(len(decode_op.inputs), 2) + # First input is the encoded tensor (original weights) + self.assertIs(decode_op.inputs[0], weights_tensor) + # Second input is ancillary tensor + self.assertEqual(decode_op.inputs[1].dtype, tflite.TensorType.UINT8) + + def test_decode_output_structure(self): + """DECODE operator output is a data-less copy of the original.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + # Snapshot what the output must look like before insertion rewrites + # the original into the encoded tensor: a copy of the original, + # renamed, with no data. + expected = weights_tensor.copy(name="weights_decoded") + expected.buffer = None + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + decode_insert.insert_decode_operators(model, compression_results) + + output = model.subgraphs[0].operators[0].outputs[0] + + # Output has no data; DECODE produces the values at runtime + self.assertIsNone(output.buffer) + self.assertIsNone(output.array) + + # Output matches the expected copy in every field, present or + # future; the new name and cleared buffer are part of the + # expectation, not exclusions + self.assertTrue(output.equal(expected)) + + def test_consumer_rewired_to_decode_output(self): + """FC operator input rewired to use DECODE output.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + fc_op = model.subgraphs[0].operators[1] + + # FC's second input (weights) should now be DECODE's output + self.assertIs(fc_op.inputs[1], decode_op.outputs[0]) + # Original weights tensor should NOT be in FC inputs + self.assertNotIn(weights_tensor, fc_op.inputs) + + def test_shared_tensor_decode_per_consumer(self): + """Tensor used by multiple ops gets separate DECODE for each consumer.""" + model = _build_shared_weights_model() + + compression_results = {(0, 0): _make_dummy_compression_result()} + + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # Should have 4 operators: 2 DECODEs + 2 FCs (DECODE before each FC) + self.assertEqual(len(sg.operators), 4) + self.assertEqual(sg.operators[0].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[1].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + self.assertEqual(sg.operators[2].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[3].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + + decode_op1 = sg.operators[0] + fc_op1 = sg.operators[1] + decode_op2 = sg.operators[2] + fc_op2 = sg.operators[3] + + # Each FC should use its own DECODE's output + self.assertIs(fc_op1.inputs[1], decode_op1.outputs[0]) + self.assertIs(fc_op2.inputs[1], decode_op2.outputs[0]) + # The two DECODEs should have different outputs + self.assertIsNot(decode_op1.outputs[0], decode_op2.outputs[0]) + # The two DECODEs should share the same encoded tensor + self.assertIs(decode_op1.inputs[0], decode_op2.inputs[0]) + # The two DECODEs should share the same ancillary tensor + self.assertIs(decode_op1.inputs[1], decode_op2.inputs[1]) + + def test_ancillary_buffer_shared_across_subgraphs(self): + """Tensors sharing a Buffer get ancillary tensors sharing a Buffer.""" + model = _build_shared_buffer_model(subgraph_count=2) + + result = _make_dummy_compression_result() + compression_results = {(0, 0): result, (1, 0): result} + + decode_insert.insert_decode_operators(model, compression_results) + + decode0 = model.subgraphs[0].operators[0] + decode1 = model.subgraphs[1].operators[0] + encoded0, ancillary0 = decode0.inputs[0], decode0.inputs[1] + encoded1, ancillary1 = decode1.inputs[0], decode1.inputs[1] + + # Tensors are per-subgraph; buffers are shared across subgraphs + self.assertIsNot(ancillary0, ancillary1) + self.assertIs(ancillary0.buffer, ancillary1.buffer) + + # The encoded tensors keep sharing their original Buffer + self.assertIsNot(encoded0, encoded1) + self.assertIs(encoded0.buffer, encoded1.buffer) + + def test_ancillary_tensor_shared_within_subgraph(self): + """Aliased tensors in one subgraph share one ancillary tensor.""" + model = _build_shared_buffer_model(subgraph_count=1) + sg = model.subgraphs[0] + + # Add a second tensor aliasing the weights buffer, with its own + # consumer; copy() shares the buffer, as converter dedup does + weights = sg.tensor_by_name("weights0") + alias = weights.copy(name="alias") + sg.tensors.append(alias) + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input_alias", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output_alias", + ) + sg.operators.append( + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, alias], + outputs=[output_t], + )) + + result = _make_dummy_compression_result() + compression_results = {(0, 0): result, (0, 1): result} + + decode_insert.insert_decode_operators(model, compression_results) + + # One DECODE before each consumer, sharing one ancillary tensor + decodes = [ + op for op in sg.operators + if op.custom_code == decode_insert.DECODE_CUSTOM_OP_NAME + ] + self.assertEqual(len(decodes), 2) + self.assertIs(decodes[0].inputs[1], decodes[1].inputs[1]) + + def test_insertion_on_read_model(self): + """Insertion works on a model read from a flatbuffer. + + Unlike the from-scratch fixtures elsewhere in this file, a read + model carries a populated buffer list, so this exercises the + pruning of the buffer orphaned when compression rewrites the + weights tensor. + """ + scratch = _build_simple_fc_model() + model = model_editor.read(bytes(scratch.build())) + weights_bytes = model.subgraphs[0].tensor_by_name("weights").buffer.data + + result = _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + decode_insert.insert_decode_operators(model, {(0, 0): result}) + + final = model_editor.read(bytes(model.build())) + sg = final.subgraphs[0] + decode_op = sg.operators[0] + self.assertEqual(decode_op.custom_code, + decode_insert.DECODE_CUSTOM_OP_NAME) + self.assertEqual(decode_op.inputs[0].buffer.data, result.encoded_data) + self.assertEqual(decode_op.inputs[1].buffer.data, result.ancillary_data) + + # The original weights buffer, orphaned by the rewrite, is pruned: + # only the conventional empty buffer, the encoded data, and the + # ancillary data remain + self.assertEqual(len(final.buffers), 3) + for buffer in final.buffers: + self.assertNotEqual(buffer.data, weights_bytes) + + def test_divergent_aliases_dissolve_sharing(self): + """Aliases whose compression results differ get separate buffers.""" + model = _build_shared_buffer_model(subgraph_count=2) + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=1, value_table=b'\x01'), + (1, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04'), + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode0 = model.subgraphs[0].operators[0] + decode1 = model.subgraphs[1].operators[0] + encoded0, ancillary0 = decode0.inputs[0], decode0.inputs[1] + encoded1, ancillary1 = decode1.inputs[0], decode1.inputs[1] + + # Each alias carries its own results; nothing is shared + self.assertIsNot(encoded0.buffer, encoded1.buffer) + self.assertIsNot(ancillary0.buffer, ancillary1.buffer) + self.assertEqual(encoded0.buffer.data, + compression_results[(0, 0)].encoded_data) + self.assertEqual(encoded1.buffer.data, + compression_results[(1, 0)].encoded_data) + self.assertEqual(ancillary0.buffer.data, + compression_results[(0, 0)].ancillary_data) + self.assertEqual(ancillary1.buffer.data, + compression_results[(1, 0)].ancillary_data) + + def test_partially_covered_buffer_not_compressed(self): + """A tensor sharing its buffer with an uncompressed tensor is + skipped.""" + model = _build_shared_buffer_model(subgraph_count=2) + weights0 = model.subgraphs[0].tensor_by_name("weights0") + original_buffer = weights0.buffer + original_data = original_buffer.data + + # Compress only one of the two tensors sharing the buffer + compression_results = {(0, 0): _make_dummy_compression_result()} + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + decode_insert.insert_decode_operators(model, compression_results) + + self.assertEqual(len(caught), 1) + self.assertIn("weights0", str(caught[0].message)) + self.assertIn("weights1", str(caught[0].message)) + + # No DECODE inserted anywhere; the tensor is untouched + for subgraph in model.subgraphs: + self.assertEqual(len(subgraph.operators), 1) + self.assertIs(weights0.buffer, original_buffer) + self.assertEqual(weights0.buffer.data, original_data) + self.assertEqual(weights0.dtype, tflite.TensorType.INT8) + + def test_output_constant_gets_decode(self): + """Compressed tensor in subgraph outputs gets DECODE appended last.""" + model = _build_output_constant_model() + sg = model.subgraphs[0] + table = sg.tensor_by_name("table") + original_first_output = sg.outputs[0] + + compression_results = {(0, 0): _make_dummy_compression_result()} + + decode_insert.insert_decode_operators(model, compression_results) + + # DECODE appended after the FC operator + self.assertEqual(len(sg.operators), 2) + decode_op = sg.operators[1] + self.assertEqual(decode_op.opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(decode_op.custom_code, + decode_insert.DECODE_CUSTOM_OP_NAME) + self.assertIs(decode_op.inputs[0], table) + + # Output list entry rewired to the decoded tensor; other entry untouched + self.assertIs(sg.outputs[0], original_first_output) + self.assertIs(sg.outputs[1], decode_op.outputs[0]) + + # Original tensor rewritten to encoded data + self.assertEqual(table.dtype, tflite.TensorType.UINT8) + + def test_consumed_and_output_tensor(self): + """Tensor both consumed and a subgraph output gets both DECODEs.""" + model = _build_simple_fc_model() + sg = model.subgraphs[0] + weights = sg.tensor_by_name("weights") + fc_output = sg.operators[0].outputs[0] + sg.outputs = [fc_output, weights] + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + decode_insert.insert_decode_operators(model, compression_results) + + # Consumer DECODE before FC, output DECODE appended last + self.assertEqual(len(sg.operators), 3) + consumer_decode = sg.operators[0] + fc_op = sg.operators[1] + output_decode = sg.operators[2] + self.assertEqual(consumer_decode.opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(fc_op.opcode, tflite.BuiltinOperator.FULLY_CONNECTED) + self.assertEqual(output_decode.opcode, tflite.BuiltinOperator.CUSTOM) + + # Each reader gets its own decoded tensor + self.assertIs(fc_op.inputs[1], consumer_decode.outputs[0]) + self.assertIs(sg.outputs[1], output_decode.outputs[0]) + self.assertIsNot(consumer_decode.outputs[0], output_decode.outputs[0]) + + # Both DECODEs share the encoded tensor + self.assertIs(consumer_decode.inputs[0], output_decode.inputs[0]) + + # Both DECODEs share the ancillary tensor + self.assertIs(consumer_decode.inputs[1], output_decode.inputs[1]) + + def test_multiple_input_tensors_share_one_decode(self): + """All compressed tensors of one consumer are decoded by one DECODE.""" + weights1 = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="weights1", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + weights2 = model_editor.Tensor( + shape=(2, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((2, 4), dtype=np.int8), + name="weights2", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + output_t = model_editor.Tensor( + shape=(6, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights1, weights2], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.CONCATENATION, + inputs=[weights1, weights2], + outputs=[output_t], + ) + ], + outputs=[output_t], + ) + ]) + sg = model.subgraphs[0] + + compression_results = { + (0, 0): _make_dummy_compression_result(element_count=16), + (0, 1): _make_dummy_compression_result(element_count=8), + } + + decode_insert.insert_decode_operators(model, compression_results) + + # One DECODE with two encoded/ancillary pairs, inserted before the + # CONCATENATION + self.assertEqual(len(sg.operators), 2) + decode_op = sg.operators[0] + concat_op = sg.operators[1] + self.assertEqual(decode_op.opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(concat_op.opcode, tflite.BuiltinOperator.CONCATENATION) + self.assertEqual(len(decode_op.inputs), 4) + self.assertEqual(len(decode_op.outputs), 2) + self.assertIs(decode_op.inputs[0], weights1) + self.assertIs(decode_op.inputs[2], weights2) + + # The consumer reads each decoded tensor from its original position + self.assertIs(concat_op.inputs[0], decode_op.outputs[0]) + self.assertIs(concat_op.inputs[1], decode_op.outputs[1]) + + def test_multiple_output_tensors_share_one_decode(self): + """All compressed subgraph outputs are decoded by a single DECODE.""" + table1 = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="table1", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + table2 = model_editor.Tensor( + shape=(2, 2), + dtype=tflite.TensorType.INT8, + data=np.ones((2, 2), dtype=np.int8), + name="table2", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[table1, table2], + operators=[], + outputs=[table1, output_t, table2], + ) + ]) + sg = model.subgraphs[0] + + compression_results = { + (0, 0): _make_dummy_compression_result(element_count=16), + (0, 1): _make_dummy_compression_result(element_count=4), + } + + decode_insert.insert_decode_operators(model, compression_results) + + # One DECODE with two encoded/ancillary pairs and two outputs + self.assertEqual(len(sg.operators), 1) + decode_op = sg.operators[0] + self.assertEqual(len(decode_op.inputs), 4) + self.assertEqual(len(decode_op.outputs), 2) + self.assertIs(decode_op.inputs[0], table1) + self.assertIs(decode_op.inputs[2], table2) + + # Output list rewired in place; uncompressed entry untouched + self.assertIs(sg.outputs[0], decode_op.outputs[0]) + self.assertIs(sg.outputs[1], output_t) + self.assertIs(sg.outputs[2], decode_op.outputs[1]) + + def test_ancillary_tensor_contains_dcm(self): + """Ancillary tensor data contains valid DCM header.""" + model = _build_simple_fc_model() + + result = _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + compression_results = {(0, 0): result} + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + ancillary_tensor = decode_op.inputs[1] + + # Ancillary tensor data should match what we provided + self.assertEqual(bytes(ancillary_tensor.array), result.ancillary_data) + + # Verify DCM header + dcm_bytes = bytes(ancillary_tensor.array[:16]) + self.assertEqual(dcm_bytes[0], decode.DecodeType.LUT) + self.assertEqual(dcm_bytes[1], 1) # DCM version + + def test_no_consumers_no_decode(self): + """Tensor with no consumers gets no DECODE operator and emits warning.""" + # Create model where compressed tensor is not used as input + unused_tensor = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="unused", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + other_weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="other_weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[unused_tensor, other_weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, other_weights], + outputs=[output_t], + ) + ], + ) + ]) + + # Compress the unused tensor + compression_results = {(0, 0): _make_dummy_compression_result()} + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + decode_insert.insert_decode_operators(model, compression_results) + + # Should emit a warning about no consumers + self.assertEqual(len(w), 1) + self.assertIn("no consumers", str(w[0].message)) + self.assertIn("unused", str(w[0].message)) + + # Should still have just 1 operator (no DECODE inserted) + self.assertEqual(len(model.subgraphs[0].operators), 1) + + def test_tensor_naming(self): + """Output and ancillary tensors get appropriate names.""" + model = _build_simple_fc_model() + + compression_results = { + (0, 0): + _make_dummy_compression_result(bitwidth=2, + value_table=b'\x01\x02\x03\x04') + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + ancillary = decode_op.inputs[1] + output = decode_op.outputs[0] + + self.assertEqual(ancillary.name, "weights_ancillary") + self.assertEqual(output.name, "weights_decoded") + + def test_encoded_tensor_rewritten(self): + """Compressed tensor is rewritten with encoded data, UINT8 type, no quant.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + encoded_data = b'\xAB\xCD\xEF' + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=encoded_data, + ancillary_data=_make_dummy_compression_result().ancillary_data, + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + # Original tensor should be rewritten + self.assertEqual(weights_tensor.shape, (len(encoded_data), )) + self.assertEqual(weights_tensor.dtype, tflite.TensorType.UINT8) + self.assertIsNone(weights_tensor.quantization) + self.assertEqual(weights_tensor.buffer.data, encoded_data) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorflow/lite/micro/compression/model_editor.py b/tensorflow/lite/micro/compression/model_editor.py index c04b51f1855..2f140a9fc13 100644 --- a/tensorflow/lite/micro/compression/model_editor.py +++ b/tensorflow/lite/micro/compression/model_editor.py @@ -16,6 +16,7 @@ Provides a clean API for creating, reading, and modifying TFLite models. """ +import copy from dataclasses import dataclass, field from typing import Optional, Union, List import numpy as np @@ -88,6 +89,25 @@ def to_tflite(self) -> tflite.QuantizationParametersT: return q +def _fields_equal(a, b) -> bool: + """Compare two values recursively, including flatbuffer objects. + + Sequences compare elementwise. Objects with fields, such as the + schema's flatbuffer classes, compare field by field, so fields + unknown to this module still participate. + """ + if isinstance(a, (list, tuple, np.ndarray)) and isinstance( + b, (list, tuple, np.ndarray)): + return np.array_equal(a, b) + if hasattr(a, '__dict__') and hasattr(b, '__dict__'): + if type(a) is not type(b): + return False + keys = vars(a).keys() | vars(b).keys() + return all( + _fields_equal(getattr(a, k, None), getattr(b, k, None)) for k in keys) + return a == b + + class Tensor: """Tensor specification wrapping a TensorT flatbuffer object. @@ -215,6 +235,62 @@ def array(self, value: np.ndarray): else: self.buffer.data = buf_data + def copy(self, name: Optional[str] = None) -> 'Tensor': + """Return a copy of this tensor, sharing this tensor's buffer. + + The copy duplicates every field, including those of the backing + TensorT, except that it references the same Buffer object as the + original and has no index until added to a subgraph. Assign None + to the copy's buffer for a tensor with no data. + + Args: + name: Optional name for the copy. If None, the copy keeps this + tensor's name. + + Returns: + A new Tensor duplicating this one. + """ + # Seed the memo so the deepcopy preserves buffer identity: sharing + # is by Buffer object, and a duplicate would compile to a duplicate + # buffer table entry. + memo = {id(self.buffer): self.buffer} + duplicate = copy.deepcopy(self, memo) + duplicate._index = None + if name is not None: + duplicate.name = name + return duplicate + + def equal(self, other: 'Tensor') -> bool: + """Return True if this tensor equals other, field by field. + + Compare the backing TensorTs recursively, so fields this module + does not manage still participate, plus quantization and buffer. + Buffers compare by identity, mirroring how the model expresses + sharing: tensors referencing distinct but byte-identical Buffers + compile to distinct buffer table entries and are not equal. The + tensors' positions in any subgraph do not participate. + + Args: + other: The tensor to compare against. + + Returns: + True if the tensors are equal. + """ + if self.buffer is not other.buffer: + return False + if self.quantization != other.quantization: + return False + # Exclude the TensorT fields mirrored by the wrapper attributes + # compared above: buffer, an index assigned at build time, and + # quantization, which build() syncs from the wrapper attribute. + # Comparing the raw fields too would misreport tensors of mixed + # provenance, read versus constructed. + excluded = ('buffer', 'quantization') + keys = vars(self._fb).keys() | vars(other._fb).keys() + return all( + _fields_equal(getattr(self._fb, k, None), getattr(other._fb, k, None)) + for k in keys if k not in excluded) + @property def index(self) -> Optional[int]: """Tensor index in the subgraph's tensor list. @@ -448,6 +524,17 @@ def tensor_by_name(self, name: str) -> Tensor: return t raise KeyError(f"No tensor named {name!r}") + def consumers_of(self, tensor: Tensor) -> List[Operator]: + """Find the operators in this subgraph that read a tensor. + + Args: + tensor: The tensor whose consumers to find. + + Returns: + The operators, in subgraph order, with tensor among their inputs. + """ + return [op for op in self.operators if tensor in op.inputs] + @property def index(self) -> Optional[int]: """Subgraph index in the model's subgraph list. @@ -521,6 +608,79 @@ def build(self) -> bytearray: return compiler.compile() +def iter_tensors(model: Model): + """Yield every tensor in the model exactly once. + + Walk the same sources the compiler collects from: each subgraph's + tensor list, inputs, outputs, and the tensors inline on operators. + + Args: + model: The model whose tensors to yield. + + Yields: + Each distinct Tensor in the model. + """ + seen = set() + for sg in model.subgraphs: + sources = [sg.tensors, sg.inputs, sg.outputs] + sources.extend(op.inputs for op in sg.operators) + sources.extend(op.outputs for op in sg.operators) + for source in sources: + for tensor in source: + if id(tensor) not in seen: + seen.add(id(tensor)) + yield tensor + + +def dedupe_buffers(model: Model) -> None: + """Merge byte-identical buffers into one shared Buffer. + + Repoint tensors whose buffers hold equal contents at a single + canonical Buffer object, the first encountered, mirroring the TfLite + converter's deduplication of identical constants. Leave tensors + marked is_variable alone: mutable data must not alias. Merged-away + buffers linger in model.buffers until pruned. + + Args: + model: The model to modify in place. + """ + canonical: dict[bytes, Buffer] = {} + for tensor in iter_tensors(model): + if tensor.buffer is None or tensor._fb.isVariable: + continue + existing = canonical.get(tensor.buffer.data) + if existing is None: + canonical[tensor.buffer.data] = tensor.buffer + else: + tensor.buffer = existing + + +def prune_buffers(model: Model) -> None: + """Drop buffers that no tensor references from model.buffers. + + Rebuild the buffer list with only the conventional empty buffer 0 + and the buffers some tensor references, renumbering indices. A model + built from scratch keeps an empty buffer list and the compiler emits + only referenced buffers, so pruning matters for models from read(), + whose buffer list the compiler preserves wholesale. + + Args: + model: The model to modify in place. + """ + if not model.buffers: + return + referenced = { + id(tensor.buffer) + for tensor in iter_tensors(model) if tensor.buffer is not None + } + survivors = _BufferList() + survivors.append(model.buffers[0]) + for buffer in model.buffers[1:]: + if id(buffer) in referenced: + survivors.append(buffer) + model.buffers = survivors + + def read(buffer: bytes) -> Model: """Read a TFLite flatbuffer and return a Model object.""" fb_model = tflite.ModelT.InitFromPackedBuf(buffer, 0) diff --git a/tensorflow/lite/micro/compression/model_editor_test.py b/tensorflow/lite/micro/compression/model_editor_test.py index 1f036b1a1e6..ff7c1da76e2 100644 --- a/tensorflow/lite/micro/compression/model_editor_test.py +++ b/tensorflow/lite/micro/compression/model_editor_test.py @@ -827,6 +827,284 @@ def test_tensor_by_name_not_found_raises(self): with self.assertRaises(KeyError): model.subgraphs[0].tensor_by_name("nonexistent") + def test_consumers_of(self): + """consumers_of finds the operators reading a tensor, in order.""" + shared = Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="shared", + ) + input1 = Tensor(shape=(1, 4), dtype=tflite.TensorType.INT8, name="input1") + input2 = Tensor(shape=(1, 4), dtype=tflite.TensorType.INT8, name="input2") + output1 = Tensor(shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output1") + output2 = Tensor(shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output2") + + fc1 = Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input1, shared], + outputs=[output1], + ) + only_produces = Operator( + opcode=tflite.BuiltinOperator.RESHAPE, + inputs=[output1], + outputs=[output2], + ) + fc2 = Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[output2, shared], + outputs=[Tensor(shape=(1, 4), dtype=tflite.TensorType.INT8)], + ) + sg = Subgraph(tensors=[shared], + operators=[fc1, only_produces, fc2], + inputs=[input1, input2]) + + self.assertEqual(sg.consumers_of(shared), [fc1, fc2]) + self.assertEqual(sg.consumers_of(input1), [fc1]) + self.assertEqual(sg.consumers_of(input2), []) + + +class TestTensorCopy(unittest.TestCase): + """Tests for Tensor.copy().""" + + def _original(self): + return Tensor( + shape=(2, 2), + dtype=tflite.TensorType.INT8, + data=np.array([[1, 2], [3, 4]], dtype=np.int8), + name="original", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + + def test_copy_shares_buffer(self): + """The copy references the same Buffer object as the original.""" + original = self._original() + duplicate = original.copy() + + self.assertIsNot(duplicate, original) + self.assertIs(duplicate.buffer, original.buffer) + + def test_copy_is_independent(self): + """Mutating the copy leaves the original untouched.""" + original = self._original() + duplicate = original.copy(name="duplicate") + self.assertEqual(duplicate.quantization, original.quantization) + + duplicate.shape = (4, ) + duplicate.dtype = tflite.TensorType.UINT8 + duplicate.quantization.scales = 2.0 + + self.assertEqual(duplicate.name, "duplicate") + self.assertEqual(original.name, "original") + self.assertEqual(original.shape, (2, 2)) + self.assertEqual(original.dtype, tflite.TensorType.INT8) + self.assertEqual(original.quantization.scales, 0.5) + + +class TestTensorEqual(unittest.TestCase): + """Tests for Tensor.equal().""" + + def _original(self): + return Tensor( + shape=(2, 2), + dtype=tflite.TensorType.INT8, + data=np.array([[1, 2], [3, 4]], dtype=np.int8), + name="original", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + + def test_copy_equals_original(self): + """A copy sharing the original's buffer compares equal.""" + original = self._original() + self.assertTrue(original.copy().equal(original)) + + def test_any_field_differing_is_unequal(self): + """A difference in any field, wrapper-managed or not, is unequal.""" + original = self._original() + + renamed = original.copy(name="renamed") + self.assertFalse(renamed.equal(original)) + + requantized = original.copy() + requantized.quantization.scales = 2.0 + self.assertFalse(requantized.equal(original)) + + dataless = original.copy() + dataless.buffer = None + self.assertFalse(dataless.equal(original)) + + # Buffers compare by identity: equal bytes in a distinct Buffer + # still make a structurally different model + rebuffered = original.copy() + rebuffered.buffer = model_editor.Buffer(data=original.buffer.data) + self.assertFalse(rebuffered.equal(original)) + + variable = original.copy() + variable._fb.isVariable = True + self.assertFalse(variable.equal(original)) + + def test_backing_quantization_details_excluded(self): + """Quantization held only by the backing TensorT does not + participate. + + equal() compares quantization at the wrapper level, where scale, + zero point, and axis live; min, max, and details carried only by + the backing TensorT are excluded by design. + """ + scratch = Model(subgraphs=[Subgraph(tensors=[self._original()])]) + fb = bytes(scratch.build()) + t1 = model_editor.read(fb).subgraphs[0].tensor_by_name("original") + + t2 = t1.copy() + t2._fb.quantization.min = [-1.0] + + self.assertTrue(t1.equal(t2)) + + +class TestDedupeBuffers(unittest.TestCase): + """Tests for dedupe_buffers().""" + + @staticmethod + def _constant(name, values): + return Tensor( + shape=(4, ), + dtype=tflite.TensorType.INT8, + data=np.array(values, dtype=np.int8), + name=name, + ) + + def test_merges_identical_buffers_across_subgraphs(self): + """Byte-identical buffers converge on one canonical Buffer.""" + c1 = self._constant("c1", [1, 2, 3, 4]) + c2 = self._constant("c2", [1, 2, 3, 4]) + model = Model(subgraphs=[ + Subgraph(tensors=[c1]), + Subgraph(tensors=[c2]), + ]) + self.assertIsNot(c1.buffer, c2.buffer) + + model_editor.dedupe_buffers(model) + + self.assertIs(c2.buffer, c1.buffer) + + # The compiled model carries one buffer for both tensors, plus the + # conventional empty buffer 0 + fb_model = tflite.ModelT.InitFromPackedBuf(bytes(model.build()), 0) + self.assertEqual(len(fb_model.buffers), 2) + + def test_leaves_distinct_and_variable_buffers(self): + """Differing contents never merge; variable tensors never alias.""" + c1 = self._constant("c1", [1, 2, 3, 4]) + c2 = self._constant("c2", [5, 6, 7, 8]) + variable = self._constant("v", [1, 2, 3, 4]) + variable._fb.isVariable = True + model = Model(subgraphs=[Subgraph(tensors=[c1, c2, variable])]) + + model_editor.dedupe_buffers(model) + + self.assertIsNot(c2.buffer, c1.buffer) + self.assertIsNot(variable.buffer, c1.buffer) + + def test_reaches_inline_tensors(self): + """Tensors on operators, absent from the tensor list, participate.""" + listed = self._constant("listed", [1, 2, 3, 4]) + inline = self._constant("inline", [1, 2, 3, 4]) + output_t = Tensor(shape=(4, ), dtype=tflite.TensorType.INT8, name="out") + model = Model(subgraphs=[ + Subgraph( + tensors=[listed], + operators=[ + Operator( + opcode=tflite.BuiltinOperator.ADD, + inputs=[listed, inline], + outputs=[output_t], + ) + ], + outputs=[output_t], + ) + ]) + + model_editor.dedupe_buffers(model) + + self.assertIs(inline.buffer, listed.buffer) + + +class TestPruneBuffers(unittest.TestCase): + """Tests for prune_buffers().""" + + @staticmethod + def _read_model_with_two_constants(metadata=None): + c1 = Tensor( + shape=(4, ), + dtype=tflite.TensorType.INT8, + data=np.array([1, 2, 3, 4], dtype=np.int8), + name="c1", + ) + c2 = Tensor( + shape=(4, ), + dtype=tflite.TensorType.INT8, + data=np.array([5, 6, 7, 8], dtype=np.int8), + name="c2", + ) + scratch = Model(subgraphs=[Subgraph(tensors=[c1, c2])], metadata=metadata) + return model_editor.read(bytes(scratch.build())) + + def test_drops_unreferenced_buffers(self): + """Orphaned buffers vanish; surviving data and indices stay right.""" + model = self._read_model_with_two_constants() + self.assertEqual(len(model.buffers), 3) + sg = model.subgraphs[0] + + # Orphan c1's buffer by repointing the tensor at c2's. The + # surviving buffer then changes position, exercising renumbering. + sg.tensor_by_name("c1").buffer = sg.tensor_by_name("c2").buffer + + model_editor.prune_buffers(model) + + self.assertEqual(len(model.buffers), 2) + roundtrip = model_editor.read(bytes(model.build())) + rt_sg = roundtrip.subgraphs[0] + np.testing.assert_array_equal( + rt_sg.tensor_by_name("c1").array, [5, 6, 7, 8]) + np.testing.assert_array_equal( + rt_sg.tensor_by_name("c2").array, [5, 6, 7, 8]) + + def test_keeps_referenced_buffers(self): + """With every buffer referenced, pruning removes nothing.""" + model = self._read_model_with_two_constants() + count_before = len(model.buffers) + + model_editor.prune_buffers(model) + + self.assertEqual(len(model.buffers), count_before) + + def test_metadata_survives_pruning(self): + """Metadata, carried outside tensor buffers, is unaffected.""" + model = self._read_model_with_two_constants(metadata={"m": b"payload"}) + + model_editor.prune_buffers(model) + + roundtrip = model_editor.read(bytes(model.build())) + self.assertEqual(roundtrip.metadata["m"], b"payload") + + def test_noop_on_scratch_model(self): + """A from-scratch model has no buffer list to prune.""" + c1 = Tensor(shape=(4, ), + dtype=tflite.TensorType.INT8, + data=np.array([1, 2, 3, 4], dtype=np.int8), + name="c1") + model = Model(subgraphs=[Subgraph(tensors=[c1])]) + + model_editor.prune_buffers(model) + + roundtrip = model_editor.read(bytes(model.build())) + np.testing.assert_array_equal( + roundtrip.subgraphs[0].tensor_by_name("c1").array, [1, 2, 3, 4]) + class TestReadEdgeCases(unittest.TestCase): """Test model_editor.read() with edge cases from real-world models. @@ -1111,6 +1389,25 @@ def test_tensor_is_variable_preserved(self): model_t2 = tflite.ModelT.InitFromPackedBuf(fb2, 0) self.assertTrue(model_t2.subgraphs[0].tensors[0].isVariable) + def test_tensor_copy_preserves_fields(self): + """Verify Tensor.copy() preserves fields model_editor doesn't manage.""" + model_t = self._create_base_model() + model_t.subgraphs[0].tensors[0].isVariable = True + + fb = self._build_model_with_schema(model_t) + + model = model_editor.read(fb) + sg = model.subgraphs[0] + duplicate = sg.tensors[0].copy(name="copy") + sg.tensors.append(duplicate) + fb2 = model.build() + + model_t2 = tflite.ModelT.InitFromPackedBuf(fb2, 0) + tensors = model_t2.subgraphs[0].tensors + self.assertEqual(tensors[-1].name, b"copy") + self.assertTrue(tensors[-1].isVariable) + self.assertEqual(tensors[-1].buffer, tensors[0].buffer) + def test_tensor_shape_signature_preserved(self): """Verify Tensor.shapeSignature is preserved through read-modify-write.""" model_t = self._create_base_model()