NXP backend: Resolve limitations of uncertain tensor formats #13942
Conversation
meta-cla
bot
added
the
CLA Signed
|
@pytorchbot label "module: nxp" "release notes: nxp" |
pytorch-bot
bot
added
module: nxp
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/13942
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Cancelled Job, 3 Unrelated FailuresAs of commit 8255593 with merge base 03333c5 ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
mergennachin
requested review from
robert-kalmar,
jirioc,
roman-janik-nxp and
digantdesai
… node format. The pass `RemoveGetItemPass` replaces a `max_pool2d_with_indices` node with a `max_pool2d` node, that doesn't require a GetItem afterward. The new operator must, however, preserve the original node format. Therefore, a copy of the pass was created in `backends/nxp/_passes`, where it was modified. The new directory was created, because the pass doesn't follow the `NeutronEdgePass` interface.
Before, the format inference was done during conversion to NeutronIR (after partitioning), so the partitioner didn't yet know the formats. Now, the partitioner has the format data, which can be used to accurately select nodes for delegation.
robert-kalmar
force-pushed
the
upstream/main-nxp/EIEX-392-resolve-limitations-of-uncertain-tensor-formats
branch
from
1d726a4 to
8255593
Compare
| node.meta = {} | ||
| if NXP_NODE_FORMAT not in node.meta: | ||
| logging.warning(f"Node `{node}` does not have inferred format.") | ||
| node.meta[NXP_NODE_FORMAT] = NodeFormat.NONE |
There was a problem hiding this comment.
As we now perform the node format inference during partition (that is on the whole Edge Program), it is likely that some nodes wont have the format determined, as the NodeFormat inference algorithm does not know them. Right?
We should make sure we stop the channel_first tag propagation on unknown operator, as we cannot determine if it propagates the channel_first or stops it. As example, the Reshape stops the propagation of channel first tag. But if we would not know the Reshape, op we were incorrectly propagate the channel_first tag behind it in the compute path.
So we must defensively stop the propagation at every unknown node. Is my thought process correct?
There was a problem hiding this comment.
Yes, you are correct, with a few caveats.
The fact that we currently propagate the format through unknown nodes, should never cause crashes, as our format handling system is quite robust. It can, however, result in unnecessary transpositions. As the "unknown" operators will inevitably not be delegated, they will split the graph, resulting in multiple delegated partitions. It is possible (and likely) that one of these partitions requires NHWC, which is propagated to the second partition (via the "unknown" node), but the second partition doesn't require NHWC. If we keep our format inference as is, unnecessary transpositions would have to be done at the inputs and outputs of the second partition.
I will update the code to not propagate the format through "unknown" operators.
| if node.op == "call_function": | ||
| if ( | ||
| node.target.__name__ == "aten.max_pool2d_with_indices.default" | ||
| or node.target.__name__ == "aten.max.dim" |
There was a problem hiding this comment.
Why do we handle the aten.max.dim too? Is it a loftover from original pass?
There was a problem hiding this comment.
Yes, it was in the original file. I wanted to make as few changes as possible, as it is not the main focus of this PR.
Summary
This PR resolves format related issues by inferring the format (NCHW/NHWC) for all nodes before partitioning. These formats are then used by the NeutronPartitioner to accurately determine which nodes are supported on Neutron.
Test plan
Unit tests provided, and correct function is tested by nearly every test in the nxp backend.
cc @robert-kalmar @roman-janik-nxp @StrycekSimon @jirioc