Skip to content

Lower mean.dim beyond 4D pooling as an XNNPACK static mean - #24741

Open
JakeStevens wants to merge 1 commit into
pytorch:mainfrom
JakeStevens:xnnpack-mean-dim-static-reduce
Open

Lower mean.dim beyond 4D pooling as an XNNPACK static mean#24741
JakeStevens wants to merge 1 commit into
pytorch:mainfrom
JakeStevens:xnnpack-mean-dim-static-reduce

Conversation

@JakeStevens

@JakeStevens JakeStevens commented Nov 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Ulterior Motive

Run RMSNorm in LeRobot policies on XNNPACK (#24667), where mean.dim fails to partition with "only supports averaging 4D tensors, got tensor of rank 3".

Rationale

What: Add an XNNStaticMean node backed by xnn_define_static_reduce(xnn_reduce_mean), and lower mean.dim that is not Global Average Pooling to it: ranks 1 to 6, any dims, and either keepdim, except small static means (see Details).

Why: XNNPACK only lowered a 4D mean over dims [2, 3] with keepdim=True. RMSNorm averages the last dim of a 3-D tensor, so its mean stayed on the portable kernel and split the delegate around every norm.

Details

Review order: the serialization schema and XNNCompiler.cpp, then op_mean_dim.py and is_global_avg_pool_mean, then ChannelsLastTaggedReshapePass, then MeanDimConfig, then the tests.

  • The 4D, dims [2, 3], keepdim=True case still lowers to Global Average Pooling on NHWC, at every size. For those programs the .pte is byte-identical to main (fp32, fp16, qs8, conv-bn-hardtanh-mean, and a CNN ending in AdaptiveAvgPool2d(1)).
  • Every other mean.dim reduces over NCHW dims, so ChannelsLastTaggedReshapePass gives it an NCHW input. After a convolution that is the same NHWC to NCHW copy main made at the delegate boundary.
  • The node uses a shared _XNNStaticReduce table, so other XNNPACK reductions can reuse it and defineGenericReduceNode. It is appended to XNodeUnion, so existing node IDs do not change. Re-exporting a model with a mean that is not Global Average Pooling, such as RMSNorm, emits the new node, and that program needs an updated runtime.
  • An empty dim list averages over every dim, as in PyTorch.
  • XNNPACK splits a reduction with more than one output into threadpool tasks, so on a multi-threaded pool a small mean costs more in XNNPACK than on the portable kernel. A 1x8x64 mean over the last dim took 4.9 us on the portable kernel and 27.9 us in XNNPACK at 16 threads; at 1 thread both took about 5 us. A static mean with more than one output and fewer than 8192 input elements therefore stays on the portable kernel and produces the same program as main. Dynamic shapes and single-output means are lowered at any size, so a small mean with a dynamic shape can be slower than on main with many threads.
  • Latency, median against main on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: RMSNorm 1x128x2048 0.48x, 3-D last-dim mean 0.45x, RMSNorm then Linear 0.77x, 4 pre-norm residual MLP blocks 0.99x for decode (1 token) and 0.80x for prefill (64 tokens), conv then mean over channels 0.88x. RMSNorm on 10 static shapes from 1x1x128 to 1x32x2048, at 1, 4 and 16 threads, is 0.43x to 1.02x. Each number is a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime built from this branch.
  • A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts: at 16 threads a bare fp32 mean over the last dim of 1x128x128 is 3.26x, 1x64x512 1.90x and 1x128x512 1.02x, and larger means are faster (1x128x2048 0.39x). At 4 threads only 1x128x128 is slower (1.14x), and at 1 thread every size is faster. The same holds for any op lowered alone, such as add (2.95x at 1x128x128 and 16 threads on main), so this is the existing single-op trade-off of the XNNPACK partitioner. Inside a model the neighbouring delegated ops already pay the wake-up, which is why the RMSNorm and block numbers above do not regress.
  • The LoRA CI scripts check the start of the 8da4w-quantized Qwen3 base model's greedy output with a regex. With RMSNorm and QK-norm now in XNNPACK, its first token can tip from :\n to : (assistant: me instead of assistant:\n<think>): unquantized, main and this branch generate the same tokens with logits within 1e-4, but 8da4w re-rounds activations and the two candidates are within about 0.5 logits. The regex now accepts any text on the assistant: line and an optional <think>, and still requires Okay, so I need to calculate 15% of 80..
  • Tests cover 1-D to 5-D inputs, all dims, keepdim=False, fp16 and qs8 over the innermost and a middle dim, dynamic kept and reduced dims, RMSNorm, convolution followed by a mean that is not Global Average Pooling, the size threshold on both sides, and the layout copies the pass inserts.

Test plan:

python -m pytest backends/xnnpack/test/ops/test_mean_dim.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.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

Authored with AI using Claude Code.

@JakeStevens JakeStevens added the release notes: xnnpack Changes to the XNNPack backend delegate label Nov 6, 2026
@pytorch-bot

pytorch-bot Bot commented Nov 6, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure, 1 Pending, 3 Unrelated Failures

As of commit 99c4063 with merge base 558c071 (image):

NEW FAILURE - The following job has failed:

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.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 6, 2026
@JakeStevens
Lower mean.dim beyond 4D pooling as an XNNPACK static mean
99c4063 
## Ulterior Motive

Run RMSNorm in LeRobot policies on XNNPACK (pytorch#24667), where `mean.dim` fails to partition with "only supports averaging 4D tensors, got tensor of rank 3".

## Rationale

**What**: Add an `XNNStaticMean` node backed by `xnn_define_static_reduce(xnn_reduce_mean)`, and lower `mean.dim` that is not Global Average Pooling to it: ranks 1 to 6, any dims, and either `keepdim`, except small static means (see Details).

**Why**: XNNPACK only lowered a 4D mean over dims [2, 3] with `keepdim=True`. RMSNorm averages the last dim of a 3-D tensor, so its mean stayed on the portable kernel and split the delegate around every norm.

## Details

Review order: the serialization schema and `XNNCompiler.cpp`, then `op_mean_dim.py` and `is_global_avg_pool_mean`, then `ChannelsLastTaggedReshapePass`, then `MeanDimConfig`, then the tests.

- The 4D, dims [2, 3], `keepdim=True` case still lowers to Global Average Pooling on NHWC, at every size. For those programs the `.pte` is byte-identical to main (fp32, fp16, qs8, conv-bn-hardtanh-mean, and a CNN ending in `AdaptiveAvgPool2d(1)`).
- Every other `mean.dim` reduces over NCHW dims, so `ChannelsLastTaggedReshapePass` gives it an NCHW input. After a convolution that is the same NHWC to NCHW copy main made at the delegate boundary.
- The node uses a shared `_XNNStaticReduce` table, so other XNNPACK reductions can reuse it and `defineGenericReduceNode`. It is appended to `XNodeUnion`, so existing node IDs do not change. Re-exporting a model with a mean that is not Global Average Pooling, such as RMSNorm, emits the new node, and that program needs an updated runtime.
- An empty dim list averages over every dim, as in PyTorch.
- XNNPACK splits a reduction with more than one output into threadpool tasks, so on a multi-threaded pool a small mean costs more in XNNPACK than on the portable kernel. A 1x8x64 mean over the last dim took 4.9 us on the portable kernel and 27.9 us in XNNPACK at 16 threads; at 1 thread both took about 5 us. A static mean with more than one output and fewer than 8192 input elements therefore stays on the portable kernel and produces the same program as main. Dynamic shapes and single-output means are lowered at any size, so a small mean with a dynamic shape can be slower than on main with many threads.
- Latency, median against main on an Apple M-series CPU (12 performance + 4 efficiency cores) with the default 16 threads: RMSNorm 1x128x2048 0.48x, 3-D last-dim mean 0.45x, RMSNorm then Linear 0.77x, 4 pre-norm residual MLP blocks 0.99x for decode (1 token) and 0.80x for prefill (64 tokens), conv then mean over channels 0.88x. RMSNorm on 10 static shapes from 1x1x128 to 1x32x2048, at 1, 4 and 16 threads, is 0.43x to 1.02x. Each number is a median over at least 1000 timed runs per program, alternating main and branch programs on one runtime built from this branch.
- A mean alone in its delegate, with nothing else lowered around it, pays the threadpool wake-up on its own and can be slower than the portable kernel at high thread counts: at 16 threads a bare fp32 mean over the last dim of 1x128x128 is 3.26x, 1x64x512 1.90x and 1x128x512 1.02x, and larger means are faster (1x128x2048 0.39x). At 4 threads only 1x128x128 is slower (1.14x), and at 1 thread every size is faster. The same holds for any op lowered alone, such as `add` (2.95x at 1x128x128 and 16 threads on main), so this is the existing single-op trade-off of the XNNPACK partitioner. Inside a model the neighbouring delegated ops already pay the wake-up, which is why the RMSNorm and block numbers above do not regress.
- The LoRA CI scripts check the start of the 8da4w-quantized Qwen3 base model's greedy output with a regex. With RMSNorm and QK-norm now in XNNPACK, its first token can tip from `:\n` to `:` (`assistant: me` instead of `assistant:\n<think>`): unquantized, main and this branch generate the same tokens with logits within 1e-4, but 8da4w re-rounds activations and the two candidates are within about 0.5 logits. The regex now accepts any text on the `assistant:` line and an optional `<think>`, and still requires `Okay, so I need to calculate 15% of 80.`.
- Tests cover 1-D to 5-D inputs, all dims, `keepdim=False`, fp16 and qs8 over the innermost and a middle dim, dynamic kept and reduced dims, RMSNorm, convolution followed by a mean that is not Global Average Pooling, the size threshold on both sides, and the layout copies the pass inserts.

Test plan:

```
python -m pytest backends/xnnpack/test/ops/test_mean_dim.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.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
```

Authored with AI using Claude Code.
@JakeStevens
JakeStevens force-pushed the xnnpack-mean-dim-static-reduce branch from 7e905e0 to 99c4063 Compare November 6, 2026 20:06
@JakeStevens JakeStevens mentioned this pull request Nov 6, 2026
Open
@JakeStevens

JakeStevens commented Nov 6, 2026

Copy link
Copy Markdown
Contributor Author

@claude review effort=max

@claude

claude Bot commented Nov 6, 2026 •
edited
Loading

Copy link
Copy Markdown

Claude finished @JakeStevens's task in 36m 9s —— View job

Needs Discussion: the 8192-element cut-off in MeanDimConfig doesn't follow from the PR's own latency data, and it still splits the delegate around small static RMSNorm and QK-norm means; the other findings are test gaps and minor items.

Detailed review

PR Review: #24741

Summary

This PR adds an XNNStaticMean node backed by xnn_define_static_reduce. Every mean.dim other than a 4D [2, 3]/keepdim=True Global Average Pooling now lowers to it on NCHW, except small static means with more than one output. The cut-off for those small means needs discussion. The rest are gaps in the tests and minor cleanups.

Code Quality

  • backends/xnnpack/utils/utils.py:96 — is_global_avg_pool_mean reads the input's meta["val"] partway through ChannelsLastTaggedReshapePass (channels_last_tagged_reshape_pass.py:320, :329).
    • Crash when the conv is visited first. Take a static per-tensor dq shared by a conv and a mean. The conv's input_to_nhwc calls insert_copy_q_dq. Its before.replace_all_uses_with(dq) (:405) moves every user of the old dq onto the new dq, including the mean. The new dq copies its meta from the _to_copy created at :591-598, which has no "val". When the pass then reaches the mean, it raises KeyError. On main, mean.dim was a plain membership check in memory_sensitive_ops_nhwc, so this graph worked for GAP means. If the graph is reachable, this is a regression.
    • Silent wrong axes when the mean is visited first. A non-GAP mean is processed, then moved onto the NHWC dq. node_visitor.py:427 ORs in the NHWC tag, so the NCHW reduction_axes index NHWC data. The re-converge at :745-766 only covers broadcast binary ops; view_copy already has the same exposure.
    • Confidence is low on reachability. PT2E's DuplicateDQPass normally gives each user its own dq.
    • Cheap partial fix. Check keepdim first, then take the rank from the mean's own node.meta["val"]. With keepdim=True it equals the input rank. That removes the KeyError half. Fix this →

Testing

  • No test asserts which XNNPACK node is serialized.
    • test_mean_dim.py only checks that mean.dim is consumed and that the outputs match. GAP and XNNStaticMean produce the same numbers.
    • So if the [2, 3]/keepdim=True case were routed through XNNStaticMean, with an extra transpose after a conv, every test would still pass. The description's claim that those .ptes are byte-identical to main would then be false without anything noticing.
    • _xnnpack_node_types (test_channels_last_tagged_reshape.py:493) already decompiles the delegate. Move it into a shared helper and assert the XNNGlobalAvgPooling2d, XNNStaticMean and XNNStaticTranspose counts. Fix this →
  • The dynamic-shape tests run a single random shape.
    • test_mean_dim.py:188-196 go through _test_mean_dim, which calls run_method_and_compare_outputs() with the default num_runs=1 (backends/test/harness/tester.py:320).
    • Each run therefore draws one seq in [1, 32]. The bounds, including a reduced dim of size 1, are only hit by chance.
    • test_add.py, test_softmax.py and test_conv1d.py pass num_runs= for this. Fix this →
  • Some cases are missing.
    • A qs8 conv followed by a non-GAP mean is the only path where the new NCHW requirement meets the quantized copy wrapper (insert_copy_q_dq). test_fp32_conv_mean_dim is fp32-only, and the qs8 tests (:241-251) have no conv and all use keepdim=True.
    • The description says ranks 1 to 6 are supported, but the highest rank tested is 5-D (:156).
    • Fix this →
  • The base-model check in .ci/scripts/test_lora.sh:52-53 and .ci/scripts/test_lora_multimethod.sh:52-53 is wider than the observed flip.
    • The regex went from assistant:\n(?:<think>\n)+ to assistant:[^\n]*\n(?:<think>\n)*. It now accepts any text after assistant:, and also no <think> at all. The one observed flip is assistant: me.
    • Pin the alternate first line, for example assistant:(?: me)?\n. Keep + unless the flipped output really drops <think>.
    • On failure, the scripts still print only EXPECTED_QUANT_PREFIX (test_lora.sh:243,247) and EXPECTED_BASE_PREFIX (test_lora_multimethod.sh:156). Someone debugging a failure doesn't see what the check actually accepts. Fix this →

Backward Compatibility

  • Practically every re-exported Llama-family .pte will now contain XNNStaticMean.
    • examples/models/llama/norm.py:35 is (x * x).mean(-1, keepdim=True). Decode RMSNorm has a single output, and dynamic-shape exports skip the cut-off.
    • Runtimes built before this PR, such as the released AAR and xcframework, reject the program in defineNotImplementedNode (XNNCompiler.cpp:1831).
    • The description does note that an updated runtime is needed. But "such as RMSNorm" undersells it: effectively every new XNNPACK LLM export is affected.
    • There is no partitioner option to keep GAP-only lowering. Removing MeanDimConfig would also drop GAP.
    • Strictly this is forward rather than backward compatibility, but it is the change users will notice most, and the release note should say so explicitly.

Performance

  • The min_multi_row_input_numel = 8192 cut-off (generic_node_configs.py:707, :735-745) doesn't follow from the PR's own numbers.

    • Measured on a desktop thread count. It was derived at 16 threads on an Apple M-series CPU. At 1 thread the description measures the 1x8x64 case as a tie (~5 µs), and at 4 threads only 1x128x128 is slower. At mobile thread counts it rejects means that cost nothing extra in XNNPACK.
    • Slow means above it still lower. The description measured lone means at 1x128x128 (16384 elements, 3.26×) and 1x64x512 (1.90×) as slower, and both are still lowered.
    • In-model means below it are rejected. The description itself says that inside a model "the neighbouring delegated ops already pay the wake-up". Yet two examples stay on portable and split the delegate between linears:
      • a static-shape RMSNorm on 1x16x256 (4096 elements, 16 outputs);
      • Qwen3-0.6B's static-decode QK-norm (16×128 and 8×128 per token).
    • Check the motivating shapes. It is worth confirming that the LeRobot shapes behind XNNPACK: LeRobot policies that do not run yet, and what blocks them #24667 clear the cut-off.
    • Static and dynamic exports diverge. is_shape_dynamic bypasses the check (:732), so the same module partitions differently under static and dynamic export.
    • No precedent, and it can't be tuned. No other XNNPACK config gates on size for performance. It is a class attribute rather than an XNNPartitionerConfig option like force_non_static_weights_for_f32_linear (xnnpack_config.py:64), so it can't be set per deployment.

    I suggest dropping it and treating a lone small mean as the "existing single-op trade-off" the description already accepts for add. If it stays, justify 8192 with in-model numbers at the shapes it rejects, and make it a partitioner option.

  • Every non-GAP 4D mean on an NHWC-tagged tensor gets a full NHWC→NCHW copy inside the delegate (channels_last_tagged_reshape_pass.py:328-329).

    • Examples are the torchvision ShuffleNetV2 and MNASNet heads (x.mean([2, 3])) and CBAM-style x.mean(1, keepdim=True).
    • test_fp32_channels_last_tagged_reshape_pass_conv_mean_channels (test_channels_last_tagged_reshape.py:930-938) locks this in by expecting two copies.
    • This isn't a regression, because main paid the same transpose at the delegate boundary.
    • Remapping the reduction axes for an NHWC input (NCHW dim 1→3, 2→1, 3→2) would avoid the copy. That works for keepdim=True, and for keepdim=False whenever the kept dims stay in order, which covers both examples. This is fine as a follow-up.

Specific Comments

  • backends/xnnpack/partition/config/generic_node_configs.py:707 — if the cut-off stays, "multi_row" is not what the condition tests (output_numel > 1). The why message at :743 ("needs at least 8192 input elements") reads like an XNNPACK capability limit rather than a performance heuristic, which will mislead anyone chasing a partition split.
  • docs/source/backends/xnnpack/op-support.csv:32 — "a static mean" reads like XNNPACK's static-mean node. "A mean with a static shape", as the MeanDimConfig docstring puts it, is clearer. The hard-coded 8192 will also drift from min_multi_row_input_numel.
  • .ci/scripts/test_lora.sh:49-50 (same in test_lora_multimethod.sh) — "can flip between builds" hides the cause. Per the description, the flip comes from RMSNorm and QK-norm moving into XNNPACK under 8da4w.
  • backends/xnnpack/runtime/XNNCompiler.cpp:2059 — the in-tree CPU provider calls xnn_define_static_reduce_v2 (with int64_t axes) against the same XNNPACK pin (backends/cpu/runtime/providers/xnnpack/XNNPACKProvider.cpp:633). If v1 is deprecated there, -Wno-deprecated-declarations (backends/xnnpack/CMakeLists.txt:40) would hide the warning. Confidence is low, because I couldn't read the pinned header.

Recommendation

Needs Discussion

The cut-off is the open design question: either drop it, or justify it with in-model data and make it a partitioner option. The node-type assertions and the tighter LoRA regex should land with this PR. Everything else is minor or can be a follow-up.

@JakeStevens JakeStevens mentioned this pull request Nov 6, 2026
Merged
JakeStevens added a commit that referenced this pull request Nov 6, 2026
@JakeStevens
[XNNPACK] Lower fp32 <-> bf16 casts to XNNPACK converts (#24745)
11821a3 
### Ulterior Motive
Improving overall BF16 support in light of
#24667

### Summary
`ToDimOrderCopyConfig` already lowers fp32 <-> fp16 casts to XNNPACK
convert nodes. This allows fp32 <-> bf16 as well, when bf16 is enabled
in the partitioner. fp16 <-> bf16 stays unsupported, because XNNPACK has
no direct convert between them. bf16 models that upcast a step to fp32
(HF softmax with `dtype=torch.float32`, Llama-style RMSNorm) now stay in
one delegate, instead of running the casts in portable kernels between
partitions. ExecuTorch's current XNNPACK pin already has the bf16 <->
fp32 convert kernels, so no runtime or pin changes are needed.

The config matches `_to_dim_order_copy`, which export emits when dim
order is kept, the `to_edge_transform_and_lower` default. Programs
exported with `skip_dim_order=True`, as `LLMEdgeManager` and
`export_llama` do by default, keep `aten._to_copy`, which no XNNPACK
partitioner config matches yet, so their casts (fp16 ones included)
still run in portable kernels. Partitioning `_to_copy` dtype casts under
the same rules would be a follow-up.

### Benchmarks
Measured with `executor_runner` (Release; XNNPACK + portable ops),
median of 50 iterations. "Before" exports the same model with the cast
lowering disabled. Shapes:
- softmax: `softmax(x, -1, dtype=float32).to(bf16)` on `[1,16,S,S]`.
- RMSNorm: Llama-style, `[1,T,2048]`.
- elementwise: `(x.float()*2+1).to(bf16)` on `[1,128,4096]`.

Single thread (Android pinned to the fastest core with `taskset`), ms,
before -> after:

| model | Apple M4 | iPhone 17 (A19) | Pixel 10 (Tensor G5, X4) | Pixel
6 (Tensor G1, X1; no BF16) |
|---|---|---|---|---|
| softmax S=128 | 1.172 -> 0.141 (8.3x) | 1.415 -> 0.403 (3.5x) | 1.785
-> 0.701 (2.5x) | 1.760 -> 0.684 (2.6x) |
| softmax S=512 | 18.64 -> 2.343 (8.0x) | 20.18 -> 3.246 (6.2x) | 31.16
-> 5.031 (6.2x) | 35.60 -> 6.484 (5.5x) |
| RMSNorm T=1 | 0.012 -> 0.002 (5.6x) | 0.021 -> 0.004 (4.7x) | 0.098 ->
0.012 (8.3x) | 0.013 -> 0.004 (3.1x) |
| RMSNorm T=128 | 1.266 -> 0.243 (5.2x) | 1.620 -> 0.652 (2.5x) | 2.502
-> 0.820 (3.1x) | 1.731 -> 0.712 (2.4x) |
| elementwise | 2.161 -> 0.118 (18.3x) | 2.244 -> 0.418 (5.4x) | 3.754
-> 0.467 (8.0x) | 4.053 -> 0.576 (7.0x) |

Default `executor_runner` threads (M4: 16, iPhone: 6, Pixel 10: 6, Pixel
6: 4; unpinned), ms, before -> after:

| model | Apple M4 | iPhone 17 (A19) | Pixel 10 (Tensor G5) | Pixel 6
(Tensor G1) |
|---|---|---|---|---|
| softmax S=128 | 1.199 -> 0.152 (7.9x) | 1.483 -> 0.294 (5.0x) | 3.858
-> 1.020 (3.8x) | 1.980 -> 0.317 (6.2x) |
| softmax S=512 | 17.11 -> 0.635 (27x) | 18.98 -> 1.956 (9.7x) | 38.88
-> 5.518 (7.0x) | 25.97 -> 5.074 (5.1x) |
| RMSNorm T=1 | 0.010 -> 0.002 (4.8x) | 0.020 -> 0.004 (4.4x) | 0.101 ->
0.101 (1.0x)* | 0.133 -> 0.025 (5.3x) |
| RMSNorm T=128 | 1.323 -> 0.347 (3.8x) | 1.616 -> 0.720 (2.2x) | 8.936
-> 1.765 (5.1x) | 2.544 -> 0.642 (4.0x) |
| elementwise | 2.154 -> 0.154 (14x) | 2.376 -> 0.355 (6.7x) | 3.241 ->
0.839 (3.9x) | 2.806 -> 0.462 (6.1x) |

\* Unpinned, this ~10 us graph runs on a core that doesn't ramp its
clock up, so the run hits a ~0.1 ms floor even with 1 thread. Pinned,
it's 0.012 ms (table above).

RMSNorm still has `mean` outside the delegate (`mean.dim` over the last
dim isn't lowered yet; #24741 lowers it). The casts no longer split it
further (3 -> 2 delegates).

### Test plan
```
python -m pytest backends/xnnpack/test/ops/test_to_copy_dtype.py backends/xnnpack/test/ops/test_to_copy.py
```
New tests cover:
- fp32 I/O around a bf16 linear;
- a bf16 upcast;
- a bf16 softmax computed in fp32 and cast back, as one delegate;
- fp32 -> bf16 rounding matching PyTorch exactly (zero tolerance),
including values halfway between two bf16 numbers;
- no lowering without `enable_bf16`;
- fp16 <-> bf16 not partitioned.

The tests that run bf16 in XNNPACK are skipped in fbcode, like the other
XNNPACK bf16 runtime tests. The delegation tests fail without the config
change.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mergennachin mergennachin Awaiting requested review from mergennachin mergennachin is a code owner
@digantdesai digantdesai Awaiting requested review from digantdesai digantdesai is a code owner
@shoumikhin shoumikhin Awaiting requested review from shoumikhin

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. release notes: xnnpack Changes to the XNNPack backend delegate

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@JakeStevens