feat(compression): add DECODE operator insertion#3624
Conversation
| if not consumers: | ||
| # Check if tensor is a subgraph output | ||
| is_output = tensor in subgraph.outputs | ||
| if is_output: |
There was a problem hiding this comment.
I guess if it_output should be outside of if not consumers? It might be consumed and an oiutput of subgraph?
There was a problem hiding this comment.
You're right: the check missed tensors that are both consumed by an operator and listed as a subgraph output, so a subgraph would have output the encoded bytes instead of the decoded values. Addressed in 3999889 by treating the output list as one more consumer, fed by a DECODE appended after the last operator.
f3a8e56 pulls forward a fix from later in the series: separate DECODEs per tensor don't work with alternate decompression memory when the decoded values are needed at the same time, and with output handling added, it's even more worth getting this right initially than rewriting it later. It's a separate commit for easy review, relying on the squash when the PR merges.
There was a problem hiding this comment.
The design doc update of 2026-07-13 settles this. DECODE may feed IF and WHILE, with the invoking operator's inputs as the consumers. A DECODE output may be a subgraph output, with the DECODE as the subgraph's final operator and the invoking operator's outputs as the consumers. The insertion code conforms to both rules. cbf22b5 adds runtime tests that invoke models with DECODE across subgraph boundaries.
There was a problem hiding this comment.
Looks really good, almost there. What is here will create a valid and usable model, but not exactly as it should be.
Maximal reduction of the model is not yet achieved due to possible duplication of the ancillary tensor Buffer. Here is the scenario: Two const-tensors, either in the same or in different subgraphs can contain the exact same data. The TfLite Converter can (and does) create two tensors in the model, but using a single shared buffer. For DECODE insertion, the existing code works just fine for encoded tensors, as they are not given a new buffer, just new buffer content. But this means that the related ancillary tensors, can also share a single buffer. However, the current insertion code always creates a new Buffer object when creating ancillary tensors.
The output tensors of DECODE must be exact copies of the original unencoded tensors, meaning when the output tensors are created they must use the original tensor's TensorT object. Only the name of the output tensors can be modified (as is currently done by the insertion code), and the Buffer object set to None (buffer in the TensorT object is zero).
For the unit tests, maybe a comment could be added stating the models are just for unit test purposes and do not represent fully valid models. Same for the encoded_data and ancillary_data of the CompressionResult objects. Something saying these are just dummy payloads for unit testing and do not correspond to the model's tensor data. Someone looking at these for the first time might mistake these for valid examples and be confused.
| sg = model.subgraphs[0] | ||
| weights = sg.tensor_by_name("shared_weights") | ||
|
|
||
| consumers = decode_insert._find_tensor_consumers(sg, weights) |
There was a problem hiding this comment.
Use of a private method (_find_tensor_consumers). Maybe have a set of public methods just for use by unit tests?
There was a problem hiding this comment.
Addressed in 16fedda: consumer lookup moved to model_editor as a public API, Subgraph.consumers_of, with its own tests, and the insertion tests no longer reach into private helpers.
| 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], | ||
| ) | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Doesn't the Subgraph inputs/outputs need to be specified? Otherwise the subgraph compiler will just assign empty vectors to the tflite.SubGraphT.
There was a problem hiding this comment.
Addressed in c2b32ce: the test models now declare their subgraph inputs and outputs.
| 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], | ||
| ) | ||
| ], | ||
| outputs=[output_t, table], | ||
| ) | ||
| ]) |
There was a problem hiding this comment.
Doesn't the Subgraph inputs need to be specified? Otherwise the subgraph compiler will just assign an empty vector in the tflite.SubGraphT.
There was a problem hiding this comment.
Addressed in c2b32ce: the test models now declare their subgraph inputs and outputs.
| encoded_data=b'\x00\x00', | ||
| ancillary_data=_make_dummy_ancillary_data(), |
There was a problem hiding this comment.
mismatch between encoded_data and the ancillary_data. Perhaps _make_dummy_ancillary_data should have arguments so that it can generate a DCM that matches the model data? And encoded_data appears to be too short in length.
There was a problem hiding this comment.
Addressed in af0fbc8: a helper now builds the whole CompressionResult, parameterized by element count, index bitwidth, and value table, so each test passes payloads consistent with the tensor it compresses and the encoded data is sized to match. The module docstring states these are structural fixtures, not decodable data, per your general comment.
| def test_insert_single_decode_operator(self): | ||
| """DECODE operator inserted before FC that uses compressed weights.""" | ||
| model = _build_simple_fc_model() | ||
| weights_tensor = model.subgraphs[0].tensor_by_name("weights") |
There was a problem hiding this comment.
unused variable weights_tensor
| 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], | ||
| ), | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Doesn't the Subgraph inputs/outputs need to be specified? Otherwise the subgraph compiler will just assign empty vectors to the tflite.SubGraphT.
There was a problem hiding this comment.
Addressed in c2b32ce: the test models now declare their subgraph inputs and outputs.
| # The two DECODEs should share the same ancillary tensor | ||
| self.assertIs(decode_op1.inputs[1], decode_op2.inputs[1]) |
There was a problem hiding this comment.
also check the two DECODES share the same encoded tensor
There was a problem hiding this comment.
Addressed in 88cdede: both tests now assert the DECODEs read the same encoded tensor object.
| self.assertIs(consumer_decode.inputs[1], output_decode.inputs[1]) | ||
|
|
There was a problem hiding this comment.
also check both DECODEs share the encoded tensor
There was a problem hiding this comment.
Addressed in 88cdede: both tests now assert the DECODEs read the same encoded tensor object.
| opcode=tflite.BuiltinOperator.FULLY_CONNECTED, | ||
| inputs=[input_t, weights1, weights2], |
There was a problem hiding this comment.
CONCATENATION would be a better example here (with output shape [6,4])
There was a problem hiding this comment.
Addressed in d633f59: the test now uses CONCATENATION with output shape (6, 4).
|
|
||
| # Verify DCM header | ||
| dcm_bytes = ancillary_tensor.array[:16] | ||
| self.assertEqual(dcm_bytes[0], 0) # decode_type = LUT |
Insert DECODE operators before consumers of compressed tensors. Each consumer gets its own DECODE operator to support alternate decompression memory, which resets allocations between DECODE invocations. After insertion, compressed tensors are rewritten to hold encoded data as UINT8 with shape matching byte count. BUG=part of tensorflow#3256
A compressed tensor can be listed in a subgraph's output list, where it is read not by an operator, but by the operator calling the subgraph (IF, WHILE), which copies subgraph outputs when the subgraph returns, or by the client, which reads model outputs after invocation. DECODE insertion previously checked the output list only for tensors with no consumers, and refused those as unsupported. A tensor both consumed and listed as an output slipped through. The pass rewired its consumers to decoded values, then rewrote the tensor to hold encoded bytes, which the output list delivered as if decoded. Treat the output list as one more consumer, one which reads its tensors only after the last operator runs. Append a DECODE after the last operator for each compressed tensor in the output list, and rewire the list entry to the decoded value. BUG=part of tensorflow#3256
A consumer reading several compressed tensors needs all their decoded values at once, but under alternate decompression memory, values produced by different DECODE operators cannot coexist. Each DECODE resets the allocation offset during Prepare, placing every DECODE's outputs at the same address, so each DECODE overwrites the outputs of the one before it. Decoding a consumer's tensors with separate DECODE operators corrupts all but the last value. Decode all compressed tensors read by one consumer with a single DECODE operator carrying one encoded/ancillary input pair and one output per tensor. Outputs of a single DECODE coexist, since the allocation reset happens between operators, not between the outputs of one. The subgraph output list, treated as one more consumer, gets the same treatment. One DECODE, appended after the last operator, decodes every compressed tensor in the list. BUG=part of tensorflow#3256
DECODE insertion sorted consumers and located insertion points with list.index, a linear scan of the operator list per lookup. Build a map of operator positions once per subgraph and consult it instead. The positions recorded before any insertion remain correct throughout, because consumers are handled in reverse position order, so each insertion falls after every consumer still to be processed. BUG=part of tensorflow#3256
Add a test suite that exercises DECODE outputs crossing subgraph boundaries on the TFLM interpreter, rather than only checking the rewritten flatbuffer structure. The tests build multi-subgraph WHILE models with model_editor, compress constants with the LUT compressor, insert DECODE operators with decode_insert, and verify inference results, in both arena and alternate decompression memory modes. The case of a DECODE output feeding a WHILE input, with a second DECODE in the cond subgraph and alternate decompression memory in use, requires the WhileEval fix from tensorflow#3633 (issue tensorflow#3632). WHILE formerly re-read its inputs after invoking the cond subgraph, picking up the value the cond subgraph's DECODE wrote over shared alternate memory.
Add Tensor.copy(), which duplicates a tensor's backing TensorT and shares the original's Buffer object. Duplicating the TensorT preserves fields model_editor does not otherwise manage, such as is_variable and shape_signature. An optional name argument gives the copy its own name. Clients that need a data-less copy, such as tooling that creates stand-in tensors, can assign None to the copy's buffer. Add Tensor.equal(), a field-wise equality over the backing TensorT, quantization, and buffer. Fields unknown to model_editor participate via recursive comparison, so clients can compare tensors without enumerating fields. Buffers compare by identity, mirroring how the model expresses buffer sharing.
Add dedupe_buffers(), which repoints tensors whose buffers hold byte-identical contents at one canonical Buffer object, mirroring the TfLite converter's deduplication of identical constants. Tensors marked is_variable are left alone, since mutable data must not alias. The walk covers every tensor the compiler collects, including tensors inline on operators that never appear in a subgraph's tensor list. Merged-away buffers linger in model.buffers until pruned.
Add prune_buffers(), which rebuilds model.buffers with only the conventional empty buffer 0 and the buffers some tensor references, renumbering indices in the process. Models built from scratch keep an empty buffer list and compile only referenced buffers, so pruning matters for models from read(), whose buffer list the compiler preserves wholesale, including entries orphaned by editing.
Create the output tensor of a DECODE operator by copying the original tensor, clearing its data, and renaming it, rather than by building a new tensor from the original's shape, dtype, and quantization. Copying preserves TensorT fields the insertion code does not otherwise handle, such as is_variable and shape_signature, so the decoded stand-in is indistinguishable from the original tensor it replaces. Verify the output against a copy of the original snapshotted before insertion rewrites it, compared with field-wise tensor equality so every field participates without the test enumerating them.
Distinct tensors can share one buffer, in the same or different subgraphs, where the converter deduplicated identical constants. Give each rewritten encoded tensor and each ancillary tensor a fresh buffer, then merge byte-identical buffers and prune unreferenced ones after insertion. Sharing survives compression wherever aliases compress to identical results, extends to any ancillary data that coincides, and dissolves where results diverge (possible for tensors sharing bytes but quantized with different structures), rather than one alias corrupting another through a shared buffer rewritten in place. Skip, with a warning, compressed tensors that share a buffer with an uncompressed tensor. The uncompressed data must remain in the model for the other tensors, so compressing such an alias cannot reduce model size.
Describe the placement of DECODE operators in terms of the operator's contract. Outputs have a lifetime limited to the very next operator in the subgraph, and DECODE trades increased latency for decreased memory usage. Remove the explanation of interpreter alternate-memory behavior that previously justified the per-consumer placement, along with a confusing aside about clients reading model outputs. Describe the output tensor as a copy of the original, matching the implementation.
Add Subgraph.consumers_of(), which returns the operators reading a given tensor, in subgraph order. Replace decode_insert's private helper _find_tensor_consumers with it. The helper's unit test called a private method of decode_insert; consumer lookup is now a public API, tested in model_editor_test.
Specify the subgraph inputs and outputs in the three test model builders. Two builders previously declared neither, and the third declared only its outputs, so the subgraph compiler emitted empty vectors for whatever was missing, making the models structurally unlike anything the converter produces.
Remove two weights_tensor assignments never read by their tests.
When one compressed tensor feeds multiple DECODE operators, assert that the operators read the same encoded tensor object, alongside the existing assertion that they share the ancillary tensor. Cover both situations in which one tensor feeds multiple DECODE operators, a tensor with two consumers and a tensor both consumed and listed as a subgraph output.
Compare the DCM's decode type byte against the DecodeType.LUT constant instead of a magic zero. Convert the DCM slice from a numpy array to bytes first, so that indexing yields a plain integer whose comparison defers to the constant's own equality.
Exercise the one-DECODE-per-consumer batching with a CONCATENATION of two compressed tensors instead of a FULLY_CONNECTED with an extra weights input, which is not a valid FC signature. CONCATENATION takes any number of inputs, so the model resembles something a converter could produce.
Replace the fixed dummy ancillary data helper with one that builds a whole CompressionResult, parameterized by element count, index bitwidth, and value table, so each test passes values consistent with the tensor it compresses and the encoded data is sized accordingly. State in the module docstring that the test models and payloads are structural fixtures, not valid runnable models or decodable data, so a reader does not mistake them for real examples.
Two tensors can share one buffer when the converter deduplicates identical constants. Add tests for the two situations in which insertion cannot preserve that sharing. In the first, both tensors are compressed but their compression results differ. Insertion dissolves the sharing, and the test verifies that each tensor receives its own encoded and ancillary buffers holding its own results. In the second, only one of the tensors is compressed. The uncompressed tensor keeps the original data in the model, so compressing its alias would grow the model rather than shrink it. Insertion declines to compress, and the test verifies that no DECODE operator is inserted, that the tensor is untouched, and that a warning explains why.
All other insertion tests build their models from scratch, and a from-scratch model keeps an empty buffer list, which makes buffer pruning a no-op. Add a test that round-trips a model through build() and read() before insertion, then verifies the packed result carries the DECODE operator with its encoded and ancillary data, and that the buffer orphaned when compression rewrote the weights tensor is pruned rather than left in the model.
50f2fcf to
726b26f
Compare
|
Everything is addressed in new commits, kept separate for easy review, relying on the squash when the PR merges. On the two structural points: DECODE outputs are now full copies of the original tensors, renamed and with no buffer (8330ba3), built by a new Tensor.copy in model_editor (d0fd38b). Buffer sharing is preserved by deduplication rather than a cache. Insertion gives every rewritten encoded tensor and every ancillary tensor a fresh buffer, then model_editor.dedupe_buffers merges byte-identical buffers across the whole model and prune_buffers drops unreferenced ones (363a71e, e1aa8e6, f62df14). That yields what a Buffer-keyed cache would, across subgraphs and for coinciding ancillary data, and covers a case a cache cannot: aliases whose compression results differ. Identical bytes can encode differently, because quantization structure shapes the LUT value table. There the sharing dissolves, rather than one alias corrupting another through a buffer rewritten in place. A compressed tensor sharing a buffer with an uncompressed tensor is skipped with a warning, because the uncompressed alias keeps the original data in the model, so compression cannot shrink it. 37505cd tests these cases; 726b26f exercises dedup and pruning end to end on a model read from a flatbuffer. |
veblush
left a comment
veblush
left a comment
There was a problem hiding this comment.
Thanks for the PR! I'll get this merged once David also approved it.
3 participants
Insert DECODE operators before consumers of compressed tensors. Each
consumer gets its own DECODE operator to support alternate decompression
memory, which resets allocations between DECODE invocations.
After insertion, compressed tensors are rewritten to hold encoded data
as UINT8 with shape matching byte count.
Compression tooling only; adds decode_insert.py and its test, no
runtime changes.
Part of the DECODE compression-tooling stack.
BUG=#3256