[JakeStevens]

Ulterior Motive

Lower the fp32 <-> fp16 and fp32 <-> bf16 casts of programs exported with skip_dim_order=True into the XNNPACK delegate, as programs exported with dim order already are.

Rationale

What: Add ToCopyConfig, which partitions aten._to_copy dtype casts under the same dtype rules as ToDimOrderCopyConfig: fp32 <-> fp16, and fp32 <-> bf16 with enable_bf16, as long as the cast keeps the dim order.

Why: to_edge only rewrites _to_copy to _to_dim_order_copy when dim order is enabled. skip_dim_order=True is the default of LLMEdgeManager (and so export_llm), examples/xnnpack/aot_compiler.py and extension/export_util, so their casts stay aten._to_copy, which no XNNPACK partitioner config matched. They ran in portable kernels between delegates, even though ToCopyVisitor already serializes a dtype-changing _to_copy as an XNNPACK convert. fp16 casts need no opt-in, so fp16 programs from these paths change as well as bf16 ones with enable_bf16.

Details

• Layout-only _to_copy nodes are not partitioned: inside the delegate, ChannelsLastTaggedReshapePass, RemoveRedundantCopyPass and ToCopyVisitor read any _to_copy that is not a cast as an NCHW <-> NHWC layout copy. Neither are casts between other dtypes, bf16 casts without enable_bf16, nor casts that also change the dim order.
• The q/dq around a cast are no longer folded into it, and RemoveRedundantCopyPass no longer removes a cast between a dq and a q. Either one dropped the requantize of dq -> cast -> q with different scales, which also affected programs exported with dim order. XNNPACKQuantizer does not annotate casts, so PT2E-quantized models do not contain this pattern.
• Programs exported with dim order and no such q/dq pattern are byte-identical with and without this change (29 fp32, fp16 and bf16 programs), since they contain no aten._to_copy before partitioning.
• Casts of lifted integer scalar constants to float, such as x.float() * 2 emits, stay in portable kernels with or without dim order, because XNNPACK does not convert integers.
• Latency on phones, after over before (below 1 is faster), median of 3 alternating rounds of 8 prefill or 40 decode runs on the Pixels (4 and 20 on the iPhone 17), first run of each launch dropped. Models are exported with export_llm, model.use_kv_cache=True, model.use_sdpa_with_kv_cache=True, export.max_seq_length=128, export.max_context_length=1024, backend.xnnpack.enabled=True and model.dtype_override set to bf16 (with backend.xnnpack.enable_bf16=True) or fp16; prefill uses model.enable_dynamic_shape=True and runs 128 tokens, decode runs 1. They run in executor_runner built with the optimized kernels and the LLM custom ops; 1 thread is pinned to the fastest core on the Pixels.

Model	dtype	Device	Casts in portable kernels	Prefill, 1 thread	Prefill, default threads	Decode, 1 thread	Decode, default threads
Qwen3-0.6B	bf16	Pixel 10	282 -> 0	0.999x	0.989x	0.996x	0.963x
Qwen3-0.6B	bf16	Pixel 6 Pro	282 -> 0	1.008x	1.005x	0.988x	0.996x
Qwen3-0.6B	fp16	Pixel 10	394 -> 0	1.005x	0.757x	0.996x	0.977x
Qwen3-0.6B	fp16	Pixel 6 Pro	394 -> 0	1.012x	0.986x	0.968x	0.929x
SmolLM2-135M	bf16	Pixel 10	542 -> 0	0.999x	1.003x	1.002x	0.997x
SmolLM2-135M	bf16	iPhone 17	542 -> 0	1.008x	1.004x	1.002x	0.998x
SmolLM2-135M	bf16	Pixel 6 Pro	542 -> 0	1.006x	1.007x	0.991x	0.973x
SmolLM2-135M	fp16	Pixel 10	662 -> 0	1.007x	1.012x	0.997x	1.069x
SmolLM2-135M	fp16	iPhone 17	662 -> 0	0.996x	1.077x	0.995x	1.067x
SmolLM2-135M	fp16	Pixel 6 Pro	662 -> 0	1.034x	1.033x	0.997x	0.955x

Large casts gain from running on the thread pool, as in Qwen3 fp16 prefill (441 ms to 334 ms). Small ones cost more in a delegate that wakes the pool than in the single-threaded portable kernel, so fp16 SmolLM2 decode at default threads is about 7% slower on the Pixel 10 and iPhone 17, though 4.5% faster on the Pixel 6 Pro, whose fp16 SmolLM2 prefill is about 3% slower; programs exported with dim order make the same trade-off today. Qwen3 does not fit the iPhone's memory limit in executor_runner, which reads the whole program into memory.

Test plan:

python -m pytest backends/xnnpack/test/ops/test_to_copy_dtype.py backends/xnnpack/test/ops/test_to_copy.py backends/xnnpack/test/passes/test_remove_redundant_copy_pass.py
python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/test_xnnpack_partitioner.py backends/xnnpack/test/serialization backends/xnnpack/test/quantizer

The cast tests in test_to_copy_dtype.py now run with and without skip_dim_order; the skip_dim_order cases fail without ToCopyConfig, and test_requantize_around_cast fails without either q/dq fix.

Authored with AI using Claude Code.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24768

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 483e267 with merge base 472953c ():

NEW FAILURE - The following job has failed:

• pull / android / run-emulator (gh)
  ##[error]Timeout waiting for emulator to boot.

This comment was automatically generated by Dr. CI and updates every 15 minutes.