WHILE operator input/output copy fix - #3633
Merged
Merged
Conversation
@tensorflow/micro Remove extraneous tensor copy operation after first invocation of condition subgraph. Move copy of operator inputs to outputs, such that it occurs before the first invocation of the condition subgraph. This preserves the operator inputs when one or more of them is the output of DECODE, and alternate decompression memory is in use. This is because the output of DECODE is for immediate consumption by the next operator in the graph (WHILE), yet it is possible for the WHILE subgraph invocations to share memory with the original DECODE output. Update the unit test for multiple invocations of the condition and body subgraphs. When copying tensors between operator inputs/outputs and subgraph inputs/outputs, check if the source and destination tensors share memory. bug=fixes tensorflow#3632
ddavis-2015
added
type:bug
ddavis-2015
temporarily deployed
to
integration-test
July 17, 2026 07:19 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
Contributor
|
I verified this fix against a runtime test suite I've been building for DECODE insertion (Python: model_editor plus the LUT compressor, DECODE insertion, and the TFLM interpreter bindings). The suite includes a reproducer for exactly the failure this PR describes: a DECODE output feeding a WHILE input, with a second DECODE in the cond subgraph and alternate decompression memory in use. Without this fix, WHILE re-reads its inputs after invoking the cond subgraph and picks up whatever the cond subgraph's DECODE wrote over the shared alternate memory, and the test fails. With this PR cherry-picked onto my decode-insert branch, the reproducer passes, along with the rest of the suite and the full kernel test suite, so I see no regressions from the kernel_util changes. The change looks right to me. The runtime suite is part of #3624, which is awaiting this fix. |
veblush
approved these changes
Jul 20, 2026
rkuester
added a commit
to rkuester/tflite-micro
that referenced
this pull request
Jul 21, 2026
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.
pull Bot
pushed a commit
to QSLee-Net/tflite-micro
that referenced
this pull request
Jul 29, 2026
* feat(compression): add DECODE operator insertion 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 * fix(compression): decode compressed subgraph outputs 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 * fix(compression): batch multiple compressed tensors per DECODE 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 * refactor(compression): precompute operator positions 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 * test(compression): add runtime tests for DECODE across subgraphs 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. * feat(compression): add tensor copying and equality to model_editor 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. * feat(compression): add buffer deduplication to model_editor 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. * feat(compression): add buffer pruning to model_editor 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. * fix(compression): make DECODE outputs full copies of their originals 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. * fix(compression): share buffers among aliases by deduplication 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. * docs(compression): reword DECODE insertion docstring 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. * feat(compression): add tensor consumer lookup to model_editor 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. * test(compression): declare subgraph inputs and outputs in test models 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. * test(compression): remove unused variables Remove two weights_tensor assignments never read by their tests. * test(compression): assert DECODEs share the encoded tensor 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. * test(compression): compare decode type against DecodeType.LUT 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. * test(compression): use CONCATENATION in multi-input DECODE test 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. * test(compression): make dummy compression payloads self-consistent 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. * test(compression): cover buffer alias divergence and partial coverage 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. * test(compression): exercise insertion on a model read from a flatbuffer 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. * test(compression): share one flatbuffer packing helper Two test classes each defined an identical method that packs a model into a flatbuffer. Replace both with one module-level function. The next commit adds fixtures at module level, which will also reuse the function. * test(compression): assert absent inputs survive The flatbuffer schema marks an absent optional operator input with an index of -1. E.g., a fully-connected operator uses that index when it has no bias. In fact, a fused LSTM leaves most of its inputs absent. The model editor does not currently preserve an absent input. Add expected-to-fail tests to expose the bug. The next commit will fix the bug and drop the expected-failure flags. * fix(compression): preserve absent optional operator inputs The flatbuffer schema marks an absent optional operator input with an index of -1. Read that index as None among an operator's inputs, and write None back out as -1. Reject any other negative index. Drop the expected-failure flags from the tests added by the previous commit, and add two more tests covering iteration and the consumer lookup. * fix(compression): reject negative tensor indices Only an operator's inputs give a negative tensor index a meaning, an index of -1 marking an absent optional input. An operator's outputs and a subgraph's inputs and outputs give none. Reject a negative index in those three. Reading one used to substitute a tensor counted from the end of the subgraph's tensor list, and writing the model back out then recorded that tensor's real index. A malformed model became a well-formed one naming a different tensor. Resolving a subgraph's inputs and outputs no longer needs a guard against an absent or empty list, since resolving nothing yields the empty list a subgraph already starts with. * test(compression): strengthen the metadata pruning test Pruning a model's buffers must not disturb its metadata. The test for that was weak and had no buffers to remove. Fix by orphaning a buffer, so pruning removes it and shifts the buffers that remain. Assert the buffer count drops. Check that the metadata still survives a roundtrip. * style(compression): drop an unused import The model editor imports dataclasses.field and never uses it. * docs(compression): say where the appended DECODE goes The docstring said a DECODE is appended, without saying to what.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
@tensorflow/micro
Remove extraneous tensor copy operation after first invocation of condition subgraph.
Move copy of operator inputs to outputs, such that it occurs before the first invocation of the condition subgraph. This preserves the operator inputs when one or more of them is the output of DECODE, and alternate decompression memory is in use. This is because the output of DECODE is for immediate consumption by the next operator in the graph (WHILE), yet it is possible for the WHILE subgraph invocations to share memory with the original DECODE output.
Update the unit test for multiple invocations of the condition and body subgraphs.
When copying tensors between operator inputs/outputs and subgraph inputs/outputs, check if the source and destination tensors share memory.
bug=fixes #3632