WHILE operator input/output copy fix#3633
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
|
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. |
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.
@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