From 3cd88970673fd70abd4e62969058d64660892d9f Mon Sep 17 00:00:00 2001
From: Puneet Matharu <puneet.matharu@arm.com>
Date: Wed, 29 Oct 2025 12:34:14 +0000
Subject: [PATCH] [CI] add aarch64 wheels.

---
 .../workflows/build_wheels_linux_aarch64.yml  | 93 +++++++++++++++++++
 ...s_linux.yml => build_wheels_linux_x86.yml} |  5 +-
 packaging/env_var_script_linux.sh             |  7 ++
 3 files changed, 103 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/build_wheels_linux_aarch64.yml
 rename .github/workflows/{build_wheels_linux.yml => build_wheels_linux_x86.yml} (96%)

diff --git a/.github/workflows/build_wheels_linux_aarch64.yml b/.github/workflows/build_wheels_linux_aarch64.yml
new file mode 100644
index 0000000000..d3c88f8230
--- /dev/null
+++ b/.github/workflows/build_wheels_linux_aarch64.yml
@@ -0,0 +1,93 @@
+name: Build Linux Wheels (AArch64)
+
+on:
+  pull_request:
+    paths:
+      - build/packaging/**
+      - packaging/**
+      - .github/workflows/build_wheels_linux_aarch64.yml
+      - setup.py
+      - .github/scripts/**
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+    tags:
+        # NOTE: Binary build pipelines should only get triggered on release candidate builds
+        # Release candidate tags look like: v1.11.0-rc1
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  schedule:
+    - cron: '0 0 * * *'  # Runs at midnight UTC every day
+  workflow_dispatch:
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux-aarch64
+      with-cpu: enable
+      with-cuda: disable
+      # Note: if free-threaded python is required add py3.13t here
+      python-versions: '["3.10"]'
+
+  build:
+    needs: generate-matrix
+    permissions:
+      id-token: write
+      contents: read
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      # Set the ref to an empty string instead of the default nightly because
+      # torchao doesn't have nightly branch setup yet, instead the build is
+      # triggered daily from main with a schedule
+      repository: pytorch/ao
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      env-var-script: packaging/env_var_script_linux.sh
+      pre-script: packaging/pre_build_script.sh
+      post-script: packaging/post_build_script.sh
+      smoke-test-script: packaging/smoke_test.py
+      package-name: torchao
+      trigger-event: ${{ github.event_name }}
+      architecture: aarch64
+    secrets:
+      PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
+  notify:
+    runs-on: ubuntu-latest
+    name: Email notification
+    needs: [generate-matrix, build]
+    if: failure() && github.event_name == 'schedule'
+    steps:
+      - uses: dawidd6/action-send-mail@v4
+        with:
+          server_address: smtp.gmail.com
+          server_port: 465
+          username: torchao.notify
+          password: ${{ secrets.TORCHAO_NOTIFY_PASSWORD }}
+          from: torchao.notify@gmail.com
+          to: ${{ secrets.TORCHAO_NOTIFY_RECIPIENT }}
+          subject: Scheduled Build Failure for TorchAO
+          body: |
+            Build Failure Notification for TorchAO
+
+            A failure occurred in the Build Linux Wheels workflow.
+
+            Run Details:
+            - Workflow: ${{ github.workflow }}
+            - Run Type: ${{ github.event_name }}
+            - Repository: ${{ github.repository }}
+            - Branch/PR: ${{ github.ref }}
+            - Commit: ${{ github.sha }}
+
+            You can view the full run details here:
+            ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}
+
+            Error Information:
+            ${{ needs.generate-matrix.result == 'failure' && 'Matrix generation failed' || '' }}
+            ${{ needs.build.result == 'failure' && 'Build job failed' || '' }}
+
+            This is an automated notification. Please check the GitHub Actions page for more details about the failure.
diff --git a/.github/workflows/build_wheels_linux.yml b/.github/workflows/build_wheels_linux_x86.yml
similarity index 96%
rename from .github/workflows/build_wheels_linux.yml
rename to .github/workflows/build_wheels_linux_x86.yml
index 2ae04b4461..c58baa6597 100644
--- a/.github/workflows/build_wheels_linux.yml
+++ b/.github/workflows/build_wheels_linux_x86.yml
@@ -1,12 +1,12 @@
 # From https://github.com/pytorch/test-infra/wiki/Using-Nova-Reusable-Build-Workflows
-name: Build Linux Wheels
+name: Build Linux Wheels (x86)
 
 on:
   pull_request:
     paths:
       - build/packaging/**
       - packaging/**
-      - .github/workflows/build_wheels_linux.yml
+      - .github/workflows/build_wheels_linux_x86.yml
       - setup.py
       - .github/scripts/**
   push:
@@ -58,6 +58,7 @@ jobs:
       trigger-event: ${{ github.event_name }}
       # This is the CUDA version to be uploaded to torchao-nightly pypi
       upload-to-pypi: cu121
+      architecture: x86_64
     secrets:
       PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }}
   notify:
diff --git a/packaging/env_var_script_linux.sh b/packaging/env_var_script_linux.sh
index 3d3394fbd5..0c84337aa5 100644
--- a/packaging/env_var_script_linux.sh
+++ b/packaging/env_var_script_linux.sh
@@ -1,4 +1,5 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright (c) 2025 Arm Limited and/or its affiliates.
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -17,3 +18,9 @@ TORCH_CUDA_ARCH_LIST="8.0;8.6"
 if [[ ${CU_VERSION:-} == "cu124" ]]; then
   TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};9.0"
 fi
+
+# Enable kleidiai for aarch64 build
+if [[ $(uname -m) == "aarch64" ]]; then
+    echo "Enabling kleidiai for aarch64 build"
+    export TORCHAO_BUILD_KLEIDIAI=1
+fi