[pytorchbot]

What is wrong today

XnnpackPartitioner(enable_bf16=True) lets XNNPACK run bf16 nodes. Without it, bf16 matrix multiplies run on the single-threaded portable kernels, which is very slow. With it, a program can fail to load:

1. The partitioner hands XNNPACK bf16 nodes that XNNPACK cannot define in bf16: softmax, mean, average and max pooling, bilinear upsampling, and batch norm that does not fold into a linear. Loading fails with xnn_status_invalid_parameter.
2. Batch matrix multiply (bmm) is defined with a bf16 output, but XNNPACK only multiplies bf16 matrices into an fp32 result. Loading fails with Failed to create BMM node.

A bf16 attention block written with matmul and softmax hits both.

What this change does

• A config can now say it does not support bf16 (supports_bf16 = False), so its bf16 nodes stay on the portable kernels and the partitioner logs why. Softmax, mean, both poolings, bilinear upsampling and convolution set it. The default is True: XNNPACK rewrites most bf16 ops to fp32 internally, and at the pinned commit only these reject bf16 when defined. A bf16 run of the other float ops found no other failure.
• A bf16 batch norm is kept in XNNPACK only when FuseBatchNormPass.can_fuse says it folds into the linear before it.
• can_fuse now also requires the batch norm to be the only user of the linear output, and the linear output to be 2-D. Without these checks, a bf16 batch norm could be kept but left unfused, so loading failed, and a batch norm after a 3-D linear output was folded along the wrong dim, which gave wrong results in fp32 and bf16.
• A bf16 bmm now writes an fp32 result and converts it to bf16, as the linear layer already did.

With this change, the SmolVLA robot policy (a vision-language-action model) in bf16 with enable_bf16=True loads, and then runs one action chunk in about 11 seconds on 16 x86 cores, instead of about 9 minutes with the default partitioner.

Tests

• test_enable_bf16_leaves_ops_xnnpack_cannot_define_in_bf16: each op, and a convolution followed by a batch norm, is delegated in fp32 but not in bf16, and the bf16 program matches eager.
• test_enable_bf16_fuses_batch_norm_into_a_linear: a bf16 linear followed by BatchNorm1d runs fully in XNNPACK. When the linear output is also used elsewhere, the batch norm stays out.
• test_enable_bf16_runs_an_attention_block: the linear layers and both bmms run in XNNPACK, the softmax stays outside.
• test_bf16_bmm in test_bmm.py, and test_fp32_linear_batch_norm_on_3d_output in test_batch_norm_fusion.py.

All of these fail before this change, except the fused batch norm case, which guards the fusion that already worked. Ran the XNNPACK partitioner, op and pass tests on Linux x86_64. The only failures, in the per-channel int4 linear tests, also happen without this change.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Investigating a potential degradation on Meta's CI clusters

⏳ No Failures, 133 Pending

As of commit 958ed92 with merge base 4db778e ():
💚 Looks good so far! There are no failures yet. 💚

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

[shoumikhin]

Closing: built on release/1.6 without the #24380 pick, so it lacks an import. Replaced by a pick stacked on the #24380 pick, which I will link here.