Add CI for Arm Linux (#2211)

* Update torchao_experimental_test.yml

Add CI for arm linux

* Update torchao_experimental_test.yml

* Update torchao_experimental_test.yml

Author: metascroy

.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:
@@ -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
@@ -39,27 +40,37 @@ jobs:
           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 .
+      - 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
           pytest torchao/experimental/tests/test_int8_dynamic_activation_intx_weight.py
           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
           sh build_and_run_tests.sh
           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
           sh build_and_run_tests.sh
           rm -rf /tmp/cmake-out
           popd
       - name: ET ops build
+        if: runner.os == 'macOS'
         run: |
           conda activate venv
           pushd torchao/experimental