[xnnpack] Fold float dtype conversions into the delegate - #21791
Open
JakeStevens wants to merge 1 commit into
Open
[xnnpack] Fold float dtype conversions into the delegate#21791JakeStevens wants to merge 1 commit into
JakeStevens wants to merge 1 commit into
Conversation
Partition and serialize fp16/bf16 <-> fp32 dtype-conversion copies into the XNNPACK delegate as xnn_define_convert, instead of leaving them as portable _to_copy ops that fragment the graph. This lets a mixed-precision model (e.g. a bf16 model whose RMSNorm computes in fp32) lower as cleanly as its fp32 counterpart: on google/gemma-3-1b-it bf16, delegate segments drop 446 -> 211 and portable dtype converts 364 -> 46, matching the fp32 structure, with identical outputs (argmax unchanged). Changes: - partition/config/generic_node_configs.py: ToDimOrderCopyConfig now allows the XNNPACK-serializable float conversions (fp32<->fp16, fp32<->bf16) for dynamic activations (constant/param converts are left for constant folding). Add ToCopyConfig (target _to_copy.default) so the same folding applies when the graph is lowered with EdgeCompileConfig(_skip_dim_order=True). - partition/config/__init__.py: register ToCopyConfig. - operators/op_to_copy.py: emit XNNConvert for a dtype-changing copy (memory format copies still emit XNNStaticTranspose). Two backend passes assumed _to_copy is always memory-format-only; partitioning dtype converts exposed both: - _passes/channels_last_tagged_reshape_pass.py: tag_node used node.kwargs['memory_format'] directly, which KeyErrors on a dtype-only copy (no memory_format kwarg). Default to contiguous (a dtype convert keeps layout). - _passes/remove_redundant_copy_pass.py: redundant-copy removal compared only memory format, so it would elide a genuine bf16<->fp32 convert as if it were a redundant layout copy, silently dropping the conversion. Skip copies that change dtype. Note: this is a graph-cleanliness/correctness change. On x86 it is latency neutral for bf16 (the bf16 cost there is intrinsic XNNPACK fp32-fallback for non-GEMM ops, not the converts); the win is a clean lowering that benefits memory-bound / ARM targets.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21791
Note: Links to docs will display an error until the docs builds have been completed. ❌ 4 New Failures, 1 Unrelated FailureAs of commit d5c2889 with merge base d92a619 ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following job 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. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
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
1 participant
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.
Partition and serialize fp16/bf16 <-> fp32 dtype-conversion copies into the XNNPACK delegate as xnn_define_convert, instead of leaving them as portable _to_copy ops that fragment the graph.