feat(compression): add DECODE operator insertion#3624
Conversation
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
| 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.
@rkuester @veblush Passing the output of DECODE to or from a subgraph simply won't work (for a number of reasons, primarily involving the WHILE operator and the tensor copying that occurs between subgraphs and the calling operator). The design document will be updated to address what to do in this situation.
| for consumer in info.consumers: | ||
| consumer_pairs.append((consumer, info)) | ||
|
|
||
| consumer_pairs.sort( |
There was a problem hiding this comment.
This could be slow because of subgraph.operators.index in the sort?
There was a problem hiding this comment.
Addressed in 50f2fcf: a position map built once per subgraph replaces the scans in the sort and insertions.
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: 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: consumers are handled in reverse position order, so each insertion falls after every consumer still to be processed. BUG=part of tensorflow#3256
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