From ad31723256ea52679e838951a50c11a4aee43a7a Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Wed, 14 May 2025 16:23:38 -0700
Subject: [PATCH 1/3] Update torchao_experimental_test.yml

Add CI for arm linux
---
 .github/workflows/torchao_experimental_test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/torchao_experimental_test.yml b/.github/workflows/torchao_experimental_test.yml
index a712e7a624..2b93ed4232 100644
--- a/.github/workflows/torchao_experimental_test.yml
+++ b/.github/workflows/torchao_experimental_test.yml
@@ -14,7 +14,7 @@ jobs:
   test-cpu-ops:
     strategy:
       matrix:
-        runner: [macos-14]
+        runner: [macos-14, linux.arm64.2xlarge]
     runs-on: ${{matrix.runner}}
     defaults:
       run:
@@ -38,7 +38,7 @@ jobs:
           pip install executorch
           pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu --force-reinstall
           pip install -r dev-requirements.txt
-          USE_CPP=1 TORCHAO_BUILD_KLEIDIAI=1 pip install .
+          USE_CPP=1 TORCHAO_BUILD_KLEIDIAI=1 BUILD_TORCHAO_EXPERIMENTAL=1 TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_BUILD_KLEIDIAI=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 pip install .
       - name: Run python tests
         run: |
           conda activate venv

From b47ed2fbf1fca0bb45d0b8baa4f21aa95ffdd044 Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Wed, 14 May 2025 16:40:58 -0700
Subject: [PATCH 2/3] Update torchao_experimental_test.yml

---
 .github/workflows/torchao_experimental_test.yml | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/torchao_experimental_test.yml b/.github/workflows/torchao_experimental_test.yml
index 2b93ed4232..4f9c2d69bb 100644
--- a/.github/workflows/torchao_experimental_test.yml
+++ b/.github/workflows/torchao_experimental_test.yml
@@ -30,7 +30,8 @@ jobs:
           python-version: "3.10"
           miniconda-version: "latest"
           activate-environment: venv
-      - name: Install requirements
+      - name: Install requirements mac
+        if: runner.os == 'macOS'
         run: |
           conda activate venv
           # Install executorch first because it installs its own version
@@ -38,7 +39,14 @@ jobs:
           pip install executorch
           pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu --force-reinstall
           pip install -r dev-requirements.txt
-          USE_CPP=1 TORCHAO_BUILD_KLEIDIAI=1 BUILD_TORCHAO_EXPERIMENTAL=1 TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_BUILD_KLEIDIAI=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 pip install .
+          USE_CPP=1 TORCHAO_BUILD_KLEIDIAI=1 pip install .
+      - name: Install requirements linux
+        if: runner.os == 'Linux'
+        run: |
+          conda activate venv
+          pip install torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu --force-reinstall
+          pip install -r dev-requirements.txt
+          BUILD_TORCHAO_EXPERIMENTAL=1 TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_BUILD_KLEIDIAI=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 TORCHAO_PARALLEL_BACKEND=OPENMP pip install .
       - name: Run python tests
         run: |
           conda activate venv
@@ -60,6 +68,7 @@ jobs:
           rm -rf /tmp/cmake-out
           popd
       - name: ET ops build
+        if: runner.os == 'macOS'
         run: |
           conda activate venv
           pushd torchao/experimental

From dacbb2ac14f188a3001e1d3ebc9eeb64bf883b2d Mon Sep 17 00:00:00 2001
From: Scott Roy <161522778+metascroy@users.noreply.github.com>
Date: Wed, 14 May 2025 17:03:06 -0700
Subject: [PATCH 3/3] Update torchao_experimental_test.yml

---
 .github/workflows/torchao_experimental_test.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/torchao_experimental_test.yml b/.github/workflows/torchao_experimental_test.yml
index 4f9c2d69bb..4c56ec0c0e 100644
--- a/.github/workflows/torchao_experimental_test.yml
+++ b/.github/workflows/torchao_experimental_test.yml
@@ -54,6 +54,7 @@ jobs:
           python torchao/experimental/tests/test_embedding_xbit_quantizer.py
           python torchao/experimental/tests/test_quant_passes.py
       - name: Run kernels/cpu/aarch64/tests
+        if: runner.os == 'macOS'
         run: |
           conda activate venv
           pushd torchao/experimental/kernels/cpu/aarch64/tests
@@ -61,6 +62,7 @@ jobs:
           rm -rf /tmp/cmake-out
           popd
       - name: Run torchao/experimental/ops/tests
+        if: runner.os == 'macOS'
         run: |
           conda activate venv
           pushd torchao/experimental/ops/tests