Skip to content

WHILE operator input/output copy fix#3633

Merged
veblush merged 1 commit into
tensorflow:mainfrom
ddavis-2015:while-fix
Jul 20, 2026
Merged

WHILE operator input/output copy fix#3633
veblush merged 1 commit into
tensorflow:mainfrom
ddavis-2015:while-fix

Conversation

@ddavis-2015

Copy link
Copy Markdown
Member

@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

@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
ddavis-2015 requested a review from a team as a code owner July 17, 2026 07:19
@ddavis-2015 ddavis-2015 added type:bug Something isn't working as intended. ci:full Triggers the comprehensive cross-platform test suite. labels Jul 17, 2026
@ddavis-2015 ddavis-2015 self-assigned this Jul 17, 2026
@ddavis-2015
ddavis-2015 deployed to integration-test July 17, 2026 07:19 — with GitHub Actions Active
@ddavis-2015
ddavis-2015 requested a review from veblush July 17, 2026 08:26
@rkuester

Copy link
Copy Markdown
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
veblush added this pull request to the merge queue Jul 20, 2026
Merged via the queue into tensorflow:main with commit 8f1f3b2 Jul 20, 2026
52 of 55 checks passed
@ddavis-2015
ddavis-2015 deleted the while-fix branch July 20, 2026 23:56
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ci:full Triggers the comprehensive cross-platform test suite. type:bug Something isn't working as intended.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

WHILE operator input/output copy fix

3 participants