[ci][aarch64] add Linux aarch64 wheels

[puneetmatharu]

Adds support for also building Linux AArch64 wheels; currently ao only releases Linux x86 wheels.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3359

• 📄 Preview Python docs built from this PR

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

❗ 2 Active SEVs

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

• Dr CI is temporarily not working due to API fairewall
• Pytorch CI is paused for the time being

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

[puneetmatharu]

@jerryzh168 I was made aware that I should contact you to take a look at this. Any chance I can get some help pushing this through?

[aditew01]

@huydhn can you please take a look and review the changes?

[puneetmatharu]

Hi @huydhn following up on this to see if we can get some traction on this PR.

[puneetmatharu]

@huydhn bump. Could you please take a look as it would be beneficial for us to have AArch64 wheels available on PyPI.

[huydhn]

CI changes LGTM! But let's wait for TorchAO maintainers to give the final stamp

[puneetmatharu]

Hi @jerryzh168 , any chance you could take a look and get this pushed through? Thanks so much.

[gottbrath]

@jerryzh168 - are you the right reviewer for this? If not can you reassign?

[jerryzh168]

I'll ask in our team group and get back

[jerryzh168]

didn't hear back concerns from the team, I think we can merge

[jerryzh168]

I think we can merge, is this tested?

[puneetmatharu]

We have tested the build with the updated env_var_script_linux.sh and can confirm the build works and includes KleidiAI. To confirm this we ran the unit tests test_accuracy_kleidiai and test_accuracy_aten:

❯ pytest -k test_accuracy_kleidiai -vv 
========================================================================= test session starts ==========================================================================
platform linux -- Python 3.12.3, pytest-8.4.2, pluggy-1.6.0 -- /home/punmat01/venv-ao/bin/python3
cachedir: .pytest_cache
hypothesis profile 'default'
rootdir: /home/punmat01/ao
configfile: pyproject.toml
plugins: hypothesis-6.148.2
collected 10592 items / 10591 deselected / 15 skipped / 1 selected                                                                                                     

quantization/test_int8_dynamic_activation_intx_weight_config_v1.py::TestInt8DynamicActivationIntxWeight::test_accuracy_kleidiai PASSED                           [100%]

❯ pytest -k test_accuracy_aten    
========================================================================= test session starts ==========================================================================
platform linux -- Python 3.12.3, pytest-8.4.2, pluggy-1.6.0
rootdir: /home/punmat01/ao
configfile: pyproject.toml
plugins: hypothesis-6.148.2
collected 10592 items / 10591 deselected / 15 skipped / 1 selected                                                                                                     

quantization/test_int8_dynamic_activation_intx_weight_config_v1.py .                                                                                             [100%]

Both completed successfully.

[nikhil-arm]

@jerryzh168 is there a wider test suite you would like us to run?

[jerryzh168]

We have tested the build with the updated env_var_script_linux.sh and can confirm the build works and includes KleidiAI. To confirm this we ran the unit tests test_accuracy_kleidiai and test_accuracy_aten:

I think that should be enough, thanks