Ship the TorchAO kernels on the Linux aarch64 wheel too

[shoumikhin]

Closes #22299.

Summary

The wheel builds the TorchAO kernels for macOS arm64 and for nothing else, so a Linux aarch64
install has no libexecutorch_kernels_torchao even though it has exactly the hardware they target.
Verified against the published nightly 1.5.0.dev20260828 cp310 manylinux_2_28_aarch64: that row
ships seven libraries and none of them is this one.

The gate was one condition. pybind.cmake enables the option inside the Darwin branch under
CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64", and Linux aarch64 reports aarch64, so it never matched.

Nothing about these kernels is Apple specific: they are selected by TORCHAO_BUILD_CPU_AARCH64 and
reach their NEON paths through TORCHAO_ENABLE_ARM_NEON_DOT, both of which hold on any aarch64
host. The in-tree comment claiming otherwise is corrected here, along with two others that repeated
it.

x86 is deliberately left out

Worth stating, because "enable it everywhere" is the tempting fix and it is wrong. In torchao's
shared_kernels/linear_8bit_act_xbit_weight/kernel_selector.h:

'return PackedWeightsFormat' inside ARM guards  : 4
'return PackedWeightsFormat' outside (portable) : 0

With no NEON define both format selectors have no reachable return and fall through to
throw std::runtime_error("No packed_weights_format was selected"). Enabling it on x86 would build
and then throw on every op, which is worse than shipping nothing.

Test plan

Built on real Linux aarch64 hardware (gcc 11.4, cmake 3.22, /proc/cpuinfo reports asimddp),
at the pinned TorchAO commit b47f1a3655, with the same options the preset now sets:

step	result
configure torchao/csrc/cpu	rc=0
build libtorchao_kernels_aarch64.a	built
kernel_selector.h under gcc -march=armv8.4-a+dotprod	compiles — and that is the header the ExecuTorch ops include

Prerequisites are already satisfied on this row: EXECUTORCH_BUILD_XNNPACK is on for the wheel and
check_required_options_on forces EXECUTORCH_BUILD_CPUINFO and EXECUTORCH_BUILD_PTHREADPOOL
with it.

Checked the new condition matches only what it should:

aarch64 ON   arm64 ON   x86_64 off   amd64 off   i686 off   armv7l off

so 32-bit ARM stays out. The macOS path is untouched, and setup.py already ships the library
whenever the flag is set, so no packaging change is needed.

One thing for reviewers to confirm

I could build and syntax-check the kernels on aarch64 Linux, but I could not produce a full aarch64
wheel and run a lowbit export end to end, since PR wheel jobs publish x86_64 only and aarch64
builds on the nightly cron. The aarch64 wheel row in this PR's CI is the real test of that.

Separately, two GCC portability bugs exist upstream in TorchAO's own aarch64 tree, found while
building its test suite on this host: embedding_lut.h uses std::memcpy without including
<cstring>, and test_bitpacking.cpp brace-initialises int8x16_t. Neither file is reached by
torchao_ops_executorch, so neither affects this build, but both are worth fixing in the submodule.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit c4a82b9 with merge base 9b558d9 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (trunk failure)

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.