From ddceb963c1599f803b5c4beca42b802de5134b44 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Wed, 29 Jan 2025 17:54:22 -0800
Subject: [PATCH] Add GitHub Actions based XLA ARM64 CPU build

PiperOrigin-RevId: 721177984
---
 third_party/xla/.github/workflows/ci.yml           | 11 +++++++++++
 third_party/xla/build_tools/ci/build.py            | 12 ++++++++++++
 third_party/xla/build_tools/ci/golden_commands.txt |  5 +++++
 3 files changed, 28 insertions(+)

diff --git a/third_party/xla/.github/workflows/ci.yml b/third_party/xla/.github/workflows/ci.yml
index 3af91a2f4049d7..966f8a2830c4ef 100644
--- a/third_party/xla/.github/workflows/ci.yml
+++ b/third_party/xla/.github/workflows/ci.yml
@@ -21,10 +21,15 @@ on:
     branches:
       - main
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
 
 jobs:
   Tests:
     strategy:
+      # Don't fail fast - want to see results for all builds even if one fails.
+      fail-fast: false
       matrix:
         job_info: [
           {
@@ -32,6 +37,12 @@ jobs:
             container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest",
             kokoro_job_name: "xla-linux-x86-cpu",
             pretty_name: "XLA Linux x86 CPU",
+          },
+          {
+            pool: "linux-arm64-c4a-16",
+            container: "us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest",
+            kokoro_job_name: "xla-linux-arm64-cpu",
+            pretty_name: "XLA Linux ARM64 CPU",
           }
         ]
     name: ${{ matrix.job_info.pretty_name }}
diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py
index e5cbe4cfa609fd..4fd6b897fb26f5 100755
--- a/third_party/xla/build_tools/ci/build.py
+++ b/third_party/xla/build_tools/ci/build.py
@@ -89,6 +89,7 @@ class BuildType(enum.Enum):
   """Enum representing all types of builds."""
   CPU_X86_SELF_HOSTED = enum.auto()
   CPU_ARM64 = enum.auto()
+  CPU_ARM64_SELF_HOSTED = enum.auto()
   GPU = enum.auto()
   GPU_CONTINUOUS = enum.auto()
 
@@ -312,6 +313,16 @@ def nvidia_gpu_build_with_compute_capability(
     build_tag_filters=cpu_arm_tag_filter,
     test_tag_filters=cpu_arm_tag_filter,
 )
+_CPU_ARM64_SELF_HOSTED_BUILD = Build(
+    type_=BuildType.CPU_ARM64_SELF_HOSTED,
+    repo="openxla/xla",
+    image_url=None,
+    configs=("warnings", "rbe_cross_compile_linux_arm64", "nonccl"),
+    target_patterns=_XLA_DEFAULT_TARGET_PATTERNS,
+    options={**_DEFAULT_BAZEL_OPTIONS, "build_tests_only": True},
+    build_tag_filters=cpu_arm_tag_filter,
+    test_tag_filters=cpu_arm_tag_filter,
+)
 # TODO(ddunleavy): Setup additional build for a100 tests once L4 RBE is ready.
 _GPU_BUILD = nvidia_gpu_build_with_compute_capability(
     type_=BuildType.GPU,
@@ -464,6 +475,7 @@ def nvidia_gpu_build_with_compute_capability(
     "tensorflow/xla/tensorflow/cpu/build_cpu": _TENSORFLOW_CPU_BUILD,
     "tensorflow/xla/tensorflow/gpu/build_gpu": _TENSORFLOW_GPU_BUILD,
     "xla-linux-x86-cpu": _CPU_X86_SELF_HOSTED_BUILD,
+    "xla-linux-arm64-cpu": _CPU_ARM64_SELF_HOSTED_BUILD,
 }
 
 
diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt
index 57f55853225271..1c50ffd0403249 100644
--- a/third_party/xla/build_tools/ci/golden_commands.txt
+++ b/third_party/xla/build_tools/ci/golden_commands.txt
@@ -7,6 +7,11 @@ docker exec xla_ci bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvi
 docker exec xla_ci bazel analyze-profile profile.json.gz
 docker stop xla_ci
 # END BuildType.CPU_ARM64
+# BEGIN BuildType.CPU_ARM64_SELF_HOSTED
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only --nobuild -- //xla/... //build_tools/... @local_tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only -- //xla/... //build_tools/... @local_tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.CPU_ARM64_SELF_HOSTED
 # BEGIN BuildType.CPU_X86_SELF_HOSTED
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --config=warnings --config=nonccl --config=rbe_linux_cpu --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/... //build_tools/... @local_tsl//tsl/...
 bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd --config=warnings --config=nonccl --config=rbe_linux_cpu --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/... //build_tools/... @local_tsl//tsl/...