chore: import upstream snapshot with attribution

third_party/xla/.bant-macros

@@ -0,0 +1,41 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+# Bant macro expansions for XLA custom rules.
+#
+# Syntax: https://github.com/hzeller/bant?tab=readme-ov-file#macros.
+#
+# See developer guide on how to use Bant tool for making sure that your targets
+# list all required dependencies:
+# https://openxla.org/xla/developer_environment#remove_unused_dependencies_from_build_files.
+
+# XLA tests
+xla_test = bant_forward_args(cc_test())
+xla_cc_test = bant_forward_args(cc_test())
+
+# tf_proto_library() wraps proto_library() and cc_proto_library().
+# It generates a proto_library with the given name and a cc_proto_library
+# with a "_cc" suffix that provides .pb.h headers.
+tf_proto_library = (
+    proto_library(
+        name = name,
+        srcs = srcs,
+        deps = deps,
+    ),
+    cc_proto_library(
+        name = name + "_cc",
+        deps = [":" + name],
+    ),
+)

third_party/xla/.bazelrc

@@ -0,0 +1,42 @@
+# Disable Bzlmod for now
+common --noenable_bzlmod --enable_workspace
+
+# Turn off empty glob check (flipped in Bazel 8).
+common --incompatible_disallow_empty_glob=false
+
+# Enable legacy external runfiles
+# This is false in Bazel 8, but broke some tests.
+common --legacy_external_runfiles=true
+
+common:bzlmod --enable_bzlmod --noenable_workspace
+
+common --incompatible_enable_cc_toolchain_resolution
+common --repo_env USE_HERMETIC_CC_TOOLCHAIN=1
+
+# TODO: Migrate for https://github.com/bazelbuild/bazel/issues/7260
+common:clang_local --noincompatible_enable_cc_toolchain_resolution
+common:clang_local --@rules_ml_toolchain//common:enable_hermetic_cc=False
+common:clang_local --repo_env USE_HERMETIC_CC_TOOLCHAIN=0
+
+common:clang-tidy-common --aspects=@bazel_clang_tidy//clang_tidy:clang_tidy.bzl%clang_tidy_aspect
+common:clang-tidy-common --output_groups=report
+common:clang-tidy-common --@bazel_clang_tidy//:clang_tidy_executable=//build_tools/ci:clang_tidy_bin
+
+common:clang-tidy --config=clang-tidy-common
+common:clang-tidy --@bazel_clang_tidy//:clang_tidy_config=//:clang_tidy_config
+
+common:clang-tidy-noerrors --config=clang-tidy-common
+common:clang-tidy-noerrors --@bazel_clang_tidy//:clang_tidy_config=//:clang_tidy_config_noerrors
+
+# Load the TensorFlow bazelrc
+import %workspace%/tensorflow.bazelrc
+import %workspace%/warnings.bazelrc
+
+# Treat warnings as errors
+build --config=warnings
+
+# Enable layering check for CI builds
+common:rbe_linux_cpu --features=layering_check
+common:bzlmod --features=-layering_check
+
+try-import %workspace%/xla_configure.bazelrc

third_party/xla/.bazelversion

@@ -0,0 +1 @@
+8.7.0

third_party/xla/.clang-format

@@ -0,0 +1,56 @@
+BasedOnStyle: Google
+Language: Cpp
+PointerBindsToType: true
+AlignTrailingComments:
+  Kind: Always
+IncludeBlocks: Regroup
+IncludeCategories:
+  # Priority 1: Standard C/C++ headers (angle brackets) - separated by blank line
+  - Regex:    '^<[a-z_]+>$'
+    Priority: 1
+    SortPriority: 1
+  # Priority 2: All other headers (no blank lines between them, sorted by SortPriority)
+  # gtest/gmock headers (angle brackets, but not standard headers)
+  - Regex:    '^<g(test|mock)/.*>$'
+    Priority: 2
+    SortPriority: 2
+  # absl headers
+  - Regex:    '^"absl/.*"$'
+    Priority: 2
+    SortPriority: 3
+  # xla/tsl/platform/status_macros.h — internally third_party/gloop, sorts after absl/
+  - Regex:    '^"xla/tsl/platform/status_macros\.h"$'
+    Priority: 2
+    SortPriority: 3
+  # llvm headers
+  - Regex:    '^"llvm/.*"$'
+    Priority: 2
+    SortPriority: 4
+  # mlir headers
+  - Regex:    '^"mlir/.*"$'
+    Priority: 2
+    SortPriority: 5
+  # third_party libraries
+  - Regex:    '^"third_party/.*"$'
+    Priority: 2
+    SortPriority: 6
+  # google (protobuf, etc.) headers
+  - Regex:    '^"google/.*"$'
+    Priority: 2
+    SortPriority: 7
+  # xla headers (including xla/tsl)
+  - Regex:    '^"xla/.*"$'
+    Priority: 2
+    SortPriority: 100
+  # tsl headers (without xla/ prefix)
+  - Regex:    '^"tsl/.*"$'
+    Priority: 2
+    SortPriority: 101
+  # triton headers
+  - Regex:    '^"triton/.*"$'
+    Priority: 2
+    SortPriority: 102
+  # Default catch-all for any other headers (IMPORTANT: ordered before XLA)
+  - Regex:    '.*'
+    Priority: 2
+    SortPriority: 8

third_party/xla/.clang-tidy

@@ -0,0 +1,165 @@
+ExtraArgs: [
+  "-Wno-everything",
+  "-Wno-error",
+  "-Wdeprecated-declarations",
+  "-D_LIBCPP_DISABLE_DEPRECATION_WARNINGS",
+  "-Wdeprecated-register",
+  "-Wexpansion-to-defined",
+  "-Wignored-attributes",
+  "-Wnon-pod-varargs",
+  "-Wshadow-field",
+  "-Wshift-sign-overflow",
+  "-Wtautological-undefined-compare",
+  "-Wthread-safety-analysis",
+  "-Wthread-safety-reference",
+  "-Wundefined-bool-conversion",
+  "-Wunreachable-code",
+  "-Wunused-const-variable",
+  "-Wunused-function",
+  "-Wunused-lambda-capture",
+  "-Wunused-local-typedef",
+  "-Wunused-private-field",
+  "-Wuser-defined-warnings",
+]
+# Get colors when outputting through `bazel build --config=clang-tidy`.
+UseColor: true
+# Necessary for using --config=clang-tidy to emit errors.
+WarningsAsErrors: "*"
+Checks: "-*,
+  abseil-duration-addition,
+  abseil-duration-comparison,
+  abseil-duration-conversion-cast,
+  abseil-duration-division,
+  abseil-duration-factory-float,
+  abseil-duration-factory-scale,
+  abseil-duration-subtraction,
+  abseil-duration-unnecessary-conversion,
+  abseil-faster-strsplit-delimiter,
+  abseil-no-internal-dependencies,
+  abseil-redundant-strcat-calls,
+  abseil-str-cat-append,
+  abseil-string-find-startswith,
+  abseil-string-find-str-contains,
+  abseil-time-comparison,
+  abseil-time-subtraction,
+  bugprone-argument-comment,
+  bugprone-assert-side-effect,
+  bugprone-bool-pointer-implicit-conversion,
+  bugprone-fold-init-type,
+  bugprone-forward-declaration-namespace,
+  bugprone-inaccurate-erase,
+  bugprone-macro-repeated-side-effects,
+  bugprone-move-forwarding-reference,
+  bugprone-multiple-statement-macro,
+  bugprone-string-constructor,
+  bugprone-stringview-nullptr,
+  bugprone-suspicious-memset-usage,
+  bugprone-undefined-memory-manipulation,
+  bugprone-undelegated-constructor,
+  bugprone-unused-raii,
+  bugprone-use-after-move,
+  clang-analyzer-optin.cplusplus.VirtualCall,
+  clang-diagnostic-c++17-compat,
+  clang-diagnostic-c++17-compat-mangling,
+  clang-diagnostic-deprecated-declarations,
+  clang-diagnostic-deprecated-increment-bool,
+  clang-diagnostic-deprecated-register,
+  clang-diagnostic-expansion-to-defined,
+  clang-diagnostic-ignored-attributes,
+  clang-diagnostic-mismatched-tags,
+  clang-diagnostic-non-pod-varargs,
+  clang-diagnostic-pre-c++20-compat,
+  clang-diagnostic-pre-c++20-compat-pedantic,
+  clang-diagnostic-shadow-field,
+  clang-diagnostic-shift-sign-overflow,
+  clang-diagnostic-tautological-undefined-compare,
+  clang-diagnostic-thread-safety*,
+  clang-diagnostic-undefined-bool-conversion,
+  clang-diagnostic-unreachable-code,
+  clang-diagnostic-unreachable-code-loop-increment,
+  clang-diagnostic-unused-const-variable,
+  clang-diagnostic-unused-function,
+  clang-diagnostic-unused-lambda-capture,
+  clang-diagnostic-unused-local-typedef,
+  clang-diagnostic-unused-private-field,
+  clang-diagnostic-user-defined-warnings,
+  darwin-avoid-spinlock,
+  google-build-explicit-make-pair,
+  google-build-namespaces,
+  google-build-using-namespace,
+  google-default-arguments,
+  google-explicit-constructor,
+  google-global-names-in-headers,
+  google-objc-avoid-nsobject-new,
+  google-objc-function-naming,
+  google-objc-global-variable-declaration,
+  google-readability-braces-around-statements,
+  google-readability-function-size,
+  google-readability-namespace-comments,
+  google-runtime-int,
+  google-runtime-operator,
+  misc-coroutine-hostile-raii,
+  misc-definitions-in-headers,
+  misc-include-cleaner,
+  misc-static-assert,
+  misc-unconventional-assign-operator,
+  misc-uniqueptr-reset-release,
+  misc-unused-alias-decls,
+  misc-unused-using-decls,
+  modernize-make-unique,
+  modernize-redundant-void-arg,
+  modernize-replace-auto-ptr,
+  modernize-shrink-to-fit,
+  modernize-use-bool-literals,
+  modernize-use-equals-default,
+  modernize-use-nullptr,
+  modernize-use-override,
+  objc-avoid-nserror-init,
+  objc-dealloc-in-category,
+  objc-forbidden-subclassing,
+  objc-nsinvocation-argument-lifetime,
+  objc-property-declaration,
+  objc-super-self,
+  performance-faster-string-find,
+  performance-for-range-copy,
+  performance-implicit-conversion-in-loop,
+  performance-inefficient-algorithm,
+  performance-inefficient-vector-operation,
+  performance-move-constructor-init,
+  portability-std-allocator-const,
+  readability-avoid-const-params-in-decls,
+  readability-const-return-type,
+  readability-container-size-empty,
+  readability-else-after-return,
+  readability-inconsistent-declaration-parameter-name,
+  readability-misleading-indentation,
+  readability-redundant-control-flow,
+  readability-redundant-smartptr-get,
+  readability-string-compare,"
+CheckOptions:
+  - key: 'bugprone-assert-side-effect.AssertMacros'
+    value: assert,DCHECK
+  - key: 'bugprone-string-constructor.WarnOnLargeLength'
+    value: 0
+  - key: 'google-readability-braces-around-statements.ShortStatementLines'
+    value: 0
+  - key: 'google-readability-function-size.ParameterThreshold'
+    value: 100
+  - key: 'modernize-make-unique.IncludeStyle'
+    value: 'google'
+  - key: 'performance-inefficient-vector-operation.VectorLikeClasses'
+    value: '::std::vector;::absl::InlinedVector'
+  - key: 'performance-inefficient-vector-operation.EnableProto'
+    value: 1
+  - key: 'abseil-string-find-startswith.IncludeStyle'
+    value: 'google'
+  - key: 'abseil-string-find-startswith.AbseilStringsMatchHeader'
+    value: 'absl/strings/match.h'
+  - key: 'abseil-string-find-startswith.StringLikeClasses'
+    value: '::std::string_view;::absl::string_view;::basic_string;::std::basic_string;'
+  - key: 'abseil-string-find-str-contains.IncludeStyle'
+    value: 'google'
+  - key: 'abseil-string-find-str-contains.AbseilStringsMatchHeader'
+    value: 'absl/strings/match.h'
+  - key: 'abseil-string-find-str-contains.StringLikeClasses'
+    value: '::std::basic_string_view;::absl::string_view;::basic_string;::std::basic_string;'

third_party/xla/.clangd

@@ -0,0 +1,27 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+CompileFlags:
+  # Use the compile_commands.json for build flags
+  CompilationDatabase: .
+
+Diagnostics:
+  # Enable strict include checking
+  UnusedIncludes: Strict
+  MissingIncludes: Strict
+
+# Improve code completion and navigation
+Index:
+  Background: Build
+  StandardLibrary: Yes

third_party/xla/.github/actionlint.yml

@@ -0,0 +1,30 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+self-hosted-runner:
+  labels:
+    - "linux-x86-n2-16" # Linux X86 runner using the 16 vcpu n2-standard-16 machine.
+    - "linux-x86-n2-64" # Linux X86 runner using the 64 vcpu n2-standard-64 machine.
+    - "linux-x86-n2-128" # Linux X86 runner using the 128 vcpu n2-standard-128 machine.
+    - "linux-x86-g2-16-l4-1gpu" # Linux X86 GPU runner using g2-standard-16 machine with 1 NVIDIA L4 GPU attached.
+    - "linux-x86-g2-48-l4-4gpu" # Linux X86 GPU runner using g2-standard-48 machine with 4 NVIDIA L4 GPUs attached.
+    - "linux-arm64-c4a-16" # Linux ARM64 CPU Runner using the 16 vcpu c4a-standard-16 machine.
+    - "linux-arm64-t2a-48" # Linux ARM64 CPU Runner using the 48 vcpu t2a-standard-48 machine.
+    - "windows-x86-n2-16" # Windows X86 runner using n2-standard-16 machine.
+    - "amd-do-linux-xla-gpu-gfx950-1" # AMD runner 1 GPU.
+    - "amd-do-linux-xla-gpu-gfx950-4" # AMD runner 4 GPU.
+    - "amd-do-linux-xla-gpu-gfx950-8" # AMD runner 8 GPUs.
+    - "linux-x86-a4-224-b200-1gpu" # Linux X86 GPU runner using 1 B200 GPU and 1/8 the resources of a a4-highgpu-8g machine
+    - "linux-x86-a3-8g-h100-8gpu" # Linux X86 GPU runner using a3-highgpu-8g machine with 8 NVIDIA H100 GPUs attached.
+    - "linux-x64-battlemage-256-1gpu-intel" # Linux X86 GPU runner using an EMR4148 machine with 2 Intel BMG GPUs attached.

third_party/xla/.github/dependabot.yml

@@ -0,0 +1,33 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: pip
+    directory: /
+    schedule:
+      interval: weekly
+    cooldown:
+      default-days: 7
+  - package-ecosystem: github-actions
+    directory: /
+    schedule:
+      interval: weekly 
+    cooldown:
+      default-days: 7

third_party/xla/.github/pull_request_template.md

@@ -0,0 +1,24 @@
+📝 Summary of Changes
+Please provide a clear and concise summary of the changes you've made.
+
+🎯 Justification
+Explain why this change is important and which workload benefits from this
+change.
+
+🚀 Kind of Contribution
+Please remove what does not apply: 🐛 Bug Fix, ⚡️ Performance Improvement,
+✨ New Feature, ♻️ Cleanup, 📚 Documentation, 🧪 Tests
+
+📊 Benchmark (for Performance Improvements)
+Please measure and include speedups for one of the public HLOs in
+`compiler/xla/tools/benchmarks/hlo/`.
+
+🧪 Unit Tests:
+What unit tests were added? For example, a new pass should be tested on minimal
+HLO. The transformation can be tested with FileCheck tests or assertions on the
+transformed HLO.
+
+🧪 Execution Tests:
+What execution tests were added? For example, a new optimization should be
+tested with an end-to-end execution test triggering the optimization and
+asserting correctness. Please provide test cases running with at most 2 GPUs.

third_party/xla/.github/workflows/actions-lint.yml

@@ -0,0 +1,39 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+name: actionlint
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+permissions: {}
+
+jobs:
+  actionlint:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+
+      - name: actionlint
+        uses: docker://rhysd/actionlint@sha256:b1934ee5f1c509618f2508e6eb47ee0d3520686341fec936f3b79331f9315667  # v1.7.12
+        with:
+          args: -color

third_party/xla/.github/workflows/bazel_dependency_violations.yml

@@ -0,0 +1,71 @@
+# Copyright 2024 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Bazel Dependency Violations
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  dependency-violations:
+    strategy:
+      matrix:
+        tag: [gpu, cuda-only, rocm-only]
+    name: no-${{ matrix.tag }}-targets-in-cpu-build
+    runs-on: "linux-x86-n2-16"
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 10
+    continue-on-error: true
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Run bazel build --nobuild //xla/... with retries"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
+      - name: "Run bazel cquery ... //xla/..."
+        run: |
+          set -euo pipefail
+          TAG_WITH_UNDERSCORES="${{ matrix.tag }}"
+          TAG_WITH_UNDERSCORES="${TAG_WITH_UNDERSCORES/-/_}"
+          if ! OUTPUT=$(bazel cquery --aspects "build_tools/dependencies/aspects.bzl%validate_${TAG_WITH_UNDERSCORES}_tag" //xla/... 2>&1); then
+            echo "Failed to run bazel cquery. Output:"
+            echo "$OUTPUT"
+            exit 1
+          fi
+          if echo "$OUTPUT" | grep 'Violation' >/dev/null; then
+            echo "The following dependency violations were found:"
+            echo "$OUTPUT" | grep 'Violation' | sed -e 's/^.*\[Violation\]/ -/'
+            echo ""
+            echo ""
+            echo "There are a couple of potential solutions for this/these violation(s):"
+            echo ""
+            echo "1. Tag the dependent target with the same tag as the dependee."
+            echo ""
+            echo "2. If unavoidable make the dependency selective using the"
+            echo "   'if_{gpu|cuda|rocm}_is_configured' macro. This is discouraged"
+            echo "   outside of stream_executor."
+            echo ""
+            exit 1
+          fi
+
+          echo "No dependency violations found for tag '${{ matrix.tag }}'."

third_party/xla/.github/workflows/bazel_query.yml

@@ -0,0 +1,47 @@
+# Copyright 2024 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Bazel Query
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  bazel-query:
+    runs-on: "linux-x86-n2-16"
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 10
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Run bazel build --nobuild //xla/... with retries"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
+      - name: "Run bazel query //xla/..."
+        run: bazel query //xla/... > /dev/null
+      - name: "Run bazel query deps(//xla/...)"
+        run: bazel query "deps(//xla/...)" > /dev/null
+      - name: "Check that TSL doesn't depend on XLA"
+        run: |
+          bazel query "somepath(//xla/tsl/..., //xla/... -//xla/tsl/...)" 2>&1 | tee query_output.log
+          grep -q "INFO: Empty results" query_output.log

third_party/xla/.github/workflows/bazel_tags.yml

@@ -0,0 +1,41 @@
+# Copyright 2024 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Bazel Tags
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  bazel-tags:
+    runs-on: "linux-x86-n2-16"
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest" # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 10
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Run bazel build --nobuild //xla/... with retries"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --nobuild //xla/...
+      - name: "Assert all tags are documented"
+        run: bazel query //xla/... --output=build | python3 build_tools/lint/tags.py

third_party/xla/.github/workflows/benchmark_postsubmit.yml

@@ -0,0 +1,279 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Postsubmit Benchmarks
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  push:
+    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: [
+          {
+            pool: "linux-arm64-t2a-48",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
+            pretty_name: "XLA Linux ARM64 CPU 48 vcpu Presubmit",
+            bazel_arch_dir: "aarch64-opt",
+            platform: "CPU",
+            runner_binary: "hlo_runner_main",
+            hlo_test_file: "xla/tools/hlo_opt/tests/cpu_hlo.hlo"
+          },
+          {
+            pool: "linux-x86-n2-128",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
+            bazel_arch_dir: "k8-opt",
+            platform: "CPU",
+            runner_binary: "hlo_runner_main",
+            hlo_test_file: "xla/tools/hlo_opt/tests/cpu_hlo.hlo"
+          },
+          {
+            pool: "linux-x86-g2-16-l4-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
+            pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
+            bazel_arch_dir: "k8-opt",
+            platform: "GPU",
+            runner_binary: "hlo_runner_main_gpu",
+            hlo_test_file: "xla/tools/hlo_opt/tests/gpu_hlo_backend.hlo"
+          },
+        ]
+    name: ${{ matrix.job_info.pretty_name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 60
+    env:
+      OUTPUT_FILE_PATH: ${{ github.workspace }}/hlo_test_output.txt
+    steps:
+      - name: Print GitHub Context
+        run: |
+          echo "GitHub SHA: ${{ github.sha }}"
+          echo "GitHub Ref: ${GITHUB_REF}"
+          echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
+          echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
+          echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
+          echo "GitHub Repository: ${{ github.repository }}"
+          echo "GitHub Run ID: ${{ github.run_id }}"
+          echo "GitHub Run Number: ${{ github.run_number }}"
+          echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
+          echo "GitHub Actor: ${GITHUB_ACTOR}"
+          echo "GitHub Event Name: ${{ github.event_name }}"
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "Pull Request Number: ${{ github.event.pull_request.number }}"
+            echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
+            echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
+          fi
+        env:
+          GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
+          GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
+      # Find the current PR number, if any, because github context doesn't have it for push events.
+      - uses: jwalton/gh-find-current-pr@f3d61b485d2801773f7a07b2aaa3306bd8f8e653 # v1.3.5
+        id: find_pr
+        with:
+          # Can be "open", "closed", or "all".  Defaults to "open".
+          state: all
+      # This will echo "Your PR is 7", or be skipped if there is no current PR.
+      - run: echo "PR_NUMBER=${STEPS_FIND_PR_OUTPUTS_PR}"
+        if: success() && steps.find_pr.outputs.number
+        env:
+          PR: ${{ steps.find_pr.outputs.pr }}
+          STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
+      - name: Checkout OpenXLA
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Wait For Connection
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Configure GPU backend
+        if: ${{ matrix.job_info.platform == 'GPU' }}
+        run: |
+          ./configure.py --backend=CUDA --cuda_compiler=nvcc
+      - name: "Run build.py"
+        run: |
+          ./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
+      # Run the corresponding HLO tests based on platform
+      - name: Run HLO tests
+        run: |
+          bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
+          binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
+          runner_binary="${{ matrix.job_info.runner_binary }}"
+          test_hlo_file="${{ matrix.job_info.hlo_test_file }}"
+          platform="${{ matrix.job_info.platform }}"
+          output_file="$GITHUB_WORKSPACE/${test_hlo_file}_output.txt"
+          xspace_file="$GITHUB_WORKSPACE/xspace.pb"
+
+          echo "Running ${platform} test with binary: $binary_dir/$runner_binary"
+          pwd #print working directory
+
+          if [[ "$platform" == "CPU" ]]; then
+            $binary_dir/multihost_hlo_runner/$runner_binary --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
+            $binary_dir/compute_xspace_stats_main --input="$xspace_file" --device_type=CPU >> "$output_file"
+          elif [[ "$platform" == "GPU" ]]; then
+            $binary_dir/multihost_hlo_runner/$runner_binary --device_type=gpu --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
+            $binary_dir/compute_xspace_stats_main_gpu --input="$xspace_file" --device_type=GPU >> "$output_file"
+          else
+            echo "Unsupported platform: $platform"
+
+            exit 1
+          fi
+          cat "$output_file"
+          echo "Output written to: $output_file"
+
+          echo "XSpace file written to: $xspace_file"
+          echo "OUTPUT_FILE=$output_file" >> "$GITHUB_ENV"
+          echo "XSPACE_FILE=$xspace_file" >> "$GITHUB_ENV"
+
+      - name: Upload HLO test output to a GCS bucket
+        run: |
+          GCS_BUCKET="gs://openxla-postsubmit-transient"
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+          DATE_FOLDER=$(date +%Y%m%d)
+          FILENAME=$(basename "$OUTPUT_FILE")
+          COMMIT_SHA="${{ github.sha }}"
+          RUN_ID="${{ github.run_id }}"
+          PR_NUMBER="${STEPS_FIND_PR_OUTPUTS_PR}"
+          PLATFORM="${{ matrix.job_info.platform }}"
+
+          POOL="${{ matrix.job_info.pool }}"
+          PR_SUFFIX=""
+
+          if [[ -n "$PR_NUMBER" ]]; then
+            PR_SUFFIX="_pr_${PR_NUMBER}"
+          fi
+
+          GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${POOL}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${PLATFORM}_${FILENAME}"
+
+          echo "Uploading $OUTPUT_FILE to $GCS_BUCKET/$GCS_OBJECT_NAME"
+          ls -l "$OUTPUT_FILE" # Verify the file exists
+          gcloud storage cp "$OUTPUT_FILE" "$GCS_BUCKET/$GCS_OBJECT_NAME"
+        env:
+          STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
+      - name: Download Gemma Hlo Files
+        run: |
+          mkdir -p tmp_hlo
+          cd tmp_hlo
+
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
+          cd ..
+      - name: Run Gemma HLO tests
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+        run: |
+          bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
+          binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
+          runner_binary="${{ matrix.job_info.runner_binary }}"
+          test_hlo_file="tmp_hlo/gemma3_1b_flax_call.hlo"
+          platform="${{ matrix.job_info.platform }}"
+          output_file="$GITHUB_WORKSPACE/${test_hlo_file}_output.txt"
+          xspace_file="$GITHUB_WORKSPACE/xspace.pb"
+
+          if [[ "$platform" == "CPU" ]]; then
+            $binary_dir/multihost_hlo_runner/$runner_binary --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
+
+            $binary_dir/compute_xspace_stats_main --input="$xspace_file" --device_type=CPU >> "$output_file"
+
+          elif [[ "$platform" == "GPU" ]]; then
+            $binary_dir/multihost_hlo_runner/$runner_binary --device_type=gpu --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$xspace_file" "$test_hlo_file" > "$output_file"
+            $binary_dir/compute_xspace_stats_main_gpu --input="$xspace_file" --device_type=GPU >> "$output_file"
+          else
+            echo "Unsupported platform: $platform"
+
+            exit 1
+          fi
+          cat "$output_file"
+          echo "Output written to: $output_file"
+          echo "XSpace file written to: $xspace_file"
+          echo "OUTPUT_FILE=$output_file" >> "$GITHUB_ENV"
+
+          echo "XSPACE_FILE=$xspace_file" >> "$GITHUB_ENV"
+      - name: Upload gemma HLO test output to a GCS bucket
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+        run: |
+          GCS_BUCKET="gs://openxla-postsubmit-transient"
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+          DATE_FOLDER=$(date +%Y%m%d)
+          FILENAME="gemma3_1b_flax_call.hlo"
+          COMMIT_SHA="${{ github.sha }}"
+          RUN_ID="${{ github.run_id }}"
+
+          PR_NUMBER="${STEPS_FIND_PR_OUTPUTS_PR}"
+          PLATFORM="${{ matrix.job_info.platform }}"
+          POOL="${{ matrix.job_info.pool }}"
+          PR_SUFFIX=""
+          if [[ -n "$PR_NUMBER" ]]; then
+            PR_SUFFIX="_pr_${PR_NUMBER}"
+          fi
+
+          GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${POOL}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${PLATFORM}_${FILENAME}"
+
+          echo "Uploading $OUTPUT_FILE to $GCS_BUCKET/$GCS_OBJECT_NAME"
+          ls -l "$OUTPUT_FILE" # Verify the file exists
+          gcloud storage cp "$OUTPUT_FILE" "$GCS_BUCKET/$GCS_OBJECT_NAME"
+        env:
+          STEPS_FIND_PR_OUTPUTS_PR: ${{ steps.find_pr.outputs.pr }}
+
+      - name: Upload XSpace artifacts
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: xspace-artifacts-${{ matrix.job_info.pool }}-${{ matrix.job_info.platform }}
+          path: ${{ env.XSPACE_FILE }}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/benchmark_presubmit.yml

@@ -0,0 +1,154 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Presubmit Benchmarks
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  pull_request:
+    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: [
+          {
+            pool: "linux-x86-n2-128",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
+            bazel_arch_dir: "k8-opt",
+            platform: "CPU"
+          },
+          {
+            pool: "linux-x86-g2-16-l4-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
+            pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
+            bazel_arch_dir: "k8-opt",
+            platform: "GPU"
+          },
+        ]
+    name: ${{ matrix.job_info.pretty_name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 15
+    env:
+      OUTPUT_DIR: ${{ github.workspace }}/output
+    steps:
+      - name: Print GitHub Context
+        run: |
+          echo "GitHub SHA: ${{ github.sha }}"
+          echo "GitHub Ref: ${GITHUB_REF}"
+          echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
+          echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
+          echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
+          echo "GitHub Repository: ${{ github.repository }}"
+          echo "GitHub Run ID: ${{ github.run_id }}"
+          echo "GitHub Run Number: ${{ github.run_number }}"
+          echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
+          echo "GitHub Actor: ${GITHUB_ACTOR}"
+          echo "GitHub Event Name: ${{ github.event_name }}"
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "Pull Request Number: ${{ github.event.pull_request.number }}"
+            echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
+            echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
+          fi
+        env:
+          GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
+          GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
+      - name: Checkout OpenXLA
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Configure GPU backend
+        if: ${{ matrix.job_info.platform == 'GPU' }}
+        run: |
+          ./configure.py --backend=CUDA --cuda_compiler=nvcc
+      - name: "Run build.py"
+        run: |
+          ./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
+      - name: Download Gemma Hlo Files
+        run: |
+          mkdir -p "$OUTPUT_DIR/tmp_hlo"
+          cd "$OUTPUT_DIR/tmp_hlo"
+
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
+          cd ..
+      # Run the corresponding HLO tests based on platform
+      - name: Run HLO tests
+
+        run: |
+          BAZEL_ARCH_DIR="${{ matrix.job_info.bazel_arch_dir }}"
+          BINARY_DIR="./bazel-out/${BAZEL_ARCH_DIR}/bin/xla/tools"
+          OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_output.txt"
+
+          TEST_HLO_FILE="$OUTPUT_DIR/tmp_hlo/gemma3_1b_flax_call.hlo"
+          XSPACE_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_xspace.pb"
+          DEVICE_TYPE=""
+
+          if [[ ${{ matrix.job_info.platform }} == "CPU" ]]; then
+            DEVICE_TYPE="host"
+            HLO_RUNNER_BINARY="$BINARY_DIR/multihost_hlo_runner/hlo_runner_main"
+            COMPUTE_XSPACE_STATS_BINARY="$BINARY_DIR/compute_xspace_stats_main"
+          elif [[ ${{ matrix.job_info.platform }} == "GPU" ]]; then
+            DEVICE_TYPE="gpu"
+
+            HLO_RUNNER_BINARY="$BINARY_DIR/multihost_hlo_runner/hlo_runner_main_gpu"
+            COMPUTE_XSPACE_STATS_BINARY="$BINARY_DIR/compute_xspace_stats_main_gpu"
+          else
+            echo "Unsupported platform: ${{ matrix.job_info.platform }}"
+            exit 1
+          fi
+
+          echo "Running test with binary: $HLO_RUNNER_BINARY"
+          pwd #print working directory
+
+          $HLO_RUNNER_BINARY --device_type=$DEVICE_TYPE --use_spmd_partitioning --num_repeats=5 --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$TEST_HLO_FILE" > "$OUTPUT_FILE_PATH"
+          $COMPUTE_XSPACE_STATS_BINARY --input="$XSPACE_FILE_PATH" --device_type="${{ matrix.job_info.platform }}" >> "$OUTPUT_FILE_PATH"
+          cat "$OUTPUT_FILE_PATH"
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/benchmarks/README.md

@@ -0,0 +1,237 @@
+# Onboard New Microbenchmarks to OpenXLA
+
+This guide provides step-by-step instructions for contributing new
+microbenchmarks to OpenXLA microbenchmarking infrastructure.
+
+## Overview
+
+The OpenXLA microbenchmarking system is designed to automatically detect
+performance regressions and track performance trends across different hardware
+backends (CPU, GPU) in presubmit, postsubmit and nightly workflows. By adding
+your microbenchmark, you help ensure XLA's performance remains robust for your
+specific use cases.
+
+The process involves:
+
+1.  **Preparing your Benchmark Artifact:** Ensuring your HLO file is OSS-friendly.
+2.  **Defining the Benchmark Configuration:** Adding an entry to a benchmark registry file.
+3.  **Establishing a Baseline:** Adding initial performance thresholds if your benchmark will run in presubmit/postsubmit/nightly jobs.
+
+## Prerequisites
+
+*   **Benchmark Artifact (HLO):** You should have your benchmark ready as either
+    an HLO text file (`.hlo`) (Note: StableHLO MLIR text file (`.mlir`) will be
+    supported later).
+    *   For small artifacts, you can place them in `xla/tools/benchmarks/hlo/`.
+    *   **[Not yet supported]** For larger artifacts, upload them to a GCS
+        bucket (e.g., `gs://xla-benchmarking-temp/your-benchmark.hlo`) and
+        ensure it's publicly readable.
+*   **GitHub Access:** You'll need to create a Pull Request (PR) to the OpenXLA repository.
+
+## Step-by-Step Guide
+
+### Step 1: Prepare Your Benchmark Artifact
+
+Ensure your HLO file is ready and accessible.
+
+*   **Store in the XLA Repository**
+    1.  Place your `.hlo` file in the `xla/xla/tools/benchmarks/hlo/` directory.
+    2.  Make sure your hlo benchmarks run < 15min/20min/30min for
+        presubmit/postsubmit/nightly workflows.
+
+### Step 2: Define the Benchmark Configuration
+
+You'll need to add a new entry to a benchmark registry YAML file. For most
+community contributions, this will be
+`xla/xla/tools/benchmarks/registries/default_registry.yml`.
+
+Each benchmark configuration is a YAML object with the following key fields:
+
+*   `name`: A unique, descriptive name for your benchmark (e.g., `"my_model_attention_layer"`).
+*   `description`: A brief explanation of what the benchmark measures.
+*   `owner`: Your GitHub handle or relevant team alias (e.g., `"your-github-username@"`).
+*   `input_artifact`:
+    *   `input_format`: Currently we support `HLO_TEXT`, and `STABLEHLO_MLIR` will be supported in the future.
+    *   `artifact_path`: (If stored in repo) Relative path from `xla`, e.g., `xla/tools/benchmarks/hlo/my_new_benchmark.hlo`.
+    *   `artifact_gcs_bucket_path`: (If stored in GCS) Full GCS URL.
+*   `model_source_info`: A list of strings describing the origin of the benchmark (e.g., `["Gemma2 2B"]`).
+*   `hardware_targets`: A list defining on which hardware configurations this
+    benchmark should run. Each target has:
+    *   `hardware_category`: e.g., `GPU_L4`, `CPU_X86`, `GPU_B200`.
+    *   `topology`:
+        *   `num_hosts`: Number of hosts (default: 1).
+        *   `num_devices_per_host`: Number of devices per host (default: 1).
+        *   `multi_host`: `true` or `false`.
+    *   `multi_device`: `true` or `false`.
+    *   `target_metrics`: A list of metrics to collect, e.g., `[GPU_DEVICE_TIME, PEAK_GPU_MEMORY]`.
+    *   `run_frequencies`: When to run this benchmark, e.g., `[PRESUBMIT, POSTSUBMIT]`, `[SCHEDULED]`.
+*   `update_frequency_policy`: How often this benchmark definition should
+    be reviewed, e.g., `QUARTERLY`.
+*   `xla_compilation_flags` (Optional): List of XLA flags, e.g.,
+    `["--xla_gpu_enable_cudnn_fusion=false"]`.
+*   `runtime_flags` (Optional): List of flags for the
+    `multihost_hlo_runner`, e.g., `["--num_repeats=5"]`.
+*   `github_labels` (Optional): GitHub labels to manually trigger this
+    specific benchmark.
+
+**Example: Adding "gemma3\_1b\_flax\_sample\_loop" to `default_registry.yml`**
+
+```yaml
+# xla/xla/tools/benchmarks/registries/default_registry.yml
+benchmarks: [
+  # ... existing benchmarks ...
+  {
+    name: "gemma3_1b_flax_sample_loop"
+    description: "Gemma3 1B in Flax Sample Loop."
+    owner: "company-A@" # Replace with your GitHub handle or team
+    input_artifact: {
+      input_format: HLO_TEXT, # Or STABLEHLO_MLIR
+      artifact_path: "xla/tools/benchmarks/hlo/gemma3_1b_flax_sample_loop.hlo"
+      # Option 2 (for large hlo):
+      #`artifact_gcs_bucket_path`: (If stored in GCS) Full GCS URL (not supported yet).
+    }
+    model_source_info: ["Gemma3 1B"] # Describe the source of your HLO
+    hardware_targets: [{
+      hardware_category: GPU_L4
+      topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
+      target_metrics: [GPU_DEVICE_TIME, GPU_DEVICE_MEMCPY_TIME]
+      run_frequencies: [PRESUBMIT, POSTSUBMIT] # Run on PRs in presubmit and postsubmit
+      runtime_flags: ["--num_repeats=5"] # Example: run 5 times to reduce noise
+    },
+    {
+      hardware_category: CPU_X86
+      topology: { num_hosts: 1, num_devices_per_host: 1, multi_host: false, multi_device: false }
+      target_metrics: [CPU_TIME, WALL_TIME]
+      run_frequencies: [PRESUBMIT] # Only run on PRs for presubmit
+      runtime_flags: ["--num_repeats=5"]
+    }]
+    update_frequency_policy: QUARTERLY # Review this benchmark definition quarterly
+  }
+]
+```
+
+### Step 3: Establish a Baseline
+
+1.  **Determine Baseline Values:**
+    *   The best way to get initial baseline values is to run your benchmark
+    manually on the target hardware or let it run once in postsubmit after your
+    initial PR (without presubmit blocking) is merged.
+    *   Promote your benchmarks from postsubmit to presubmit once you get
+        stable results for baseline values.
+    *   Run the benchmark multiple times (e.g., using `--num_repeats=5` or
+        more) and take the median or a stable average.
+    *   Benchmarks must run < 15min for presubmit, < 20min for postsubmit and < 30min for nightly.
+
+2.  **Add to `presubmit_baseline.yml`:**
+    Edit the file `xla/xla/tools/benchmarks/baseline/presubmit_baseline.yml`. The
+    key for each entry is the `config_id`.
+
+    **Note on `config_id` generation:**
+    `config_id` follows the below pattern:
+    `"{benchmark_name}_{hardware_category_simplified}_{topology_simplified}_{workflow_type}"`.
+    *   `hardware_category_simplified`: e.g., `l4` (for `GPU_L4`), `b200` (for `GPU_B200`), `x86` (for `CPU_X86`).
+    *   `topology_simplified`: e.g., `1h1d` for 1 host, 1 device. <!-- disableFinding(LINE_OVER_80) -->
+    *   `workflow_type`: e.g., `presubmit`, `postsubmit`, `scheduled`.
+
+    If unsure, you can check the GitHub Actions workflow logs for the
+    `generate_benchmark_matrices.py` script output, which will show the generated
+    `config_id`s.
+
+    For each metric you want to track in presubmit (must be in `target_metrics`
+    in the registry):
+
+    *   `baseline_ms`: The baseline performance in milliseconds.
+    *   `threshold`: The maximum allowed regression percentage (e.g.,
+        `0.30` for 30%).
+    *   **Note on metrics:** Currently, we support `GPU_DEVICE_TIME`,
+        `GPU_DEVICE_MEMCPY_TIME` for GPU, and `CPU_TIME`, `WALL_TIME` for CPU.
+
+**Example: Adding baseline for "gemma3\_1b\_flax\_sample\_loop"**
+
+Assuming the `name` is `"gemma3_1b_flax_sample_loop"`:
+
+*   For `GPU_L4`, `1` host, `1` device: `config_id` becomes
+    `gemma3_1b_flax_sample_loop_l4_1h1d_presubmit`
+*   For `CPU_X86`, `1` host, `1` device: `config_id` becomes
+    `gemma3_1b_flax_sample_loop_x86_1h1d_presubmit`
+
+```yaml
+# xla/xla/tools/benchmarks/baseline/presubmit_baseline.yml
+{
+  # ... existing baselines ...
+
+  "gemma3_1b_flax_sample_loop_l4_1h1d_presubmit": {
+    "GPU_DEVICE_TIME": {
+      "baseline_ms": 4,  # Your measured baseline
+      "threshold": 0.30
+    },
+    "GPU_DEVICE_MEMCPY_TIME": {
+      "baseline_ms": 10, # Your measured baseline
+      "threshold": 0.30
+    }
+  },
+  "gemma3_1b_flax_sample_loop_x86_1h1d_presubmit": {
+    "CPU_TIME": {
+      "baseline_ms": 8000, # Your measured baseline
+      "threshold": 0.30
+    },
+    "WALL_TIME": {
+      "baseline_ms": 1300, # Your measured baseline
+      "threshold": 0.30
+    }
+  }
+}
+```
+
+### Step 4: Create a Pull Request
+
+1.  Commit your changes:
+    *   The HLO file (if added to the repo).
+    *   The updated benchmark registry file (e.g., `default_registry.yml`).
+    *   The updated `presubmit_baseline.yml` (if applicable).
+2.  Push your branch and open a Pull Request against the `openxla/xla` main branch.
+3.  A member of the OpenXLA repository or organization will need to review your
+    PR for safety before the CI system is invoked.
+     * **Note**: This step happens
+    automatically for organization members and most Googlers, but require
+    manual review for external contributors.
+4.  Once approved, the CI system will pick up your new benchmark configuration.
+    *   If it's a `PRESUBMIT` benchmark, it will run against your PR and check for regressions based on the baseline you provided.
+    *   If it's `POSTSUBMIT` or `SCHEDULED`, it will run after your PR is merged.
+5.  Monitor the CI checks. If the presubmit check fails due to your new
+    benchmark (e.g., performance is significantly different from your initial
+    baseline), you might need to adjust the baseline values in
+    `presubmit_baseline.yml` and update your PR.
+
+## Best Practices
+
+*   **Establish Baselines First:** Since, a baseline value per metric is required, always add the benchmark with only
+    `POSTSUBMIT` or `SCHEDULED` frequency first to establish stable baseline
+    values. Once it runs a few times and you have stable performance data, you
+    can add `PRESUBMIT` and the corresponding baseline entry in a follow-up PR.
+*   **Meaningful Names and Descriptions:** Make it easy for others to
+    understand what your benchmark does.
+*   **Targeted Metrics:** Only include relevant metrics in `target_metrics`.
+*   **Noise Reduction:** Use `runtime_flags: ["--num_repeats=X"]` (e.g., X=5 or
+    10) to run the benchmark multiple times within a single execution, which
+    helps in getting more stable measurements. The runner typically reports the
+    median or average.
+*   **Keep Baselines Updated:** If your benchmark's performance characteristics
+    change significantly (due to XLA improvements or changes in the benchmark
+    itself), the baseline values in `presubmit_baseline.yml` will need to be
+    updated. This is usually done by the benchmark owner or XLA maintainers.
+
+## Troubleshooting
+
+*   **Workflow Failures:** Check the GitHub Actions logs for detailed error
+    messages. The logs for the "Compare Benchmarks" step are particularly useful
+    for presubmit issues.
+*   **Incorrect `config_id`:** If your presubmit benchmark isn't being picked up
+    or matched to a baseline, double-check the `config_id` format in
+    `presubmit_baseline.yml`.
+*   **Performance Fluctuations:** Microbenchmarks can be sensitive to noise.
+    Ensure you're using `--num_repeats` and that your baseline reflects typical
+    performance.
+
+If you encounter issues, feel free to ask for help on the OpenXLA communication
+channels or tag the juliagmt-google@ on your PR.

third_party/xla/.github/workflows/benchmarks/build_binaries.sh

@@ -0,0 +1,135 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/build_binaries.sh
+# TODO(juliagmt): convert this to a python script.
+# Builds HLO runner and stats computation binaries using build.py.
+# Expects HARDWARE_CATEGORY (matching enum values like CPU_X86, GPU_L4, etc.)
+# and OUTPUT_DIR environment variables to be set.
+# Outputs: runner_binary, stats_binary, device_type_flag to GITHUB_OUTPUT
+#!/bin/bash
+set -e # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+# set -o pipefail # Causes pipelines to fail if any command fails (see Run script)
+
+echo "--- Configuring and Building Binaries ---"
+echo "Workspace: $(pwd)"
+echo "Hardware Category: $HARDWARE_CATEGORY"
+echo "Output Directory (for profile.json.gz): $OUTPUT_DIR"
+
+# Sanitize HARDWARE_CATEGORY for use in filenames.
+# Replaces non-alphanumeric with underscore, converts to lowercase.
+HW_CATEGORY_SLUG=$(echo "${HARDWARE_CATEGORY:-UNSPECIFIED}" | tr '[:upper:]' '[:lower:]' | sed 's/[^A-Z0-9_]/_/g')
+
+# --- 1. Configure Backend (using ./configure.py if available) ---
+# This part can remain, as build.py doesn't currently handle `configure.py`
+configure_backend() {
+  echo "Configuring backend using ./configure.py if present..."
+  if [ ! -f "./configure.py" ]; then
+    echo "INFO: ./configure.py not found. Skipping configuration step."
+    return
+  fi
+
+  local hw_category_upper_for_configure
+  hw_category_upper_for_configure=$(echo "${HARDWARE_CATEGORY:-UNSPECIFIED}" | tr '[:lower:]' '[:upper:]')
+
+  case "$hw_category_upper_for_configure" in
+    CPU_X86 | CPU_ARM64)
+      echo "Running: ./configure.py --backend=CPU"
+      ./configure.py --backend=CPU || echo "INFO: CPU Configure script failed or is not applicable."
+      ;;
+    GPU_L4 | GPU_B200)
+      echo "Running: ./configure.py --backend=CUDA --cuda_compiler=nvcc"
+      ./configure.py --backend=CUDA --cuda_compiler=nvcc || echo "INFO: GPU Configure script failed or is not applicable."
+      ;;
+    *)
+      echo "INFO: Unknown hardware category '$hw_category_upper_for_configure'"
+      ;;
+  esac
+  echo "Configuration step finished."
+}
+
+# --- 2. Main Build Logic using build.py ---
+declare BAZEL_BIN_DIR="bazel-bin"
+declare runner_binary_path=""
+declare stats_binary_path=""
+declare device_type_flag_value=""
+
+configure_backend
+
+echo "Building with build.py for HARDWARE_CATEGORY: $HARDWARE_CATEGORY"
+
+BUILD_TYPE=""
+case "$HARDWARE_CATEGORY" in
+  CPU_X86)
+    BUILD_TYPE="XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main"
+    device_type_flag_value="host"
+    ;;
+  CPU_ARM64)
+    BUILD_TYPE="XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main"
+    device_type_flag_value="host"
+    ;;
+  GPU_L4)
+    BUILD_TYPE="XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS" # Or _48_VCPU if that's the more common
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main_gpu"
+    device_type_flag_value="gpu"
+    ;;
+  GPU_B200)
+    BUILD_TYPE="XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS"
+    runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu"
+    stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main_gpu"
+    device_type_flag_value="gpu"
+    ;;
+  *)
+    echo "::error::Unsupported HARDWARE_CATEGORY: '$HARDWARE_CATEGORY'. This script is configured to handle specific values from the HardwareCategory enum (CPU_X86, CPU_ARM64, GPU_L4, GPU_B200)."
+    exit 1
+    ;;
+esac
+
+echo "Executing build with build.py for build type: $BUILD_TYPE"
+# Run build.py with the determined build type
+python3 build_tools/ci/build.py --build="$BUILD_TYPE" || {
+  echo "::error::build.py failed for $BUILD_TYPE!"
+  exit 1
+}
+
+
+# --- 3. Verify Binaries and Set Outputs ---
+echo "Verifying binary existence..."
+if [ -z "$runner_binary_path" ] || [ ! -f "$runner_binary_path" ]; then
+  echo "::error::Runner binary path not set or binary '$runner_binary_path' not found after build for $HARDWARE_CATEGORY!"
+  exit 1
+fi
+if [ -z "$stats_binary_path" ] || [ ! -f "$stats_binary_path" ]; then
+  echo "::error::Stats binary path not set or binary '$stats_binary_path' not found after build for $HARDWARE_CATEGORY!"
+  exit 1
+fi
+echo "Binaries verified: $runner_binary_path, $stats_binary_path"
+
+echo "Setting step outputs..."
+{
+  echo "runner_binary=$runner_binary_path"
+  echo "stats_binary=$stats_binary_path"
+  echo "device_type_flag=$device_type_flag_value"
+} >> "$GITHUB_OUTPUT"
+
+echo "  Runner binary: $runner_binary_path"
+echo "  Stats binary: $stats_binary_path"
+echo "  Device type flag: $device_type_flag_value"
+echo "--- Build Script Finished ---"

third_party/xla/.github/workflows/benchmarks/compare_with_baseline.py

@@ -0,0 +1,340 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Compares benchmark results to predefined baselines and checks for regressions."""
+import argparse
+import json
+import os
+import sys
+import yaml
+
+
+def load_results_data(results_json_file):
+  """Loads and parses the results JSON file."""
+  try:
+    with open(results_json_file, "r") as f:
+      results_data = json.load(f)
+  except json.JSONDecodeError as e:
+    print(
+        f"::error file={results_json_file}::Failed to parse results JSON: {e}"
+    )
+    sys.exit(1)
+  except OSError as e:  # Catch other potential file reading errors
+    print(f"::error file={results_json_file}::Error reading results JSON: {e}")
+    sys.exit(1)
+  return results_data
+
+
+def load_baseline_data(baseline_yaml_file):
+  """Loads and parses the baseline YAML file."""
+  try:
+    with open(baseline_yaml_file, "r") as f:
+      baseline_data_full = yaml.safe_load(f)
+  except yaml.YAMLError as e:
+    print(
+        f"::error file={baseline_yaml_file}::Failed to parse baseline YAML: {e}"
+    )
+    sys.exit(1)
+  except OSError as e:  # Catch other potential file reading errors
+    print(
+        f"::error file={baseline_yaml_file}::Error reading baseline YAML: {e}"
+    )
+    sys.exit(1)
+  return baseline_data_full
+
+
+def validate_baseline_data(baseline_data_full, config_id, baseline_yaml_file):
+  """Validates the loaded baseline data and extracts config-specific baselines."""
+  if not baseline_data_full or not isinstance(baseline_data_full, dict):
+    print(
+        f"::warning file={baseline_yaml_file}::Baseline YAML is empty or"
+        " not a dictionary. Skipping comparison."
+    )
+    sys.exit(0)
+
+  if config_id not in baseline_data_full:
+    print(
+        f"::notice::No baseline found for config_id '{config_id}' in"
+        f" {baseline_yaml_file}. Skipping comparison."
+    )
+    sys.exit(0)
+
+  config_baselines = baseline_data_full[config_id]
+  if not isinstance(config_baselines, dict):
+    print(
+        f"::warning file={baseline_yaml_file},title=Invalid Baseline"
+        f" Structure::Baseline entry for '{config_id}' is not a"
+        " dictionary. Skipping."
+    )
+    sys.exit(0)
+  return config_baselines
+
+
+def compare_metrics(
+    results_data, config_baselines, results_json_file, config_id
+):
+  """Compares metrics from results to baselines and reports findings."""
+  regressions_found = False
+  summary_messages = ["\n--- Comparison Summary ---"]
+
+  # --- IMPORTANT: Metric Extraction Logic ---
+  # This assumes your results.json (produced by compute_xspace_stats_main or
+  # fallback)
+  # has a structure like:
+  # {
+  #   ...,
+  #   "metrics": {
+  #     "GPU_DEVICE_TIME": { "value": 150.0, "unit": "ms", ... },
+  #     "GPU_DEVICE_MEMCPY_TIME": { "value": 1.2, "unit": "ms", ... },
+  #     "PEAK_GPU_MEMORY": { "value": 12.45, "unit": "GB", ... },
+  #     ...
+  #   }
+  # }
+  # If your actual results.json structure is different, you MUST adapt this
+  # section.
+  actual_metrics_container = results_data.get("metrics")
+  metric_to_baseline_key = {
+      "WALL_TIME": "baseline_ms",
+      "GPU_DEVICE_TIME": "baseline_ms",
+      "GPU_DEVICE_MEMCPY_TIME": "baseline_ms",
+      "CPU_TIME": "baseline_ms",
+      "PEAK_CPU_MEMORY": "baseline_gb",
+      "PEAK_GPU_MEMORY": "baseline_gb",
+  }
+
+  if not actual_metrics_container or not isinstance(
+      actual_metrics_container, dict
+  ):
+    summary_messages.append(
+        "::warning title=Missing Metrics in Results::'metrics' key not found"
+        f" or not a dictionary in '{results_json_file}'. Cannot perform"
+        " comparison."
+    )
+    # Depending on your policy, if metrics are always expected,
+    # this could be sys.exit(1)
+    # For now, it will skip comparisons and pass if no metrics are found.
+    print("\n".join(summary_messages))
+    sys.exit(0)  # Exit cleanly if no metrics to compare
+
+  for metric_name, baseline_info in config_baselines.items():
+    if (
+        not isinstance(baseline_info, dict)
+        or "threshold" not in baseline_info
+        or not {"baseline_ms", "baseline_gb"}.intersection(baseline_info.keys())
+    ):
+      summary_messages.append(
+          f"::warning title=Malformed Baseline::Metric '{metric_name}' in"
+          f" baseline for '{config_id}' is missing 'threshold', or is missing"
+          " both 'baseline_ms' and 'baseline_gb', or is not structured as a"
+          " dictionary. Skipping."
+      )
+      continue
+
+    baseline_key = metric_to_baseline_key.get(metric_name)
+    if not baseline_key:
+      summary_messages.append(
+          f"::warning title=Unsupported Metric::Metric '{metric_name}' is not"
+          " supported by this script. Skipping."
+      )
+      continue
+
+    try:
+      baseline_value = float(baseline_info[baseline_key])
+      threshold_percentage = float(baseline_info["threshold"])
+    except ValueError:
+      summary_messages.append(
+          f"::warning title=Invalid Baseline Value::Metric '{metric_name}' in"
+          f" baseline for '{config_id}' has non-numeric '{baseline_key}' or"
+          " 'threshold'. Skipping."
+      )
+      continue
+
+    # Extract the actual metric value from results.json
+    actual_metric_entry = actual_metrics_container.get(metric_name)
+
+    if not actual_metric_entry or "value" not in actual_metric_entry:
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value or 'value' key not found in"
+          " results, or not a dictionary. Skipping."
+      )
+      # For debugging:
+      # available_keys = (
+      #     list(actual_metrics_container.keys())
+      #     if actual_metrics_container
+      #     else "None"
+      # )
+      # summary_messages.append(
+      #   f"  Available metric keys in results 'metrics' object:"
+      #   f" {available_keys}"
+      # )
+      continue
+
+    try:
+      actual_value = float(actual_metric_entry["value"])
+    except (ValueError, TypeError):
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value"
+          f" '{actual_metric_entry['value']}' is not a valid number."
+          " Skipping."
+      )
+      continue
+
+    actual_unit = actual_metric_entry.get("unit")
+    if not actual_unit:
+      summary_messages.append(
+          f"Metric '{metric_name}': Actual value unit is not specified."
+          " Skipping."
+      )
+      continue
+
+    summary_messages.append(f"\nComparing metric: {metric_name}")
+    summary_messages.append(f"  Actual Value: {actual_value:.3f} {actual_unit}")
+    summary_messages.append(
+        f"  Baseline Value: {baseline_value:.3f} {actual_unit}"
+    )
+    summary_messages.append(
+        f"  Allowed Threshold: {threshold_percentage*100:.1f}%"
+    )
+
+    # Higher value is worse for time-based metrics
+    allowed_upper_bound = baseline_value * (1.0 + threshold_percentage)
+    summary_messages.append(
+        "  Allowed Upper Bound (Baseline * (1 + Threshold)):"
+        f" {allowed_upper_bound:.3f} {actual_unit}"
+    )
+
+    if actual_value > allowed_upper_bound:
+      percentage_diff = 0.0
+      if (
+          abs(baseline_value) > 1e-9
+      ):  # Avoid division by zero for very small baselines
+        percentage_diff = (
+            (actual_value - baseline_value) / baseline_value
+        ) * 100.0
+      elif (
+          actual_value > 0
+      ):  # If baseline is effectively zero, any positive value is infinitely
+        # worse.
+        percentage_diff = float("inf")
+
+      # Use GitHub Actions error annotation for better visibility
+      error_title = f"REGRESSION: {metric_name}"
+      error_details = (
+          f"Value {actual_value:.3f} {actual_unit} is {percentage_diff:.2f}%"
+          f" worse than baseline {baseline_value:.3f} {actual_unit}. Exceeds"
+          f" threshold of {threshold_percentage*100:.1f}% (max allowed:"
+          f" {allowed_upper_bound:.3f} {actual_unit})."
+      )
+      summary_messages.append(
+          "  ::error"
+          f" file={results_json_file},title={error_title}::{error_details}"
+      )
+      regressions_found = True
+    else:
+      summary_messages.append(
+          f"  Metric '{metric_name}' is within threshold. PASSED."
+      )
+  return regressions_found, summary_messages
+
+
+def main():
+  parser = argparse.ArgumentParser(
+      description=(
+          "Compare benchmark results to baselines and fail on regression."
+      ),
+      formatter_class=argparse.RawTextHelpFormatter,
+  )  # For better help text formatting
+  parser.add_argument(
+      "--results-json-file",
+      required=True,
+      help=(
+          "Path to the benchmark results JSON file (e.g., output/results.json)."
+      ),
+  )
+  parser.add_argument(
+      "--baseline-yaml-file",
+      required=True,
+      help=(
+          "Path to the baseline YAML file (e.g.,"
+          " xla/tools/benchmarks/baseline/presubmit_baseline.yml)."
+      ),
+  )
+  parser.add_argument(
+      "--config-id",
+      required=True,
+      help=(
+          "The configuration ID for the current benchmark run. \nThis ID must"
+          " exactly match a top-level key in the baseline YAML file. \nExample:"
+          " 'gemma3_1b_flax_call_gpu_b200_1_host_1_device'"
+      ),
+  )
+
+  args = parser.parse_args()
+
+  print("--- Benchmark Baseline Comparison ---")
+  print(f"Results JSON: {args.results_json_file}")
+  print(f"Baseline YAML: {args.baseline_yaml_file}")
+  print(f"Config ID for Baseline Lookup: {args.config_id}")
+
+  if not os.path.exists(args.results_json_file):
+    print(
+        f"::error file={args.results_json_file}::Results JSON file not found."
+    )
+    sys.exit(1)
+
+  if not os.path.exists(args.baseline_yaml_file):
+    print(
+        f"::notice file={args.baseline_yaml_file}::Baseline YAML file not"
+        " found. Skipping comparison."
+    )
+    # Exiting 0 because no baseline means no regression to detect.
+    # If a baseline is mandatory, this could be sys.exit(1).
+    sys.exit(0)
+
+  results_data = load_results_data(args.results_json_file)
+  print("\nLoaded Results Data:")
+  print(json.dumps(results_data, indent=2))
+
+  baseline_data_full = load_baseline_data(args.baseline_yaml_file)
+  config_baselines = validate_baseline_data(
+      baseline_data_full, args.config_id, args.baseline_yaml_file
+  )
+  print("\nLoaded Baseline Data for this config_id:")
+  print(json.dumps(config_baselines, indent=2))
+
+  regressions_found, summary_messages = compare_metrics(
+      results_data,
+      config_baselines,
+      args.results_json_file,
+      args.config_id,
+  )
+
+  print("\n".join(summary_messages))
+
+  if regressions_found:
+    print(
+        "\n::error::One or more benchmark metrics regressed beyond the allowed"
+        " threshold. Failing the check."
+    )
+    sys.exit(1)  # Exit with error code 1 to fail the GitHub Actions step
+  else:
+    print(
+        "\nAll benchmark metrics are within allowed thresholds (or no"
+        " applicable baselines found)."
+    )
+    sys.exit(0)  # Exit with success
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/.github/workflows/benchmarks/prepare_artifact.sh

@@ -0,0 +1,86 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/prepare_artifact.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+set -e # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+
+echo "--- prepare_artifact.sh (Self-creating directory version) ---"
+echo "SCRIPT: Current PWD: $(pwd)"
+echo "SCRIPT: GITHUB_WORKSPACE is: $GITHUB_WORKSPACE"
+echo "SCRIPT: Intended OUTPUT_DIR is: $OUTPUT_DIR"
+
+# Create the directory HERE, inside this script, right before using it.
+echo "SCRIPT: Ensuring directory '$OUTPUT_DIR' exists by creating it with mkdir -p."
+mkdir -p "$OUTPUT_DIR"
+
+# Verify creation immediately
+echo "SCRIPT: Verifying directory '$OUTPUT_DIR' after mkdir with 'ls -ld':"
+ls -ld "$OUTPUT_DIR" || echo "SCRIPT: 'ls -ld ""$OUTPUT_DIR""' FAILED even after mkdir in script!"
+
+# Now, check with [ -d ... ]
+if [ ! -d "$OUTPUT_DIR" ]; then
+  echo "::error::SCRIPT: Output directory '$OUTPUT_DIR' STILL NOT found with [ -d ... ] even after mkdir in this script."
+  echo "SCRIPT: Listing parent directory '$(dirname "$OUTPUT_DIR")' using 'ls -la':"
+  ls -la "$(dirname "$OUTPUT_DIR")" || echo "SCRIPT: Failed to list parent directory."
+  exit 1
+else
+  echo "SCRIPT: Output directory '$OUTPUT_DIR' IS now found with [ -d ... ]."
+fi
+
+# --- Original script logic from here ---
+echo "--- Preparing Artifact (main logic) ---"
+
+ARTIFACT_FILE_NAME=$(basename "$ARTIFACT_LOCATION")
+LOCAL_ARTIFACT_PATH="$OUTPUT_DIR/$ARTIFACT_FILE_NAME"
+
+echo "Target local path: ${LOCAL_ARTIFACT_PATH}"
+
+if [ "$IS_GCS_ARTIFACT" == "true" ]; then
+  echo "Downloading GCS artifact from: $ARTIFACT_LOCATION"
+   if ! command -v wget &> /dev/null; then
+     echo "::error::wget command not found in container. Cannot download GCS artifact."
+     exit 1
+   fi
+
+   wget -q -nv -O "$LOCAL_ARTIFACT_PATH" "$ARTIFACT_LOCATION"
+   WGET_EXIT_CODE=$?
+    if [ $WGET_EXIT_CODE -ne 0 ]; then
+      echo "::error::wget failed to download GCS artifact from $ARTIFACT_LOCATION (Exit code: $WGET_EXIT_CODE)"
+      rm -f "$LOCAL_ARTIFACT_PATH" # Clean up partial file
+      exit $WGET_EXIT_CODE
+    fi
+   echo "GCS artifact downloaded."
+else
+   REPO_ARTIFACT_PATH="$GITHUB_WORKSPACE/$ARTIFACT_LOCATION" # ARTIFACT_LOCATION is the relative repo path here
+   echo "Copying local artifact from workspace path: $REPO_ARTIFACT_PATH (IS_GCS_ARTIFACT was false)"
+    if [ ! -f "$REPO_ARTIFACT_PATH" ]; then
+       echo "::error::Local artifact not found at repository path: $REPO_ARTIFACT_PATH"
+       exit 1
+    fi
+    cp -v "$REPO_ARTIFACT_PATH" "$LOCAL_ARTIFACT_PATH" || exit 1 # Exit if copy fails
+   echo "Local artifact copied successfully."
+fi
+
+# Verify the final destination file exists
+if [ ! -f "$LOCAL_ARTIFACT_PATH" ]; then
+   echo "::error::Final artifact file not found at destination: $LOCAL_ARTIFACT_PATH"
+   exit 1
+fi
+echo "Artifact successfully prepared at $LOCAL_ARTIFACT_PATH."
+
+echo "artifact_local_path=$LOCAL_ARTIFACT_PATH" >> "$GITHUB_OUTPUT"
+echo "--- Artifact Prep Finished ---"

third_party/xla/.github/workflows/benchmarks/run_benchmark.sh

@@ -0,0 +1,276 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/run_benchmark.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+
+# This script encapsulates the logic to run a benchmark and generate results.json.
+
+set -u # Treat unset variables as an error when substituting.
+# IMPORTANT: pipefail is handled specifically around the runner command.
+set -e # Exit on errors, EXCEPT where explicitly handled.
+
+echo "--- Running Benchmark Script ---"
+echo "Benchmark Name: $BENCHMARK_NAME"
+echo "Config ID: $CONFIG_ID"
+echo "Hardware Category: $HARDWARE_CATEGORY"
+echo "Output Directory: $OUTPUT_DIR"
+echo "Runner Binary: $RUNNER_BINARY"
+echo "Stats Binary: $STATS_BINARY"
+echo "Device Type Flag: $DEVICE_TYPE_FLAG"
+echo "Local Artifact Path: $LOCAL_ARTIFACT_PATH"
+echo "Input Format: $INPUT_FORMAT"
+echo "XLA Flags JSON: $XLA_FLAGS_JSON"
+echo "Runtime Flags JSON: $RUNTIME_FLAGS_JSON"
+echo "Commit SHA: $COMMIT_SHA"
+echo "Workflow Run ID: $WORKFLOW_RUN_ID"
+
+
+# --- Validate Inputs ---
+if [ -z "$LOCAL_ARTIFACT_PATH" ] || [ ! -f "$LOCAL_ARTIFACT_PATH" ]; then echo "::error::LOCAL_ARTIFACT_PATH path is invalid or file not found: '$LOCAL_ARTIFACT_PATH'"; exit 1; fi
+if [ -z "$RUNNER_BINARY" ] || [ ! -x "$RUNNER_BINARY" ]; then echo "::error::RUNNER_BINARY path is invalid or file not executable: '$RUNNER_BINARY'"; exit 1; fi
+if [ -z "$DEVICE_TYPE_FLAG" ]; then echo "::error::DEVICE_TYPE_FLAG is empty"; exit 1; fi
+if [ -z "$STATS_BINARY" ] || [ ! -x "$STATS_BINARY" ]; then echo "::error::STATS_BINARY path is invalid or file not executable: '$STATS_BINARY'"; exit 1; fi
+if ! command -v jq &> /dev/null; then echo "::error::jq command not found."; exit 1; fi
+
+RUNNER_STDOUT_FILE="$OUTPUT_DIR/runner_stdout.txt"
+XSPACE_FILE_PATH="$OUTPUT_DIR/xspace.pb"
+RESULTS_JSON_FILE="$OUTPUT_DIR/results.json"
+
+# --- Prepare flags ---
+declare -a xla_flags_array=()
+declare -a runtime_flags_array=()
+
+# Use JQ to safely parse JSON and populate bash arrays
+if echo "$XLA_FLAGS_JSON" | jq -e '. | arrays and length > 0' > /dev/null; then
+    mapfile -t xla_flags_array < <(echo "$XLA_FLAGS_JSON" | jq -r '.[]')
+fi
+if echo "$RUNTIME_FLAGS_JSON" | jq -e '. | arrays and length > 0' > /dev/null; then
+   mapfile -t runtime_flags_array < <(echo "$RUNTIME_FLAGS_JSON" | jq -r '.[]')
+fi
+
+needs_xspace_dump_flag=true # By default, we enable profiler and xspace dump.
+
+# --- Build Runner Command ---
+declare -a runner_command_array=(
+    "$RUNNER_BINARY"
+    "--device_type=$DEVICE_TYPE_FLAG"
+)
+if [ ${#runtime_flags_array[@]} -gt 0 ]; then runner_command_array+=("${runtime_flags_array[@]}"); fi
+if [ ${#xla_flags_array[@]} -gt 0 ]; then runner_command_array+=("${xla_flags_array[@]}"); fi
+if $needs_xspace_dump_flag; then
+   runner_command_array+=("--xla_gpu_dump_xspace_to=$XSPACE_FILE_PATH")
+fi
+runner_command_array+=("$LOCAL_ARTIFACT_PATH")
+
+# --- Execute Runner ---
+echo "Executing HLO Runner command:"
+printf "%q " "${runner_command_array[@]}"; echo
+
+set +e # Disable exit-on-error temporarily to capture exit code
+set -o pipefail # Ensure tee doesn't mask the runner's exit code
+"${runner_command_array[@]}" 2>&1 | tee "$RUNNER_STDOUT_FILE"
+RUNNER_EXIT_CODE=${PIPESTATUS[0]}
+set +o pipefail
+set -e # Re-enable exit-on-error
+
+echo "Runner stdout/stderr saved to $RUNNER_STDOUT_FILE"
+echo "Runner exited with code: $RUNNER_EXIT_CODE"
+
+
+# --- Process Stats and Generate results.json ---
+STATS_RUN_SUCCESSFUL=false
+METRICS_JSON_CONTENT="{}" # Initialize as an empty JSON object
+
+if [ -f "$XSPACE_FILE_PATH" ] && [ $RUNNER_EXIT_CODE -eq 0 ]; then
+  echo "XSpace file found. Running compute_xspace_stats_main..."
+  STATS_PLATFORM_TYPE=$([[ "$HARDWARE_CATEGORY" == GPU* ]] && echo "GPU" || echo "CPU")
+
+  # Capture the output of compute_xspace_stats_main to parse it
+  # Do not write its output directly to results.json yet
+  echo "Executing Stats command and capturing its output:"
+
+  set +e # Temporarily disable exit-on-error for stats command
+  STATS_OUTPUT=$("$STATS_BINARY" --input="$XSPACE_FILE_PATH" --device_type="$STATS_PLATFORM_TYPE" 2>&1)
+  STATS_EXIT_CODE=$?
+  set -e
+
+  echo "compute_xspace_stats_main output:"
+  echo "$STATS_OUTPUT"
+  echo "compute_xspace_stats_main exited with code: $STATS_EXIT_CODE"
+
+  # Append stats tool's raw output to the main runner log for complete record
+  echo -e "\n--- compute_xspace_stats_main Raw Output ---" >> "$RUNNER_STDOUT_FILE"
+  echo "$STATS_OUTPUT" >> "$RUNNER_STDOUT_FILE"
+  echo "--- End compute_xspace_stats_main Raw Output ---" >> "$RUNNER_STDOUT_FILE"
+
+
+  if [ $STATS_EXIT_CODE -eq 0 ]; then
+    STATS_RUN_SUCCESSFUL=true
+    metrics_obj_str="{"
+    first_metric=true
+
+    # Process each line of STATS_OUTPUT
+    while IFS=':' read -r key value; do
+        # Trim leading/trailing whitespace from key and value
+        key=$(echo "$key" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+        value=$(echo "$value" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
+
+        # Sanitize base metric key (e.g., "Device Time" -> "DEVICE_TIME")
+        base_metric_key=$(echo "$key" | tr ' ' '_' | tr '[:lower:]' '[:upper:]')
+        final_metric_key=""
+
+        # Determine the final metric key based on HARDWARE_CATEGORY for baseline matching
+        if [[ "$HARDWARE_CATEGORY" == GPU* ]]; then
+            if [[ "$base_metric_key" == "DEVICE_TIME" ]]; then
+                final_metric_key="GPU_DEVICE_TIME"
+            elif [[ "$base_metric_key" == "DEVICE_MEMCPY_TIME" ]]; then
+                final_metric_key="GPU_DEVICE_MEMCPY_TIME"
+            elif  [[ "$base_metric_key" == "PEAK_MEMORY" ]]; then
+                final_metric_key="PEAK_GPU_MEMORY"
+            # Add other specific GPU mappings here if needed
+            # else
+            #    final_metric_key="GPU_${base_metric_key}" # Generic prefix
+            else
+                final_metric_key="$base_metric_key" # If no specific GPU mapping, use base
+            fi
+        elif [[ "$HARDWARE_CATEGORY" == CPU* ]]; then
+            if [[ "$base_metric_key" == "CPU_TIME" ]] || [[ "$base_metric_key" == "TIME" ]]; then # Handle "CPU Time" or just "Time"
+                final_metric_key="CPU_TIME"
+            elif [[ "$base_metric_key" == "WALL_TIME" ]]; then
+                final_metric_key="WALL_TIME" # Wall time is generic
+            # Add other specific CPU mappings here if needed
+            # else
+            #    final_metric_key="CPU_${base_metric_key}" # Generic prefix
+            else
+                final_metric_key="$base_metric_key" # If no specific CPU mapping, use base
+            fi
+        else
+            final_metric_key="$base_metric_key" # For unknown/other categories
+        fi
+
+        # Expecting lines like "Metric Name: 123.45 us" or "Metric Name: 12345 bytes"
+        read number unit <<< $(echo "$value" | sed -E 's/([0-9]+\.?[0-9]*)\s*([a-zA-Z]+).*/\1 \2/')
+        # Convert microseconds to milliseconds
+        if [[ "$unit" == "us" ]]; then
+            final_metric_value=$(LC_ALL=C awk -v num="$number" 'BEGIN { printf "%.3f", num / 1000 }')
+            final_unit="ms"
+        elif [[ "$unit" == "bytes" ]]; then
+            # Convert bytes to GB
+            final_metric_value=$(echo "$number" | awk '{printf "%.2f", $1/1024^3}')
+            final_unit="GB"
+        else
+          echo "::warning::Skipping unsupported unit: $unit"
+          continue
+        fi
+        if ! $first_metric; then metrics_obj_str+=","; fi
+        metrics_obj_str+="\"$final_metric_key\": {\"value\": $final_metric_value, \"unit\": \"$final_unit\"}"
+        first_metric=false
+        echo "INFO: Parsed metric: OriginalKey='$key', BaseKey='$base_metric_key', FinalKey='$final_metric_key', Value='$final_metric_value $final_unit'"
+    done <<< "$STATS_OUTPUT"
+    metrics_obj_str+="}"
+
+    if echo "$metrics_obj_str" | jq -e . > /dev/null 2>&1; then
+        METRICS_JSON_CONTENT=$(echo "$metrics_obj_str" | jq '.')
+        echo "Successfully parsed metrics from stats output."
+    else
+        echo "::warning::Could not construct valid JSON from stats output. Metrics object will be empty."
+        echo "Problematic metrics string constructed: $metrics_obj_str"
+        METRICS_JSON_CONTENT="{}"
+        STATS_RUN_SUCCESSFUL=false 
+    fi
+  else
+    echo "::warning::compute_xspace_stats_main failed with code $STATS_EXIT_CODE. No metrics will be parsed from its output."
+  fi
+else
+   if [ $RUNNER_EXIT_CODE -ne 0 ]; then 
+      echo "::warning::Runner failed (Exit Code: $RUNNER_EXIT_CODE), skipping stats processing."
+   else 
+     echo "::warning::XSpace file missing at $XSPACE_FILE_PATH, skipping stats processing."
+   fi
+fi
+
+# --- Construct Final results.json ---
+RUN_STATUS_MSG=""
+ERROR_MSG_CONTENT=""
+
+if [ $RUNNER_EXIT_CODE -ne 0 ]; then
+    RUN_STATUS_MSG="FAILURE"
+    ERROR_MSG_CONTENT="Runner failed with code $RUNNER_EXIT_CODE"
+elif [ ! -f "$XSPACE_FILE_PATH" ]; then
+    RUN_STATUS_MSG="SUCCESS_NO_PROFILE"
+    ERROR_MSG_CONTENT="XSpace file not generated by successful run."
+elif [ $STATS_EXIT_CODE -ne 0 ] || [ "$STATS_RUN_SUCCESSFUL" = false ] ; then
+    RUN_STATUS_MSG="STATS_FAILURE"
+    ERROR_MSG_CONTENT="compute_xspace_stats_main failed (code $STATS_EXIT_CODE) or metrics parsing failed. Runner was successful."
+else
+    RUN_STATUS_MSG="SUCCESS"
+    ERROR_MSG_CONTENT=""
+fi
+
+# Use jq to build the final JSON, incorporating the parsed metrics
+jq -n \
+  --arg bn "$BENCHMARK_NAME" \
+  --arg cid "$CONFIG_ID" \
+  --arg hc "$HARDWARE_CATEGORY" \
+  --arg rs "$RUN_STATUS_MSG" \
+  --arg em "$ERROR_MSG_CONTENT" \
+  --arg cs "$COMMIT_SHA" \
+  --arg wrid "$WORKFLOW_RUN_ID" \
+  --argjson metrics "$METRICS_JSON_CONTENT" \
+  '{
+     benchmark_name: $bn,
+     config_id: $cid,
+     hardware_category: $hc,
+     run_status: $rs,
+     error_message: $em,
+     commit_sha: $cs,
+     workflow_run_id: $wrid,
+     metrics: $metrics
+   }' > "$RESULTS_JSON_FILE"
+
+if [ $? -eq 0 ]; then
+    echo "Final results.json created at $RESULTS_JSON_FILE."
+else
+    echo "::error::FATAL: Failed to create final results.json using jq."
+    echo "FATAL JQ ERROR. Benchmark Name: $BENCHMARK_NAME, Run Status: $RUN_STATUS_MSG, Error: $ERROR_MSG_CONTENT" > "$RESULTS_JSON_FILE.txt"
+    exit 1 
+fi
+
+# --- Debug: Verify file creation ---
+echo "DEBUG: Listing contents of OUTPUT_DIR ($OUTPUT_DIR):"
+ls -la "$OUTPUT_DIR"
+echo "DEBUG: Checking for RESULTS_JSON_FILE ($RESULTS_JSON_FILE):"
+if [ -f "$RESULTS_JSON_FILE" ]; then
+  echo "DEBUG: RESULTS_JSON_FILE exists. Content (first 20 lines):"
+  head -n 20 "$RESULTS_JSON_FILE"
+else
+  echo "DEBUG: RESULTS_JSON_FILE does NOT exist."
+  if [ -f "${RESULTS_JSON_FILE}.txt" ]; then
+    echo "DEBUG: RESULTS_JSON_FILE.txt exists. Content:"
+    cat "${RESULTS_JSON_FILE}.txt"
+  else
+    echo "DEBUG: RESULTS_JSON_FILE.txt also does NOT exist."
+  fi
+fi
+echo "DEBUG: End of file check."
+
+# --- Final Exit Status of the script ---
+if [ $RUNNER_EXIT_CODE -ne 0 ]; then
+  echo "::error::Benchmark run failed (HLO Runner Exit Code: $RUNNER_EXIT_CODE)."
+  exit $RUNNER_EXIT_CODE
+fi
+
+echo "--- Run Benchmark Script Finished ---"

third_party/xla/.github/workflows/benchmarks/run_comparison.sh

@@ -0,0 +1,59 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/benchmarks/run_comparison.sh
+# TODO(juliagmt): convert this to a python script.
+#!/bin/bash
+
+# This script encapsulates the logic to compare benchmark results against a baseline.
+
+set -e  # Exit immediately if a command exits with a non-zero status.
+set -u # Treat unset variables as an error when substituting.
+
+echo "--- Running Comparison Script ---"
+echo "CONFIG_ID for baseline lookup: $CONFIG_ID" # This is now the comprehensive ID
+echo "Results Directory: $RESOLVED_OUTPUT_DIR"
+echo "Baseline File: $RESOLVED_BASELINE_YAML"
+echo "Comparison Python Script: $RESOLVED_COMPARISON_SCRIPT"
+
+ACTUAL_RESULTS_JSON_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+
+if [ ! -f "$ACTUAL_RESULTS_JSON_PATH" ]; then
+  echo "::warning::Primary results file '$ACTUAL_RESULTS_JSON_PATH' not found. Cannot perform baseline comparison."
+  # Check for fallback .txt file for more info if primary is missing
+  if [ -f "${ACTUAL_RESULTS_JSON_PATH}.txt" ]; then
+    echo "Fallback results.json.txt found:"
+    cat "${ACTUAL_RESULTS_JSON_PATH}.txt"
+  fi
+  exit 0 # Exiting cleanly to not block if results are missing; comparison script handles this too
+fi
+SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP="${CONFIG_ID}"
+
+echo "Using Config ID for baseline lookup: $SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+echo "Constructed Config ID for baseline lookup: $SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+python3 "$RESOLVED_COMPARISON_SCRIPT" \
+  --results-json-file="$ACTUAL_RESULTS_JSON_PATH" \
+  --baseline-yaml-file="$RESOLVED_BASELINE_YAML" \
+  --config-id="$SCRIPT_CONFIG_ID_FOR_BASELINE_LOOKUP"
+
+COMPARISON_EXIT_CODE=$?
+if [ $COMPARISON_EXIT_CODE -ne 0 ]; then
+  echo "::error::Baseline comparison script failed or regressions detected (Exit Code: $COMPARISON_EXIT_CODE)."
+  exit $COMPARISON_EXIT_CODE
+fi
+
+echo "Baseline comparison successful: No regressions detected or no applicable baselines found."
+echo "--- Comparison Script Finished ---"

third_party/xla/.github/workflows/buildifier.yml

@@ -0,0 +1,43 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Buildifier
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+env:
+  # Have `go install` place binaries in $PATH
+  GOBIN: "/usr/local/bin"
+
+jobs:
+  buildifier-lint:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 6
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Install buildifier"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
+      - name: "Run buildifier"
+        run: buildifier --lint=warn --warnings=-out-of-order-load -r xla/

third_party/xla/.github/workflows/check_contents.yml

@@ -0,0 +1,140 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+# The goal of these checks is to externalize all check_contents checks that run
+# internally. For now, no automation will try to keep these in sync as many
+# METADATA files are doing different checks, and we'll have to consolidate those
+# files once XLA moves out of Tensorflow internally.
+# TODO(ddunleavy): Update this after METADATA files are consolidated.
+name: Check Contents
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+env:
+  # A bit tricky here: this lets us invoke python files outside of `bazel run`
+  # as long as the python file has imports as if it expects bazel to invoke it,
+  # and only depends on other python files in this workspace or the standard
+  # library.
+  PYTHONPATH: "."
+
+jobs:
+  check-contents:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 6
+    if: |
+      github.event.sender.type == 'User' ||
+      contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Fetch HEAD of main branch"
+        run: git fetch origin main --depth=1
+
+      # ----------------C++ checks----------------
+      - name: "Check for tsl::{Status,StatusOr}"
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'tsl::Status(Or)?' --suppression_regex 'TENSORFLOW_STATUS_OK' --failure_message 'Use plain Status(Or) rather than tsl::Status(Or) within XLA.'"
+      - name: "Check for tsl::Status::OK()"
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'tsl::Status::OK' --failure_message 'Use OkStatus instead of Status::OK.'"
+      - name: "Check for std::call_once"
+
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex 'std::call_once' --failure_message 'Use absl::call_once instead of std::call_once'"
+      - name: "Check for Abseil compatibility API"
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex '\\b(absl::(any|get|get_if|make_any|make_unique|nullopt|nullopt_t|MakeUnique|make_optional|optional|underlying_type_t|variant|visit))\\b' --suppression_regex '(?i)non-std ok' --failure_message 'Use standard library version instead of the Abseil compatibility API.'"
+      - name: "Use absl instead of Tensorflow/TSL functions/types"
+
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex_exclusion '.github/workflows/check_contents\\.yml' --prohibited_regex '\\b(gtl::FlatMap|gtl::FlatSet|gtl::InlinedVector|gtl::optional|strings::Appendf|strings::Printf|strings::safe_strto64|strings::safe_strtof|strings::StrAppend|strings::StrCat|str_util::|tensorflow::StringPiece)\\b' --suppression_regex '(?i)non-absl ok' --failure_message 'Use Abseil types and functions instead of those in Tensorflow/TSL.'"
+      - name: "Check for Header Guards"
+        if: always()
+        run: "python3 build_tools/lint/check_header_guards.py"
+      - name: "Check for TF Status Macros"
+        if: always()
+        run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.(h|cc|cpp)' --path_regex_exclusion 'xla/tsl/.*' 'xla/python/.*' 'xla/pjrt/.*' --prohibited_regex '\\bTF_(RETURN_IF_ERROR|ASSIGN_OR_RETURN)\\b' --failure_message 'Within XLA, you can now use RETURN_IF_ERROR/ASSIGN_OR_RETURN from xla/tsl/platform/status_macros.h rather than their TSL TF_ equivalents.'"
+      # ----------------Python checks----------------
+      - name: "Check for python print()"
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.py' --prohibited_regex '(^|\\n|\\s)print\\(.*\\)' --suppression_regex 'DISABLE_DEBUG_PRINT_CHECK' --failure_message 'Found debug print() statement left in the code. Add DISABLE_DEBUG_PRINT_CHECK to disable'"
+
+      - name: "Check for mock.patch"
+        if: always() # We want to run all checks, not just til first failure
+        run: "python3 build_tools/lint/check_contents.py --path_regex '.*\\.py' --prohibited_regex 'mock\\.patch\\(' --failure_message 'Please use match.patch.object instead of mock.patch.'"
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/ci.yml

@@ -0,0 +1,224 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: CI
+permissions:
+  contents: read
+on:
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  pull_request:
+  push:
+    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: [
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux x86 CPU",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "windows-x86-n2-16",
+            container: null,
+            name: "XLA Windows x86 CPU",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux x86 CPU Bzlmod",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux X86 GPU ONEAPI",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux x86 GPU Hermetic ROCm",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-x86-g2-16-l4-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux x86 GPU L4",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-arm64-c4a-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
+            name: "XLA Linux ARM64 CPU",
+            repo: "openxla/xla",
+          },
+          {
+            pool: "linux-x86-n2-16",
+            container: "gcr.io/tensorflow-sigs/build:latest-python3.11",
+            name: "JAX Linux x86 CPU Bzlmod",
+            repo: "jax-ml/jax",
+          },
+          {
+            pool: "windows-x86-n2-16",
+            container: null,
+            name: "JAX Windows x86 CPU",
+            repo: "jax-ml/jax",
+          },
+          {
+            pool: "linux-x86-g2-16-l4-1gpu",
+            container: "gcr.io/tensorflow-sigs/build:latest-python3.12",
+            name: "JAX Linux x86 GPU L4",
+            repo: "jax-ml/jax",
+          },
+          {
+            pool: "linux-x86-n2-16",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "TensorFlow Linux x86 CPU",
+            repo: "tensorflow/tensorflow",
+          },
+          {
+            pool: "linux-x86-g2-16-l4-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "TensorFlow Linux x86 GPU L4",
+            repo: "tensorflow/tensorflow",
+          },
+        ]
+    name: ${{ matrix.job_info.name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: "Checking out openxla/xla"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          path: "openxla/xla"
+          persist-credentials: false
+      - name: Checking out ${{ matrix.job_info.repo }}
+        if: ${{ matrix.job_info.repo != 'openxla/xla' }}
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          repository: ${{ matrix.job_info.repo }}
+          path: ${{ matrix.job_info.repo }}
+          persist-credentials: false
+      - name: Checking out rules_ml_toolchain # TODO(alekstheod):  remove once jax is migrated to the latest rules_ml_toolchain changes
+        run: |
+          git clone https://github.com/google-ml-infra/rules_ml_toolchain.git openxla/rules_ml_toolchain
+          cd openxla/rules_ml_toolchain
+          git checkout b11745590f513ec55b32e2d126073576fde18c71
+      - name: "Wait For Connection"
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: "Run build.py"
+        working-directory: ${{ matrix.job_info.repo }}
+        timeout-minutes: 60
+        run: |
+          # Shellcheck doesn't understand GitHub Actions expressions and incorrectly
+          # thinks this comparison will always fail.
+          # shellcheck disable=SC2193
+          if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
+            python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
+          else
+            "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
+          fi
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/ci_multi_device.yml

@@ -0,0 +1,85 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Multi-Device CI
+permissions:
+  contents: read
+on:
+  workflow_dispatch:  # Allows manual triggering
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  schedule:
+    - cron: "0 0 * * *"  # Run at midnight every day
+
+jobs:
+  Tests:
+    strategy:
+      # Don't fail fast - want to see results for all builds even if one fails.
+      fail-fast: false
+      matrix:
+        job_info: [
+          {
+            pool: "linux-x86-a3-8g-h100-8gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            name: "XLA Linux x86 GPU 8X H100",
+            repo: "openxla/xla",
+          },
+        ]
+    name: ${{ matrix.job_info.name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 60
+    steps:
+      - name: "Checking out openxla/xla"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          path: "openxla/xla"
+          persist-credentials: false
+      - name: Checking out ${{ matrix.job_info.repo }}
+        if: ${{ matrix.job_info.repo != 'openxla/xla' }}
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          repository: ${{ matrix.job_info.repo }}
+          path: ${{ matrix.job_info.repo }}
+          persist-credentials: false
+      - name: "Wait For Connection"
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: "Run build.py"
+        working-directory: ${{ matrix.job_info.repo }}
+        run: |
+          # shellcheck disable=SC2193
+          if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
+            python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
+          else
+            "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
+          fi
+
+
+
+
+
+
+

third_party/xla/.github/workflows/clang_format.yml

@@ -0,0 +1,48 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+name: Clang Format
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  clang-format:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 6
+    if: |
+      github.event.sender.type == 'User' ||
+      contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          fetch-depth: '0'
+          persist-credentials: false
+      - name: "Fetch HEAD of main branch"
+        run: git fetch origin main --depth=1
+      - name: Install uv dependencies
+        uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
+      - name: "Run clang-format check"
+        env:
+          TARGET_REF: "origin/${{ github.base_ref || 'main' }}"
+        run: ./build_tools/ci/run_clang_format.sh

third_party/xla/.github/workflows/clang_tidy.yml

@@ -0,0 +1,63 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Clang Tidy
+permissions:
+  contents: read
+on:
+  pull_request:
+jobs:
+  clang-tidy-cuda:
+    runs-on: "linux-x86-n2-16"
+    container:
+      image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"  # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 30
+    continue-on-error: true
+    if: |
+      github.event.sender.type == 'User' ||
+      contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          # Fetch more commits to ensure we can find the merge base for clang-tidy.
+          # This is better than 0 since we don't want to fetch everything as well.
+          fetch-depth: '100'
+          persist-credentials: false
+      - name: "Update base branch"
+        run: |
+          git config --global --add safe.directory "$GITHUB_WORKSPACE"
+          BASE_REF="${GITHUB_BASE_REF:-main}"
+          # Keep this fetch unshallow so clang-tidy can find the merge base even
+          # when the PR branch is more than checkout fetch-depth commits behind.
+          git fetch origin "refs/heads/${BASE_REF}:refs/remotes/origin/${BASE_REF}"
+      - name: Configure Workspace
+      # The actual compute capability does not matter. This allows deviceless run for this pipeline.
+      # The common configs ensure that bazel query and bazel build don't refetch repo mappings.
+        run: |
+          set -euo pipefail
+          ./configure.py --backend=CUDA --cuda_compute_capabilities="9.0"
+          echo "common --config cuda_nvcc" >> xla_configure.bazelrc
+          echo "common --config rbe_linux_cpu" >> xla_configure.bazelrc
+      - name: "Run Clang Tidy"
+        env:
+          TARGET_REF: "origin/${{ github.base_ref || 'main' }}"
+          TAGS_TO_IGNORE: "no_oss|rocm-only|oneapi-only"
+          BAZEL_CMD: "bazel"
+        run: |
+          set -euo pipefail
+          ./build_tools/ci/run_clang_tidy.sh

third_party/xla/.github/workflows/cpu_benchmarks_nightly.yml

@@ -0,0 +1,248 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Nightly CPU Benchmarks
+permissions:
+  contents: read
+on:
+  workflow_dispatch: # Allows manual triggering
+  schedule:
+    - cron: "0 0 * * *" # Run at midnight every day
+
+jobs:
+  Tests:
+    strategy:
+      # Don't fail fast - want to see results for all builds even if one fails.
+      fail-fast: false
+      matrix:
+        job_info: [
+           {
+            pool: "linux-arm64-t2a-48",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
+            pretty_name: "XLA Linux ARM64 CPU 48 vcpu Presubmit",
+            bazel_arch_dir: "aarch64-opt",
+            platform: "CPU",
+            output_dir: "cpu_48_output"
+          },
+          {
+            pool: "linux-x86-n2-128",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            pretty_name: "XLA Linux x86 CPU 128 vcpu Presubmit",
+            bazel_arch_dir: "k8-opt",
+            platform: "CPU",
+            output_dir: "cpu_128_output"
+          }
+        ]
+    name: ${{ matrix.job_info.pretty_name }}
+    runs-on: ${{ matrix.job_info.pool }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 540
+    env:
+      OUTPUT_DIR: ${{ github.workspace }}/${{ matrix.job_info.output_dir }}
+    steps:
+      - name: Print machine specs
+        run: |
+          lscpu
+          free -h  # Memory information
+          df -h    # Disk space information
+          uname -a # Kernel information
+      - name: Check Python Version in Container
+        run: python3 --version
+      - name: Install Python 3.10 if not present (IN CONTAINER)
+        run: |
+          if ! python3 --version > /dev/null 2>&1; then # check for python3
+            echo "Python 3 not found, installing..."
+            apt-get update
+            apt-get install -y python3.10 python3-pip
+          else
+            echo "Python 3 found."
+          fi
+      - name: Checkout OpenXLA
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Create results directory
+        run: mkdir results
+
+      - name: Configure CPU backend
+        run: |
+
+          ./configure.py --backend=CPU
+      - name: Download and run Gemma HLO files and collect runtime and compile time stats
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+        run: |
+          mkdir -p "$OUTPUT_DIR/tmp_hlo"
+          cd "$OUTPUT_DIR/tmp_hlo"
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
+          cd ..
+
+      - name: Print GitHub Context
+        run: |
+          echo "GitHub SHA: ${{ github.sha }}"
+          echo "GitHub Ref: ${GITHUB_REF}"
+
+          echo "GitHub Ref Name: ${GITHUB_REF_NAME}"
+          echo "GitHub Head Ref: ${GITHUB_HEAD_REF})"
+          echo "GitHub Base Ref: ${GITHUB_BASE_REF})"
+
+          echo "GitHub Repository: ${{ github.repository }}"
+          echo "GitHub Run ID: ${{ github.run_id }}"
+          echo "GitHub Run Number: ${{ github.run_number }}"
+
+          echo "GitHub Workflow: ${GITHUB_WORKFLOW}"
+          echo "GitHub Actor: ${GITHUB_ACTOR}"
+          echo "GitHub Event Name: ${{ github.event_name }}"
+          if [[ "${{ github.event_name }}" == "pull_request" ]]; then
+            echo "Pull Request Number: ${{ github.event.pull_request.number }}"
+            echo "Pull Request Head Ref: ${GITHUB_EVENT_PULL_REQUEST_HEAD_REF}"
+            echo "Pull Request Base Ref: ${GITHUB_EVENT_PULL_REQUEST_BASE_REF}"
+          fi
+
+        env:
+          GITHUB_EVENT_PULL_REQUEST_HEAD_REF: ${{ github.event.pull_request.head.ref }}
+          GITHUB_EVENT_PULL_REQUEST_BASE_REF: ${{ github.event.pull_request.base.ref }}
+      # Find the current PR number, if any, because github context doesn't have it for push events.
+      - uses: jwalton/gh-find-current-pr@f3d61b485d2801773f7a07b2aaa3306bd8f8e653 # v1.3.5
+        id: findPr
+        with:
+          # Can be "open", "closed", or "all".  Defaults to "open".
+          state: all
+      # This will echo "Your PR is 7", or be skipped if there is no current PR.
+      - run: echo "PR_NUMBER=${STEPS_FINDPR_OUTPUTS_PR}"
+        if: success() && steps.findPr.outputs.number
+        env:
+          PR: ${{ steps.findPr.outputs.pr }}
+          STEPS_FINDPR_OUTPUTS_PR: ${{ steps.findPr.outputs.pr }}
+      - name: "Run build.py"
+        run: |
+          ./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
+      - name: Run hlo_runner_main on cpu_hlo.hlo and collect profile stats
+        run: |
+          bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
+
+          binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
+          mkdir -p "$OUTPUT_DIR"
+          OUTPUT_FILE_PATH="$OUTPUT_DIR/cpu_hlo_output.txt"
+          XSPACE_FILE_PATH="$OUTPUT_DIR/cpu_hlo_xspace.pb"
+          test_hlo_file="xla/tools/hlo_opt/tests/cpu_hlo.hlo"
+          echo "Running CPU test with binary: $binary_dir"
+          pwd #print working directory
+          $binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
+          $binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
+          cat "$OUTPUT_FILE_PATH"
+      - name: Run hlo_runner_main on gemma2_2b_keras_jax.hlo and collect profile stats
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+
+        run: |
+          bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
+          binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
+
+
+          OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma2_2b_keras_jax_output.txt"
+          XSPACE_FILE_PATH="$OUTPUT_DIR/gemma2_2b_xspace.pb"
+          test_hlo_file="$OUTPUT_DIR/tmp_hlo/gemma2_2b_keras_jax.hlo"
+          echo "Running CPU test with binary: $binary_dir"
+          pwd #print working directory
+          $binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
+          $binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
+          cat "$OUTPUT_FILE_PATH"
+      - name: Run hlo_runner_main on gemma3_1b_flax_call.hlo and collect profile stats
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+        run: |
+          bazel_arch_dir="${{ matrix.job_info.bazel_arch_dir }}"
+          binary_dir="./bazel-out/${bazel_arch_dir}/bin/xla/tools"
+
+          OUTPUT_FILE_PATH="$OUTPUT_DIR/gemma3_1b_flax_call_output.txt"
+          XSPACE_FILE_PATH="$OUTPUT_DIR/gemma3_1b_xspace.pb"
+          test_hlo_file="$OUTPUT_DIR/tmp_hlo/gemma3_1b_flax_call.hlo"
+          echo "Running CPU test with binary: $binary_dir"
+          pwd #print working directory
+          $binary_dir/multihost_hlo_runner/hlo_runner_main --device_type=host --num_repeats=5 --use_spmd_partitioning --profile_execution=True --xla_gpu_dump_xspace_to="$XSPACE_FILE_PATH" "$test_hlo_file" > "$OUTPUT_FILE_PATH"
+          $binary_dir/compute_xspace_stats_main --input="$XSPACE_FILE_PATH" --device_type=CPU >> "$OUTPUT_FILE_PATH"
+          cat "$OUTPUT_FILE_PATH"
+      - name: Upload HLO test output to a GCS bucket
+        if: ${{ matrix.job_info.pool != 'linux-arm64-t2a-48' }}
+        run: |
+          GCS_BUCKET="gs://openxla-nightly-transient"
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+
+          DATE_FOLDER=$(date +%Y%m%d)
+          FILENAME_GEMMA="gemma2_2b_keras_jax_output.txt"
+          FILENAME_CPU="cpu_hlo_output.txt"
+          FILENAME_GEMMA3="gemma3_1b_flax_call_output.txt"
+          COMMIT_SHA="${{ github.sha }}"
+          RUN_ID="${{ github.run_id }}"
+          PR_NUMBER="${STEPS_FINDPR_OUTPUTS_PR}"
+          OUTPUT_SUBDIR="${{ matrix.job_info.output_dir }}"
+          PR_SUFFIX=""
+          if [[ -n "$PR_NUMBER" ]]; then
+            PR_SUFFIX="_pr_${PR_NUMBER}"
+          fi
+
+
+          GEMMA_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_GEMMA}"
+          CPU_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_CPU}"
+          GEMMA3_GCS_OBJECT_NAME="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.pool }}_run_${RUN_ID}_commit_${COMMIT_SHA}${PR_SUFFIX}_${OUTPUT_SUBDIR}_${FILENAME_GEMMA3}"
+
+          echo "Uploading $OUTPUT_DIR/$FILENAME_GEMMA to $GCS_BUCKET/$GEMMA_GCS_OBJECT_NAME"
+          ls -l "$OUTPUT_DIR/$FILENAME_GEMMA" # Verify the file exists
+          gcloud storage cp "$OUTPUT_DIR/$FILENAME_GEMMA" "$GCS_BUCKET/$GEMMA_GCS_OBJECT_NAME"
+
+          echo "Uploading $OUTPUT_DIR/$FILENAME_CPU to $GCS_BUCKET/$CPU_GCS_OBJECT_NAME"
+          ls -l "$OUTPUT_DIR/$FILENAME_CPU" # Verify the file exists
+          gcloud storage cp "$OUTPUT_DIR/$FILENAME_CPU" "$GCS_BUCKET/$CPU_GCS_OBJECT_NAME"
+
+          echo "Uploading $OUTPUT_DIR/$FILENAME_GEMMA3 to $GCS_BUCKET/$GEMMA3_GCS_OBJECT_NAME"
+          ls -l "$OUTPUT_DIR/$FILENAME_GEMMA3" # Verify the file exists
+          gcloud storage cp "$OUTPUT_DIR/$FILENAME_GEMMA3" "$GCS_BUCKET/$GEMMA3_GCS_OBJECT_NAME"
+        env:
+          STEPS_FINDPR_OUTPUTS_PR: ${{ steps.findPr.outputs.pr }}
+
+      - name: Upload XSpace artifacts
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+
+          name: cpu-xla-benchmarks-xspace-${{ matrix.job_info.pool }}
+          path: ${{ github.workspace }}/${{ matrix.job_info.output_dir }}/*_xspace.pb
+      - name: Run E2E benchmarks flax_2b
+
+        timeout-minutes: 60
+        run: |
+          cd xla/backends/cpu/benchmarks/e2e/gemma2/flax_2b
+
+          bash setup.sh
+          bash run.sh | tee -a ../../../../../../../results/flax_2b.log
+
+# TODO(juliagmt): Add CPU XSpace collection support to hlo_runner_main.
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/dwyu.yml

@@ -0,0 +1,64 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: DWYU (Depend on What You Use)
+permissions:
+  contents: read
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+env:
+  GOBIN: "/usr/local/bin"
+  PYTHONPATH: "."
+  BANT_VERSION: "0.2.8"
+
+jobs:
+  dwyu:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 15
+    continue-on-error: true
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          fetch-depth: '0'
+          persist-credentials: false
+      - name: "Fetch HEAD of main branch"
+        run: git fetch origin main --depth=1
+      - name: "Install bant"
+        run: |
+          curl -fSL "https://github.com/hzeller/bant/releases/download/v${BANT_VERSION}/bant-v${BANT_VERSION}-linux-static-x86_64.tar.gz" -o /tmp/bant.tar.gz
+          tar xzf /tmp/bant.tar.gz -C /tmp
+          sudo install "/tmp/bant-v${BANT_VERSION}-linux-static-x86_64/bin/bant" /usr/local/bin/bant
+      - name: "Install bazelisk"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- go install github.com/bazelbuild/bazelisk@24651ab # v1.20.0
+      - name: "Run bazel build --nobuild //xla/... with retries"
+        run: parallel --ungroup --retries 3 --delay 15 --nonall -- bazelisk build --nobuild //xla/...
+      - name: "Find affected targets and run bant dwyu"
+        run: |
+          set -euo pipefail
+          TARGETS=$(python3 build_tools/lint/check_dwyu.py)
+          if [ -z "$TARGETS" ]; then
+            echo "No targets to check."
+            exit 0
+          fi
+          echo "Running: bant dwyu $TARGETS"
+          # shellcheck disable=SC2086
+          bant dwyu $TARGETS

third_party/xla/.github/workflows/generate_benchmark_matrix.yml

@@ -0,0 +1,187 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/generate_benchmark_matrix.yml
+name: Generate Benchmark Matrix
+
+permissions:
+  contents: read
+
+on:
+  workflow_call:
+    inputs:
+      registry_file:
+        description: 'Path to the benchmark registry file (relative to repository root)'
+        required: false
+        type: string
+        default: 'xla/tools/benchmarks/registries/default_registry.yml'
+      workflow_type:
+        description: 'The type of workflow triggering this generation (e.g. PRESUBMIT, POSTSUBMIT, SCHEDULED, MANUAL, or aliases like nightly)'
+        required: true
+        type: string
+      checkout_ref:
+        description: 'The Git ref (branch, tag, or SHA) to checkout'
+        required: false
+        type: string
+        default: ''
+    outputs:
+      matrix_include_json:
+        description: 'JSON string representing the array for the "include" matrix strategy'
+        value: ${{ jobs.generate.outputs.matrix_json_output }}
+
+jobs:
+  generate:
+    name: Generate Matrix (${{ inputs.workflow_type }})
+    runs-on: linux-x86-n2-64
+    container: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest # zizmor: ignore[unpinned-images]
+    outputs:
+      matrix_json_output: ${{ steps.run_generator.outputs.matrix_json }}
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 60
+
+    steps:
+      - name: Checkout OpenXLA
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          # Use inputs.checkout_ref if provided, otherwise default to the event's ref
+          # (e.g., PR's HEAD SHA or caller's commit SHA)
+          ref: ${{ inputs.checkout_ref || github.ref }}
+          fetch-depth: 0 # Fetch all history, might be needed by build/configure scripts
+          persist-credentials: false
+
+      - name: Configure OpenXLA
+        run: "echo \"Configuring OpenXLA for CPU to build the generator tool...\"\nif [ -f \"./configure.py\" ]; then\n  ./configure.py --backend=CPU \nelse\n  echo \"::warning::configure.py not found. Assuming C++ tool build doesn't require it or is pre-configured.\"\nfi\n"
+      # TODO(juliagmt): Use build.py to build the binary.
+      - name: Build generate_benchmark_matrices_main
+        run: |
+          echo "Building generator..."
+          if ! 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 \
+                --color=yes \
+                --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/tools/benchmarks/utils:generate_benchmark_matrices_main; then
+             echo "::error::Failed to build generate_benchmark_matrices_main"
+             exit 1
+          fi
+      - name: Run generate_benchmark_matrices_main
+        id: run_generator
+        run: |
+          BINARY_PATH="./bazel-bin/xla/tools/benchmarks/utils/generate_benchmark_matrices_main"
+          REGISTRY_PATH="${INPUTS_REGISTRY_FILE}"
+          # Convert workflow_type input to uppercase for the C++ binary, 
+          # as the binary's GetWorkflowTypeFromStr expects uppercase or known aliases.
+          WORKFLOW_TYPE_ARG_UPPER=$(echo "${INPUTS_WORKFLOW_TYPE}" | tr '[:lower:]' '[:upper:]')
+
+
+          if [ ! -f "$BINARY_PATH" ]; then 
+            echo "::error::Generator binary not found at $BINARY_PATH after build."
+            # List bazel-bin for debugging if binary not found
+            echo "Listing bazel-bin directory:"
+            ls -R bazel-bin || echo "bazel-bin directory not found."
+            exit 1
+          fi
+
+          echo "Generating matrix for workflow type: $WORKFLOW_TYPE_ARG_UPPER (original input: ${INPUTS_WORKFLOW_TYPE}) using registry: $REGISTRY_PATH"
+
+          # Execute and capture output
+          JSON_ARRAY_STRING=$("$BINARY_PATH" --registry_file="$REGISTRY_PATH" --workflow_type="$WORKFLOW_TYPE_ARG_UPPER")
+          BINARY_EXIT_CODE=$?
+
+          if [ $BINARY_EXIT_CODE -ne 0 ]; then
+            echo "::error::generate_benchmark_matrices_main failed with exit code $BINARY_EXIT_CODE."
+            echo "Output from binary was:"
+            echo "$JSON_ARRAY_STRING" # This might contain error messages from the binary's LOG(QFATAL)
+            exit 1
+          fi
+
+          # Basic validation: must be an array (even if empty "[]")
+          if ! echo "$JSON_ARRAY_STRING" | jq -e '. | type == "array"' > /dev/null; then
+             echo "::error::Generator output is not a valid JSON array."
+             echo "Output was: $JSON_ARRAY_STRING"
+             exit 1
+          fi
+
+          echo "Generated matrix JSON array string:"
+          echo "$JSON_ARRAY_STRING"
+
+          {
+            echo "matrix_json<<EOF_MATRIX_JSON"
+            echo "$JSON_ARRAY_STRING"
+            echo "EOF_MATRIX_JSON"
+          } >> "$GITHUB_OUTPUT"
+        env:
+          INPUTS_REGISTRY_FILE: ${{ inputs.registry_file }}
+          INPUTS_WORKFLOW_TYPE: ${{ inputs.workflow_type }}
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+

third_party/xla/.github/workflows/github_api.py

@@ -0,0 +1,160 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Provides a Python interface to various parts of the GitHub API.
+
+NOTE: not using PyGithub https://github.com/PyGithub/PyGithub as it doesn't
+have full API coverage and it's easy enough to use the endpoints we need like
+this (with zero dependencies as a bonus!)
+"""
+from typing import Any, Optional
+import urllib
+
+import requests
+
+
+class GitHubAPI:
+  """Wraps the GitHub REST API."""
+
+  def __init__(self, token: Optional[str] = None):
+    self._session = requests.Session()
+    self._session.headers["Accept"] = "application/vnd.github+json"
+    if token:
+      self._session.headers["Authorization"] = f"token {token}"
+
+  def _make_request(
+      self, verb: str, endpoint: str, **kwargs: dict[str, Any]
+  ) -> requests.Response:
+    """Helper method to make a request and raise an HTTPError if one occurred.
+
+    Arguments:
+      verb: The HTTP verb to use
+      endpoint: The endpoint to make the request to
+      **kwargs: The json that will be sent as the body of the request.
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    res = self._session.request(
+        verb,
+        urllib.parse.urljoin("https://api.github.com", endpoint),
+        json=kwargs,
+    )
+    res.raise_for_status()
+    return res.json()
+
+  def get_commit(self, repo: str, commit_id: str) -> requests.Response:
+    """Gets a commit by it's SHA-1 hash.
+
+    https://docs.github.com/en/rest/commits/commits?apiVersion=2022-11-28#get-a-
+    commit
+
+    Arguments:
+      repo: a string of the form `owner/repo_name`, e.g. openxla/xla.
+      commit_id: a string describing the commit to get, e.g. `deadbeef` or
+        `HEAD`.
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    endpoint = f"repos/{repo}/commits/{commit_id}"
+    return self._make_request("GET", endpoint)
+
+  def write_issue_comment(
+      self, repo: str, issue_number: int, body: str
+  ) -> requests.Response:
+    """Writes a comment on an issue (or PR).
+
+    https://docs.github.com/en/rest/issues/comments?apiVersion=2022-11-
+    28#create-an-issue-comment
+
+    Arguments:
+      repo: a string of the form `owner/repo_name`, e.g. openxla/xla
+      issue_number: the issue (or PR) to comment on
+      body: the body of the comment
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    endpoint = f"repos/{repo}/issues/{issue_number}/comments"
+    return self._make_request("POST", endpoint, body=body)
+
+  def set_issue_status(
+      self, repo: str, issue_number: int, status: str
+  ) -> requests.Response:
+    """Sets the status of an issue (or PR).
+
+    https://docs.github.com/en/rest/issues/issues?apiVersion=2022-11-28#update-
+    an-issue
+
+    Arguments:
+      repo: a string of the form `owner/repo_name`, e.g. openxla/xla
+      issue_number: the issue (or PR) to set the status of
+      status: the status to set
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    endpoint = f"repos/{repo}/issues/{issue_number}"
+    return self._make_request("POST", endpoint, status=status)
+
+  def add_issue_labels(
+      self, repo: str, issue_number: int, labels: list[str]
+  ) -> requests.Response:
+    """Adds labels to an issue (or PR).
+
+    https://docs.github.com/en/actions/managing-issues-and-pull-requests/adding-labels-to-issues
+
+    Arguments:
+      repo: a string of the form `owner/repo_name`, e.g. openxla/xla
+      issue_number: the issue (or PR) to set the status of
+      labels: the labels to add to the issue
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    endpoint = f"repos/{repo}/issues/{issue_number}/labels"
+    return self._make_request("POST", endpoint, labels=labels)
+
+  def get_user_orgs(self, username: str) -> requests.Response:
+    """Gets all public org memberships for a user.
+
+    https://docs.github.com/en/rest/orgs/orgs?apiVersion=2022-11-28#list-organizations-for-a-user
+
+    Arguments:
+      username: The user's GitHub username as a string.
+
+    Returns:
+      a requests.Response object containing the response from the API.
+
+    Raises:
+      requests.exceptions.HTTPError
+    """
+    endpoint = f"users/{username}/orgs"
+    return self._make_request("GET", endpoint, username=username)

third_party/xla/.github/workflows/gpu_benchmarks_nightly.yml

@@ -0,0 +1,211 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Nightly GPU Benchmarks
+permissions:
+  contents: read
+on:
+  workflow_dispatch: # Allows manual triggering
+  schedule:
+    - cron: "0 0 * * *" # Run at midnight every day
+
+
+jobs:
+  Tests:
+    strategy:
+      # Don't fail fast - want to see results for all builds even if one fails.
+      fail-fast: false
+      matrix:
+       job_info: [
+          {
+            os: "linux-x86-g2-48-l4-4gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            pretty_name: "XLA Linux x86 GPU L4 48 vcpu Presubmit",
+          },
+          {
+            os: "linux-x86-g2-16-l4-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
+            pretty_name: "XLA Linux x86 GPU L4 16 vcpu Presubmit",
+          },
+          {
+            os: "linux-x86-a4-224-b200-1gpu",
+            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cuda12.8-cudnn9.8:latest",
+            pretty_name: "XLA Linux x86 GPU A4 224 vcpu Presubmit",
+          },
+          # Expect more GPU types in the future.
+        ]
+    name: ${{ matrix.job_info.pretty_name }}
+    runs-on: ${{ matrix.job_info.os }}
+    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 240
+    env:
+      OUTPUT_DIR: ${{ github.workspace }}/output
+    steps:
+      - name: Checkout XLA
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Download Gemma Hlo Files
+        run: |
+          mkdir -p tmp_hlo
+          cd tmp_hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma2_2b_keras_jax.hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_call.hlo
+          wget https://storage.googleapis.com/xla-benchmarking-temp/gemma3_1b_flax_sample_loop.hlo
+          cd ..
+      - name: Configure GPU backend
+        run: |
+          ./configure.py --backend=CUDA --cuda_compiler=nvcc
+      - name: "Run build.py"
+        run: |
+          ./build_tools/ci/build.py --build="${{ matrix.job_info.pretty_name }}_github_actions"
+      - name: Run HLO tests and collect data
+
+        run: |
+          binary_dir="./bazel-out/k8-opt/bin/xla/tools"
+          mkdir -p "$OUTPUT_DIR"
+
+          # Run gpu_hlo_backend.hlo
+          HLO_FILE_GB="xla/tools/hlo_opt/tests/gpu_hlo_backend.hlo"
+          OUTPUT_PREFIX_GB="$OUTPUT_DIR/gpu_hlo_backend"
+          echo "Running GPU test: $HLO_FILE_GB"
+
+          $binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
+            --device_type=gpu \
+            --num_repeats=5 \
+            --use_spmd_partitioning \
+            --profile_execution=True \
+            --xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GB}_xspace.pb" \
+            "$HLO_FILE_GB" \
+            > "${OUTPUT_PREFIX_GB}.txt"
+
+          $binary_dir/compute_xspace_stats_main_gpu \
+            --input="${OUTPUT_PREFIX_GB}_xspace.pb" \
+            --device_type=GPU \
+            >> "${OUTPUT_PREFIX_GB}.txt"
+
+          cat "${OUTPUT_PREFIX_GB}.txt"
+
+          # Run gemma2_2b_keras_jax.hlo
+          HLO_FILE_GEMMA="tmp_hlo/gemma2_2b_keras_jax.hlo"
+          OUTPUT_PREFIX_GEMMA="$OUTPUT_DIR/gemma2_2b_keras_jax"
+          echo "Running GPU test: $HLO_FILE_GEMMA"
+          $binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
+            --device_type=gpu \
+            --num_repeats=5 \
+            --use_spmd_partitioning \
+            --profile_execution=True \
+            --xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA}_xspace.pb" \
+            "$HLO_FILE_GEMMA" \
+            > "${OUTPUT_PREFIX_GEMMA}.txt"
+
+          $binary_dir/compute_xspace_stats_main_gpu \
+            --input="${OUTPUT_PREFIX_GEMMA}_xspace.pb" \
+            --device_type=GPU \
+            >> "${OUTPUT_PREFIX_GEMMA}.txt"
+
+          cat "${OUTPUT_PREFIX_GEMMA}.txt"
+
+          echo "Output written to: ${OUTPUT_PREFIX_GB}.txt and ${OUTPUT_PREFIX_GEMMA}.txt"
+
+          # Run gemma3_1b_flax_call.hlo
+
+          HLO_FILE_GEMMA3_CALL="tmp_hlo/gemma3_1b_flax_call.hlo"
+          OUTPUT_PREFIX_GEMMA3_CALL="$OUTPUT_DIR/gemma3_1b_flax_call"
+          echo "Running GPU test: $HLO_FILE_GEMMA3_CALL"
+          $binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
+            --device_type=gpu \
+            --num_repeats=5 \
+            --use_spmd_partitioning \
+            --profile_execution=True \
+            --xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA3_CALL}_xspace.pb" \
+            "$HLO_FILE_GEMMA3_CALL" \
+            > "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
+
+            $binary_dir/compute_xspace_stats_main_gpu \
+            --input="${OUTPUT_PREFIX_GEMMA3_CALL}_xspace.pb" \
+            --device_type=GPU \
+            >> "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
+
+            cat "${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
+
+            echo "Output written to: ${OUTPUT_PREFIX_GEMMA3_CALL}.txt"
+
+            # Run gemma3_1b_flax_sample_loop.hlo
+          HLO_FILE_GEMMA3_SAMPLE_LOOP="tmp_hlo/gemma3_1b_flax_sample_loop.hlo"
+
+          OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP="$OUTPUT_DIR/gemma3_1b_flax_sample_loop"
+
+          echo "Running GPU test: $HLO_FILE_GEMMA3_SAMPLE_LOOP"
+
+          $binary_dir/multihost_hlo_runner/hlo_runner_main_gpu \
+            --device_type=gpu \
+            --num_repeats=5 \
+            --use_spmd_partitioning \
+            --profile_execution=True \
+            --xla_gpu_dump_xspace_to="${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}_xspace.pb" \
+            "$HLO_FILE_GEMMA3_SAMPLE_LOOP" \
+            > "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
+
+            $binary_dir/compute_xspace_stats_main_gpu \
+            --input="${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}_xspace.pb" \
+            --device_type=GPU \
+            >> "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
+
+            cat "${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
+
+
+            echo "Output written to: ${OUTPUT_PREFIX_GEMMA3_SAMPLE_LOOP}.txt"
+
+      - name: Upload HLO test output to GCS
+
+        run: |
+          GCS_BUCKET="gs://openxla-nightly-transient"
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+          DATE_FOLDER=$(date +%Y%m%d)
+          COMMIT_SHA="${{ github.sha }}"
+          RUN_ID="${{ github.run_id }}"
+
+
+          upload_to_gcs() {
+            local base_name="$1"
+            local gcs_file_name="${DATE_FOLDER}/${TIMESTAMP}_${{ matrix.job_info.os }}_run_${RUN_ID}_commit_${COMMIT_SHA}_${base_name}.txt"
+            echo "Uploading $OUTPUT_DIR/${base_name}.txt to $GCS_BUCKET/$gcs_file_name"
+            gcloud storage cp "$OUTPUT_DIR/${base_name}.txt" "$GCS_BUCKET/$gcs_file_name"
+          }
+
+          # Upload output for gpu_hlo_backend
+          GB_BASE_NAME="gpu_hlo_backend"
+          upload_to_gcs "$GB_BASE_NAME"
+
+          # Upload output for gemma2_2b_keras_jax
+          GEMMA_BASE_NAME="gemma2_2b_keras_jax"
+          upload_to_gcs "$GEMMA_BASE_NAME"
+
+          # Upload output for gemma3_1b_flax_call
+          GEMMA3_CALL_BASE_NAME="gemma3_1b_flax_call"
+          upload_to_gcs "$GEMMA3_CALL_BASE_NAME"
+
+          # Upload output for gemma3_1b_flax_sample_loop
+          GEMMA3_SAMPLE_LOOP_BASE_NAME="gemma3_1b_flax_sample_loop"
+          upload_to_gcs "$GEMMA3_SAMPLE_LOOP_BASE_NAME"
+
+      - name: Upload XSpace artifacts
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: gpu-xla-benchmarks-xspace-${{ matrix.job_info.os }}
+          path: ${{ github.workspace }}/output/*_xspace.pb

third_party/xla/.github/workflows/license_check.yml

@@ -0,0 +1,43 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: License Check
+
+permissions: {}
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  license-check:
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+        with:
+          persist-credentials: false
+      - name: "Set up Go"
+        uses: actions/setup-go@924ae3a1cded613372ab5595356fb5720e22ba16 # v6.5.0
+        with:
+          go-version: 'stable'
+      - name: "Run license check"
+        uses: google-ml-infra/actions/ci_check_license@432d36c6ac62d76fa99b3679758ccfbc10140fb6
+        with:
+          copyright-holder: "The XLA Authors"
+          license-type: "apache"
+          exclude-paths: "third_party/"
+          exclude-files: "BUILD _toc.yaml MHLOConfig.cmake.in"

third_party/xla/.github/workflows/nightly_benchmarks.yml

@@ -0,0 +1,202 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/nightly_benchmarks.yml
+name: Nightly Benchmarks
+permissions:
+  contents: read
+on:
+  workflow_dispatch: # Allows manual triggering
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  schedule:
+    - cron: "0 0 * * *" # Run at midnight every day
+
+jobs:
+  # =================================================================
+  # Job 1: Generate the matrix specifically for SCHEDULED benchmarks
+  # =================================================================
+  generate_matrix:
+    name: Generate Scheduled Matrix
+    uses: ./.github/workflows/generate_benchmark_matrix.yml
+    with:
+      workflow_type: 'SCHEDULED'
+      registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
+      checkout_ref: ${{ github.sha }} # On push/dispatch to main, github.sha is the commit SHA
+  run_benchmarks:
+    name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # config_id will indicate the workflow type, e.g., '_scheduled'
+    needs: generate_matrix
+    # Also check if the generate_matrix job itself was skipped due to its `if` condition
+    if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
+    strategy:
+      fail-fast: false # Continue running other benchmarks even if one fails
+      matrix:
+        benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
+    runs-on: ${{ matrix.benchmark_entry.runner_label }}
+    container: ${{ matrix.benchmark_entry.container_image }}
+    defaults:
+      run:
+        shell: bash
+    timeout-minutes: 60
+    env:
+      # Variables from the matrix
+      BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
+      CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
+      RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
+      CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
+      ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
+      IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
+      INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
+      XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
+      RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
+      TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
+      TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
+      HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }} # e.g., NVIDIA_GPU_B200
+      # Workflow context variables
+      CHECKOUT_REF: ${{ github.sha }}
+      COMMIT_SHA: ${{ github.sha }}
+      WORKFLOW_RUN_ID: ${{ github.run_id }}
+      # Paths and script locations
+      OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }} # config_id includes workflow type
+      SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
+      BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/nightly_baseline.yml
+      COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
+    steps:
+      - name: "Wait For Connection"
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+      - name: Print Job Info & Set Full Paths in ENV
+        run: |
+          # Resolve full paths based on GITHUB_WORKSPACE and relative paths defined in env
+          FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
+          RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
+          RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}" # Will use nightly_baseline.yml
+          RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
+          RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
+          echo "--- Benchmark Job Info (Nightly) ---"
+          echo "Config ID (from matrix): $CONFIG_ID"
+          echo "Benchmark Name: $BENCHMARK_NAME"
+          echo "Runner Label: $RUNNER_LABEL"
+          echo "Hardware Category: $HARDWARE_CATEGORY"
+          echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
+          echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
+          echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
+          echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH" # This will point to nightly_baseline.yml
+          echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
+          echo "Full Run Comparison Script Path: $RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
+          echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
+          echo "Current PWD: $(pwd)"
+          echo "---------------------------------------"
+          # Make resolved paths available to subsequent steps via GITHUB_ENV
+          {
+            echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
+            echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
+            echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
+            echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
+            echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
+          } >> "$GITHUB_ENV"
+          # Create the output directory; scripts will assume it exists
+          mkdir -p "$FULL_OUTPUT_DIR_PATH"
+          if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
+            echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
+            exit 1
+          fi
+      - name: Checkout OpenXLA Repository
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ env.CHECKOUT_REF }}
+          persist-credentials: false
+      - name: Build Binaries
+        id: build_binaries
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+      - name: Prepare Benchmark Artifact
+        id: prep_artifact
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+      - name: Run Benchmark and Generate Stats
+        id: run_hlo
+        env:
+          RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
+          STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
+          DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
+          LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+          # Other job-level env vars are automatically inherited
+        run: |
+          # The run_benchmark.sh script is expected to create $RESOLVED_OUTPUT_DIR/results.json
+          bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
+          # Print the content of results.json for debugging purposes
+          RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+          echo "--- Content of results.json (from workflow) ---"
+          if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
+            if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
+              jq '.' "$RESULTS_JSON_FILE_PATH"
+            else
+              echo "results.json may not be valid JSON or jq failed, printing with cat:"
+              cat "$RESULTS_JSON_FILE_PATH"
+            fi
+          elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then # Fallback for non-JSON or renamed output
+            echo "results.json not found, printing fallback .txt file:"
+            cat "${RESULTS_JSON_FILE_PATH}.txt"
+          else
+            echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
+          fi
+          echo "---------------------------------------------"
+      - name: Compare Results to Baseline
+        if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
+        run: |
+          echo "Starting baseline comparison (against nightly baseline)..."
+          bash "$RESOLVED_RUN_COMPARISON_SCRIPT" # This script calls the Python comparison script
+          echo "Baseline comparison finished."
+          echo "---------------------------------------------"
+      - name: Upload results.json directly to GCS
+        run: |
+          GCS_BUCKET="gs://openxla-nightly-transient"
+          RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+
+          # Check if the results file exists
+          if [ ! -f "$RESULTS_JSON_FILE_PATH" ]; then
+            echo "::error::results.json not found at $RESULTS_JSON_FILE_PATH"
+            exit 1
+          fi
+
+          # Construct a GCS object name
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+          DATE_FOLDER=$(date +%Y%m%d)
+          COMMIT_SHA_SHORT=$(echo "${{ github.sha }}" | cut -c1-8)
+          GCS_OBJECT_NAME="${BENCHMARK_NAME}/${DATE_FOLDER}/${TIMESTAMP}_run_${WORKFLOW_RUN_ID}_commit_${COMMIT_SHA_SHORT}.json"
+
+          echo "Uploading $RESULTS_JSON_FILE_PATH to $GCS_BUCKET/$GCS_OBJECT_NAME"
+          gcloud storage cp "$RESULTS_JSON_FILE_PATH" "$GCS_BUCKET/$GCS_OBJECT_NAME"
+      - name: Upload Benchmark Artifacts
+        if: always()
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: results-${{ env.CONFIG_ID }}
+          path: ${{ env.RESOLVED_OUTPUT_DIR }}
+          retention-days: 14
+          if-no-files-found: error

third_party/xla/.github/workflows/oneAPI_Postsubmit.yml

@@ -0,0 +1,53 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+name: Postsubmit CI
+
+permissions:
+  contents: read
+
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+jobs:
+  Tests:
+    name: XLA Linux X86 GPU ONEAPI POSTSUBMIT
+    runs-on: linux-x64-battlemage-256-1gpu-intel
+    container: 
+      image: us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest  # zizmor: ignore[unpinned-images]
+      options: --device=/dev/dri
+  
+    timeout-minutes: 120
+
+    defaults:
+      run:
+        shell: bash
+
+    steps:
+      - name: Checking out openxla/xla
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0
+        with:
+          path: openxla/xla
+          persist-credentials: false
+
+      - name: Run oneAPI CI
+        working-directory: openxla/xla
+        run: bash build_tools/sycl/ci_docker_test.sh

third_party/xla/.github/workflows/postsubmit_benchmark.yml

@@ -0,0 +1,235 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/postsubmit_benchmarks.yml
+name: Postsubmit - Run Benchmarks
+
+permissions:
+  contents: read
+
+on:
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  push:
+    branches:
+      - main
+
+concurrency:
+  # Run every push to main and do not cancel in-progress jobs; the timeout is 60 minutes.
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: false
+
+jobs:
+  # =================================================================
+  # Job 1: Generate the matrix specifically for POSTSUBMIT benchmarks
+  # =================================================================
+  generate_matrix:
+    name: Generate Postsubmit Matrix
+    # Condition: Run if manually dispatched OR if it's a push event to the main branch.
+    if: |
+      github.event_name == 'workflow_dispatch' ||
+      (github.event_name == 'push' && github.ref == 'refs/heads/main')
+    uses: ./.github/workflows/generate_benchmark_matrix.yml
+    with:
+      workflow_type: 'POSTSUBMIT'
+      registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
+      checkout_ref: ${{ github.sha }} # On push/dispatch to main, github.sha is the commit SHA
+
+  run_benchmarks:
+    name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # config_id will indicate the workflow type, e.g., '_postsubmit'
+    needs: generate_matrix
+    # Also check if the generate_matrix job itself was skipped due to its `if` condition
+    if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
+
+    strategy:
+      fail-fast: false # Continue running other benchmarks even if one fails
+      matrix:
+        benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
+
+    runs-on: ${{ matrix.benchmark_entry.runner_label }}
+    container: ${{ matrix.benchmark_entry.container_image }}
+
+    defaults:
+      run:
+        shell: bash
+
+    timeout-minutes: 60
+
+    env:
+      # Variables from the matrix
+      BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
+      CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
+      RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
+      CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
+      ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
+      IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
+      INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
+      XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
+      RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
+      TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
+      TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
+      HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }} # e.g., NVIDIA_GPU_B200
+
+      # Workflow context variables
+      CHECKOUT_REF: ${{ github.sha }} # For postsubmit, use the commit SHA from the push/dispatch event
+      COMMIT_SHA: ${{ github.sha }} # For postsubmit, use the commit SHA
+      WORKFLOW_RUN_ID: ${{ github.run_id }}
+
+      # Paths and script locations
+      OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }} # config_id includes workflow type
+      SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
+      BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/postsubmit_baseline.yml
+      COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
+
+    steps:
+      - name: "Wait For Connection"
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      - name: Print Job Info & Set Full Paths in ENV
+        run: |
+          # Resolve full paths based on GITHUB_WORKSPACE and relative paths defined in env
+          FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
+          RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
+          RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}" # Will use postsubmit_baseline.yml
+          RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
+          RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
+
+          echo "--- Benchmark Job Info (Postsubmit) ---"
+          echo "Config ID (from matrix): $CONFIG_ID"
+          echo "Benchmark Name: $BENCHMARK_NAME"
+          echo "Runner Label: $RUNNER_LABEL"
+          echo "Hardware Category: $HARDWARE_CATEGORY"
+          echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
+          echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
+          echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
+          echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH" # This will point to postsubmit_baseline.yml
+          echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
+          echo "Full Run Comparison Script Path: $RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
+          echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
+          echo "Current PWD: $(pwd)"
+          echo "---------------------------------------"
+
+          # Make resolved paths available to subsequent steps via GITHUB_ENV
+          {
+            echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
+            echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
+            echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
+            echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
+            echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
+          } >> "$GITHUB_ENV"
+
+          # Create the output directory; scripts will assume it exists
+          mkdir -p "$FULL_OUTPUT_DIR_PATH"
+          if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
+            echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
+            exit 1
+          fi
+
+      - name: Checkout OpenXLA Repository
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ env.CHECKOUT_REF }}
+          persist-credentials: false
+
+      - name: Build Binaries
+        id: build_binaries
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+
+      - name: Prepare Benchmark Artifact
+        id: prep_artifact
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+
+      - name: Run Benchmark and Generate Stats
+        id: run_hlo
+        env:
+          RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
+          STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
+          DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
+          LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+          # Other job-level env vars are automatically inherited
+        run: |
+          # The run_benchmark.sh script is expected to create $RESOLVED_OUTPUT_DIR/results.json
+          bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
+
+          # Print the content of results.json for debugging purposes
+          RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+          echo "--- Content of results.json (from workflow) ---"
+          if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
+            if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
+              jq '.' "$RESULTS_JSON_FILE_PATH"
+            else
+              echo "results.json may not be valid JSON or jq failed, printing with cat:"
+              cat "$RESULTS_JSON_FILE_PATH"
+            fi
+          elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then # Fallback for non-JSON or renamed output
+            echo "results.json not found, printing fallback .txt file:"
+            cat "${RESULTS_JSON_FILE_PATH}.txt"
+          else
+            echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
+          fi
+          echo "---------------------------------------------"
+
+      - name: Compare Results to Baseline
+        if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
+        run: |
+          echo "Starting baseline comparison (against postsubmit baseline)..."
+          bash "$RESOLVED_RUN_COMPARISON_SCRIPT" # This script calls the Python comparison script
+          echo "Baseline comparison finished."
+          echo "---------------------------------------------"
+
+      - name: Upload results.json directly to GCS
+        run: |
+          GCS_BUCKET="gs://openxla-postsubmit-transient"
+          RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+
+          # Check if the results file exists
+          if [ ! -f "$RESULTS_JSON_FILE_PATH" ]; then
+            echo "::error::results.json not found at $RESULTS_JSON_FILE_PATH"
+            exit 1
+          fi
+
+          # Construct a GCS object name
+          TIMESTAMP=$(date +%Y%m%d_%H%M%S)
+          DATE_FOLDER=$(date +%Y%m%d)
+          COMMIT_SHA_SHORT=$(echo "${{ github.sha }}" | cut -c1-8)
+          GCS_OBJECT_NAME="${BENCHMARK_NAME}/${DATE_FOLDER}/${TIMESTAMP}_run_${WORKFLOW_RUN_ID}_commit_${COMMIT_SHA_SHORT}.json"
+
+          echo "Uploading $RESULTS_JSON_FILE_PATH to $GCS_BUCKET/$GCS_OBJECT_NAME"
+          gcloud storage cp "$RESULTS_JSON_FILE_PATH" "$GCS_BUCKET/$GCS_OBJECT_NAME"
+
+      - name: Upload Benchmark Artifacts
+        if: always()
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: results-${{ env.CONFIG_ID }}
+          path: ${{ env.RESOLVED_OUTPUT_DIR }}
+          retention-days: 14
+          if-no-files-found: error

third_party/xla/.github/workflows/pr_size_check.py

@@ -0,0 +1,85 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Checks PR size and posts a comment if it's too large."""
+
+import os
+import sys
+
+import github_api
+
+THRESHOLD_MESSAGES = {
+    1000: (
+        "🔴 This PR has a very large delta of over 1000. In order to enable an"
+        " effective code reivew, please break the PR down into smaller and more"
+        " focused PRs.\nSee the [Small"
+        " CLs](https://google.github.io/eng-practices/review/developer/small-cls.html)"
+        " Google Eng practice for more details on how to write compact PRs."
+    ),
+    500: (
+        "⚠️ This PR has a large delta of over 500. Consider breaking the PR"
+        " down into smaller PRs for a faster code review.\nSee the [Small"
+        " CLs](https://google.github.io/eng-practices/review/developer/small-cls.html)"
+        " Google Eng practice for more details on how to write compact PRs."
+    ),
+}
+
+
+def main():
+  pr_number = os.getenv("PR_NUMBER")
+  if not pr_number:
+    print("PR_NUMBER is not set.")
+    sys.exit(1)
+
+  additions = int(os.getenv("ADDITIONS", "0"))
+  deletions = int(os.getenv("DELETIONS", "0"))
+  total_delta = additions + deletions
+
+  print(
+      f"PR #{pr_number} additions: {additions}, deletions: {deletions}, total:"
+      f" {total_delta}"
+  )
+
+  comment = (
+      f"PR #{pr_number} additions: {additions}, deletions: {deletions}, total:"
+      f" {total_delta}"
+  )
+  exceeded = False
+
+  for threshold in sorted(THRESHOLD_MESSAGES.keys(), reverse=True):
+    if total_delta >= threshold:
+      comment += f"\n\n{THRESHOLD_MESSAGES[threshold]}"
+      exceeded = True
+      break
+
+  if exceeded:
+    token = os.getenv("GH_TOKEN")
+    repo = os.getenv("GITHUB_REPOSITORY")
+    if not repo:
+      print("GITHUB_REPOSITORY is not set.")
+      sys.exit(1)
+
+    try:
+      api = github_api.GitHubAPI(token)
+      api.write_issue_comment(repo, int(pr_number), comment)
+      print("Comment posted.")
+    except Exception as e:  # pylint: disable=broad-except
+      print(f"Failed to post comment: {e}")
+      sys.exit(1)
+  else:
+    print("PR size is within limits.")
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/.github/workflows/pr_size_check.yml

@@ -0,0 +1,47 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: PR Size Check
+permissions:
+  pull-requests: write
+on:
+  # zizmor: ignore[dangerous-triggers]
+  pull_request_target:
+    branches:
+      - main
+
+jobs:
+  check-size:
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    env:
+      GH_TOKEN: ${{ github.token }}
+      PR_NUMBER: ${{ github.event.pull_request.number }}
+      ADDITIONS: ${{ github.event.pull_request.additions }}
+      DELETIONS: ${{ github.event.pull_request.deletions }}
+    timeout-minutes: 5
+    if: |
+      github.repository == 'openxla/xla' && (
+        github.event.sender.type == 'User' ||
+        contains(github.event.pull_request.body, 'FORCE_TEST_ACTIONS')
+      )
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Check PR Size"
+        run: python3 .github/workflows/pr_size_check.py

third_party/xla/.github/workflows/presubmit_benchmark.yml

@@ -0,0 +1,210 @@
+# Copyright 2025 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# .github/workflows/presubmit_benchmarks.yml
+name: Presubmit - Run Benchmarks
+
+permissions:
+  contents: read
+
+on:
+  workflow_dispatch:
+    inputs:
+      halt-for-connection:
+        description: 'Should this workflow run wait for a remote connection?'
+        type: choice
+        required: true
+        default: 'no'
+        options:
+          - 'yes'
+          - 'no'
+  pull_request:
+    branches:
+      - main
+
+concurrency:
+  # Group all jobs for a given PR together, and cancel presubmit if a new commit is pushed.
+  group: ${{ github.workflow }}-${{ github.head_ref }}
+  cancel-in-progress: true
+
+jobs:
+  # ================================================================
+  # Job 1: Generate the matrix specifically for PRESUBMIT benchmarks
+  # ================================================================
+  generate_matrix:
+    name: Generate Presubmit Matrix
+    uses: ./.github/workflows/generate_benchmark_matrix.yml
+    with:
+      workflow_type: 'PRESUBMIT'
+      registry_file: 'xla/tools/benchmarks/registries/default_registry.yml'
+      checkout_ref: ${{ github.event.pull_request.head.sha || github.sha }}
+
+  run_benchmarks:
+    name: Run Benchmark (${{ matrix.benchmark_entry.config_id }}) # Use the full config_id: {name}_{hw-short}_{topo}_{workflow} e.g., gemma3_1b_flax_call_l4_1h1d_presubmit
+    needs: generate_matrix
+    # Also check if the generate_matrix job itself was skipped due to its `if` condition
+    if: success() && needs.generate_matrix.result != 'skipped' && needs.generate_matrix.outputs.matrix_include_json != '[]' && needs.generate_matrix.outputs.matrix_include_json != ''
+
+    strategy:
+      fail-fast: false
+      matrix:
+        benchmark_entry: ${{ fromJson(needs.generate_matrix.outputs.matrix_include_json || '[]') }}
+
+    runs-on: ${{ matrix.benchmark_entry.runner_label }}
+    container: ${{ matrix.benchmark_entry.container_image }}
+
+    defaults:
+      run:
+        shell: bash
+
+    timeout-minutes: 60
+
+    env:
+      BENCHMARK_NAME: ${{ matrix.benchmark_entry.benchmark_name }}
+      CONFIG_ID: ${{ matrix.benchmark_entry.config_id }}
+      RUNNER_LABEL: ${{ matrix.benchmark_entry.runner_label }}
+      CONTAINER_IMAGE: ${{ matrix.benchmark_entry.container_image }}
+      ARTIFACT_LOCATION: ${{ matrix.benchmark_entry.artifact_location }}
+      IS_GCS_ARTIFACT: ${{ matrix.benchmark_entry.is_gcs_artifact }}
+      INPUT_FORMAT: ${{ matrix.benchmark_entry.input_format }}
+      XLA_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.xla_compilation_flags) }}
+      RUNTIME_FLAGS_JSON: ${{ toJson(matrix.benchmark_entry.runtime_flags) }}
+      TARGET_METRICS_JSON: ${{ toJson(matrix.benchmark_entry.target_metrics) }}
+      TOPOLOGY_JSON: ${{ toJson(matrix.benchmark_entry.topology) }}
+      HARDWARE_CATEGORY: ${{ matrix.benchmark_entry.hardware_category }}
+
+
+      CHECKOUT_REF: ${{ github.event.pull_request.head.sha || github.sha }}
+      COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
+      WORKFLOW_RUN_ID: ${{ github.run_id }}
+
+      OUTPUT_DIR_NAME: benchmark_output_${{ matrix.benchmark_entry.config_id }}
+      SCRIPT_DIR_RELATIVE: .github/workflows/benchmarks
+      BASELINE_YAML_FILE_RELATIVE: xla/tools/benchmarks/baseline/presubmit_baseline.yml
+      COMPARISON_SCRIPT_RELATIVE: .github/workflows/benchmarks/compare_with_baseline.py
+
+    steps:
+      - name: "Wait For Connection"
+        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+        with:
+          halt-dispatch-input: ${{ inputs.halt-for-connection }}
+
+      - name: Print Job Info & Set Full Paths in ENV
+        run: |
+          FULL_OUTPUT_DIR_PATH="${GITHUB_WORKSPACE}/${OUTPUT_DIR_NAME}"
+          RESOLVED_SCRIPT_DIR_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}"
+          RESOLVED_BASELINE_YAML_PATH="${GITHUB_WORKSPACE}/${BASELINE_YAML_FILE_RELATIVE}"
+          RESOLVED_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${COMPARISON_SCRIPT_RELATIVE}"
+          RESOLVED_RUN_COMPARISON_SCRIPT_PATH="${GITHUB_WORKSPACE}/${SCRIPT_DIR_RELATIVE}/run_comparison.sh"
+
+
+          echo "--- Benchmark Job Info ---"
+          echo "Config ID (from matrix, used for dir): $CONFIG_ID"
+          echo "Benchmark Name: $BENCHMARK_NAME"
+          echo "Runner Label: $RUNNER_LABEL"
+          echo "Hardware Category: $HARDWARE_CATEGORY"
+          echo "Output Directory Name (relative): $OUTPUT_DIR_NAME"
+          echo "Full Output Directory Path: $FULL_OUTPUT_DIR_PATH"
+          echo "Full Script Directory Path: $RESOLVED_SCRIPT_DIR_PATH"
+          echo "Full Baseline YAML Path: $RESOLVED_BASELINE_YAML_PATH"
+          echo "Full Comparison Script Path: $RESOLVED_COMPARISON_SCRIPT_PATH"
+          echo "Full Compare Results Script Path: $RESOLVED_COMPARE_RESULTS_SCRIPT_PATH"
+          echo "GITHUB_WORKSPACE: ${GITHUB_WORKSPACE}"
+          echo "Current PWD: $(pwd)"
+          echo "--------------------------"
+
+          # These are now available to subsequent steps via env context
+          {
+            echo "RESOLVED_OUTPUT_DIR=$FULL_OUTPUT_DIR_PATH"
+            echo "RESOLVED_SCRIPT_DIR=$RESOLVED_SCRIPT_DIR_PATH"
+            echo "RESOLVED_BASELINE_YAML=$RESOLVED_BASELINE_YAML_PATH"
+            echo "RESOLVED_COMPARISON_SCRIPT=$RESOLVED_COMPARISON_SCRIPT_PATH"
+            echo "RESOLVED_RUN_COMPARISON_SCRIPT=$RESOLVED_RUN_COMPARISON_SCRIPT_PATH"
+          } >> "$GITHUB_ENV"
+
+          # Create the output directory - scripts will assume it exists
+          mkdir -p "$FULL_OUTPUT_DIR_PATH"
+          if [ ! -d "$FULL_OUTPUT_DIR_PATH" ]; then
+            echo "::error::Failed to create output directory: $FULL_OUTPUT_DIR_PATH"
+            exit 1
+          fi
+
+      - name: Checkout OpenXLA Repository
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          ref: ${{ env.CHECKOUT_REF }}
+          persist-credentials: false
+
+      - name: Build Binaries
+        id: build_binaries
+        # Use the RESOLVED_SCRIPT_DIR and RESOLVED_OUTPUT_DIR from GITHUB_ENV
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/build_binaries.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+
+      - name: Prepare Benchmark Artifact
+        id: prep_artifact
+        run: |
+          bash "${RESOLVED_SCRIPT_DIR}/prepare_artifact.sh"
+        env:
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+
+      - name: Run Benchmark and Generate Stats
+        id: run_hlo
+        env:
+          RUNNER_BINARY: "${{ steps.build_binaries.outputs.runner_binary }}"
+          STATS_BINARY: "${{ steps.build_binaries.outputs.stats_binary }}"
+          DEVICE_TYPE_FLAG: "${{ steps.build_binaries.outputs.device_type_flag }}"
+          LOCAL_ARTIFACT_PATH: "${{ steps.prep_artifact.outputs.artifact_local_path }}"
+          OUTPUT_DIR: ${{ env.RESOLVED_OUTPUT_DIR }}
+          # Other job-level env vars are automatically inherited
+        run: |
+          # The run_benchmark.sh script is expected to create $OUTPUT_DIR/results.json
+          bash "${RESOLVED_SCRIPT_DIR}/run_benchmark.sh"
+
+          # Print the content of results.json for debugging
+          RESULTS_JSON_FILE_PATH="${RESOLVED_OUTPUT_DIR}/results.json"
+          echo "--- Content of results.json (from workflow) ---"
+          if [ -f "$RESULTS_JSON_FILE_PATH" ]; then
+            if command -v jq &> /dev/null && jq '.' "$RESULTS_JSON_FILE_PATH" > /dev/null 2>&1; then
+              jq '.' "$RESULTS_JSON_FILE_PATH"
+            else
+              echo "results.json may not be valid JSON or jq failed, printing with cat:"
+              cat "$RESULTS_JSON_FILE_PATH"
+            fi
+          elif [ -f "${RESULTS_JSON_FILE_PATH}.txt" ]; then
+            echo "results.json not found, printing fallback .txt file:"
+            cat "${RESULTS_JSON_FILE_PATH}.txt"
+          else
+            echo "::warning::Neither results.json nor results.json.txt found at $RESULTS_JSON_FILE_PATH"
+          fi
+          echo "---------------------------------------------"
+
+      - name: Compare Results to Baseline
+        if: (steps.run_hlo.outcome == 'success' || (steps.run_hlo.outcome == 'failure' && (hashFiles(format('{0}/results.json', env.RESOLVED_OUTPUT_DIR)) != '' || hashFiles(format('{0}/results.json.txt', env.RESOLVED_OUTPUT_DIR)) != ''))) && env.skip_comparison_due_to_yaml_install_failure != 'true'
+        run: |
+          echo "Starting baseline comparison..."
+          bash "$RESOLVED_RUN_COMPARISON_SCRIPT"
+          echo "Baseline comparison finished."
+          echo "---------------------------------------------"
+
+      - name: Upload Benchmark Artifacts
+        if: always()
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: results-${{ env.CONFIG_ID }}
+          path: ${{ env.RESOLVED_OUTPUT_DIR }}
+          retention-days: 7
+          if-no-files-found: error

third_party/xla/.github/workflows/rocm_ci.yml

@@ -0,0 +1,175 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+name: CI ROCm
+permissions:
+  contents: read
+on:
+  pull_request:
+  workflow_dispatch:
+  workflow_call:
+    inputs:
+      rocm_xla_branch:
+        description: 'Branch name to fetch ROCm XLA scripts from'
+        required: false
+        default: 'rocm-dev-infra'
+        type: string
+      runner_label:
+        description: 'Runner label for GPU jobs'
+        required: false
+        default: 'amd-do-linux-xla-gpu-gfx950-1'
+        type: string
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: ${{ github.ref != 'main' }}
+
+jobs:
+  rocm-config:
+    runs-on: ubuntu-latest
+    outputs:
+      docker-image: ${{ steps.out.outputs.docker-image }}
+    steps:
+      - id: out
+        shell: bash
+        run: |
+          # hermetic llvm based on  ghcr.io/rocm/jax-ubu22.rocm7.2.4:latest
+          echo "docker-image=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT"
+
+  jax:
+    needs: rocm-config
+    if: ${{ github.event_name != 'pull_request' || github.base_ref == 'main' }}
+    name: JAX Linux x86 AMD Instinct GPU
+    runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
+    timeout-minutes: 80
+    env:
+      DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
+    container: &rocm_container
+      image: ${{ needs.rocm-config.outputs.docker-image }} # zizmor: ignore[unpinned-images] image is pinned in the rocm-config job.
+      volumes:
+        - /data/ci-cert.crt:/data/ci-cert.crt
+        - /data/ci-cert.key:/data/ci-cert.key
+        - /data:/data
+      options: >-
+        --device=/dev/dri
+        --device=/dev/kfd
+        --ipc=host
+        --shm-size=64G
+        --cap-add=SYS_PTRACE
+        --security-opt=seccomp=unconfined
+        --tmpfs /root/.cache/bazel:rw,exec,size=80g
+        --group-add video
+        --env-file /etc/podinfo/gha-gpu-isolation-settings
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Checking out openxla/xla
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          repository: openxla/xla
+          persist-credentials: false
+
+      - name: Checkout JAX
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          repository: jax-ml/jax
+          path: jax
+          persist-credentials: false
+
+      - name: CPU Info
+        run: lscpu
+
+      - name: ROCm info
+        run: /opt/rocm/bin/rocminfo
+
+      - name: Run JAX Unit Tests
+        timeout-minutes: 60
+        working-directory: jax
+        run: |
+          ./ci/run_bazel_test_rocm_rbe.sh \
+            --override_repository=xla="${GITHUB_WORKSPACE}" \
+            --override_module=xla="${GITHUB_WORKSPACE}" \
+            --config=single_gpu \
+            --//jax:build_jaxlib=wheel \
+            --//jax:build_jax=true \
+            --local_test_jobs=1 \
+            --action_env=JAX_ENABLE_X64=0 \
+            --repo_env=HERMETIC_PYTHON_VERSION=3.14 \
+            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
+            --bes_keywords=xla \
+            --bes_keywords=jax \
+            --bes_keywords=upstream \
+            --bes_keywords=gpu \
+            --bes_keywords=pull_request
+  xla:
+    needs: rocm-config
+    name: XLA Linux x86 AMD Instinct GPU
+    runs-on: ${{ inputs.runner_label || 'amd-do-linux-xla-gpu-gfx950-1' }}
+    timeout-minutes: 220
+    env:
+      DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
+      EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh
+    container: *rocm_container
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Checking out openxla/xla
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          repository: openxla/xla
+          persist-credentials: false
+
+      - name: Fetch ROCm XLA Upstream Script from ROCm/xla
+        run: |
+          wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}"
+          chmod +x build_tools/rocm/execute_ci_build_upstream.sh
+
+      - name: CPU Info
+        run: lscpu
+
+      - name: ROCm info
+        run: /opt/rocm/bin/rocminfo
+
+      - name: Test XLA [single_gpu]
+        timeout-minutes: 120
+        run: |
+          build_tools/rocm/execute_ci_build_upstream.sh \
+            --config=rocm_ci \
+            --config=rocm_rbe \
+            --config=ci_single_gpu \
+            --local_test_jobs=1 \
+            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
+            --local_test_jobs=1 \
+            --internal_spawn_scheduler \
+            --strategy=TestRunner=dynamic \
+            --bes_keywords=xla \
+            --bes_keywords=upstream \
+            --bes_keywords=gpu \
+            --bes_keywords=pull_request
+      - name: Test XLA [rocm_cpu]
+        timeout-minutes: 80
+        run: |
+          build_tools/rocm/execute_ci_build_upstream.sh \
+            --config=rocm_ci \
+            --config=rocm_rbe \
+            --config=ci_rocm_cpu \
+            --local_test_jobs=200 \
+            --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
+            --bes_keywords=xla \
+            --bes_keywords=upstream \
+            --bes_keywords=cpu \
+            --bes_keywords=pull_request

third_party/xla/.github/workflows/rollback_notification.py

@@ -0,0 +1,77 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Notifies PR authors of rollbacks on committed PRs."""
+import os
+import re
+from typing import Generator, Sequence
+
+import github_api
+
+
+def get_reverted_commit_hashes(message: str) -> list[str]:
+  """Searches a commit message for `reverts <commit hash>` and returns the found SHAs.
+
+  Arguments:
+    message: the commit message to search
+
+  Returns:
+    A list of SHAs as strings.
+  """
+  print("Head commit message:", message, sep="\n")
+  regex = re.compile(r"reverts ([0-9a-f]{5,40})", flags=re.IGNORECASE)
+  commit_hashes = regex.findall(message)
+  print(f"Found commit hashes reverted in this commit: {commit_hashes}")
+  return commit_hashes
+
+
+def get_associated_prs(
+    api: github_api.GitHubAPI,
+    commit_hashes: Sequence[str],
+) -> Generator[int, None, None]:
+  """Finds PRs associated with commits.
+
+  Arguments:
+    api: GitHubAPI object which will be used to make requests
+    commit_hashes: A sequence of SHAs which may have PRs associated with them
+
+  Yields:
+    Associated pairs of (PR number, SHA), both as strings
+  """
+  regex = re.compile(r"PR #(\d+)")
+  for commit_hash in commit_hashes:
+    response = api.get_commit("openxla/xla", commit_hash)
+    message = response["commit"]["message"]
+    if maybe_match := regex.match(message):
+      pr_number = maybe_match.group(1)
+      print(f"Found PR #{pr_number} associated with commit hash {commit_hash}")
+      yield int(pr_number)
+  print(f"Didn't find any PRs associated with commit hashes: {commit_hashes}")
+
+
+def main():
+  api = github_api.GitHubAPI(os.getenv("GH_TOKEN"))
+  head_commit = api.get_commit("openxla/xla", "HEAD")
+  commit_hashes = get_reverted_commit_hashes(head_commit["commit"]["message"])
+
+  for pr_number in get_associated_prs(api, commit_hashes):
+    sha = head_commit["sha"]
+    api.write_issue_comment(
+        "openxla/xla", pr_number, f"This PR was rolled back in {sha}!"
+    )
+    api.set_issue_status("openxla/xla", pr_number, "open")
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/.github/workflows/rollback_notification.yml

@@ -0,0 +1,41 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Rollback Notification
+permissions:
+  pull-requests: write
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  rollback-notification:
+    if: contains(github.event.head_commit.message, 'revert')
+    runs-on: ubuntu-22.04
+    defaults:
+      run:
+        shell: bash
+    env:
+      GH_TOKEN: ${{ github.token }}
+      HEAD_COMMIT_MESSAGE: ${{ github.event.head_commit.message }}
+    timeout-minutes: 6
+    steps:
+      - name: "Checking out repository"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: "Check if PR was rolled back"
+        run: python3 .github/workflows/rollback_notification.py
+

third_party/xla/.github/workflows/scorecards-analysis.yml

@@ -0,0 +1,72 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+# This workflow uses actions that are not certified by GitHub. They are provided
+# by a third-party and are governed by separate terms of service, privacy
+# policy, and support documentation.
+
+name: Scorecards supply-chain security
+on:
+  # For Branch-Protection check. Only the default branch is supported. See
+  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
+  branch_protection_rule:
+  # To guarantee Maintained check is occasionally updated. See
+  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
+  schedule:
+    - cron: '26 3 * * 2'
+  push:
+    branches: ["main"]
+
+# Declare default permissions as read only.
+permissions: read-all
+
+jobs:
+  analysis:
+    if: github.repository == 'openxla/xla' # Don't do this in forks
+    name: Scorecards analysis
+    runs-on: ubuntu-latest
+    permissions:
+      # Needed to upload the results to code-scanning dashboard.
+      security-events: write
+      # Needed to publish results and get a badge (see publish_results below).
+      id-token: write
+
+    steps:
+      - name: "Checkout code"
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+
+      - name: "Run analysis"
+        uses: ossf/scorecard-action@4eaacf0543bb3f2c246792bd56e8cdeffafb205a # v2.4.3
+        with:
+          results_file: results.sarif
+          results_format: sarif
+          publish_results: true
+
+      # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
+      # format to the repository Actions tab.
+      - name: "Upload artifact"
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: SARIF file
+          path: results.sarif
+          retention-days: 5
+
+      # Upload the results to GitHub's code scanning dashboard (optional).
+      # Commenting out will disable upload of results to your repo's Code Scanning dashboard
+      - name: "Upload to code-scanning"
+        uses: github/codeql-action/upload-sarif@8aad20d150bbac5944a9f9d289da16a4b0d87c1e # v4.36.2
+        with:
+          sarif_file: results.sarif

third_party/xla/.github/workflows/zizmor.yml

@@ -0,0 +1,41 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+name: Zizmor
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: read
+
+jobs:
+  zizmor:
+    name: GitHub Actions Security Analysis
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+
+      - name: Install uv dependencies
+        uses: astral-sh/setup-uv@fac544c07dec837d0ccb6301d7b5580bf5edae39 # v8.2.0
+
+      - name: Run zizmor
+        run:  uvx zizmor --format=github .github xla
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

third_party/xla/.gitignore

@@ -0,0 +1,38 @@
+# Copyright 2023 The OpenXLA Authors
+#
+
+
+# Ignore the bazel directories when this directory is used as a local repository.
+bazel-*
+bazel-bin
+bazel-out
+bazel-testlogs
+
+# Ignore Bazel lock file
+MODULE.bazel.lock
+
+# Ignore files produced by `configure`
+.tf_configure.bazelrc
+xla_configure.bazelrc
+tools/python_bin_path.sh
+
+# Emacs autosaves
+*~
+\#*\#
+
+# Visual Studio files
+.vs/
+.vscode/
+*.sdf
+*.opensdf
+*.VC.opendb
+*.suo
+*.user
+
+# Ignore clangd files and directories: https://openxla.org/xla/lsp
+.cache
+compile_commands.json
+external
+
+# Ignore python cache files.
+__pycache__/

third_party/xla/.kokoro/generate_index_html.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+# Copyright 2022 Google LLC All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Generates a handy index.html with a bunch of Kokoro links for GitHub
+# presubmits.
+# Usage: generate_index_html.sh /path/to/output/index.html
+
+tee "$1" <<EOF
+<html>
+<head>
+<title>#$KOKORO_GITHUB_PULL_REQUEST_NUMBER_xla | $(basename "$KOKORO_JOB_NAME")</title>
+</head>
+<body>
+<h1>XLA Job Logs and Links</h1>
+<h2>Job Details</h2>
+<ul>
+<li>Job name: $KOKORO_JOB_NAME</li>
+<li>Job pool: $KOKORO_JOB_POOL</li>
+<li>Job ID: $KOKORO_BUILD_ID</li>
+<li>Current HEAD Piper Changelist (may be empty): cl/${KOKORO_PIPER_CHANGELIST:-not available}</li>
+<li>Pull Request Number: $KOKORO_GITHUB_PULL_REQUEST_NUMBER_xla</li>
+<li>Pull Request Link: <a href="$KOKORO_GITHUB_PULL_REQUEST_URL_xla">$KOKORO_GITHUB_PULL_REQUEST_URL_xla</a></li>
+<li>Commit: $KOKORO_GIT_COMMIT_xla</li>
+</ul>
+<h2>Googlers-Only Links</h2>
+<ul>
+<li><a href="http://sponge2/$KOKORO_BUILD_ID">Sponge2</a></li>
+<li><a href="http://sponge/target:$KOKORO_JOB_NAME">Sponge - recent jobs</a></li>
+</ul>
+<h2>Non-Googler Links</h2>
+<ul>
+<li><a href="https://source.cloud.google.com/results/invocations/$KOKORO_BUILD_ID">ResultStore</a></li>
+</ul>
+</body></html>
+EOF

third_party/xla/.kokoro/macos/build.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Copyright 2022 Google LLC All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
+# -o history: record shell history
+set -euox pipefail -o history
+
+# Generate a templated results file to make output accessible to everyone
+"$KOKORO_ARTIFACTS_DIR"/github/xla/.kokoro/generate_index_html.sh "$KOKORO_ARTIFACTS_DIR"/index.html
+
+# TODO(ddunleavy) figure out how to best move this into build.py
+cd "${KOKORO_ARTIFACTS_DIR}/github/xla"
+
+"$KOKORO_ARTIFACTS_DIR"/github/xla/build_tools/ci/build.py --build=XLA_MACOS_X86_CPU_KOKORO

third_party/xla/.kokoro/windows/build.bat

@@ -0,0 +1,17 @@
+:: Copyright 2026 Google LLC. All Rights Reserved.
+::
+:: Licensed under the Apache License, Version 2.0 (the "License");
+:: you may not use this file except in compliance with the License.
+:: You may obtain a copy of the License at
+::
+::     http://www.apache.org/licenses/LICENSE-2.0
+::
+:: Unless required by applicable law or agreed to in writing, software
+:: distributed under the License is distributed on an "AS IS" BASIS,
+:: WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+:: See the License for the specific language governing permissions and
+:: limitations under the License.
+:: ==============================================================================
+
+bash -l %0/../windows_build.sh %*
+exit /b %ERRORLEVEL%

third_party/xla/.kokoro/windows/windows_build.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright 2022 Google LLC All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# -e: abort script if one command fails
+# -u: error if undefined variable used
+# -o pipefail: entire command fails if pipe fails. watch out for yes | ...
+# Note: set -x <code> +x around anything you want to have logged.
+set -euox pipefail
+
+# Generate a templated results file to make output accessible to everyone
+"$KOKORO_ARTIFACTS_DIR"/github/xla/.kokoro/generate_index_html.sh "$KOKORO_ARTIFACTS_DIR"/index.html
+
+cd "${KOKORO_ARTIFACTS_DIR}/github/xla"
+
+export PATH="$PATH:/c/Python38"
+
+TARGET_FILTER="-//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/..."
+TAGS_FILTER="-no_oss,-oss_excluded,-gpu,-no_windows,-windows_excluded"
+/c/tools/bazel.exe test \
+  --keep_going \
+  --build_tag_filters=$TAGS_FILTER  --test_tag_filters=$TAGS_FILTER \
+  -- //xla/... $TARGET_FILTER |& grep -v "violates visibility of" \
+  || { exit 1; }
+
+exit 0

third_party/xla/AUTHORS

@@ -0,0 +1,9 @@
+# This is the list of OpenXLA's significant contributors.
+#
+# This does not necessarily list everyone who has contributed code,
+# especially since many employees of one corporation may be contributing.
+# To see the full list of contributors, see the revision history in
+# source control.
+Arm Limited
+Google LLC
+NVIDIA Corporation

third_party/xla/BUILD.bazel

@@ -0,0 +1,28 @@
+load("@rules_license//rules:license.bzl", "license")
+
+package(
+    default_applicable_licenses = [":license"],
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])
+
+license(
+    name = "license",
+    package_name = "xla",
+    license_kinds = ["@rules_license//licenses/spdx:Apache-2.0"],
+)
+
+filegroup(
+    name = "clang_tidy_config",
+    srcs = [".clang-tidy"],
+    visibility = ["//visibility:public"],
+)
+
+genrule(
+    name = "clang_tidy_config_noerrors",
+    srcs = ["//:.clang-tidy"],
+    outs = [".clang-tidy.noerrors"],
+    cmd = "grep -v 'WarningsAsErrors' $< > $@",
+    visibility = ["//visibility:public"],
+)

third_party/xla/CONTRIBUTING.md

@@ -0,0 +1,4 @@
+# Contributing to OpenXLA
+
+For information on how to contribute to OpenXLA, see
+[Contributing to OpenXLA](docs/contributing.md)

third_party/xla/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

third_party/xla/MODULE.bazel

@@ -0,0 +1,312 @@
+module(name = "xla")
+
+##############################################################
+# Bazel module dependencies
+
+# go/keep-sorted start
+bazel_dep(name = "abseil-cpp", version = "20260526.0", repo_name = "com_google_absl")
+bazel_dep(name = "abseil-py", version = "2.1.0", repo_name = "absl_py")
+bazel_dep(name = "bazel_features", version = "1.36.0")
+bazel_dep(name = "bazel_skylib", version = "1.8.2")
+bazel_dep(name = "boringssl", version = "0.20250818.0")
+bazel_dep(name = "curl", version = "8.12.0.bcr.1")
+bazel_dep(name = "google_benchmark", version = "1.8.5", repo_name = "com_google_benchmark")
+bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest")
+bazel_dep(name = "grpc", version = "1.81.0", repo_name = "com_github_grpc_grpc")
+bazel_dep(name = "gutil", version = "20260309.0.bcr.1", repo_name = "com_google_gutil")
+bazel_dep(name = "highway", version = "1.3.0", repo_name = "com_google_highway")
+bazel_dep(name = "jsoncpp", version = "1.9.6", repo_name = "jsoncpp_git")
+bazel_dep(name = "or-tools", version = "9.12", repo_name = "com_google_ortools")
+bazel_dep(name = "platforms", version = "1.0.0")
+bazel_dep(name = "protobuf", version = "32.1", repo_name = "com_google_protobuf")
+bazel_dep(name = "pybind11_abseil", version = "202402.0")
+bazel_dep(name = "pybind11_bazel", version = "3.0.0")
+bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7")
+bazel_dep(name = "re2", version = "2025-11-05.bcr.1", repo_name = "com_googlesource_code_re2")
+bazel_dep(name = "rules_cc", version = "0.2.17")
+bazel_dep(name = "rules_java", version = "8.16.1")
+bazel_dep(name = "rules_license", version = "1.0.0")
+bazel_dep(name = "rules_python", version = "1.8.5")
+bazel_dep(name = "rules_shell", version = "0.6.1")
+bazel_dep(name = "snappy", version = "1.2.1")
+bazel_dep(name = "zlib", version = "1.3.2")
+# go/keep-sorted end
+
+# Only for compatibility, not directly used, change repo_name to None after upgrading Bazel to latest 7.x
+bazel_dep(name = "eigen", version = "4.0.0-20241125.bcr.3", repo_name = "DO_NOT_USE_eigen")
+bazel_dep(name = "grpc-java", version = "1.75.0", repo_name = "DO_NOT_USE_grpc_java")
+
+# TODO: publish an official version of rules_ml_toolchain to BCR
+bazel_dep(name = "rules_ml_toolchain")
+
+# To calculate integrity:
+# wget -O temp_module_archive.tar.gz <archive URL>
+# HASH=$(openssl dgst -sha256 -binary temp_module_archive.tar.gz | openssl base64 -A)
+# echo "sha256-${HASH}"
+archive_override(
+    module_name = "rules_ml_toolchain",
+    integrity = "sha256-axspTL7LmKCLz+WWQUYfcuET2zD6kbRPzEnZgJm1umU=",
+    strip_prefix = "rules_ml_toolchain-b11745590f513ec55b32e2d126073576fde18c71",
+    urls = ["https://github.com/google-ml-infra/rules_ml_toolchain/archive/b11745590f513ec55b32e2d126073576fde18c71.tar.gz"],
+)
+
+# TODO: Upstream the patch?
+single_version_override(
+    module_name = "grpc",
+    patch_strip = 1,
+    patches = [
+        "//third_party/grpc:grpc.patch",
+    ],
+    version = "1.81.0",
+)
+
+# TODO: Upstream those patch?
+single_version_override(
+    module_name = "abseil-cpp",
+    version = "20260526.0",
+    patch_strip = 1,
+    patches = [
+        "//third_party/absl:btree.patch",
+        "//third_party/absl:build_dll.patch",
+        "//third_party/absl:endian.patch",
+    ],
+)
+
+single_version_override(
+    module_name = "protobuf",
+    patch_strip = 1,
+    patches = [
+        "//third_party/protobuf:protobuf_arena.patch",
+        "//third_party/protobuf:fix_message_lite_incomplete_type.patch",
+        # The following patch is needed to fix package loading for `bazel query "deps(//xla/...)"`.
+        "//third_party/protobuf:fix_python_dist_package.patch",
+    ],
+    version = "32.1",
+)
+
+# Use an unreleased version of googletest
+archive_override(
+    module_name = "googletest",
+    patch_strip = 1,
+    patches = [
+        "//third_party/googletest:0001-Add-ASSERT_OK-EXPECT_OK-ASSERT_OK_AND_ASSIGN-macros.patch",
+    ],
+    strip_prefix = "googletest-28e9d1f26771c6517c3b4be10254887673c94018",
+    urls = ["https://github.com/google/googletest/archive/28e9d1f26771c6517c3b4be10254887673c94018.zip"],
+)
+
+##############################################################
+# C++ dependencies
+
+# TODO: most of them can be Bazel modules, but we need a release strategy for them
+third_party = use_extension("//third_party/extensions:third_party.bzl", "third_party_ext")
+use_repo(
+    third_party,
+    "FXdiv",
+    "XNNPACK",
+    "cpuinfo",
+    "cudnn_frontend_archive",
+    "cutlass_archive",
+    "dlpack",
+    "ducc",
+    "eigen_archive",
+    "farmhash_archive",
+    "farmhash_gpu_archive",
+    "fmt",
+    "gloo",
+    "highwayhash",
+    "hwloc",
+    "implib_so",
+    "llvm-raw",
+    "llvm_openmp",
+    "mkl_dnn_acl_compatible",
+    "ml_dtypes_py",
+    "mpitrampoline",
+    "nanobind",
+    "nccl_archive",
+    "net_zstd",
+    "nvshmem",
+    "nvtx_archive",
+    "onednn",
+    "onednn_async",
+    "org_sqlite",
+    "pthreadpool",
+    "raft",
+    "rapids_logger",
+    "riegeli",
+    "rmm",
+    "robin_map",
+    "rocm_device_libs",
+    "shardy",
+    "slinky",
+    "spdlog",
+    "stablehlo",
+    "transformer_engine",
+    "triton",
+    "uv",
+)
+
+##############################################################
+# Python toolchain and pypi dependencies
+
+single_version_override(
+    module_name = "rules_python",
+    patch_strip = 1,
+    patches = [
+        "//third_party/py:rules_python_pip_version.patch",
+        "//third_party/py:rules_python_freethreaded.patch",
+        "//third_party/py:rules_python_versions.patch",
+        "//third_party/py:rules_python_scope.patch",
+    ],
+    version = "1.8.5",
+)
+
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.defaults(python_version = "3.11")
+python.toolchain(python_version = "3.11")
+use_repo(python, "pythons_hub")
+
+pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
+pip.whl_mods(
+    additive_build_content = """
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+cc_library(
+    name = "numpy_headers_2",
+    hdrs = glob(["site-packages/numpy/_core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/_core/include/",
+)
+cc_library(
+    name = "numpy_headers_1",
+    hdrs = glob(["site-packages/numpy/core/include/**/*.h"]),
+    strip_include_prefix="site-packages/numpy/core/include/",
+)
+cc_library(
+    name = "numpy_headers",
+    deps = [":numpy_headers_2", ":numpy_headers_1"],
+)
+""",
+    hub_name = "pypi_mods",
+    whl_name = "numpy",
+)
+pip.parse(
+    extra_hub_aliases = {
+        "numpy": ["numpy_headers"],
+    },
+    hub_name = "xla_pypi",
+    python_version = "3.11",
+    requirements_lock = "//:requirements_lock_3_11.txt",
+    whl_modifications = {
+        "@pypi_mods//:numpy.json": "numpy",
+    },
+)
+use_repo(pip, "pypi_mods", pypi = "xla_pypi")
+
+python_version_ext = use_extension("@rules_ml_toolchain//extensions:python_version.bzl", "python_version_ext")
+use_repo(python_version_ext, "python_version_repo")
+
+##############################################################
+# Other dependencies via module extensions
+
+### TSL
+tsl_extension = use_extension("//third_party/extensions:tsl.bzl", "tsl_extension")
+use_repo(tsl_extension, "tsl")
+
+### LLVM
+llvm = use_extension("//third_party/extensions:llvm.bzl", "llvm_extension")
+use_repo(llvm, "llvm-project")
+
+### pybind
+pybind11_internal_configure = use_extension(
+    "@pybind11_bazel//:internal_configure.bzl",
+    "internal_configure_extension",
+)
+use_repo(pybind11_internal_configure, "pybind11")
+
+### RBE
+rbe_config = use_extension("//third_party/extensions:rbe_config.bzl", "rbe_config_ext")
+use_repo(rbe_config, "ml_build_config_platform")
+
+remote_execution_configure = use_extension("//third_party/extensions:remote_execution_configure.bzl", "remote_execution_configure_ext")
+use_repo(remote_execution_configure, "local_config_remote_execution")
+
+### From @rules_ml_toolchain
+cuda_configure = use_extension("@rules_ml_toolchain//extensions:cuda_configure.bzl", "cuda_configure_ext")
+use_repo(cuda_configure, "local_config_cuda")
+
+nccl_configure = use_extension("@rules_ml_toolchain//extensions:nccl_configure.bzl", "nccl_configure_ext")
+use_repo(nccl_configure, "local_config_nccl")
+
+nccl_redist = use_extension("@rules_ml_toolchain//extensions:nccl_redist.bzl", "nccl_redist_ext")
+use_repo(nccl_redist, "cuda_nccl")
+
+sycl_configure = use_extension("@rules_ml_toolchain//extensions:sycl_configure.bzl", "sycl_configure_ext")
+use_repo(sycl_configure, "level_zero", "local_config_sycl", "oneapi", "zero_loader")
+
+local_sysroot = use_extension("@rules_ml_toolchain//extensions:local_sysroot_configure.bzl", "local_sysroot_ext")
+use_repo(local_sysroot, "local_sysroot_config")
+
+cuda_redist_init_ext = use_extension("@rules_ml_toolchain//extensions:cuda_redist_init.bzl", "cuda_redist_init_ext")
+use_repo(
+    cuda_redist_init_ext,
+    "cuda_cublas",
+    "cuda_cudart",
+    "cuda_cudnn",
+    "cuda_cufft",
+    "cuda_cupti",
+    "cuda_curand",
+    "cuda_cusolver",
+    "cuda_cusparse",
+    "cuda_driver",
+    "cuda_nvcc",
+    "cuda_nvdisasm",
+    "cuda_nvjitlink",
+    "cuda_nvml",
+    "cuda_nvrtc",
+    "cuda_nvtx",
+    "cuda_nvvm",
+)
+
+# Override cuda_cccl to use a newer version defined in the cccl_extension.
+cccl_extension = use_extension("//third_party/cccl:cccl_extension.bzl", "cccl_extension")
+use_repo(cccl_extension, "cuda_cccl")
+
+override_repo(
+    cuda_redist_init_ext,
+    cuda_cccl = "cuda_cccl",
+)
+
+nvshmem_redist = use_extension("@rules_ml_toolchain//extensions:nvshmem_redist.bzl", "nvshmem_redist_ext")
+use_repo(nvshmem_redist, "nvidia_nvshmem")
+
+toolchain_ext = use_extension("@rules_ml_toolchain//extensions:toolchain.bzl", "toolchain_ext")
+use_repo(toolchain_ext, "llvm18_linux_x86_64", "llvm_linux_x86_64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_sycl")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda")
+
+### Other local config repos
+local_clang_configure = use_extension("@rules_ml_toolchain//extensions:local_clang_configure.bzl", "local_clang_configure_ext")
+use_repo(local_clang_configure, "local_config_clang")
+
+rocm_configure = use_extension("//third_party/extensions:rocm_configure.bzl", "rocm_configure_ext")
+use_repo(rocm_configure, "config_rocm_hipcc", "local_config_rocm")
+
+tensorrt_configure = use_extension("//third_party/extensions:tensorrt_configure.bzl", "tensorrt_configure_ext")
+use_repo(tensorrt_configure, "local_config_tensorrt")
+
+pjrt_nightly_timestamp = use_extension("//build_tools/pjrt_wheels:nightly.bzl", "nightly_timestamp_repo_bzlmod")
+use_repo(pjrt_nightly_timestamp, "nightly_timestamp")
+
+pjrt_rc_number = use_extension("//build_tools/pjrt_wheels:release_candidate.bzl", "rc_number_repo_bzlmod")
+use_repo(pjrt_rc_number, "rc_number")
+
+### libdrm (dedicated extension to avoid Bzlmod merging issues)
+libdrm_ext = use_extension("//third_party/libdrm:extension.bzl", "libdrm_ext")
+use_repo(libdrm_ext, "libdrm")

third_party/xla/README.md

@@ -0,0 +1,49 @@
+# XLA
+
+XLA (Accelerated Linear Algebra) is an open-source machine learning (ML)
+compiler for GPUs, CPUs, and ML accelerators.
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="docs/images/openxla_dark.svg">
+  <img alt="OpenXLA Ecosystem" src="docs/images/openxla.svg">
+</picture>
+
+The XLA compiler takes models from popular ML frameworks such as PyTorch,
+TensorFlow, and JAX, and optimizes them for high-performance execution across
+different hardware platforms including GPUs, CPUs, and ML accelerators.
+
+[openxla.org](https://openxla.org/) is the project's website.
+
+## Get started
+
+If you want to use XLA to compile your ML project, refer to the corresponding
+documentation for your ML framework:
+
+* [PyTorch](https://pytorch.org/xla)
+* [TensorFlow](https://www.tensorflow.org/xla)
+* [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html)
+
+If you're not contributing code to the XLA compiler, you don't need to clone and
+build this repo. Everything here is intended for XLA contributors who want to
+develop the compiler and XLA integrators who want to debug or add support for ML
+frontends and hardware backends.
+
+## Contribute
+
+If you'd like to contribute to XLA, review
+[How to Contribute](docs/contributing.md) and then see the
+[developer guide](docs/developer_guide.md).
+
+## Contacts
+
+*   For questions, contact the maintainers - maintainers at openxla.org
+
+## Resources
+
+*   [Community Resources](https://github.com/openxla/community)
+
+## Code of Conduct
+
+While under TensorFlow governance, all community spaces for SIG OpenXLA are
+subject to the
+[TensorFlow Code of Conduct](https://github.com/tensorflow/tensorflow/blob/master/CODE_OF_CONDUCT.md).

third_party/xla/WORKSPACE

@@ -0,0 +1,159 @@
+# buildifier: disable=load-on-top
+workspace(name = "xla")
+
+# Initialize the XLA repository and all dependencies.
+#
+# The cascade of load() statements and xla_workspace?() calls works around the
+# restriction that load() statements need to be at the top of .bzl files.
+# E.g. we can not retrieve a new repository with http_archive and then load()
+# a macro from that repository in the same file.
+
+load(":workspace4.bzl", "xla_workspace4")
+
+xla_workspace4()
+
+load(":workspace3.bzl", "xla_workspace3")
+
+xla_workspace3()
+
+load("@bazel_features//:deps.bzl", "bazel_features_deps")
+
+bazel_features_deps()
+
+# Initialize hermetic C++
+load("@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl", "cc_toolchain_deps")
+
+cc_toolchain_deps()
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_cuda")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_sycl")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64_rocm")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64")
+
+register_toolchains("@rules_ml_toolchain//cc:linux_aarch64_linux_aarch64_cuda")
+
+# Initialize hermetic Python
+load("//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("@rules_ml_toolchain//py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    requirements = {
+        "3.11": "//:requirements_lock_3_11.txt",
+        "3.12": "//:requirements_lock_3_12.txt",
+    },
+)
+
+load("@rules_ml_toolchain//py:python_register_toolchain.bzl", "python_register_toolchain")
+
+python_register_toolchain()
+
+load("@rules_ml_toolchain//py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
+load(":workspace2.bzl", "xla_workspace2")
+
+xla_workspace2()
+
+load(":workspace1.bzl", "xla_workspace1")
+
+xla_workspace1()
+
+load(":workspace0.bzl", "xla_workspace0")
+
+xla_workspace0()
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_json_init_repository.bzl",
+    "cuda_json_init_repository",
+)
+
+cuda_json_init_repository()
+
+load(
+    "@cuda_redist_json//:distributions.bzl",
+    "CUDA_REDISTRIBUTIONS",
+    "CUDNN_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_redist_init_repositories.bzl",
+    "cuda_redist_init_repositories",
+    "cudnn_redist_init_repository",
+)
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_redist_versions.bzl",
+    "REDIST_VERSIONS_TO_BUILD_TEMPLATES",
+)
+load("//third_party/cccl:workspace.bzl", "CCCL_3_2_0_DIST_DICT", "CCCL_GITHUB_VERSIONS_TO_BUILD_TEMPLATES")
+
+cuda_redist_init_repositories(
+    cuda_redistributions = CUDA_REDISTRIBUTIONS | CCCL_3_2_0_DIST_DICT,
+    redist_versions_to_build_templates = REDIST_VERSIONS_TO_BUILD_TEMPLATES | CCCL_GITHUB_VERSIONS_TO_BUILD_TEMPLATES,
+)
+
+cudnn_redist_init_repository(
+    cudnn_redistributions = CUDNN_REDISTRIBUTIONS,
+)
+
+load(
+    "@rules_ml_toolchain//gpu/cuda:cuda_configure.bzl",
+    "cuda_configure",
+)
+
+cuda_configure(name = "local_config_cuda")
+
+load(
+    "@rules_ml_toolchain//gpu/nccl:nccl_redist_init_repository.bzl",
+    "nccl_redist_init_repository",
+)
+
+nccl_redist_init_repository()
+
+load(
+    "@rules_ml_toolchain//gpu/nccl:nccl_configure.bzl",
+    "nccl_configure",
+)
+
+nccl_configure(name = "local_config_nccl")
+
+load(
+    "@rules_ml_toolchain//gpu/nvshmem:nvshmem_json_init_repository.bzl",
+    "nvshmem_json_init_repository",
+)
+
+nvshmem_json_init_repository()
+
+load(
+    "@nvshmem_redist_json//:distributions.bzl",
+    "NVSHMEM_REDISTRIBUTIONS",
+)
+load(
+    "@rules_ml_toolchain//gpu/nvshmem:nvshmem_redist_init_repository.bzl",
+    "nvshmem_redist_init_repository",
+)
+
+nvshmem_redist_init_repository(
+    nvshmem_redistributions = NVSHMEM_REDISTRIBUTIONS,
+)
+
+# This is used for building nightly PJRT wheels.
+load("//build_tools/pjrt_wheels:nightly.bzl", "nightly_timestamp_repo")
+
+nightly_timestamp_repo(name = "nightly_timestamp")
+
+# This is used for building release candidate PJRT wheels.
+load("//build_tools/pjrt_wheels:release_candidate.bzl", "rc_number_repo")
+
+rc_number_repo(name = "rc_number")

third_party/xla/build_tools/BUILD

@@ -0,0 +1,28 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+load("//xla:pytype.bzl", "pytype_strict_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "test_utils",
+    testonly = True,
+    srcs = ["test_utils.py"],
+    visibility = ["//visibility:public"],
+)

third_party/xla/build_tools/ci/BUILD

@@ -0,0 +1,142 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+load("@bazel_skylib//rules:diff_test.bzl", "diff_test")
+load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
+
+# copybara:comment_begin(oss-only)
+load("@llvm_linux_x86_64//:version.bzl", _llvm_version = "VERSION")
+load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
+load("@xla//third_party/rules_python/python:py_test.bzl", "py_test")
+load("//xla:pytype.bzl", "pytype_strict_binary", "pytype_strict_library")
+load("//xla/tsl:tsl.bzl", "if_oss")
+
+_LLVM_REPO_NAME = "llvm{}_linux_x86_64".format(_llvm_version)
+# copybara:comment_end_and_uncomment_begin
+# _LLVM_REPO_NAME = ""
+# copybara:uncomment_end
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_binary(
+    name = "build",
+    srcs = ["build.py"],
+)
+
+pytype_strict_library(
+    name = "clang_tidy_diff_lib",
+    srcs = ["clang_tidy_diff.py"],
+    deps = [
+        "//build_tools/lint:diff_parser",
+    ],
+)
+
+pytype_strict_binary(
+    name = "clang_tidy_diff",
+    srcs = ["clang_tidy_diff.py"],
+    data = if_oss([
+        ":clang_apply_replacements_bin",
+    ]),
+    main = "clang_tidy_diff.py",
+    deps = [
+        ":clang_tidy_diff_lib",
+    ],
+)
+
+py_test(
+    name = "clang_tidy_diff_test",
+    srcs = ["clang_tidy_diff_test.py"],
+    deps = [
+        ":clang_tidy_diff_lib",
+        "@absl_py//absl/testing:absltest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
+genrule(
+    name = "generated_build_commands",
+    outs = ["generated_commands.txt"],
+    cmd = "$(location //build_tools/ci:build) --dump_commands > $(OUTS)",
+    tags = ["not_run:arm"],
+    tools = [":build"],
+)
+
+diff_test(
+    name = "build_command_golden_test",
+    failure_message = """Regenerate with `PYTHONDONTWRITEBYTECODE=1 python3 build.py --dump_commands > golden_commands.txt`.""",
+    file1 = "golden_commands.txt",
+    file2 = ":generated_build_commands",
+    tags = ["not_run:arm"],
+)
+
+sh_binary(
+    name = "parallel_gpu_execute",
+    srcs = ["parallel_gpu_execute.sh"],
+)
+
+# TODO: When upstream rules expose clang-tidy use that instead of explicitly creating one.
+expand_template(
+    name = "generate_wrapper_script",
+    out = "clang_tidy_wrapper.sh",
+    substitutions = {
+        "%LLVM_REPO_NAME%": _LLVM_REPO_NAME,
+    },
+    template = "clang_tidy_wrapper.sh.tpl",
+)
+
+sh_binary(
+    name = "clang_tidy_bin",
+    srcs = [":generate_wrapper_script"],
+    data = [
+        # This not for g3. Only for running clang-tidy on CI/locally for OSS workflows.
+        "@{}//:all".format(_LLVM_REPO_NAME),  # buildifier: disable=platform-specific-binaries
+    ],
+    tags = [
+        "manual",  # This binary only runs via explicit targeting.
+        "notap",  # Not for g3.
+    ],
+    visibility = ["//visibility:public"],
+)
+
+expand_template(
+    name = "generate_apply_replacements_wrapper_script",
+    out = "clang_apply_replacements_wrapper.sh",
+    substitutions = {
+        "%LLVM_REPO_NAME%": _LLVM_REPO_NAME,
+    },
+    template = "clang_apply_replacements_wrapper.sh.tpl",
+)
+
+sh_binary(
+    name = "clang_apply_replacements_bin",
+    srcs = [":generate_apply_replacements_wrapper_script"],
+    data = [
+        # This not for g3. Only for running clang-tidy on CI/locally for OSS workflows.
+        "@{}//:all".format(_LLVM_REPO_NAME),  # buildifier: disable=platform-specific-binaries
+    ],
+    tags = [
+        "manual",  # This binary only runs via explicit targeting.
+        "notap",  # Not for g3.
+    ],
+    visibility = ["//visibility:public"],
+)
+
+sh_binary(
+    name = "run_clang_tidy",
+    srcs = ["run_clang_tidy.sh"],
+)

third_party/xla/build_tools/ci/clang_apply_replacements_wrapper.sh.tpl

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -e
+
+echoerr() {
+  RED='\033[1;31m'
+  NOCOLOR='\033[0m'
+  printf "${RED}ERROR:${NOCOLOR} %s\n" "$*" >&2
+}
+
+REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-apply-replacements"
+if [ ! -f "$REAL_BIN" ]; then
+  echoerr "Failed to locate clang-apply-replacements binary at: $REAL_BIN"
+  exit 1
+fi
+echo "Using clang-apply-replacements at: " $REAL_BIN
+REAL_LIB_DIR="$(dirname "$REAL_BIN")/../lib"
+export LD_LIBRARY_PATH="${REAL_LIB_DIR}:${LD_LIBRARY_PATH}"
+exec "$REAL_BIN" "$@"

third_party/xla/build_tools/ci/clang_tidy_diff.py

@@ -0,0 +1,804 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Filters Clang-Tidy errors based on modified lines in a Git patch.
+
+This script reads a unified diff (patch file) to determine which lines of which
+files have been modified. It then parses the Bazel Build Event Protocol (BEP)
+JSON file to find the paths of all generated Clang-Tidy YAML reports. For each
+report, it checks if any reported errors fall within the modified line ranges
+and prints the matching errors.
+"""
+
+from __future__ import annotations
+
+import argparse
+import bisect
+import collections
+from collections.abc import Callable, Sequence
+import dataclasses
+import functools
+import itertools
+import json
+import logging
+import os
+import pathlib
+import subprocess
+import sys
+import tempfile
+from typing import IO, Protocol, TypedDict
+
+from build_tools.lint import diff_parser
+
+
+@dataclasses.dataclass(frozen=True)
+class AppConfig:
+  """Configuration for the ClangTidyDiffFilter application.
+
+  Attributes:
+    patch: Path to the patch file containing the diff or None if no diff
+      filtering should be applied.
+    repo_root: Absolute path to the repository root.
+    bep_file: Path to the Bazel Build Event Protocol JSON file.
+    warnings_as_errors: If True, treat Clang-Tidy warnings as errors.
+    fix: If True, apply fixes to the source files.
+  """
+
+  patch: str | None
+  repo_root: str
+  bep_file: str
+  warnings_as_errors: bool
+  fix: bool = False
+
+
+@dataclasses.dataclass(frozen=True)
+class Diagnostic:
+  """Represents a single Clang-Tidy diagnostic.
+
+  Attributes:
+    file_path: The path to the file where the diagnostic was found, relative to
+      the repo root.
+    line_num: The 1-based line number of the diagnostic.
+    col_num: The 1-based column number of the diagnostic.
+    level: The severity level of the diagnostic (e.g., "warning", "error").
+    name: The name of the Clang-Tidy check (e.g., "clang-diagnostic-error").
+    message: The diagnostic message.
+    yaml_file: The path to the .clang-tidy.yaml file where this diagnostic was
+      reported.
+    has_replacements: Whether the diagnostic has replacements available. For
+      some diagnostics, clang-tidy can automatically fix the issue, and this
+      attribute will be True if there are one or more replacements available.
+  """
+
+  file_path: str
+  line_num: int
+  col_num: int
+  level: str
+  name: str
+  message: str
+  yaml_file: str
+  has_replacements: bool
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class DiagnosticSummary:
+  """Summary for all Clang-Tidy diagnostics for a file.
+
+  Attributes:
+    file_path: The path to the file where the diagnostics were found, relative
+      to the repo root.
+    was_skipped: Whether the file was skipped due to not being in the git diff.
+    total: The total number of diagnostics found in the report.
+    matched: The number of diagnostics that matched the diff.
+  """
+
+  file_path: str
+  was_skipped: bool
+  total: int
+  matched: int
+
+
+# Can't use TypedDict with classes because linting will complain about fields
+# starting with capital letters.
+ClangTidyDiagnostic = TypedDict(
+    "ClangTidyDiagnostic",
+    {
+        "DiagnosticName": str,
+        "Message": str,
+        "FilePath": str,
+        "FileOffset": int,
+        "Level": str,
+        "HasReplacements": bool,
+    },
+    total=False,
+)
+
+
+ClangTidyReport = TypedDict(
+    "ClangTidyReport",
+    {
+        "MainSourceFile": str,
+        "Diagnostics": list[ClangTidyDiagnostic],
+    },
+    total=False,
+)
+
+
+class ApplyFixes(Protocol):
+  """A callable for applying fixes from clang-tidy diagnostics."""
+
+  def __call__(self, temp_dir: pathlib.Path) -> None:
+    """Applies fixes to the source files in the given temporary directory.
+
+    Args:
+      temp_dir: The temporary directory containing the source files to apply
+        fixes to.
+    """
+    ...
+
+
+def clang_apply_fixes(bin_path: pathlib.Path, temp_dir: pathlib.Path) -> None:
+  """Applies fixes to the source files in the given temporary directory.
+
+  Args:
+    bin_path: Path to the clang-apply-replacements binary.
+    temp_dir: The temporary directory containing the source files to apply fixes
+      to.
+  """
+  cmd = [
+      bin_path.as_posix(),
+      "-remove-change-desc-files",
+      temp_dir.as_posix(),
+  ]
+  _logger().info("Running command: %r", " ".join(cmd))
+  subprocess.check_output(cmd)
+
+
+def _logger() -> logging.Logger:
+  """Returns the logger for this module."""
+  return logging.getLogger(__name__)
+
+
+def _set_log_level(log_level: str) -> None:
+  """Sets the log level for the application."""
+  logging.basicConfig(
+      format=(
+          "[%(asctime)s] [%(levelname)s][%(filename)s:%(funcName)s:%(lineno)d]"
+          " %(message)s"
+      ),
+      datefmt="%H:%M:%S",
+  )
+
+  def _set_level(level: int) -> None:
+    _logger().setLevel(level)
+
+  match log_level:
+    case "DEBUG":
+      _set_level(logging.DEBUG)
+    case "INFO":
+      _set_level(logging.INFO)
+    case "WARNING":
+      _set_level(logging.WARNING)
+    case "ERROR":
+      _set_level(logging.ERROR)
+    case _:
+      raise ValueError(f"Unsupported log level: {log_level}")
+
+
+def _resolve_clang_apply_replacements(
+    clang_apply_replacements_bin: pathlib.Path | None,
+) -> pathlib.Path | None:
+  """Resolve path to the clang-apply-replacements binary.
+
+  If the path is explicitly provided, use that. Otherwise, try to find it using
+  bazel-runfiles.
+
+  Args:
+    clang_apply_replacements_bin: The path to the clang-apply-replacements
+      binary, or None if it should be resolved from the environment.
+
+  Returns:
+    The path to the clang-apply-replacements binary, or None if it could not be
+    resolved.
+  """
+  if clang_apply_replacements_bin:
+    return clang_apply_replacements_bin
+  runfiles_dir = os.environ.get("RUNFILES_DIR")
+  if not runfiles_dir:
+    return None
+  runfiles_path = (
+      pathlib.Path(runfiles_dir)
+      / "xla/build_tools/ci/clang_apply_replacements_bin"
+  )
+  if runfiles_path.exists():
+    return runfiles_path
+  _logger().warning(
+      "clang-apply-replacements binary not found in RUNFILES_DIR: %s",
+      runfiles_path,
+  )
+  return None
+
+
+def parse_diff(diff_path: str) -> dict[str, set[int]]:
+  """Parses a unified diff file using diff_parser and returns a dictionary mapping filenames to a set of modified line numbers."""
+  with open(diff_path, "r") as f:
+    diff_str = f.read()
+  hunks = diff_parser.parse_hunks(diff_str)
+  file_to_lines: dict[str, set[int]] = collections.defaultdict(set)
+  for hunk in hunks:
+    for line_no, _ in hunk.added_lines():
+      file_to_lines[hunk.file].add(line_no)
+  return file_to_lines
+
+
+def get_line_offsets(file_path: str) -> tuple[int, ...]:
+  """Returns a list of byte offsets for the start of each line in the file."""
+  offsets = [0]
+  with open(file_path, "rb") as f:
+    while f.readline():
+      offsets.append(f.tell())
+  return tuple(offsets)
+
+
+def offset_to_line(offsets: Sequence[int], offset: int) -> int:
+  """Converts a byte offset to a 1-based line number using binary search."""
+  if not offsets:
+    return -1
+  # bisect_right returns the index where the offset would be inserted after
+  # existing entries. Since offsets contains start of lines, bisect_right - 1
+  # gives the line index (0-based).
+  return bisect.bisect_right(offsets, offset)
+
+
+def normalize_path(path: str, repo_root: str) -> str:
+  """Normalizes a path to be relative to the repo root.
+
+  This is not foolproof for all possible path formats,
+  but it handles common cases seen locally and in CI.
+
+  Args:
+    path: The path to normalize.
+    repo_root: The absolute path to the repository root.
+
+  Returns:
+      The normalized path as a string.
+  """
+  if not path:
+    return ""
+  p = pathlib.Path(path)
+  # Handle local absolute paths under repo_root (CI Runner paths)
+  # This handles /__w/xla/xla by removing the prefix.
+  if p.is_absolute() and p.is_relative_to(repo_root):
+    return p.relative_to(repo_root).as_posix()
+  # Handle bazel execroot paths
+  if "execroot" in p.parts:
+    idx = p.parts.index("execroot")
+    if idx + 2 < len(p.parts):
+      return pathlib.Path(*p.parts[idx + 2 :]).as_posix()
+  # Handle remote execution paths
+  # p is like "/b/f/w/xla/..."
+  # NB: We don't quite know the top level directory to look for in the remote
+  # path, but since all CPP sources live mostly under "xla/" we use it as the
+  # anchor. We also include "third_party" as an anchor since some files may
+  # be under that directory.
+  _top_level_pkgs = ("xla", "third_party")
+  for pkg in _top_level_pkgs:
+    if pkg in p.parts:
+      parts = list(p.parts)
+      idx = parts.index(pkg)  # Find FIRST occurrence in remote path
+      return pathlib.Path(*parts[idx:]).as_posix()
+  return path
+
+
+def parse_bep(bep_path: str, repo_root: str) -> list[str]:
+  """Parses a Bazel BEP JSON file and returns a list of paths to .clang-tidy.yaml files.
+
+  Args:
+    bep_path: Path to the Bazel BEP JSON file.
+    repo_root: Absolute path to the repository root.
+
+  Returns:
+    A list of paths to .clang-tidy.yaml files.
+
+  Raises:
+    ValueError: If a file entry in the BEP is missing 'name' or 'pathPrefix'
+      fields.
+  """
+  yaml_files: list[str] = []
+  with open(bep_path, "r") as f:
+    for line in f:
+      try:
+        event = json.loads(line)
+        if "namedSetOfFiles" not in event:
+          continue
+        files = event["namedSetOfFiles"].get("files", [])
+        for file_info in files:
+          name = file_info.get("name")
+          prefix = file_info.get("pathPrefix")
+          if name is None:
+            raise ValueError("File entry in BEP is missing 'name' field.")
+          if not name.endswith(".clang-tidy.yaml"):
+            continue
+          if prefix is None:
+            raise ValueError("File entry in BEP is missing 'pathPrefix' field.")
+          path = (
+              pathlib.Path(repo_root) / pathlib.Path(*prefix) / name
+          ).as_posix()
+          yaml_files.append(path)
+      except json.JSONDecodeError:
+        _logger().warning(
+            "Skipping invalid JSON line in BEP file: %s", line.strip()
+        )
+        continue
+
+  return yaml_files
+
+
+def parse_clang_tidy_yaml(yaml_path: str) -> ClangTidyReport:
+  """A simple, specialized parser for clang-tidy YAML reports to avoid PyYAML dependency."""
+
+  def extract(s: str) -> str:
+    _, value = s.split(":", 1)
+    return value.strip().strip("'\"")
+
+  result: ClangTidyReport = {"Diagnostics": []}
+
+  current_diag: ClangTidyDiagnostic | None = None
+  in_diag_message = False
+
+  with open(yaml_path, "r") as f:
+    for line in f:
+      stripped = line.strip()
+      if stripped.startswith("MainSourceFile:"):
+        result["MainSourceFile"] = extract(stripped)
+      elif stripped.startswith("- DiagnosticName:"):
+        current_diag = {"DiagnosticName": extract(stripped)}
+        result["Diagnostics"].append(current_diag)
+        in_diag_message = False
+      elif stripped.startswith("DiagnosticMessage:"):
+        in_diag_message = True
+      elif in_diag_message and stripped.startswith("Message:"):
+        if current_diag is not None:
+          current_diag["Message"] = extract(stripped)
+      elif in_diag_message and stripped.startswith("FilePath:"):
+        if current_diag is not None:
+          current_diag["FilePath"] = extract(stripped)
+      elif in_diag_message and stripped.startswith("FileOffset:"):
+        if current_diag is not None:
+          current_diag["FileOffset"] = int(extract(stripped))
+      elif stripped.startswith("Level:"):
+        if current_diag is not None:
+          current_diag["Level"] = extract(stripped)
+      elif stripped.startswith("Replacements:"):
+        in_diag_message = False
+        if current_diag is not None:
+          current_diag["HasReplacements"] = not stripped.endswith("[]")
+  return result
+
+
+RESET = "\033[0m"
+BOLD = "\033[1m"
+RED = "\033[31m"
+YELLOW = "\033[33m"
+GREEN = "\033[32m"
+CYAN = "\033[36m"
+
+
+def ansiformat(text: str, color: str = "", bold: bool = False) -> str:
+  """Returns text wrapped in ANSI color codes."""
+  bold_str = BOLD if bold else ""
+  return f"{bold_str}{color}{text}{RESET}"
+
+
+def print_diagnostic(
+    diag: Diagnostic,
+    repo_root: str,
+    *,
+    warnings_as_errors: bool = False,
+    stream: IO[str] = sys.stderr,
+) -> None:
+  """Prints a diagnostic message with code snippet and color."""
+
+  # Clang-tidy format: file:line:col: level: message [name]
+  def get_level_str() -> str:
+    level_str = diag.level.lower() if diag.level else "warning"
+    if warnings_as_errors and level_str == "warning":
+      return "error"
+    return level_str
+
+  level_str = get_level_str()
+  diag_line = " ".join((
+      ansiformat(
+          ":".join((diag.file_path, str(diag.line_num), str(diag.col_num))),
+          bold=True,
+      ),
+      ansiformat(
+          f"{level_str}:",
+          color=RED if level_str == "error" else YELLOW,
+          bold=True,
+      ),
+      ansiformat(diag.message, bold=True),
+      ansiformat(f"[{diag.name}]", color=CYAN, bold=True),
+  ))
+  stream.write(f"{diag_line}\n")
+  abs_path = pathlib.Path(repo_root) / diag.file_path
+  try:
+    with open(abs_path, "r") as f:
+      lines = f.readlines()
+  except FileNotFoundError:
+    _logger().warning(
+        "Could not read file %r to print diagnostic snippet.",
+        abs_path.as_posix(),
+    )
+    return
+
+  context_lines = 5
+  start = max(0, diag.line_num - context_lines - 1)
+  end = min(len(lines), diag.line_num + context_lines)
+  for linenum, line in enumerate(lines[start:end], start=start + 1):
+    line_content = line.rstrip("\n")
+    prefix = f"{linenum:5d} | "
+    if linenum == diag.line_num:
+      stream.write(f"{ansiformat(prefix, bold=True)}{line_content}\n")
+      # Print caret
+      spaces = "".join(
+          "\t" if ch == "\t" else " " for ch in line_content[: diag.col_num - 1]
+      )
+      stream.write(
+          f"{' ' * len(prefix)}{spaces}{ansiformat('^', color=GREEN, bold=True)}\n"
+      )
+    else:
+      stream.write(f"{prefix}{line_content}\n")
+
+
+def _find_with_prefix(
+    parts: Sequence[str], prefix: str
+) -> tuple[int, str | None]:
+  """Finds the index and value of the first part starting with the prefix.
+
+  Args:
+    parts: The path parts to search.
+    prefix: The prefix to search for.
+
+  Returns:
+    A tuple of (index, dir) where dir is the directory starting with the
+    prefix or None if not found.
+  """
+  for idx, part in enumerate(parts):
+    if part.startswith(prefix):
+      return idx, part
+  return 0, None
+
+
+def _get_report_source(yaml_file_path: pathlib.Path, repo_root: str) -> str:
+  """Returns the source file path for a Clang-Tidy YAML report file."""
+  # yaml_path looks like:
+  # x/y/z/bazel_clang_tidy/path/to/file.cc.<target>.clang-tidy.yml.
+  expected_prefix = "bazel_clang_tidy_"
+  idx, aspect_dir = _find_with_prefix(yaml_file_path.parts, expected_prefix)
+  if aspect_dir is not None:
+    toplevel = [aspect_dir.removeprefix(expected_prefix)]
+    rel_parts = yaml_file_path.parts[idx + 1 :]
+    # Remove <target>.clang-tidy.yml suffix.
+    filename = rel_parts[-1].rsplit(".", 3)[0]
+  else:
+    toplevel = []
+    rel_parts = yaml_file_path.parts
+    # Remove .clang-tidy.yml without the <target>.
+    filename = rel_parts[-1].rsplit(".", 2)[0]
+  return normalize_path(
+      pathlib.Path(*toplevel, *rel_parts[:-1], filename).as_posix(),
+      repo_root,
+  )
+
+
+class ClangTidyDiffFilter:
+  """Filters Clang-Tidy diagnostics based on a diff."""
+
+  def __init__(
+      self,
+      config: AppConfig,
+      *,
+      offset_provider: Callable[[str], Sequence[int]] = get_line_offsets,
+      apply_fixes_fn: ApplyFixes | None = None,
+  ):
+    """Initializes the ClangTidyDiffFilter.
+
+    Args:
+      config: An AppConfig object containing the application configuration.
+      offset_provider: A callable that takes a file path and returns a list of
+        byte offsets for the start of each line in the file. Defaults to
+        `get_line_offsets`.
+      apply_fixes_fn: Method to call to apply fixes, if --fix is enabled. It
+        will be called with the path to the temp directory where the relevant
+        YAML files have been staged. If None, fixes will not be applied.
+    """
+    self.has_diff = config.patch is not None
+    self.diff_ranges = parse_diff(config.patch) if config.patch else {}
+    self.yaml_files = parse_bep(config.bep_file, config.repo_root)
+    self.repo_root = config.repo_root
+    self.warnings_as_errors = config.warnings_as_errors
+    self.offset_provider = offset_provider
+    self.file_offsets_cache: dict[str, Sequence[int]] = {}
+    self.seen_files: set[str] = set()
+    self.apply_fixes_fn = apply_fixes_fn
+
+  def process_file(
+      self, yaml_file: str
+  ) -> tuple[Sequence[Diagnostic], DiagnosticSummary]:
+    """Processes a single Clang-Tidy YAML report file.
+
+    Args:
+      yaml_file: The path to the Clang-Tidy YAML report file.
+
+    Returns:
+      A tuple of (matched_diagnostics, summary).
+    """
+    matched_diagnostics: list[Diagnostic] = []
+    data = parse_clang_tidy_yaml(yaml_file)
+    _logger().debug(
+        "Processing clang-tidy report file: %s with content:\n%s",
+        yaml_file,
+        data,
+    )
+    yaml_file_path = pathlib.Path(yaml_file)
+    report_source = _get_report_source(yaml_file_path, self.repo_root)
+    if not data:
+      return [], DiagnosticSummary(
+          file_path=report_source,
+          was_skipped=self.has_diff,
+          total=0,
+          matched=0,
+      )
+
+    main_source = data.get("MainSourceFile")
+    if main_source:
+      norm_main_source = normalize_path(main_source, self.repo_root)
+      self.seen_files.add(norm_main_source)
+    if self.has_diff and report_source in self.diff_ranges:
+      self.seen_files.add(report_source)
+    if "Diagnostics" not in data:
+      return [], DiagnosticSummary(
+          file_path=report_source,
+          was_skipped=self.has_diff,
+          total=0,
+          matched=0,
+      )
+
+    for diag in data["Diagnostics"]:
+      file_path = diag.get("FilePath")
+      offset = diag.get("FileOffset")
+
+      if not file_path or offset is None:
+        continue
+
+      norm_path = normalize_path(file_path, self.repo_root)
+      if self.has_diff and norm_path not in self.diff_ranges:
+        continue
+
+      abs_path = pathlib.Path(self.repo_root) / norm_path
+      if norm_path not in self.file_offsets_cache:
+        self.file_offsets_cache[norm_path] = self.offset_provider(
+            abs_path.as_posix()
+        )
+
+      offsets = self.file_offsets_cache[norm_path]
+      if not offsets:
+        _logger().warning(
+            "Could not read file %r to calculate line number.",
+            abs_path.as_posix(),
+        )
+        continue
+
+      line_num = offset_to_line(offsets, offset)
+
+      line_start_offset = offsets[line_num - 1]
+      col_num = offset - line_start_offset + 1
+
+      lines = self.diff_ranges[norm_path] if self.has_diff else set()
+      if not self.has_diff or line_num in lines:
+        matched_diagnostics.append(
+            Diagnostic(
+                file_path=norm_path,
+                line_num=line_num,
+                col_num=col_num,
+                level=diag.get("Level") or "",
+                name=diag.get("DiagnosticName") or "",
+                message=diag.get("Message") or "",
+                yaml_file=yaml_file,
+                has_replacements=diag.get("HasReplacements", False),
+            )
+        )
+    return matched_diagnostics, DiagnosticSummary(
+        file_path=report_source,
+        was_skipped=self.has_diff and report_source not in self.diff_ranges,
+        total=len(data["Diagnostics"]),
+        matched=len(matched_diagnostics),
+    )
+
+  def report_missing(self) -> None:
+    """Reports any touched files that were not processed using the logger."""
+    if not self.has_diff:
+      return
+    touched_files = set(self.diff_ranges.keys())
+    missing_files = [
+        f for f in touched_files - self.seen_files if f.endswith((".h", ".cc"))
+    ]
+    if missing_files:
+      _logger().warning(
+          "No Clang-Tidy reports were processed for the following modified"
+          " files:"
+      )
+      for f in sorted(missing_files):
+        _logger().warning("  - %s", f)
+
+  def _copy_and_normalize_yaml(
+      self, src: pathlib.Path, dest: pathlib.Path
+  ) -> None:
+    """Copies a YAML report and normalizes all remote paths to local absolute paths."""
+    with src.open("r") as f_source, dest.open("w") as f_dest:
+      for line in f_source:
+        stripped = line.strip()
+        # Paths in clang-tidy YAMLs only appear in these two fields.
+        if not stripped.startswith(("MainSourceFile:", "FilePath:")):
+          f_dest.write(line)
+          continue
+        _, val = line.split(":", 1)
+        raw_path = val.strip()
+        if not raw_path:
+          continue
+        norm_rel_path = normalize_path(raw_path, self.repo_root)
+        local_abs_path = (
+            pathlib.Path(self.repo_root) / norm_rel_path
+        ).as_posix()
+        line = line.replace(raw_path, local_abs_path)
+        f_dest.write(line)
+
+  def apply_fixes(self, yaml_files: Sequence[str]) -> None:
+    """Stages YAML files in temp directory and calls apply_fixes_fn."""
+    if not self.apply_fixes_fn or not yaml_files:
+      return
+    _logger().info("Applying fixes automatically.")
+    with tempfile.TemporaryDirectory() as temp_dir_str:
+      temp_dir = pathlib.Path(temp_dir_str)
+      counter = itertools.count()
+      for y in yaml_files:
+        dest_path = (
+            pathlib.Path(temp_dir) / f"{next(counter)}_{pathlib.Path(y).name}"
+        )
+        self._copy_and_normalize_yaml(pathlib.Path(y), dest_path)
+        _logger().info(
+            "Copied YAML report to temp dir for fixes: %s", dest_path
+        )
+      self.apply_fixes_fn(temp_dir)
+
+  def run(self) -> bool:
+    """Runs the Clang-Tidy diff filter.
+
+    Returns:
+      True if the check was successful (no errors found), False if errors were
+      found or running the check failed.
+    """
+    if self.has_diff and not self.diff_ranges:
+      _logger().error("No modified files found in patch.")
+      return False
+    if not self.yaml_files:
+      _logger().error("No YAML files provided or found in BEP.")
+      return False
+
+    found_diagnostics = False
+    yaml_files_to_fix: set[str] = set()
+    for y in self.yaml_files:
+      diagnostics, summary = self.process_file(y)
+      if summary.was_skipped and summary.total > 0:
+        _logger().info(
+            "Skipping %r with %d diagnostics because it is not in the diff.",
+            summary.file_path,
+            summary.total,
+        )
+      if summary.total > summary.matched and not summary.was_skipped:
+        _logger().info(
+            "Not all diagnostics will be reported for %r because they are not"
+            " part of the diff. Found %d diagnostics but reported %d.",
+            summary.file_path,
+            summary.total,
+            summary.matched,
+        )
+      for d in diagnostics:
+        if d.has_replacements and self.apply_fixes_fn is not None:
+          yaml_files_to_fix.add(y)
+        else:
+          print_diagnostic(
+              d, self.repo_root, warnings_as_errors=self.warnings_as_errors
+          )
+          found_diagnostics = True
+
+    self.report_missing()
+    self.apply_fixes(sorted(yaml_files_to_fix))
+    return not found_diagnostics
+
+
+def main() -> None:
+  """Main entry point for the Clang-Tidy diff filter."""
+
+  parser = argparse.ArgumentParser(
+      description="Filter Clang-Tidy errors by Git diff."
+  )
+  parser.add_argument("--patch", default=None, help="Path to the patch file.")
+  parser.add_argument(
+      "--repo-root", required=True, help="Absolute path to the repo root."
+  )
+  parser.add_argument(
+      "--bep-file",
+      required=True,
+      help="Path to Bazel Build Event Protocol JSON file.",
+  )
+  parser.add_argument(
+      "--warnings-as-errors",
+      default="true",
+      choices=["true", "false"],
+      help="Treat warnings as errors.",
+  )
+  parser.add_argument(
+      "--log-level",
+      default="INFO",
+      choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+      help="Set the log level.",
+  )
+  parser.add_argument(
+      "--fix",
+      action="store_true",
+      help="Apply automatic fixes to touched files.",
+  )
+  parser.add_argument(
+      "--clang-apply-replacements",
+      type=pathlib.Path,
+      default=None,
+      help="Path to clang-apply-replacements binary.",
+  )
+
+  args = parser.parse_args()
+  _set_log_level(args.log_level)
+
+  config = AppConfig(
+      patch=args.patch,
+      repo_root=args.repo_root,
+      bep_file=args.bep_file,
+      warnings_as_errors=args.warnings_as_errors == "true",
+      fix=args.fix,
+  )
+  apply_replacements_bin_path = (
+      _resolve_clang_apply_replacements(args.clang_apply_replacements)
+      if args.fix
+      else None
+  )
+  if args.fix and apply_replacements_bin_path is None:
+    _logger().error(
+        "--clang-apply-replacements is required when --fix is enabled "
+        "and the binary could not be resolved automatically in runfiles."
+    )
+    sys.exit(1)
+  apply_fixes_fn = (
+      functools.partial(clang_apply_fixes, apply_replacements_bin_path)
+      if apply_replacements_bin_path
+      else None
+  )
+  filterer = ClangTidyDiffFilter(config, apply_fixes_fn=apply_fixes_fn)
+  if not filterer.run():
+    sys.exit(1)
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/build_tools/ci/clang_tidy_diff_test.py

@@ -0,0 +1,649 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+from collections.abc import Sequence
+import io
+import json
+import pathlib
+import tempfile
+import textwrap
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from build_tools.ci import clang_tidy_diff
+
+
+class TestClangTidyDiff(parameterized.TestCase):
+
+  @parameterized.parameters(
+      (0, 1),
+      (5, 1),
+      (10, 2),
+      (15, 2),
+      (20, 3),
+      (25, 3),
+      (30, 4),
+      (35, 4),
+  )
+  def test_offset_to_line(self, offset, expected_line):
+    offsets = [0, 10, 20, 30]
+    self.assertEqual(
+        clang_tidy_diff.offset_to_line(offsets, offset), expected_line
+    )
+
+  def test_offset_to_line_empty(self):
+    self.assertEqual(clang_tidy_diff.offset_to_line([], 10), -1)
+
+  def test_normalize_path_relative(self):
+    self.assertEqual(
+        clang_tidy_diff.normalize_path("foo/bar.cc", "/root"), "foo/bar.cc"
+    )
+
+  def test_normalize_path_absolute_in_repo(self):
+    self.assertEqual(
+        clang_tidy_diff.normalize_path("/root/foo/bar.cc", "/root"),
+        "foo/bar.cc",
+    )
+
+  def test_normalize_path_absolute_outside_repo(self):
+    self.assertEqual(
+        clang_tidy_diff.normalize_path("/other/foo/bar.cc", "/root"),
+        "/other/foo/bar.cc",
+    )
+
+  def test_normalize_path_execroot(self):
+    path = "/usr/local/google/home/user/.cache/bazel/_bazel_user/a708d4fc59660ccd295a76cce84d113c/execroot/xla/xla/stream_executor/cuda/cuda_status.h"
+    self.assertEqual(
+        clang_tidy_diff.normalize_path(path, "/root"),
+        "xla/stream_executor/cuda/cuda_status.h",
+    )
+
+  def test_normalize_path_remote_worker(self):
+    path = "/b/f/w/xla/backends/gpu/codegen/triton/transforms/lowering_utils.h"
+    # repo_root is /__w/xla/xla, so workspace name is 'xla'
+    self.assertEqual(
+        clang_tidy_diff.normalize_path(path, "/__w/xla/xla"),
+        "xla/backends/gpu/codegen/triton/transforms/lowering_utils.h",
+    )
+
+  def test_normalize_path_local_ci_runner(self):
+    path = "/__w/xla/xla/xla/backends/gpu/codegen/triton/transforms/lowering_utils.h"
+    self.assertEqual(
+        clang_tidy_diff.normalize_path(path, "/__w/xla/xla"),
+        "xla/backends/gpu/codegen/triton/transforms/lowering_utils.h",
+    )
+
+  def test_normalize_path_remote_worker_third_party(self):
+    path = "/b/f/w/third_party/gpus/cuda/include/cuda.h"
+    self.assertEqual(
+        clang_tidy_diff.normalize_path(path, "/__w/xla/xla"),
+        "third_party/gpus/cuda/include/cuda.h",
+    )
+
+  def test_parse_diff(self):
+    tmpdir = self.create_tempdir().full_path
+    diff_path = pathlib.Path(tmpdir) / "test.diff"
+    with open(diff_path, "w") as f:
+
+      f.write(textwrap.dedent("""\
+                  diff --git a/file1.cc b/file1.cc
+                  index 123456..789012 100644
+                  --- a/file1.cc
+                  +++ b/file1.cc
+                  @@ -1,2 +1,3 @@
+                   line1
+                  +line2
+                   line3
+                  """))
+    ranges = clang_tidy_diff.parse_diff(str(diff_path))
+    self.assertEqual(ranges, {"file1.cc": {2}})
+
+  def test_parse_bep(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+      bep_path = pathlib.Path(tmpdir) / "test.bep"
+      with open(bep_path, "w") as f:
+        f.write(
+            '{"namedSetOfFiles": {"files": [{"name": "file1.clang-tidy.yaml",'
+            ' "pathPrefix": ["bazel-out", "k8-opt", "bin"]}]}}\n'
+        )
+      yaml_files = clang_tidy_diff.parse_bep(str(bep_path), "/root")
+      self.assertEqual(
+          yaml_files, ["/root/bazel-out/k8-opt/bin/file1.clang-tidy.yaml"]
+      )
+
+  def test_parse_clang_tidy_yaml(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+      yaml_path = pathlib.Path(tmpdir) / "test.yaml"
+      with open(yaml_path, "w") as f:
+
+        f.write(textwrap.dedent("""\
+                    ---
+                    MainSourceFile:  '/root/file1.cc'
+                    Diagnostics:
+                      - DiagnosticName:  misc-unused
+                        DiagnosticMessage:
+                          Message:         'unused variable'
+                          FilePath:        '/root/file1.cc'
+                          FileOffset:      15
+                    ...
+                    """))
+
+      data = clang_tidy_diff.parse_clang_tidy_yaml(str(yaml_path))
+
+      self.assertEqual(data.get("MainSourceFile"), "/root/file1.cc")
+      diagnostics = data.get("Diagnostics", [])
+      with self.subTest("Diagnostics"):
+        self.assertLen(diagnostics, 1)
+        self.assertEqual(diagnostics[0].get("DiagnosticName"), "misc-unused")
+        self.assertEqual(diagnostics[0].get("Message"), "unused variable")
+        self.assertEqual(diagnostics[0].get("FilePath"), "/root/file1.cc")
+        self.assertEqual(diagnostics[0].get("FileOffset"), 15)
+
+  def test_process_file(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+      yaml_path = pathlib.Path(tmpdir) / "file1.cc.clang-tidy.yaml"
+      with open(yaml_path, "w") as f:
+        f.write(textwrap.dedent(f"""\
+                    ---
+                    MainSourceFile:  '{tmpdir}/file1.cc'
+                    Diagnostics:
+                      - DiagnosticName:  misc-unused
+                        DiagnosticMessage:
+                          Message:         'unused variable'
+                          FilePath:        '{tmpdir}/file1.cc'
+                          FileOffset:      15
+                        Level:           Error
+                    ...
+                    """))
+
+      diff_path = pathlib.Path(tmpdir) / "test.diff"
+      with open(diff_path, "w") as f:
+
+        f.write(textwrap.dedent("""\
+                    diff --git a/file1.cc b/file1.cc
+                    index 123456..789012 100644
+                    --- a/file1.cc
+                    +++ b/file1.cc
+                    @@ -1,2 +1,3 @@
+                     line1
+                    +line2
+                     line3
+                    """))
+
+      bep_path = pathlib.Path(tmpdir) / "test.bep"
+      with open(bep_path, "w") as f:
+        f.write(
+            json.dumps({
+                "namedSetOfFiles": {
+                    "files": [{
+                        "name": "file1.cc.clang-tidy.yaml",
+                        "pathPrefix": [],
+                    }]
+                }
+            })
+            + "\n"
+        )
+
+      config = clang_tidy_diff.AppConfig(
+          patch=str(diff_path),
+          repo_root=str(tmpdir),
+          bep_file=str(bep_path),
+          warnings_as_errors=True,
+      )
+
+      def mock_offset_provider(_: str) -> list[int]:
+        return [0, 10, 20, 30]
+
+      filterer = clang_tidy_diff.ClangTidyDiffFilter(
+          config, offset_provider=mock_offset_provider
+      )
+      diagnostics, summary = filterer.process_file(str(yaml_path))
+
+      with self.subTest("Diagnostics"):
+        self.assertLen(diagnostics, 1)
+        self.assertEqual(diagnostics[0].file_path, "file1.cc")
+        self.assertEqual(diagnostics[0].line_num, 2)
+        self.assertEqual(
+            diagnostics[0].col_num, 6
+        )  # Offset 15 - Line 2 start 10 + 1 = 6
+        self.assertEqual(diagnostics[0].level, "Error")
+        self.assertEqual(diagnostics[0].name, "misc-unused")
+        self.assertEqual(diagnostics[0].message, "unused variable")
+
+      with self.subTest("DiagnosticSummary"):
+        self.assertEqual(
+            summary,
+            clang_tidy_diff.DiagnosticSummary(
+                file_path="file1.cc",
+                was_skipped=False,
+                total=1,
+                matched=1,
+            ),
+        )
+
+  def test_process_file_no_substring_false_positives(self):
+    """Tests that we don't get false positives from diff file paths being substrings of other file paths."""
+    tmpdir = self.create_tempdir().full_path
+    yaml_path = pathlib.Path(tmpdir) / "xla/long_util.cc.clang-tidy.yaml"
+    yaml_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(yaml_path, "w") as f:
+      f.write(textwrap.dedent(f"""\
+                  ---
+                  MainSourceFile:  '{tmpdir}/xla/long_util.cc'
+                  Diagnostics: []
+                  """))
+    # util and long_util are both in the diff but only long_util has a report.
+    diff_path = pathlib.Path(tmpdir) / "test.diff"
+    with open(diff_path, "w") as f:
+      f.write(textwrap.dedent("""\
+                  diff --git a/util.cc b/util.cc
+                  index 123456..789012 100644
+                  --- a/util.cc
+                  +++ b/util.cc
+                  @@ -1,1 +1,2 @@
+                   line1
+                  +line2
+                  diff --git a/xla/long_util.cc b/xla/long_util.cc
+                  index 123456..789012 100644
+                  --- a/xla/long_util.cc
+                  +++ b/xla/long_util.cc
+                  @@ -1,1 +1,2 @@
+                   line1
+                  +line2
+                  """))
+    bep_path = pathlib.Path(tmpdir) / "test.bep"
+    with open(bep_path, "w") as f:
+      f.write(
+          json.dumps({
+              "namedSetOfFiles": {
+                  "files": [{
+                      "name": "xla/long_util.cc.clang-tidy.yaml",
+                      "pathPrefix": [],
+                  }]
+              }
+          })
+          + "\n"
+      )
+    config = clang_tidy_diff.AppConfig(
+        patch=diff_path.as_posix(),
+        repo_root=tmpdir,
+        bep_file=bep_path.as_posix(),
+        warnings_as_errors=True,
+    )
+    filterer = clang_tidy_diff.ClangTidyDiffFilter(config)
+    _, summary = filterer.process_file(str(yaml_path))
+    self.assertIn("xla/long_util.cc", filterer.seen_files)
+    self.assertFalse(summary.was_skipped)
+    self.assertNotIn("util.cc", filterer.seen_files)
+
+  def test_process_file_empty_yaml_aspect_path(self):
+    tmpdir = self.create_tempdir().full_path
+    yaml_path = pathlib.Path(tmpdir) / (
+        "bazel-out/k8-opt/bin/xla/backends/bazel_clang_tidy_xla/"
+        "backends/source.cc.target.clang-tidy.yaml"
+    )
+    yaml_path.parent.mkdir(parents=True, exist_ok=True)
+    self.create_tempfile(yaml_path.as_posix(), content="")
+    diff_path = pathlib.Path(tmpdir) / "test.diff"
+    self.create_tempfile(
+        diff_path.as_posix(),
+        content=textwrap.dedent("""\
+        diff --git a/xla/backends/source.cc b/xla/backends/source.cc
+        index 123456..789012 100644
+        --- a/xla/backends/source.cc
+        +++ b/xla/backends/source.cc
+        @@ -1,1 +1,2 @@
+         line1
+        +line2
+        """),
+    )
+    bep_path = pathlib.Path(tmpdir) / "test.bep"
+    self.create_tempfile(
+        bep_path.as_posix(),
+        content=json.dumps({
+            "namedSetOfFiles": {
+                "files": [{
+                    "name": "xla/backends/source.cc.target.clang-tidy.yaml",
+                    "pathPrefix": [
+                        "bazel-out",
+                        "k8-opt",
+                        "bin",
+                        "xla",
+                        "backends",
+                        "bazel_clang_tidy_xla",
+                    ],
+                }]
+            }
+        })
+        + "\n",
+    )
+    config = clang_tidy_diff.AppConfig(
+        patch=diff_path.as_posix(),
+        repo_root=tmpdir,
+        bep_file=bep_path.as_posix(),
+        warnings_as_errors=False,
+    )
+    filterer = clang_tidy_diff.ClangTidyDiffFilter(config)
+    _, summary = filterer.process_file(str(yaml_path))
+    # File should be marked as seen despite the empty YAML.
+    self.assertIn(
+        "xla/backends/source.cc",
+        filterer.seen_files,
+    )
+    with self.subTest("DiagnosticSummary"):
+      self.assertEqual(
+          summary,
+          clang_tidy_diff.DiagnosticSummary(
+              file_path="xla/backends/source.cc",
+              was_skipped=False,
+              total=0,
+              matched=0,
+          ),
+      )
+
+  def test_run(self):
+    with tempfile.TemporaryDirectory() as tmpdir:
+      yaml_path = pathlib.Path(tmpdir) / "test.clang-tidy.yaml"
+      with open(yaml_path, "w") as f:
+        f.write(textwrap.dedent(f"""\
+                    ---
+                    MainSourceFile:  '{tmpdir}/file1.cc'
+                    Diagnostics:
+                      - DiagnosticName:  misc-unused
+                        DiagnosticMessage:
+                          Message:         'unused variable'
+                          FilePath:        '{tmpdir}/file1.cc'
+                          FileOffset:      15
+                        Level:           Error
+                    ...
+                    """))
+
+      diff_path = pathlib.Path(tmpdir) / "test.diff"
+      with open(diff_path, "w") as f:
+
+        f.write(textwrap.dedent("""\
+                    diff --git a/file1.cc b/file1.cc
+                    index 123456..789012 100644
+                    --- a/file1.cc
+                    +++ b/file1.cc
+                    @@ -1,2 +1,3 @@
+                     line1
+                    +line2
+                     line3
+                    """))
+
+      bep_path = pathlib.Path(tmpdir) / "test.bep"
+      with open(bep_path, "w") as f:
+        f.write(
+            json.dumps({
+                "namedSetOfFiles": {
+                    "files": [{
+                        "name": "test.clang-tidy.yaml",
+                        "pathPrefix": [],
+                    }]
+                }
+            })
+            + "\n"
+        )
+
+      config = clang_tidy_diff.AppConfig(
+          patch=str(diff_path),
+          repo_root=str(tmpdir),
+          bep_file=str(bep_path),
+          warnings_as_errors=True,
+      )
+
+      def mock_offset_provider(_: str) -> list[int]:
+        return [0, 10, 20, 30]
+
+      filterer = clang_tidy_diff.ClangTidyDiffFilter(
+          config, offset_provider=mock_offset_provider
+      )
+
+      self.assertFalse(filterer.run())
+
+  def test_print_diagnostic_sanity(self):
+    diag = clang_tidy_diff.Diagnostic(
+        file_path="file1.cc",
+        line_num=2,
+        col_num=3,
+        level="warning",
+        name="misc-unused",
+        message="unused variable",
+        yaml_file="test.clang-tidy.yaml",
+        has_replacements=False,
+    )
+    with tempfile.TemporaryDirectory() as tmpdir:
+      src_file = pathlib.Path(tmpdir) / "file1.cc"
+      with open(src_file, "w") as f:
+        f.write("line1\nline2\nline3\n")
+      captured_stderr = io.StringIO()
+      # Run it with warnings_as_errors=True to test that path too
+      clang_tidy_diff.print_diagnostic(
+          diag,
+          repo_root=tmpdir,
+          warnings_as_errors=True,
+          stream=captured_stderr,
+      )
+      output = captured_stderr.getvalue()
+      with self.subTest("diagnostic_string_sanity"):
+        self.assertIn("file1.cc:2:3", output)
+        self.assertIn("error:", output)
+        self.assertIn("unused variable", output)
+        self.assertIn("[misc-unused]", output)
+        self.assertIn("  2 |", output)  # Snippet line
+        self.assertIn("line2", output)
+        self.assertIn("^", output)  # Caret
+
+  @parameterized.parameters(
+      dict(replacement_text="", expected=None),
+      dict(replacement_text="Replacements: []", expected=False),
+      dict(
+          replacement_text=textwrap.dedent("""\
+            Replacements:
+              - FilePath:        '/root/file1.cc'
+                Offset:          15
+                Length:          6
+                ReplacementText: 'blehblehbleh'
+            """),
+          expected=True,
+      ),
+  )
+  def test_parse_clang_tidy_yaml_with_replacements(
+      self, replacement_text: str, expected: bool | None
+  ):
+    tmpdir = self.create_tempdir().full_path
+    yaml_path = pathlib.Path(tmpdir) / "test.yaml"
+    self.create_tempfile(
+        yaml_path.as_posix(),
+        content=textwrap.dedent(f"""\
+            ---
+            MainSourceFile:  '/root/file1.cc'
+            Diagnostics:
+              - DiagnosticName:  misc-unused
+                DiagnosticMessage:
+                  Message:         'unused variable'
+                  FilePath:        '/root/file1.cc'
+                  FileOffset:      15
+                {replacement_text}
+            ...
+            """),
+    )
+    data = clang_tidy_diff.parse_clang_tidy_yaml(str(yaml_path))
+    diagnostics = data.get("Diagnostics", [])
+    self.assertLen(diagnostics, 1)
+    self.assertEqual(diagnostics[0].get("HasReplacements"), expected)
+
+  def _setup_temp_files(
+      self,
+  ) -> tuple[str, pathlib.Path, pathlib.Path, pathlib.Path]:
+    """Setup for a simple yaml, diff, and BEP file in a temporary directory.
+
+    Returns:
+      A tuple of (tmpdir, yaml_path, diff_path, bep_path).
+    """
+    tmpdir = self.create_tempdir().full_path
+    yaml_path = pathlib.Path(tmpdir) / "file1.cc.clang-tidy.yaml"
+    self.create_tempfile(
+        yaml_path.as_posix(),
+        content=textwrap.dedent(f"""\
+            ---
+            MainSourceFile:  '{tmpdir}/file1.cc'
+            Diagnostics:
+              - DiagnosticName:  misc-unused
+                DiagnosticMessage:
+                  Message:         'unused variable'
+                  FilePath:        '{tmpdir}/file1.cc'
+                  FileOffset:      15
+                Level:           Warning
+                Replacements:
+                  - FilePath:        '{tmpdir}/file1.cc'
+                    Offset:          15
+                    Length:          6
+                    ReplacementText: ''
+            ...
+            """),
+    )
+
+    diff_path = pathlib.Path(tmpdir) / "test.diff"
+    self.create_tempfile(
+        diff_path.as_posix(),
+        content=textwrap.dedent("""\
+            diff --git a/file1.cc b/file1.cc
+            index 123456..789012 100644
+            --- a/file1.cc
+            +++ b/file1.cc
+            @@ -1,2 +1,3 @@
+             line1
+            +line2
+             line3
+            """),
+    )
+
+    bep_path = pathlib.Path(tmpdir) / "test.bep"
+    self.create_tempfile(
+        bep_path.as_posix(),
+        content=json.dumps({
+            "namedSetOfFiles": {
+                "files": [{
+                    "name": "file1.cc.clang-tidy.yaml",
+                    "pathPrefix": [],
+                }]
+            }
+        })
+        + "\n",
+    )
+    return tmpdir, yaml_path, diff_path, bep_path
+
+  def test_process_file_detects_replacements(self):
+    tmpdir, yaml_path, diff_path, bep_path = self._setup_temp_files()
+    config = clang_tidy_diff.AppConfig(
+        patch=str(diff_path),
+        repo_root=tmpdir,
+        bep_file=str(bep_path),
+        warnings_as_errors=False,
+        fix=False,
+    )
+
+    def mock_offset_provider(_: str) -> Sequence[int]:
+      return [0, 10, 20, 30]
+
+    filterer = clang_tidy_diff.ClangTidyDiffFilter(
+        config, offset_provider=mock_offset_provider
+    )
+    diagnostics, _ = filterer.process_file(yaml_path.as_posix())
+
+    self.assertLen(diagnostics, 1)
+    self.assertTrue(diagnostics[0].has_replacements)
+
+  def test_run_collects_files_to_fix(self):
+    tmpdir, _, diff_path, bep_path = self._setup_temp_files()
+    config = clang_tidy_diff.AppConfig(
+        patch=str(diff_path),
+        repo_root=tmpdir,
+        bep_file=str(bep_path),
+        warnings_as_errors=False,
+        fix=True,
+    )
+    staged_dirs = []
+
+    def fake_apply_fixes(temp_dir: pathlib.Path):
+      staged_dirs.append(temp_dir)
+      copied_files = list(
+          pathlib.Path(temp_dir).glob("*_file1.cc.clang-tidy.yaml")
+      )
+      self.assertLen(copied_files, 1)
+      self.assertTrue(copied_files[0].exists())
+
+    def mock_offset_provider(_: str) -> Sequence[int]:
+      return [0, 10, 20, 30]
+
+    filterer = clang_tidy_diff.ClangTidyDiffFilter(
+        config,
+        offset_provider=mock_offset_provider,
+        apply_fixes_fn=fake_apply_fixes,
+    )
+    filterer.run()
+    self.assertLen(staged_dirs, 1)
+
+  def test_apply_fixes_stages_files_correctly(self):
+    staged_dirs = []
+
+    def fake_apply_fixes(temp_dir: pathlib.Path) -> None:
+      staged_dirs.append(temp_dir.as_posix())
+      copied_files = list(pathlib.Path(temp_dir).glob("*_file1.yaml"))
+      self.assertLen(copied_files, 1)
+      copied_file = copied_files[0]
+      self.assertTrue(copied_file.exists())
+      expected_normalized = textwrap.dedent(f"""\
+            MainSourceFile: {tmpdir}/xla/file1.cc
+            SomeOtherField: some_value
+            FilePath: {tmpdir}/third_party/file1.cc
+            """)
+      self.assertEqual(copied_file.read_text(), expected_normalized)
+
+    tmpdir = self.create_tempdir().full_path
+    src_yaml = pathlib.Path(tmpdir) / "file1.yaml"
+    src_yaml_content = textwrap.dedent("""\
+        MainSourceFile: /b/f/w/xla/file1.cc
+        SomeOtherField: some_value
+        FilePath: /somepath/_bazel_user/1234/execroot/xla/third_party/file1.cc
+        """)
+    self.create_tempfile(src_yaml.as_posix(), content=src_yaml_content)
+    diff_path = pathlib.Path(tmpdir) / "test.diff"
+    self.create_tempfile(diff_path.as_posix(), content="")
+    bep_path = pathlib.Path(tmpdir) / "test.bep"
+    self.create_tempfile(bep_path.as_posix(), content="")
+
+    config = clang_tidy_diff.AppConfig(
+        patch=str(diff_path),
+        repo_root=tmpdir,
+        bep_file=str(bep_path),
+        warnings_as_errors=False,
+        fix=True,
+    )
+    filterer = clang_tidy_diff.ClangTidyDiffFilter(
+        config,
+        apply_fixes_fn=fake_apply_fixes,
+    )
+
+    filterer.apply_fixes([src_yaml])
+
+    self.assertLen(staged_dirs, 1)
+
+
+if __name__ == "__main__":
+  absltest.main()

third_party/xla/build_tools/ci/clang_tidy_wrapper.sh.tpl

@@ -0,0 +1,21 @@
+#!/bin/bash
+set -e
+
+echoerr() {
+  RED='\033[1;31m'
+  NOCOLOR='\033[0m'
+  printf "${RED}ERROR:${NOCOLOR} %s\n" "$*" >&2
+}
+
+REAL_BIN="$PWD/external/%LLVM_REPO_NAME%/bin/clang-tidy"
+if [ ! -f "$REAL_BIN" ]; then
+  echoerr "Failed to locate clang-tidy binary at: $REAL_BIN"
+  exit 1
+fi
+echo "Using clang-tidy at: " $REAL_BIN
+REAL_LIB_DIR="$(dirname "$REAL_BIN")/../lib"
+export LD_LIBRARY_PATH="${REAL_LIB_DIR}:${LD_LIBRARY_PATH}"
+# Intentional unquoted $@ expansion.
+# Word-splitting is required here to resolve composite tokens passed by Bazel
+# (e.g., "-include file.h" -> "-include" "file.h").
+exec "$REAL_BIN" $@

third_party/xla/build_tools/ci/golden_commands.txt

@@ -0,0 +1,128 @@
+# BEGIN BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --config=bzlmod --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=rules_ml_toolchain=$GITHUB_WORKSPACE/openxla/rules_ml_toolchain --color=yes --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 -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel test --build_tag_filters= --test_tag_filters= --config=rbe_linux_x86_64 --config=bzlmod --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=rules_ml_toolchain=$GITHUB_WORKSPACE/openxla/rules_ml_toolchain --color=yes --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 -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel analyze-profile profile.json.gz
+# END BuildType.JAX_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS
+# BEGIN BuildType.JAX_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --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 --@local_config_cuda//cuda:override_include_cuda_libs --nobuild -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel test --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --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 --@local_config_cuda//cuda:override_include_cuda_libs -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel analyze-profile profile.json.gz
+# END BuildType.JAX_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+# BEGIN BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS
+bazel --output_base=$GITHUB_WORKSPACE\bazel_output_base build --build_tag_filters= --test_tag_filters= --config=rbe_windows_amd64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE\openxla\xla --override_module=xla=$GITHUB_WORKSPACE\openxla\xla --color=yes --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_runfile_links=False -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel analyze-profile profile.json.gz
+# END BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS
+# BEGIN BuildType.TENSORFLOW_LINUX_X86_CPU_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-gpu --test_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-gpu --config=release_cpu_linux --config=rbe_linux_cpu --repo_env=USE_PYWRAP_RULES=True --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --verbose_failures --test_output=errors --profile=profile.json.gz --test_lang_filters=cc,py --color=yes --nobuild -- //tensorflow/compiler/... -//tensorflow/compiler/tf2tensorrt/... //tensorflow/python/... -//tensorflow/python/distribute/... -//tensorflow/python/kernel_tests/... -//tensorflow/python/data/... -//tensorflow/python/compiler/tensorrt/...
+bazel test --build_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-gpu --test_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-gpu --config=release_cpu_linux --config=rbe_linux_cpu --repo_env=USE_PYWRAP_RULES=True --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --verbose_failures --test_output=errors --profile=profile.json.gz --test_lang_filters=cc,py --color=yes -- //tensorflow/compiler/... -//tensorflow/compiler/tf2tensorrt/... //tensorflow/python/... -//tensorflow/python/distribute/... -//tensorflow/python/kernel_tests/... -//tensorflow/python/data/... -//tensorflow/python/compiler/tensorrt/...
+bazel analyze-profile profile.json.gz
+# END BuildType.TENSORFLOW_LINUX_X86_CPU_GITHUB_ACTIONS
+# BEGIN BuildType.TENSORFLOW_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-no_gpu,-no_gpu_presubmit,-no_cuda11,+gpu --test_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-no_gpu,-no_gpu_presubmit,-no_cuda11,+gpu --config=release_gpu_linux --config=rbe_linux_cuda --config=hermetic_cuda_umd --repo_env=USE_PYWRAP_RULES=True --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --verbose_failures --test_output=errors --profile=profile.json.gz --test_lang_filters=cc,py --color=yes --nobuild -- //tensorflow/compiler/... -//tensorflow/compiler/tf2tensorrt/... //tensorflow/python/... -//tensorflow/python/distribute/... -//tensorflow/python/kernel_tests/... -//tensorflow/python/data/... -//tensorflow/python/compiler/tensorrt/...
+bazel test --build_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-no_gpu,-no_gpu_presubmit,-no_cuda11,+gpu --test_tag_filters=-no_oss,-tf_tosa,-oss_excluded,-oss_serial,-tpu,-benchmark-test,-v1only,-no_gpu,-no_gpu_presubmit,-no_cuda11,+gpu --config=release_gpu_linux --config=rbe_linux_cuda --config=hermetic_cuda_umd --repo_env=USE_PYWRAP_RULES=True --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --verbose_failures --test_output=errors --profile=profile.json.gz --test_lang_filters=cc,py --color=yes -- //tensorflow/compiler/... -//tensorflow/compiler/tf2tensorrt/... //tensorflow/python/... -//tensorflow/python/distribute/... -//tensorflow/python/kernel_tests/... -//tensorflow/python/data/... -//tensorflow/python/compiler/tensorrt/...
+bazel analyze-profile profile.json.gz
+# END BuildType.TENSORFLOW_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --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=False --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main
+bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --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=False --//xla/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --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/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_ARM64_CPU_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --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/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main
+bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --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/tsl:ci_build -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=bzlmod --color=yes --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/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --config=bzlmod --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_CPU_BZLMOD_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --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/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_linux_cpu --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --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 --test_sharding_strategy=disabled --test_strategy=exclusive --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any --test_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nvgpu_any,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=9.0 --repo_env=REMOTE_GPU_TESTING=0 --//xla/tsl:ci_build --color=yes --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 --test_sharding_strategy=disabled --test_strategy=exclusive -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --repo_env=HERMETIC_CUDA_VERSION=12.8.0 --repo_env=HERMETIC_CUDNN_VERSION=9.8.0 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --repo_env=HERMETIC_CUDA_VERSION=12.8.0 --repo_env=HERMETIC_CUDNN_VERSION=9.8.0 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --repo_env=HERMETIC_CUDA_VERSION=12.8.0 --repo_env=HERMETIC_CUDNN_VERSION=9.8.0 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --repo_env=HERMETIC_CUDA_VERSION=12.8.0 --repo_env=HERMETIC_CUDNN_VERSION=9.8.0 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --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/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+nvidia-smi
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/... @tsl//tsl/...
+bazel test --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --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/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_L4_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --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/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
+bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS
+# BEGIN BuildType.XLA_MACOS_ARM64_CPU_KOKORO
+df -h
+bazel --version
+mkdir -p /tmpfs/bazel_output
+bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --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 --macos_minimum_os=10.15 --test_tmpdir=/tmpfs/bazel_output --test_size_filters=small,medium --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_MACOS_ARM64_CPU_KOKORO
+# BEGIN BuildType.XLA_MACOS_X86_CPU_KOKORO
+sudo wget --no-verbose -O /usr/local/bin/bazel https://github.com/bazelbuild/bazelisk/releases/download/v1.11.0/bazelisk-darwin-amd64
+chmod +x /usr/local/bin/bazel
+bazel --version
+mkdir -p /Volumes/BuildData/bazel_output
+bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --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 --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/...
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_MACOS_X86_CPU_KOKORO
+# BEGIN BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
+bazel --output_user_root=C:/x build --build_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_xla_windows_x86_cpu_2022 --color=yes --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/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -//xla/tpu/... -//xla/backends/cpu/collectives:gloo_collectives_test -//xla/backends/cpu/collectives:mpi_collectives -//xla/backends/cpu/collectives:mpi_communicator -//xla/python/transfer/... -//xla/backends/profiler/subprocess:subprocess_profiling_session -//xla/backends/profiler/subprocess:subprocess_profiling_session_test -//xla/backends/profiler/subprocess:subprocess_registry -//xla/backends/profiler/subprocess:subprocess_registry_test -//xla/tools/benchmarks/utils:generate_benchmark_matrices_cc -//xla/tools/benchmarks/utils:generate_benchmark_matrices_main -//xla/tools/benchmarks/utils:generate_benchmark_matrices_test -//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk -//xla/backends/cpu/runtime/ynnpack:ynn_interop -//xla/backends/cpu/runtime/ynnpack:ynn_threadpool -//xla/backends/gpu/... -//xla/codegen/emitters/tests/... -//xla/service/gpu/... -//xla/codegen/intrinsic/cpp:eigen_unary_test
+bazel analyze-profile profile.json.gz
+# END BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS

third_party/xla/build_tools/ci/parallel_gpu_execute.sh

@@ -0,0 +1,83 @@
+#!/usr/bin/env bash
+# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+#
+#
+# A script to run multiple GPU tests in parallel controlled with an environment
+# variable.
+#
+# Required environment variables:
+#     TF_GPU_COUNT = Number of GPUs available.
+
+TF_GPU_COUNT=${TF_GPU_COUNT:-4}
+TF_TESTS_PER_GPU=${TF_TESTS_PER_GPU:-8}
+
+# This function is used below in rlocation to check that a path is absolute
+function is_absolute {
+  [[ "$1" = /* ]] || [[ "$1" =~ ^[a-zA-Z]:[/\\].* ]]
+}
+
+export TF_PER_DEVICE_MEMORY_LIMIT_MB=${TF_PER_DEVICE_MEMORY_LIMIT_MB:-4096}
+
+# *******************************************************************
+#         This section of the script is needed to
+#         make things work on windows under msys.
+# *******************************************************************
+RUNFILES_MANIFEST_FILE="${TEST_SRCDIR}/MANIFEST"
+function rlocation() {
+  if is_absolute "$1" ; then
+    # If the file path is already fully specified, simply return it.
+    echo "$1"
+  elif [[ -e "$TEST_SRCDIR/$1" ]]; then
+    # If the file exists in the $TEST_SRCDIR then just use it.
+    echo "$TEST_SRCDIR/$1"
+  elif [[ -e "$RUNFILES_MANIFEST_FILE" ]]; then
+    # If a runfiles manifest file exists then use it.
+    echo "$(grep "^$1 " "$RUNFILES_MANIFEST_FILE" | sed 's/[^ ]* //')"
+  fi
+}
+
+TEST_BINARY="$(rlocation $TEST_WORKSPACE/${1#./})"
+shift
+# *******************************************************************
+
+mkdir -p /var/lock
+# Try to acquire any of the TF_GPU_COUNT * TF_TESTS_PER_GPU
+# slots to run a test at.
+#
+# Prefer to allocate 1 test per GPU over 4 tests on 1 GPU.
+# So, we iterate over TF_TESTS_PER_GPU first.
+for j in `seq 0 $((TF_TESTS_PER_GPU-1))`; do
+  for i in `seq 0 $((TF_GPU_COUNT-1))`; do
+    exec {lock_fd}>/var/lock/gpulock${i}_${j} || exit 1
+    if flock -n "$lock_fd";
+    then
+      (
+        # This export only works within the brackets, so it is isolated to one
+        # single command.
+        export CUDA_VISIBLE_DEVICES=$i
+        export HIP_VISIBLE_DEVICES=$i
+        echo "Running test $TEST_BINARY $* on GPU $CUDA_VISIBLE_DEVICES"
+        "$TEST_BINARY" $@
+      )
+      return_code=$?
+      flock -u "$lock_fd"
+      exit $return_code
+    fi
+  done
+done
+
+echo "Cannot find a free GPU to run the test $* on, exiting with failure..."
+exit 1

third_party/xla/build_tools/ci/readme.md

@@ -0,0 +1,39 @@
+This folder contains scripts and commands executed by GitHub actions in OpenXLA,
+TensorFlow and JAX OSS repos. The GitHub actions replaced
+KokoroPresubmit-tensorflow* workflows. They run under
+`Copybara_XLA{Presubmit,Submit}` presubmit chip (category) at Google internally.
+The tests run on several target OS configurations/GCP containers such as Linux
+x86 with GPU or Linux ARM64. They assure that XLA/Tensor Flow/JAX compiles and
+runs on those platforms. The tests uses released OSS c++ clang compiler which
+has some differences in supporting c++ standards compared Google's internal
+version.
+
+#### How it works
+
+Repo specific GitHub actions call `build.py --build="build_name"`. E.g. OpenXLA
+uses https://github.com/openxla/xla/blob/main/.github/workflows/ci.yml
+
+The build here is a set of shell script commands executing the test targets or
+doing compile only testing. Each GitHub action call translates into compile only
+test:
+
+1.  dryrun `bazel build --nobuild ... test_targets`
+1.  actual compile `bazel build ... test_targets`
+1.  analyse results `bazel analyze-profile ...`
+
+or compile and run test commands:
+
+1.  dry run `bazel build --nobuild ... test_targets`
+1.  actual test `bazel test ... test_targets`
+1.  analyse results `bazel analyze-profile profile.json.gz`
+
+Checking in changes to `build.py` regenerates `golden_commands.txt` which lets
+us see how commands are changing. `golden_commands.txt` are not called as part
+of the continuous integration process.
+
+#### When does it run?
+
+The GitHub actions are automatically called as part of Google presubmit. The
+GitHub actions are automatically called on GitHub PR commit if the committer is
+part of OpenXLA GitHub org. Committer which are not part of OpenXLA org need an
+approval from OpenXLA org member to run the actions.

third_party/xla/build_tools/ci/run_clang_format.sh

@@ -0,0 +1,94 @@
+#!/bin/bash
+
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+set -xe
+
+source "$(dirname "${BASH_SOURCE[0]}")/shell_common.sh"
+
+CLANG_FORMAT_VERSION="${CLANG_FORMAT_VERSION:-17.0.6}"
+CLANG_FORMAT_PKG="clang-format==${CLANG_FORMAT_VERSION}"
+
+BUILD_WORKSPACE_DIRECTORY="${BUILD_WORKSPACE_DIRECTORY:-$(pwd)}"
+cd "$BUILD_WORKSPACE_DIRECTORY"
+
+if ! command -v uv >/dev/null 2>&1; then
+  echoerr "uv is required to run the clang-format check."
+  echoerr "Install uv from https://docs.astral.sh/uv/"
+  exit 1
+fi
+
+FIX=false
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --fix)
+      FIX=true
+      shift
+      ;;
+    *)
+      set +x
+      echo "Unknown option: $1" >&2
+      echo "Usage: $0 [--fix]" >&2
+      exit 1
+      ;;
+  esac
+done
+
+get_merge_base() {
+  local REFERENCE
+  local REMOTE="${REMOTE:-$(git remote -v | awk '/openxla\/xla/ { print $1; exit }')}"
+
+  if ! git rev-parse --is-inside-work-tree >/dev/null 2>&1; then
+    echoerr "This script must be run inside a Git repository."
+    exit 1
+  fi
+
+  if [ -n "${TARGET_REF}" ]; then
+    REFERENCE="$TARGET_REF"
+  else
+    if [ -z "$REMOTE" ]; then
+      echoerr "Could not find a git remote pointing to openxla/xla. Please add it as a remote."
+      echoerr "Example: git remote add upstream https://github.com/openxla/xla.git"
+      exit 1
+    fi
+    REFERENCE="${REMOTE}/main"
+  fi
+
+  MERGE_BASE=$(git merge-base "$REFERENCE" HEAD || true)
+  if [ -z "$MERGE_BASE" ]; then
+    echoerr "Could not find a common ancestor with $REFERENCE. Please fetch and rebase on main."
+    echoerr "Example: git fetch ${REMOTE:-origin} main && git rebase ${REMOTE:-origin}/main"
+    exit 1
+  fi
+}
+
+get_merge_base
+
+EXTRA_ARGS=(-p1 -style=file)
+if [ "$FIX" = true ]; then
+  EXTRA_ARGS+=(-i)
+fi
+
+# Run diff against the merge base.
+# -U0: Context of 0 lines (ignore surrounding code)
+# clang-format-diff: Checks only lines present in the diff
+DIFF=$(git diff -U0 --no-color "$MERGE_BASE" -- '*.cc' '*.h' |
+  uvx --from "$CLANG_FORMAT_PKG" clang-format-diff.py "${EXTRA_ARGS[@]}")
+
+if [ -n "$DIFF" ]; then
+  echoerr "Clang-format failed on the following changes:"
+  echo "$DIFF"
+  exit 1
+fi

third_party/xla/build_tools/ci/run_clang_tidy.sh

@@ -0,0 +1,125 @@
+#!/bin/bash
+
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+set -xe
+
+# Get targets from git diff if not specified on the command nline.
+get_targets_from_diff() {
+  if ! git rev-parse --is-inside-work-tree > /dev/null 2>&1; then
+    set +x
+    echo "Error: This script must be run inside a Git repository." >&2
+    exit 1
+  fi
+  REMOTE=${REMOTE:-$(git remote -v | awk '/openxla\/xla/ { print $1; exit }')}
+  if [ -z "$REMOTE" ]; then
+    set +x
+    echo "Could not find a git remote pointing to openxla/xla. Please add it as a remote." >&2
+    echo "Example: git remote add upstream https://github.com/openxla/xla.git" >&2
+    exit 1
+  fi
+  TARGET_REF=${TARGET_REF:-${REMOTE}/main}
+  MERGE_BASE=$(git merge-base "$TARGET_REF" HEAD || true)
+  if [ -z "$MERGE_BASE" ]; then
+    set +x
+    echo "Could not find a common ancestor with $TARGET_REF. Please rebase on $REMOTE main." >&2
+    echo "Example: git pull --rebase $REMOTE main" >&2
+    exit 1
+  fi
+  CHANGED_FILES=$(git diff --name-only --diff-filter=d "$MERGE_BASE" |
+    grep -E '\.(cc|h)$' || true)
+  if [ -z "$CHANGED_FILES" ]; then
+    set +x
+    echo "No C++ files changed."
+    exit 0
+  fi
+  PACKAGES=$(echo "$CHANGED_FILES" | while read -r file; do
+      echo "//$(dirname "$file"):all"
+  done | sort -u | tr '\n' ' ')
+  BASE_QUERY="kind('cc_(library|binary|test)', rdeps(set($PACKAGES), set($CHANGED_FILES), 1))"
+  if [ -n "$TAGS_TO_IGNORE" ]; then
+    QUERY="(${BASE_QUERY} except attr('tags', '${TAGS_TO_IGNORE}', (${BASE_QUERY})))"
+  else
+    QUERY="${BASE_QUERY}"
+  fi
+  
+  # Populate global TARGETS variable
+  # We use readarray to handle multi-line output from bazel cquery
+  readarray -t TARGETS < <($BAZEL_CMD cquery --output=starlark --starlark:expr="target.label" --config="$CONFIG" "$QUERY")
+  
+  if [ ${#TARGETS[@]} -eq 0 ]; then
+    set +x
+    echo "No relevant targets found for changed files."
+    exit 0
+  fi
+}
+
+BUILD_WORKSPACE_DIRECTORY=${BUILD_WORKSPACE_DIRECTORY:-$(pwd)}
+cd "$BUILD_WORKSPACE_DIRECTORY"
+BAZEL_CMD=${BAZEL_CMD:-bazelisk}
+CONFIG="clang-tidy-noerrors"
+
+FIX=false
+# Explicitly specify targets to run clang-tidy on.
+# TODO(sohaibiftikhar): Move the bash orchestration to python.
+TARGETS=()
+while [[ $# -gt 0 ]]; do
+  case $1 in
+    --fix)
+      FIX=true
+      shift
+      ;;
+    --target)
+      TARGETS+=("$2")
+      shift 2
+      ;;
+    *)
+      echo "Unknown option: $1"
+      echo "Usage: $0 [--fix] [--target <target> ...]"
+      exit 1
+      ;;
+  esac
+done
+
+if [ ${#TARGETS[@]} -eq 0 ]; then
+  USE_DIFF=true
+  get_targets_from_diff
+else
+  USE_DIFF=false
+fi
+
+# Create temporary files for the patch and BEP
+PATCH_FILE=$(mktemp)
+BEP_FILE=$(mktemp)
+# Ensure cleanup on exit
+trap 'rm -f "$PATCH_FILE" "$BEP_FILE"' EXIT
+
+
+$BAZEL_CMD build --config="$CONFIG" --build_event_json_file="$BEP_FILE" --keep_going \
+  "${TARGETS[@]}"
+EXTRA_ARGS=()
+if [ "$FIX" = true ]; then
+  EXTRA_ARGS+=("--fix")
+fi
+if [ "$USE_DIFF" = true ]; then
+  # Generate the patch file for changed lines
+  git diff "$MERGE_BASE" > "$PATCH_FILE"
+  EXTRA_ARGS+=("--patch" "$PATCH_FILE")
+fi
+
+$BAZEL_CMD run //build_tools/ci:clang_tidy_diff -- \
+  --repo-root "$BUILD_WORKSPACE_DIRECTORY" \
+  --bep-file "$BEP_FILE" \
+  "${EXTRA_ARGS[@]}"

third_party/xla/build_tools/ci/shell_common.sh

@@ -0,0 +1,21 @@
+# Copyright 2026 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+echoerr() {
+  set +x
+  RED='\033[1;31m'
+  NOCOLOR='\033[0m'
+  printf "${RED}ERROR:${NOCOLOR} %s\n" "$*" >&2
+}

third_party/xla/build_tools/configure/BUILD

@@ -0,0 +1,90 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+load("@local_config_cuda//cuda:build_defs.bzl", "cuda_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@xla//third_party/rules_python/python:py_test.bzl", "py_test")
+load("//xla:pytype.bzl", "pytype_strict_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "configure",
+    srcs = ["configure.py"],
+)
+
+py_test(
+    name = "configure_test",
+    srcs = ["configure_test.py"],
+    data = [
+        "testdata/clang_local.bazelrc",
+        "testdata/cuda_clang.bazelrc",
+        "testdata/cuda_clang_local.bazelrc",
+        "testdata/default_cuda_clang.bazelrc",
+        "testdata/gcc.bazelrc",
+        "testdata/nvcc_clang.bazelrc",
+        "testdata/nvcc_clang_local.bazelrc",
+        "testdata/nvcc_gcc.bazelrc",
+    ],
+    # After https://github.com/openxla/xla/commit/7d3043283, this test is no
+    # longer hermetic. This works in OSS tests because the docker container has
+    # clang-17 and gcc, but it's a little sketchy.
+    tags = ["notap"],
+    deps = [
+        ":configure",
+        "//build_tools:test_utils",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+# Below targets are just for checking if the host/CUDA compiler are configured
+# as expected.
+cc_library(
+    name = "assert_clang",
+    srcs = ["assert_clang.cc"],
+    tags = ["manual"],
+)
+
+cc_library(
+    name = "assert_gcc",
+    srcs = ["assert_gcc.cc"],
+    tags = ["manual"],
+)
+
+cuda_library(
+    name = "assert_cuda_clang",
+    srcs = ["assert_cuda_clang.cu.cc"],
+    tags = [
+        "gpu",
+        "manual",
+    ],
+    deps = ["@local_config_cuda//cuda:cuda_headers"],
+)
+
+cuda_library(
+    name = "assert_nvcc",
+    srcs = ["assert_nvcc.cu.cc"],
+    tags = [
+        "gpu",
+        "manual",
+    ],
+    # Notably, this builds fine in OSS without this dependency. Apparently,
+    # NVCC can give targets access to CUDA headers without letting Bazel know,
+    # while CUDA clang cannot.
+    deps = ["@local_config_cuda//cuda:cuda_headers"],
+)

third_party/xla/build_tools/configure/assert_clang.cc

@@ -0,0 +1,18 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef __clang__
+#error "__clang__ not defined!"
+#endif  // #ifdef __clang__

third_party/xla/build_tools/configure/assert_cuda_clang.cu.cc

@@ -0,0 +1,18 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#if !defined(__clang__) || !defined(__CUDA__)
+#error "__clang__ or __CUDA__ not defined!"
+#endif  // #if !defined(__clang__) || !defined(__CUDA__)

third_party/xla/build_tools/configure/assert_gcc.cc

@@ -0,0 +1,21 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// Notably, clang will define `__GNUC__`, so need to make sure __clang__ is not
+// defined to detect GCC (or, most correctly, some compiler that supports GNU
+// extensions that is not clang).
+#if !defined(__GNUC__) || defined(__clang__)
+#error "__GNUC__ is not defined independently of __clang__!"
+#endif  // #if !defined(__GNUC__) || defined(__clang__)

third_party/xla/build_tools/configure/assert_nvcc.cu.cc

@@ -0,0 +1,17 @@
+/* Copyright 2024 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef __NVCC__
+#error "__NVCC__ not defined!"
+#endif  // #ifdef __NVCC__

third_party/xla/build_tools/configure/configure.py

@@ -0,0 +1,624 @@
+#!/usr/bin/env python3
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Configure script to get build parameters from user.
+
+This script populates a bazelrc file that tells Bazel where to look for
+cuda versions and compilers. Note: that a configuration is possible to request,
+does not mean that it is supported (e.g. building with gcc). That being said,
+if this stops working for you on an unsupported build and you have a fix, please
+send a PR!
+
+Example usage:
+  `./configure.py --backend=cpu --host_compiler=clang`
+  Will write a bazelrc to the root of the repo with the lines required to find
+  the clang in your path. If that isn't the correct clang, you can override like
+  `./configure.py --backend=cpu --clang_path=<PATH_TO_YOUR_CLANG>`.
+
+TODO(ddunleavy): add more thorough validation.
+"""
+import argparse
+import dataclasses
+import enum
+import logging
+import os
+import pathlib
+import platform
+import shutil
+import subprocess
+import sys
+from typing import Optional
+
+_DEFAULT_BUILD_AND_TEST_TAG_FILTERS = ("-no_oss",)
+# Assume we are being invoked from the symlink at the root of the repo
+_XLA_SRC_ROOT = pathlib.Path(__file__).absolute().parent
+_XLA_BAZELRC_NAME = "xla_configure.bazelrc"
+_KW_ONLY_IF_PYTHON310 = {"kw_only": True} if sys.version_info >= (3, 10) else {}
+
+
+def _find_executable(executable: str) -> Optional[str]:
+  logging.info("Trying to find path to %s...", executable)
+  # Resolving the symlink is necessary for finding system headers.
+  if unresolved_path := shutil.which(executable):
+    return str(pathlib.Path(unresolved_path).resolve())
+  return None
+
+
+def _find_executable_or_die(
+    executable_name: str, executable_path: Optional[str] = None
+) -> str:
+  """Finds executable and resolves symlinks or raises RuntimeError.
+
+  Resolving symlinks is sometimes necessary for finding system headers.
+
+  Args:
+    executable_name: The name of the executable that we want to find.
+    executable_path: If not None, the path to the executable.
+
+  Returns:
+    The path to the executable we are looking for, after symlinks are resolved.
+  Raises:
+    RuntimeError: if path to the executable cannot be found.
+  """
+  if executable_path:
+    return str(pathlib.Path(executable_path).resolve(strict=True))
+  resolved_path_to_exe = _find_executable(executable_name)
+  if resolved_path_to_exe is None:
+    raise RuntimeError(
+        f"Could not find executable `{executable_name}`! "
+        "Please change your $PATH or pass the path directly like"
+        f"`--{executable_name}_path=path/to/executable."
+    )
+  logging.info("Found path to %s at %s", executable_name, resolved_path_to_exe)
+
+  return resolved_path_to_exe
+
+
+def _get_cuda_compute_capabilities_or_die() -> list[str]:
+  """Finds compute capabilities via nvidia-smi or rasies exception.
+
+  Returns:
+    list of unique, sorted strings representing compute capabilities:
+  Raises:
+    RuntimeError: if path to nvidia-smi couldn't be found.
+    subprocess.CalledProcessError: if nvidia-smi process failed.
+  """
+  try:
+    nvidia_smi = _find_executable_or_die("nvidia-smi")
+    nvidia_smi_proc = subprocess.run(
+        [nvidia_smi, "--query-gpu=compute_cap", "--format=csv,noheader"],
+        capture_output=True,
+        check=True,
+        text=True,
+    )
+    # Command above returns a newline separated list of compute capabilities
+    # with possible repeats. So we should unique them and sort the final result.
+    capabilities = sorted(set(nvidia_smi_proc.stdout.strip().split("\n")))
+    logging.info("Found CUDA compute capabilities: %s", capabilities)
+    return capabilities
+  except (RuntimeError, subprocess.CalledProcessError) as e:
+    logging.info(
+        "Could not find nvidia-smi, or nvidia-smi command failed. Please pass"
+        " capabilities directly using --cuda_compute_capabilities."
+    )
+    raise e
+
+
+def _get_clang_major_version(path_to_clang: str) -> int:
+  """Gets the major version of the clang at `path_to_clang`.
+
+  Args:
+    path_to_clang: Path to a clang executable
+
+  Returns:
+    The major version.
+  """
+  logging.info("Running echo __clang_major__ | %s -E -P -", path_to_clang)
+  clang_version_proc = subprocess.run(
+      [path_to_clang, "-E", "-P", "-"],
+      input="__clang_major__",
+      check=True,
+      capture_output=True,
+      text=True,
+  )
+  major_version = int(clang_version_proc.stdout)
+  logging.info("%s reports major version %s.", path_to_clang, major_version)
+
+  return major_version
+
+
+def _get_gcc_major_version(path_to_gcc: str) -> int:
+  """Gets the major version of the gcc at `path_to_gcc`.
+
+  Args:
+    path_to_gcc: Path to a gcc executable
+
+  Returns:
+    The major version.
+  """
+  logging.info("Running echo __GNUC__ | %s -E -P -", path_to_gcc)
+  gcc_version_proc = subprocess.run(
+      [path_to_gcc, "-E", "-P", "-"],
+      input="__GNUC__",
+      check=True,
+      capture_output=True,
+      text=True,
+  )
+  major_version = int(gcc_version_proc.stdout)
+  logging.info("%s reports major version %s.", path_to_gcc, major_version)
+
+  return major_version
+
+
+class ArgparseableEnum(enum.Enum):
+  """Enum base class with helper methods for working with argparse.
+
+  Example usage:
+  ```
+  class Fruit(ArgparseableEnum):
+    APPLE = enum.auto()
+
+  # argparse setup
+  parser.add_argument("--fruit", type=Fruit.from_str, choices=list(Fruit))
+  ```
+  Users can pass strings like `--fruit=apple` with nice error messages and the
+  parser will get the corresponding enum value.
+
+  NOTE: PyType gets confused when this class is used to create Enums in the
+  functional style like `ArgparseableEnum("Fruit", ["APPLE", "BANANA"])`.
+  """
+
+  def __str__(self):
+    return self.name
+
+  @classmethod
+  def from_str(cls, s):
+    s = s.upper()
+    try:
+      return cls[s]
+    except KeyError:
+      # Sloppy looking exception handling, but argparse will catch ValueError
+      # and give a pleasant error message. KeyError would not work here.
+      raise ValueError  # pylint: disable=raise-missing-from
+
+
+class Backend(ArgparseableEnum):
+  CPU = enum.auto()
+  CUDA = enum.auto()
+  ROCM = enum.auto()
+  SYCL = enum.auto()
+
+
+class HostCompiler(ArgparseableEnum):
+  CLANG = enum.auto()
+  GCC = enum.auto()
+
+
+class CudaCompiler(ArgparseableEnum):
+  CLANG = enum.auto()
+  NVCC = enum.auto()
+
+
+class RocmCompiler(ArgparseableEnum):
+  HIPCC = enum.auto()
+
+
+class SyclCompiler(ArgparseableEnum):
+  ICPX = enum.auto()
+
+
+class OS(ArgparseableEnum):
+  """Modeled after the values returned by `platform.system()`."""
+  LINUX = enum.auto()
+  DARWIN = enum.auto()
+  WINDOWS = enum.auto()
+
+
+@dataclasses.dataclass(**_KW_ONLY_IF_PYTHON310)
+class DiscoverablePathsAndVersions:
+  """Paths to various tools and libraries needed to build XLA.
+
+  This class is where all 'stateful' activity should happen, like trying to read
+  environment variables or looking for things in the $PATH. An instance that has
+  all fields set should not try to do any of these things though, so that this
+  file can remain unit testable.
+  """
+
+  clang_path: Optional[str] = None
+  clang_major_version: Optional[int] = None
+  gcc_path: Optional[str] = None
+  gcc_major_version: Optional[int] = None
+  lld_path: Optional[str] = None
+  ld_library_path: Optional[str] = None
+
+  # CUDA specific
+  cuda_version: Optional[str] = None
+  cuda_compute_capabilities: Optional[list[str]] = None
+  cudnn_version: Optional[str] = None
+  local_cuda_path: Optional[str] = None
+  local_cudnn_path: Optional[str] = None
+  local_nccl_path: Optional[str] = None
+
+  def get_relevant_paths_and_versions(self, config: "XLAConfigOptions"):
+    """Gets paths and versions as needed by the config.
+
+    Args:
+      config: XLAConfigOptions instance that determines what paths and versions
+        to try to autoconfigure.
+    """
+    if self.ld_library_path is None:
+      self.ld_library_path = os.environ.get("LD_LIBRARY_PATH", None)
+
+    if config.host_compiler == HostCompiler.CLANG:
+      if self.clang_path or not is_hermetic_build(config.backend, config.os):
+        self.clang_path = _find_executable_or_die("clang", self.clang_path)
+        self.clang_major_version = (
+            self.clang_major_version
+            or _get_clang_major_version(self.clang_path)
+        )
+
+        # Notably, we don't use `_find_executable_or_die` for lld, as it changes
+        # which commands it accepts based on its name! ld.lld is symlinked to a
+        # different executable just called lld, which should not be invoked
+        # directly.
+        self.lld_path = self.lld_path or shutil.which("ld.lld")
+      else:
+        # TODO: b/443091874 - set the version of Clang when it will be
+        # available outside of rules_ml_toolchain. Current hermetic Clang
+        # version is 18
+        self.clang_major_version = 18   # Hermetic toolchain
+
+    elif config.host_compiler == HostCompiler.GCC:
+      self.gcc_path = _find_executable_or_die("gcc", self.gcc_path)
+      self.gcc_major_version = self.gcc_major_version or _get_gcc_major_version(
+          self.gcc_path
+      )
+
+    if config.backend == Backend.CUDA:
+      if config.cuda_compiler == CudaCompiler.CLANG and (
+          self.clang_path or not is_hermetic_build(config.backend, config.os)
+      ):
+        self.clang_path = _find_executable_or_die("clang", self.clang_path)
+
+      if not self.cuda_compute_capabilities:
+        self.cuda_compute_capabilities = _get_cuda_compute_capabilities_or_die()
+
+
+@dataclasses.dataclass(frozen=True, **_KW_ONLY_IF_PYTHON310)
+class XLAConfigOptions:
+  """Represents XLA configuration options."""
+
+  backend: Backend
+  os: OS
+  python_bin_path: str
+  host_compiler: HostCompiler
+  compiler_options: list[str]
+
+  # CUDA specific
+  cuda_compiler: CudaCompiler
+  using_nccl: bool
+
+  # ROCM specific
+  rocm_compiler: RocmCompiler
+
+  # SYCL specific
+  sycl_compiler: SyclCompiler
+
+  def to_bazelrc_lines(
+      self,
+      dpav: DiscoverablePathsAndVersions,
+  ) -> list[str]:
+    """Creates a bazelrc given an XLAConfigOptions.
+
+    Necessary paths are provided by the user, or retrieved via
+    `self._get_relevant_paths`.
+
+    Args:
+      dpav: DiscoverablePathsAndVersions that may hold user-specified paths and
+        versions. The dpav will then read from `self` to determine what to try
+        to auto-configure.
+
+    Returns:
+      The lines of a bazelrc.
+    """
+    dpav.get_relevant_paths_and_versions(self)
+    rc = []
+    build_and_test_tag_filters = list(_DEFAULT_BUILD_AND_TEST_TAG_FILTERS)
+
+    if self.os == OS.DARWIN:
+      build_and_test_tag_filters.append("-no_mac")
+
+    # Platform independent options based on host compiler
+    if self.host_compiler == HostCompiler.GCC:
+      rc.append(f"build --action_env GCC_HOST_COMPILER_PATH={dpav.gcc_path}")
+    elif self.host_compiler == HostCompiler.CLANG:
+      if dpav.clang_path:
+        rc.append("build --config clang_local")
+        rc.append(f"build --action_env CLANG_COMPILER_PATH={dpav.clang_path}")
+        rc.append(f"build --repo_env CC={dpav.clang_path}")
+        rc.append(f"build --repo_env BAZEL_COMPILER={dpav.clang_path}")
+      self.compiler_options.append("-Wno-error=unused-command-line-argument")
+      if dpav.lld_path:
+        rc.append(f"build --linkopt --ld-path={dpav.lld_path}")
+
+    if self.backend == Backend.CPU:
+      build_and_test_tag_filters.append("-gpu")
+
+    elif self.backend == Backend.CUDA:
+      build_and_test_tag_filters.append("-rocm-only")
+      build_and_test_tag_filters.append("-oneapi-only")
+
+      compiler_pair = self.cuda_compiler, self.host_compiler
+
+      if compiler_pair == (CudaCompiler.CLANG, HostCompiler.CLANG):
+        if not dpav.clang_path:
+          rc.append("build --config cuda_clang")
+        else:
+          rc.append("build --config cuda_clang_local")
+          rc.append(
+              f"build --action_env CLANG_CUDA_COMPILER_PATH={dpav.clang_path}"
+          )
+      elif compiler_pair == (CudaCompiler.NVCC, HostCompiler.CLANG):
+        if not dpav.clang_path:
+          rc.append("build --config cuda_nvcc")
+        else:
+          rc.append("build --config cuda_nvcc_clang_local")
+          # This is demanded by cuda_configure.bzl
+          rc.append(
+              f"build --action_env CLANG_CUDA_COMPILER_PATH={dpav.clang_path}"
+          )
+      elif compiler_pair == (CudaCompiler.NVCC, HostCompiler.GCC):
+        rc.append("build --config cuda")
+      else:
+        raise NotImplementedError(
+            "CUDA clang with host compiler gcc not supported"
+        )
+
+      # Lines needed for CUDA backend regardless of CUDA/host compiler
+      if dpav.cuda_version:
+        rc.append(
+            f"build:cuda --repo_env HERMETIC_CUDA_VERSION={dpav.cuda_version}"
+        )
+      rc.append(
+          "build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES="
+          f"{','.join(dpav.cuda_compute_capabilities)}"
+      )
+      if dpav.cudnn_version:
+        rc.append(
+            f"build:cuda --repo_env HERMETIC_CUDNN_VERSION={dpav.cudnn_version}"
+        )
+      if dpav.local_cuda_path:
+        rc.append(
+            f"build:cuda --repo_env LOCAL_CUDA_PATH={dpav.local_cuda_path}"
+        )
+      if dpav.local_cudnn_path:
+        rc.append(
+            f"build:cuda --repo_env LOCAL_CUDNN_PATH={dpav.local_cudnn_path}"
+        )
+      if dpav.local_nccl_path:
+        rc.append(
+            f"build:cuda --repo_env LOCAL_NCCL_PATH={dpav.local_nccl_path}"
+        )
+      if not self.using_nccl:
+        rc.append("build --config nonccl")
+    elif self.backend == Backend.ROCM:
+      build_and_test_tag_filters.append("-cuda-only")
+      build_and_test_tag_filters.append("-oneapi-only")
+
+      compiler_pair = self.rocm_compiler, self.host_compiler
+
+      if compiler_pair == (RocmCompiler.HIPCC, HostCompiler.CLANG):
+        rc.append("build --config rocm")
+        # This is demanded by rocm_configure.bzl.
+        rc.append(f"build --action_env CLANG_COMPILER_PATH={dpav.clang_path}")
+      elif compiler_pair == (RocmCompiler.HIPCC, HostCompiler.GCC):
+        rc.append("build --config rocm")
+      else:
+        raise NotImplementedError("ROCm clang with host compiler not supported")
+    elif self.backend == Backend.SYCL:
+      build_and_test_tag_filters.append("-cuda-only")
+      build_and_test_tag_filters.append("-rocm-only")
+      build_and_test_tag_filters.append("-no-oneapi")
+
+      compiler_pair = self.sycl_compiler, self.host_compiler
+
+      if compiler_pair == (SyclCompiler.ICPX, HostCompiler.CLANG):
+        rc.append("build --config sycl")
+        rc.append("build --config icpx_clang")
+      else:
+        raise NotImplementedError(" Sycl with host compiler not supported")
+
+    # Lines that are added for every backend
+    if dpav.ld_library_path:
+      rc.append(f"build --action_env LD_LIBRARY_PATH={dpav.ld_library_path}")
+
+    # Needed due to error in @upb//:upb which is a dep of @com_github_grpc_grpc
+    # error: defining a type within 'offsetof' is a Clang extension
+    if dpav.clang_major_version in (16, 17, 18):
+      self.compiler_options.append("-Wno-gnu-offsetof-extensions")
+    # error: defining a type within 'offsetof' is a C23 extension
+    if dpav.clang_major_version and dpav.clang_major_version >= 19:
+      self.compiler_options.append("-Wno-c23-extensions")
+
+    rc.append(f"build --action_env PYTHON_BIN_PATH={self.python_bin_path}")
+    rc.append(f"build --python_path {self.python_bin_path}")
+    rc.append("test --test_env LD_LIBRARY_PATH")
+    rc.append("test --test_size_filters small,medium")
+
+    rc.extend([
+        f"build --copt {compiler_option}"
+        for compiler_option in self.compiler_options
+    ])
+
+    # Add build and test tag filters
+    build_and_test_tag_filters = ",".join(build_and_test_tag_filters)
+    rc.append(f"build --build_tag_filters {build_and_test_tag_filters}")
+    rc.append(f"build --test_tag_filters {build_and_test_tag_filters}")
+    rc.append(f"test --build_tag_filters {build_and_test_tag_filters}")
+    rc.append(f"test --test_tag_filters {build_and_test_tag_filters}")
+
+    return rc
+
+
+def _parse_args():
+  """Creates an argparse.ArgumentParser and parses arguments."""
+  # pylint: disable=C3001
+  comma_separated_list = lambda l: [s.strip() for s in l.split(",")]
+
+  parser = argparse.ArgumentParser(allow_abbrev=False)
+  parser.add_argument(
+      "--backend",
+      type=Backend.from_str,
+      choices=list(Backend),
+      required=True,
+  )
+  parser.add_argument(
+      "--os", type=OS.from_str, choices=list(OS), default=platform.system()
+  )
+  parser.add_argument(
+      "--host_compiler",
+      type=HostCompiler.from_str,
+      choices=list(HostCompiler),
+      default="clang",
+  )
+  parser.add_argument(
+      "--cuda_compiler",
+      type=CudaCompiler.from_str,
+      choices=list(CudaCompiler),
+      default="nvcc",
+  )
+  parser.add_argument(
+      "--rocm_compiler",
+      type=RocmCompiler.from_str,
+      choices=list(RocmCompiler),
+      default="hipcc",
+  )
+  parser.add_argument(
+      "--sycl_compiler",
+      type=SyclCompiler.from_str,
+      choices=list(SyclCompiler),
+      default="icpx",
+  )
+  parser.add_argument(
+      "--cuda_compute_capabilities",
+      type=comma_separated_list,
+      default=None,
+  )
+  parser.add_argument("--python_bin_path", default=sys.executable)
+  parser.add_argument(
+      "--compiler_options",
+      type=comma_separated_list,
+      default="-Wno-sign-compare",
+  )
+  parser.add_argument("--nccl", action="store_true")
+
+  # Path and version overrides
+  path_help = "Optional: will be found on PATH if possible."
+  parser.add_argument("--clang_path", help=path_help)
+  parser.add_argument("--gcc_path", help=path_help)
+  parser.add_argument(
+      "--ld_library_path",
+      help=(
+          "Optional: will be automatically taken from the current environment"
+          " if flag is not set"
+      ),
+  )
+  parser.add_argument("--lld_path", help=path_help)
+
+  # CUDA specific
+  parser.add_argument(
+      "--cuda_version",
+      help="Optional: CUDA will be downloaded by Bazel if the flag is set",
+  )
+  parser.add_argument(
+      "--cudnn_version",
+      help="Optional: CUDNN will be downloaded by Bazel if the flag is set",
+  )
+  parser.add_argument(
+      "--local_cuda_path",
+      help=(
+          "Optional: Local CUDA dir will be used in dependencies if the flag"
+          " is set"
+      ),
+  )
+  parser.add_argument(
+      "--local_cudnn_path",
+      help=(
+          "Optional: Local CUDNN dir will be used in dependencies if the flag"
+          " is set"
+      ),
+  )
+  parser.add_argument(
+      "--local_nccl_path",
+      help=(
+          "Optional: Local NCCL dir will be used in dependencies if the flag"
+          " is set"
+      ),
+  )
+
+  return parser.parse_args()
+
+
+def main():
+  # Setup logging
+  logging.basicConfig()
+  logging.getLogger().setLevel(logging.INFO)
+
+  args = _parse_args()
+
+  config = XLAConfigOptions(
+      backend=args.backend,
+      os=args.os,
+      host_compiler=args.host_compiler,
+      cuda_compiler=args.cuda_compiler,
+      python_bin_path=args.python_bin_path,
+      compiler_options=args.compiler_options,
+      using_nccl=args.nccl,
+      rocm_compiler=args.rocm_compiler,
+      sycl_compiler=args.sycl_compiler,
+  )
+
+  bazelrc_lines = config.to_bazelrc_lines(
+      DiscoverablePathsAndVersions(
+          clang_path=args.clang_path,
+          gcc_path=args.gcc_path,
+          lld_path=args.lld_path,
+          ld_library_path=args.ld_library_path,
+          cuda_version=args.cuda_version,
+          cudnn_version=args.cudnn_version,
+          cuda_compute_capabilities=args.cuda_compute_capabilities,
+          local_cuda_path=args.local_cuda_path,
+          local_cudnn_path=args.local_cudnn_path,
+          local_nccl_path=args.local_nccl_path,
+      )
+  )
+
+  bazelrc_path = _XLA_SRC_ROOT / _XLA_BAZELRC_NAME
+  bazelrc_contents = "\n".join(bazelrc_lines) + "\n"
+
+  with (bazelrc_path).open("w") as f:
+    logging.info("Writing bazelrc to %s...", bazelrc_path)
+    f.write(bazelrc_contents)
+
+
+def is_hermetic_build(backend: Backend, os_host: OS):
+  return (
+      backend != Backend.ROCM
+      and os_host == OS.LINUX
+  )
+
+
+if __name__ == "__main__":
+  raise SystemExit(main())

third_party/xla/build_tools/configure/configure_test.py

@@ -0,0 +1,262 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import os
+
+from absl.testing import absltest
+
+from build_tools import test_utils
+from build_tools.configure import configure
+
+
+XLAConfigOptions = configure.XLAConfigOptions
+DiscoverablePathsAndVersions = configure.DiscoverablePathsAndVersions
+Backend = configure.Backend
+HostCompiler = configure.HostCompiler
+CudaCompiler = configure.CudaCompiler
+RocmCompiler = configure.RocmCompiler
+SyclCompiler = configure.SyclCompiler
+OS = configure.OS
+
+_PYTHON_BIN_PATH = "/usr/bin/python3"
+_CLANG_PATH = "/usr/lib/llvm-18/bin/clang"
+_GCC_PATH = "/usr/bin/gcc"
+_COMPILER_OPTIONS = ("-Wno-sign-compare",)
+
+# CUDA specific paths and versions
+_CUDA_SPECIFIC_PATHS_AND_VERSIONS = {
+    "cuda_version": '"12.8.0"',
+    "cuda_compute_capabilities": ["7.5"],
+    "cudnn_version": '"9.8.0"',
+    "ld_library_path": "/usr/local/nvidia/lib:/usr/local/nvidia/lib64",
+}
+_CUDA_COMPUTE_CAPABILITIES_AND_LD_LIBRARY_PATH = {
+    "cuda_compute_capabilities": [
+        "sm_50",
+        "sm_60",
+        "sm_70",
+        "sm_80",
+        "compute_90",
+    ],
+    "ld_library_path": "/usr/local/nvidia/lib:/usr/local/nvidia/lib64",
+}
+
+
+class ConfigureTest(absltest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+
+    testdata = (
+        test_utils.xla_src_root() / "build_tools" / "configure" / "testdata"
+    )
+
+    with (testdata / "clang_local.bazelrc").open() as f:
+      cls.clang_local_bazelrc_lines = [line.strip() for line in f.readlines()]
+
+    with (testdata / "gcc.bazelrc").open() as f:
+      resolved_gcc_path = os.path.realpath(_GCC_PATH)
+      cls.gcc_bazelrc_lines = [
+          line.strip().replace(_GCC_PATH, resolved_gcc_path)
+          for line in f.readlines()
+      ]
+
+    with (testdata / "cuda_clang.bazelrc").open() as f:
+      cls.cuda_clang_bazelrc_lines = [line.strip() for line in f.readlines()]
+
+    with (testdata / "cuda_clang_local.bazelrc").open() as f:
+      cls.cuda_clang_local_bazelrc_lines = [
+          line.strip() for line in f.readlines()
+      ]
+
+    with (testdata / "default_cuda_clang.bazelrc").open() as f:
+      cls.default_cuda_clang_bazelrc_lines = [
+          line.strip() for line in f.readlines()
+      ]
+
+    with (testdata / "nvcc_clang_local.bazelrc").open() as f:
+      cls.nvcc_clang_local_bazelrc_lines = [
+          line.strip() for line in f.readlines()
+      ]
+
+    with (testdata / "nvcc_gcc.bazelrc").open() as f:
+      resolved_gcc_path = os.path.realpath(_GCC_PATH)
+      cls.nvcc_gcc_bazelrc_lines = [
+          line.strip().replace(_GCC_PATH, resolved_gcc_path)
+          for line in f.readlines()
+      ]
+
+  def test_clang_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CPU,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.CLANG,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.NVCC,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            clang_path=_CLANG_PATH,
+            ld_library_path="",
+            clang_major_version=18,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.clang_local_bazelrc_lines)
+
+  def test_gcc_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CPU,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.GCC,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.NVCC,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            gcc_path=_GCC_PATH,
+            ld_library_path="",
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.gcc_bazelrc_lines)
+
+  def test_cuda_clang_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CUDA,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.CLANG,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.CLANG,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            clang_path=None,
+            clang_major_version=None,
+            **_CUDA_SPECIFIC_PATHS_AND_VERSIONS,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.cuda_clang_bazelrc_lines)
+
+  def test_cuda_clang_local_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CUDA,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.CLANG,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.CLANG,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            clang_path=_CLANG_PATH,
+            clang_major_version=18,
+            **_CUDA_SPECIFIC_PATHS_AND_VERSIONS,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.cuda_clang_local_bazelrc_lines)
+
+  def test_default_cuda_clang_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CUDA,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.CLANG,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.CLANG,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            clang_path=_CLANG_PATH,
+            clang_major_version=17,
+            **_CUDA_COMPUTE_CAPABILITIES_AND_LD_LIBRARY_PATH,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.default_cuda_clang_bazelrc_lines)
+
+  def test_nvcc_clang_local_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CUDA,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.CLANG,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.NVCC,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            clang_path=_CLANG_PATH,
+            clang_major_version=18,
+            **_CUDA_SPECIFIC_PATHS_AND_VERSIONS,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.nvcc_clang_local_bazelrc_lines)
+
+  def test_nvcc_gcc_bazelrc(self):
+    config = XLAConfigOptions(
+        backend=Backend.CUDA,
+        os=OS.LINUX,
+        python_bin_path=_PYTHON_BIN_PATH,
+        host_compiler=HostCompiler.GCC,
+        compiler_options=list(_COMPILER_OPTIONS),
+        cuda_compiler=CudaCompiler.NVCC,
+        using_nccl=False,
+        rocm_compiler=RocmCompiler.HIPCC,
+        sycl_compiler=SyclCompiler.ICPX,
+    )
+
+    bazelrc_lines = config.to_bazelrc_lines(
+        DiscoverablePathsAndVersions(
+            gcc_path=_GCC_PATH,
+            **_CUDA_SPECIFIC_PATHS_AND_VERSIONS,
+        )
+    )
+
+    self.assertEqual(bazelrc_lines, self.nvcc_gcc_bazelrc_lines)
+
+
+if __name__ == "__main__":
+  absltest.main()

third_party/xla/build_tools/configure/testdata/clang_local.bazelrc

@@ -0,0 +1,15 @@
+build --config clang_local
+build --action_env CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build --repo_env CC=/usr/lib/llvm-18/bin/clang
+build --repo_env BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-gpu
+build --test_tag_filters -no_oss,-gpu
+test --build_tag_filters -no_oss,-gpu
+test --test_tag_filters -no_oss,-gpu

third_party/xla/build_tools/configure/testdata/cuda_clang.bazelrc

@@ -0,0 +1,17 @@
+build --config cuda_clang
+build:cuda --repo_env HERMETIC_CUDA_VERSION="12.8.0"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=7.5
+build:cuda --repo_env HERMETIC_CUDNN_VERSION="9.8.0"
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/configure/testdata/cuda_clang_local.bazelrc

@@ -0,0 +1,22 @@
+build --config clang_local
+build --action_env CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build --repo_env CC=/usr/lib/llvm-18/bin/clang
+build --repo_env BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang
+build --config cuda_clang_local
+build --action_env CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build:cuda --repo_env HERMETIC_CUDA_VERSION="12.8.0"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=7.5
+build:cuda --repo_env HERMETIC_CUDNN_VERSION="9.8.0"
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/configure/testdata/default_cuda_clang.bazelrc

@@ -0,0 +1,20 @@
+build --config clang_local
+build --action_env CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build --repo_env CC=/usr/lib/llvm-18/bin/clang
+build --repo_env BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang
+build --config cuda_clang_local
+build --action_env CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=sm_50,sm_60,sm_70,sm_80,compute_90
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/configure/testdata/gcc.bazelrc

@@ -0,0 +1,10 @@
+build --action_env GCC_HOST_COMPILER_PATH=/usr/bin/gcc
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --build_tag_filters -no_oss,-gpu
+build --test_tag_filters -no_oss,-gpu
+test --build_tag_filters -no_oss,-gpu
+test --test_tag_filters -no_oss,-gpu

third_party/xla/build_tools/configure/testdata/nvcc_clang.bazelrc

@@ -0,0 +1,17 @@
+build --config cuda_nvcc
+build:cuda --repo_env HERMETIC_CUDA_VERSION="12.8.0"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=7.5
+build:cuda --repo_env HERMETIC_CUDNN_VERSION="9.8.0"
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/configure/testdata/nvcc_clang_local.bazelrc

@@ -0,0 +1,22 @@
+build --config clang_local
+build --action_env CLANG_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build --repo_env CC=/usr/lib/llvm-18/bin/clang
+build --repo_env BAZEL_COMPILER=/usr/lib/llvm-18/bin/clang
+build --config cuda_nvcc_clang_local
+build --action_env CLANG_CUDA_COMPILER_PATH=/usr/lib/llvm-18/bin/clang
+build:cuda --repo_env HERMETIC_CUDA_VERSION="12.8.0"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=7.5
+build:cuda --repo_env HERMETIC_CUDNN_VERSION="9.8.0"
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --copt -Wno-error=unused-command-line-argument
+build --copt -Wno-gnu-offsetof-extensions
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/configure/testdata/nvcc_gcc.bazelrc

@@ -0,0 +1,16 @@
+build --action_env GCC_HOST_COMPILER_PATH=/usr/bin/gcc
+build --config cuda
+build:cuda --repo_env HERMETIC_CUDA_VERSION="12.8.0"
+build:cuda --repo_env HERMETIC_CUDA_COMPUTE_CAPABILITIES=7.5
+build:cuda --repo_env HERMETIC_CUDNN_VERSION="9.8.0"
+build --config nonccl
+build --action_env LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+build --action_env PYTHON_BIN_PATH=/usr/bin/python3
+build --python_path /usr/bin/python3
+test --test_env LD_LIBRARY_PATH
+test --test_size_filters small,medium
+build --copt -Wno-sign-compare
+build --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+build --test_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --build_tag_filters -no_oss,-rocm-only,-oneapi-only
+test --test_tag_filters -no_oss,-rocm-only,-oneapi-only

third_party/xla/build_tools/dependencies/BUILD

@@ -0,0 +1,19 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)

third_party/xla/build_tools/dependencies/aspects.bzl

@@ -0,0 +1,90 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""A collection of Bazel aspects that can help detecting dependency violations
+
+The dependency violation detection works by iterating through all targets in XLA
+and comparing the applied tags of each target to the tags of all its dependencies.
+
+If a target is tagged `gpu` it means it can only be used in an XLA build with the
+GPU backend enabled. Hence all targets that are NOT tagged `gpu` may never depend
+on a target that IS tagged `gpu` if we are building XLA with only the CPU backend
+enabled.
+
+The Bazel aspect runs after Bazel's analysis phase. That means all `select` expressions
+(and its derivatives like the `if_gpu_is_configured` macro) have been evaluated and
+the actual build configuration is taken into account.
+
+The easiest way to run the aspect is during a build:
+`bazel build --aspects build_tools/dependencies/aspects.bzl%validate_gpu_tag //xla/...`
+
+But a cquery expression also works:
+`bazel cquery --aspects build_tools/dependencies/aspects.bzl%validate_gpu_tag //xla/...`
+
+The results are reported as debug prints and need to be fished out of stderr. There
+are ways to make it less hacky but the complexity of the aspect would also increase
+quite a bit.
+"""
+
+DependencyViolationInfo = provider(
+    "Internal provider needed by the dependency violation check",
+    fields = {
+        # We can't access the tags of a dependency through the context, so instead we
+        # "send" the tags to the dependee through this provider.
+        "tags": "Tags of the dependecy",
+    },
+)
+
+def _dependency_violation_aspect_impl(_, ctx, tag):
+    if not hasattr(ctx.rule.attr, "deps"):
+        return [DependencyViolationInfo(tags = ctx.rule.attr.tags)]
+
+    for dep in ctx.rule.attr.deps:
+        if DependencyViolationInfo not in dep:
+            continue
+        dep_tags = dep[DependencyViolationInfo].tags
+        if tag in dep_tags and tag not in ctx.rule.attr.tags:
+            print("[Violation] {} (not tagged {}) depends on {} (tagged {})".format(
+                ctx.label,
+                tag,
+                dep.label,
+                tag,
+            ))  # buildifier: disable=print
+
+    return [DependencyViolationInfo(tags = ctx.rule.attr.tags)]
+
+def _gpu_tag_violation_aspect_impl(target, ctx):
+    return _dependency_violation_aspect_impl(target, ctx, "gpu")
+
+validate_gpu_tag = aspect(
+    implementation = _gpu_tag_violation_aspect_impl,
+    attr_aspects = ["deps"],
+)
+
+def _cuda_only_tag_violation_aspect_impl(target, ctx):
+    return _dependency_violation_aspect_impl(target, ctx, "cuda-only")
+
+validate_cuda_only_tag = aspect(
+    implementation = _cuda_only_tag_violation_aspect_impl,
+    attr_aspects = ["deps"],
+)
+
+def _rocm_only_tag_violation_aspect_impl(target, ctx):
+    return _dependency_violation_aspect_impl(target, ctx, "rocm-only")
+
+validate_rocm_only_tag = aspect(
+    implementation = _rocm_only_tag_violation_aspect_impl,
+    attr_aspects = ["deps"],
+)

third_party/xla/build_tools/dependencies/gen_disable_layering_check_patch.sh

@@ -0,0 +1,61 @@
+#!/bin/bash
+# Copyright 2025 The OpenXLA Authors.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# Generates a patch file that disables the layering check for all cc_library
+# targets in the archive. Both BUILD and BUILD.bazel files are taken into account.
+#
+# The script takes one argument: the URL of the .tar.gz archive to download.
+#
+# The following tools are needed (need to be installed on the machine):
+# - curl
+# - git
+# - buildozer (from Bazel buildtools)
+#
+# The tool has originally been written for ortools but should work for similarly structured
+# projects as well.
+#
+# Example:
+# build_tools/dependencies/gen_disable_layering_check_patch.sh \
+# https://github.com/google/or-tools/archive/v9.11.tar.gz \
+# > third_party/ortools/layering_check.patch
+
+set -euo pipefail
+
+readonly TMP_DIR=$(mktemp -d)
+trap 'rm -rf -- $TMP_DIR' EXIT
+
+echo "Downloading archive $1..." >&2
+curl -Lqo "$TMP_DIR/archive.tar.gz" "$1" 1>&2
+
+echo "Extracting archive..." >&2
+mkdir -p "$TMP_DIR/extracted" 1>&2
+tar  -x -C "$TMP_DIR/extracted" -f "$TMP_DIR/archive.tar.gz" --strip-components=1 1>&2
+
+echo "Initialzing temporary git repo..." >&2
+git -C "$TMP_DIR/extracted" init 1>&2
+git -C "$TMP_DIR/extracted" add . 1>&2
+git -C "$TMP_DIR/extracted" commit --no-verify -m "original state" -q 1>&2
+
+echo "Patching build targets..." >&2
+find $TMP_DIR/extracted -name BUILD.bazel -or -name BUILD | while read f; do
+   buildozer 'add features "-layering_check"' $(dirname $f):%cc_library 1>&2 || exit_code=$?
+   if [[ $exit_code -ne 0 && $exit_code -ne 3 ]]; then
+     echo "Buildozer command failed with exit code: $exit_code" >&2
+     exit $exit_code
+   fi
+done
+
+echo "Generating diff..." >&2
+git -C "$TMP_DIR/extracted" --no-pager diff

third_party/xla/build_tools/lint/BUILD

@@ -0,0 +1,97 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+load("@xla//third_party/rules_python/python:py_test.bzl", "py_test")
+load("//xla:pytype.bzl", "pytype_strict_binary", "pytype_strict_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "check_contents",
+    srcs = ["check_contents.py"],
+    deps = [":diff_parser"],
+)
+
+pytype_strict_library(
+    name = "check_header_guards",
+    srcs = ["check_header_guards.py"],
+)
+
+pytype_strict_library(
+    name = "diff_parser",
+    srcs = ["diff_parser.py"],
+    visibility = ["//visibility:public"],
+)
+
+pytype_strict_library(
+    name = "generate_compile_commands",
+    srcs = ["generate_compile_commands.py"],
+)
+
+py_test(
+    name = "check_contents_test",
+    srcs = ["check_contents_test.py"],
+    data = [
+        "testdata/bad_cc.diff",
+        "testdata/important_cc.diff",
+    ],
+    deps = [
+        ":check_contents",
+        ":diff_parser",
+        "//build_tools:test_utils",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+py_test(
+    name = "diff_parser_test",
+    srcs = ["diff_parser_test.py"],
+    data = [
+        "testdata/bad_cc.diff",
+        "testdata/crosstool.diff",
+        "testdata/important_cc.diff",
+    ],
+    deps = [
+        ":diff_parser",
+        "//build_tools:test_utils",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+py_test(
+    name = "generate_compile_commands_test",
+    srcs = ["generate_compile_commands_test.py"],
+    deps = [
+        ":generate_compile_commands",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+py_test(
+    name = "check_header_guards_test",
+    srcs = ["check_header_guards_test.py"],
+    deps = [
+        ":check_header_guards",
+        "@absl_py//absl/testing:absltest",
+    ],
+)
+
+pytype_strict_binary(
+    name = "tags",
+    srcs = ["tags.py"],
+)

third_party/xla/build_tools/lint/check_contents.py

@@ -0,0 +1,179 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Command line tool for checking for regexes in diffs.
+
+Filters `git diff` by path, then checks to make sure no lines matching a regex
+have been added in the diff.
+"""
+import argparse
+import dataclasses
+import logging  # Intended to run on vanilla Github Actions runner
+import re
+import sys
+from typing import Iterable, Sequence
+
+from build_tools.lint import diff_parser
+
+
+@dataclasses.dataclass
+class RegexLocation:
+  """Path and line where a prohibited regex was found.
+
+  Attributes:
+    path: Path of the file which has the prohibited regex.
+    line_number: The number of the offending line.
+    line_contents: The text of the offending line.
+    matched_text: The exact string matched by the regex.
+  """
+
+  path: str
+  line_number: int
+  line_contents: str
+  matched_text: str
+
+
+def filter_hunks_by_path(
+    hunks: Iterable[diff_parser.Hunk],
+    *,
+    path_regexes: list[str],
+    path_regex_exclusions: list[str],
+) -> list[diff_parser.Hunk]:
+  """Filters files according to path_regexes.
+
+  If a file matches both a path_regex and a path_regex_exclusion, then
+  it will be filtered out.
+
+  Arguments:
+    hunks: A sequence of Hunk objects representing the hunks of the diff in the
+      change.
+    path_regexes: A list of regexes. Paths matching these will pass through the
+      filter. By default, every path is matched.
+    path_regex_exclusions: A list of regexes. Paths that match both a path_regex
+      and a path_regex_exclusion won't pass through the filter.
+
+  Returns:
+    A list of FileDiffs whose paths match a path_regex and don't match
+      any path_regex_exclusions.
+  """
+
+  if not path_regexes:
+    path_regexes = [".*"]  # by default match everything
+
+  path_regexes = [re.compile(regex) for regex in path_regexes]
+
+  def should_include(path: str) -> bool:
+    return any(regex.search(path) for regex in path_regexes)
+
+  path_regex_exclusions = [re.compile(regex) for regex in path_regex_exclusions]
+
+  def should_exclude(path: str) -> bool:
+    return any(regex.search(path) for regex in path_regex_exclusions)
+
+  return [
+      hunk
+      for hunk in hunks
+      if should_include(hunk.file) and not should_exclude(hunk.file)
+  ]
+
+
+def check_diffs(
+    hunks: Iterable[diff_parser.Hunk],
+    *,
+    prohibited_regex: str,
+    suppression_regex: str | None = None,
+) -> list[RegexLocation]:
+  """Checks FileDiffs for prohibited regexes.
+
+  Arguments:
+    hunks: A sequence of Hunk objects representing the hunks of the diff.
+    prohibited_regex: The regex that isn't allowed in the diff.
+    suppression_regex: A regex used as an escape hatch to allow the prohibited
+      regex in the diff. If this is found on the same line as prohibited_regex,
+      there is no error.
+
+  Returns:
+    A list of RegexLocations where the prohibited_regex is found.
+  """
+
+  prohibited_regex = re.compile(prohibited_regex)
+  if suppression_regex is not None:
+    suppression_regex = re.compile(suppression_regex)
+
+  def should_not_suppress(line) -> bool:
+    if suppression_regex:
+      return not suppression_regex.search(line)
+    return True
+
+  regex_locations = []
+  for hunk in hunks:
+    for line_no, line in hunk.added_lines():
+      if should_not_suppress(line):
+        regex_locations.extend(
+            [
+                RegexLocation(hunk.file, line_no, line, regex_match.group())
+                for regex_match in prohibited_regex.finditer(line)
+            ]
+        )
+
+  return regex_locations
+
+
+def main(argv: Sequence[str]):
+  parser = argparse.ArgumentParser(
+      description="Check `git diff` for prohibited regexes."
+  )
+  parser.add_argument("--path_regex", nargs="*", default=[])
+  parser.add_argument("--path_regex_exclusion", nargs="*", default=[])
+  parser.add_argument("--prohibited_regex", required=True)
+  parser.add_argument("--suppression_regex")
+  parser.add_argument("--failure_message", required=True)
+
+  # We don't want to include path/to/check_contents.py as an argument
+  args = parser.parse_args(argv[1:])
+
+  file_diffs = filter_hunks_by_path(
+      diff_parser.parse_hunks(diff_parser.get_git_diff_stdout()),
+      path_regexes=args.path_regex,
+      path_regex_exclusions=args.path_regex_exclusion,
+  )
+
+  regex_locations = check_diffs(
+      file_diffs,
+      prohibited_regex=args.prohibited_regex,
+      suppression_regex=args.suppression_regex,
+  )
+
+  if regex_locations:
+    for loc in regex_locations:
+      logging.error(
+          "Found `%s` in %s:%s",
+          args.prohibited_regex,
+          loc.path,
+          loc.line_number,
+      )
+      logging.error(
+          "Matched `%s` in line `%s`", loc.matched_text, loc.line_contents
+      )
+      logging.error("Failure message: %s", args.failure_message)
+    sys.exit(1)
+  else:
+    logging.info(
+        "Prohibited regex `%s` not found in diff!", args.prohibited_regex
+    )
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+  main(sys.argv)

third_party/xla/build_tools/lint/check_contents_test.py

@@ -0,0 +1,114 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from absl.testing import absltest
+
+from build_tools import test_utils
+from build_tools.lint import check_contents
+from build_tools.lint import diff_parser
+
+
+class CheckDiffsTest(absltest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+
+    testdata = test_utils.xla_src_root() / "build_tools" / "lint" / "testdata"
+    with (testdata / "bad_cc.diff").open() as f:
+      cls.bad_cc_hunks = diff_parser.parse_hunks(f.read())
+
+    with (testdata / "important_cc.diff").open() as f:
+      cls.important_cc_hunks = diff_parser.parse_hunks(f.read())
+
+  def test_check_good_diff(self):
+
+    locs = check_contents.check_diffs(
+        self.bad_cc_hunks,
+        prohibited_regex="Make_Unique",
+        suppression_regex="OK",
+    )
+    self.assertEmpty(locs, 0)
+
+  def test_check_suppressed_diff_without_suppressions(self):
+
+    locs = check_contents.check_diffs(
+        self.bad_cc_hunks, prohibited_regex="Make_Unique"
+    )
+
+    expected_locs = [
+        check_contents.RegexLocation(
+            path="src/dir/bad.cc",
+            line_number=3,
+            line_contents="using Make_Unique = std::make_unique; // OK",
+            matched_text="Make_Unique",
+        ),
+        check_contents.RegexLocation(
+            path="src/dir/bad.cc",
+            line_number=6,
+            line_contents="  return Make_Unique<int>(a + b); // OK. Fixed now!",
+            matched_text="Make_Unique",
+        ),
+    ]
+
+    self.assertEqual(locs, expected_locs)
+
+  def test_check_suppressed_diff_with_path_regexes(self):
+    filtered_hunks = check_contents.filter_hunks_by_path(
+        self.bad_cc_hunks,
+        path_regexes=["src/important\\..*"],
+        path_regex_exclusions=[],
+    )
+
+    self.assertLen(filtered_hunks, 1)
+
+    locs = check_contents.check_diffs(
+        filtered_hunks, prohibited_regex="Make_Unique"
+    )
+
+    self.assertEmpty(locs)
+
+  def test_check_suppressed_diff_with_exclusions(self):
+    filtered_hunks = check_contents.filter_hunks_by_path(
+        self.bad_cc_hunks,
+        path_regexes=[],
+        path_regex_exclusions=["src/dir/.*"],
+    )
+
+    self.assertLen(filtered_hunks, 1)
+
+    locs = check_contents.check_diffs(
+        filtered_hunks, prohibited_regex="Make_Unique"
+    )
+
+    self.assertEmpty(locs)
+
+  def test_check_suppressed_diff_with_suppression(self):
+
+    filtered_hunks = check_contents.filter_hunks_by_path(
+        self.bad_cc_hunks, path_regexes=[], path_regex_exclusions=[]
+    )
+
+    # filtering without path_regex(_exclusions) is a noop
+    self.assertEqual(self.bad_cc_hunks, filtered_hunks)
+
+    locs = check_contents.check_diffs(
+        filtered_hunks, prohibited_regex="Make_Unique", suppression_regex="OK"
+    )
+
+    self.assertEmpty(locs)
+
+
+if __name__ == "__main__":
+  absltest.main()

third_party/xla/build_tools/lint/check_dwyu.py

@@ -0,0 +1,250 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Find modified Bazel targets of allowed types for DWYU checking.
+
+Finds Bazel targets whose source files (srcs/hdrs) were modified in the
+current diff, and prints their labels. Intended to be used with `bant dwyu`
+in CI to check that modified targets depend on what they use.
+
+Only targets that directly include a modified file are checked, rather than
+all targets in affected packages.
+
+Usage:
+  python3 build_tools/lint/check_dwyu.py --allowed_rules cc_library xla_test
+"""
+
+import argparse
+import logging
+import os
+import re
+import subprocess
+import sys
+from typing import Sequence
+
+from build_tools.lint import diff_parser
+
+_RULE_START = re.compile(r"^(\w+)\s*\(", re.MULTILINE)
+_NAME_ATTR = re.compile(r'name\s*=\s*"([^"]+)"')
+_STRING_LITERAL = re.compile(r'"([^"]+)"')
+
+DEFAULT_ALLOWED_RULES = ("cc_library", "xla_test", "xla_cc_test")
+
+
+def get_diff(base_ref: str) -> str:
+  """Run git diff against base_ref and return stdout."""
+  proc = subprocess.run(
+      ["git", "diff", base_ref, "HEAD"],
+      capture_output=True,
+      check=True,
+      text=True,
+  )
+  return proc.stdout
+
+
+def changed_files_from_diff(diff: str) -> list[str]:
+  """Extract unique file paths from a parsed diff."""
+  hunks = diff_parser.parse_hunks(diff)
+  return sorted(set(h.file for h in hunks))
+
+
+def find_packages(changed_files: list[str]) -> set[str]:
+  """Find Bazel packages containing the changed files."""
+  packages = set()
+  for filepath in changed_files:
+    dirpath = os.path.dirname(filepath)
+    while dirpath:
+      if os.path.isfile(os.path.join(dirpath, "BUILD")) or os.path.isfile(
+          os.path.join(dirpath, "BUILD.bazel")
+      ):
+        packages.add(dirpath)
+        break
+      dirpath = os.path.dirname(dirpath)
+  return packages
+
+
+def _find_rule_end(lines: list[str], start: int) -> int:
+  """Find the line index where the rule's parentheses are balanced."""
+  depth = 0
+  for i in range(start, len(lines)):
+    depth += lines[i].count("(") - lines[i].count(")")
+    if depth <= 0:
+      return i
+  return len(lines) - 1
+
+
+def _extract_string_list(block: str, attr: str) -> list[str]:
+  """Extract string literals from a list-valued attribute in a rule block."""
+  pattern = re.compile(rf"{attr}\s*=\s*\[([^\]]*)\]", re.DOTALL)
+  m = pattern.search(block)
+  if not m:
+    return []
+  return _STRING_LITERAL.findall(m.group(1))
+
+
+def extract_targets(
+    build_content: str, allowed_rules: set[str]
+) -> list[tuple[str, set[str]]]:
+  """Extract target names and their source files from allowed rule types.
+
+  Args:
+    build_content: the contents of a BUILD file.
+    allowed_rules: set of rule types to consider.
+
+  Returns:
+    A list of (target_name, source_files) tuples, where source_files
+    is the set of filenames referenced in srcs and hdrs attributes.
+  """
+  targets = []
+  lines = build_content.split("\n")
+  for i, line in enumerate(lines):
+    m = _RULE_START.match(line.strip())
+    if not m or m.group(1) not in allowed_rules:
+      continue
+    # Find the target name.
+    name = None
+    for j in range(i, min(i + 5, len(lines))):
+      nm = _NAME_ATTR.search(lines[j])
+      if nm:
+        name = nm.group(1)
+        break
+    if name is None:
+      continue
+    # Extract the full rule block to find srcs/hdrs.
+    end = _find_rule_end(lines, i)
+    block = "\n".join(lines[i : end + 1])
+    source_files = set()
+    source_files.update(_extract_string_list(block, "srcs"))
+    source_files.update(_extract_string_list(block, "hdrs"))
+    targets.append((name, source_files))
+  return targets
+
+
+def find_affected_targets(
+    packages: set[str],
+    allowed_rules: set[str],
+    changed_basenames: dict[str, set[str]],
+    build_files_changed: set[str],
+) -> list[str]:
+  """Find targets whose source files were modified.
+
+  Args:
+    packages: set of Bazel package paths to scan.
+    allowed_rules: set of rule types to consider.
+    changed_basenames: mapping from package path to set of changed file
+      basenames within that package.
+    build_files_changed: set of package paths whose BUILD files were modified.
+
+  Returns:
+    list of Bazel target labels that include modified files.
+  """
+  targets = []
+  for package in sorted(packages):
+    for build_name in ("BUILD", "BUILD.bazel"):
+      build_path = os.path.join(package, build_name)
+      if not os.path.isfile(build_path):
+        continue
+      with open(build_path) as f:
+        content = f.read()
+      pkg_changed = changed_basenames.get(package, set())
+      build_changed = package in build_files_changed
+      for target_name, source_files in extract_targets(content, allowed_rules):
+        # Include target if: (1) any of its srcs/hdrs were modified, or
+        # (2) the BUILD file itself was modified (deps may have changed).
+        if build_changed or (source_files & pkg_changed):
+          targets.append(f"//{package}:{target_name}")
+      break
+  return targets
+
+
+def _group_changed_files_by_package(
+    changed_files: list[str], packages: set[str]
+) -> tuple[dict[str, set[str]], set[str]]:
+  """Group changed file basenames by their Bazel package.
+
+  Args:
+    changed_files: list of changed file paths.
+    packages: set of Bazel package paths to consider.
+
+  Returns:
+    A tuple of (changed_basenames, build_files_changed) where
+    changed_basenames maps package path to set of changed file basenames,
+    and build_files_changed is a set of package paths whose BUILD file
+    was modified.
+  """
+  changed_basenames: dict[str, set[str]] = {}
+  build_files_changed: set[str] = set()
+  for filepath in changed_files:
+    basename = os.path.basename(filepath)
+    dirpath = os.path.dirname(filepath)
+    # Walk up to find which package this file belongs to.
+    while dirpath:
+      if dirpath in packages:
+        changed_basenames.setdefault(dirpath, set()).add(basename)
+        if basename in ("BUILD", "BUILD.bazel"):
+          build_files_changed.add(dirpath)
+        break
+      dirpath = os.path.dirname(dirpath)
+  return changed_basenames, build_files_changed
+
+
+def main(argv: Sequence[str]):
+  parser = argparse.ArgumentParser(
+      description="Find modified Bazel targets for DWYU checking."
+  )
+  parser.add_argument(
+      "--allowed_rules",
+      nargs="+",
+      default=list(DEFAULT_ALLOWED_RULES),
+      help="Rule types to include (default: %(default)s)",
+  )
+  parser.add_argument(
+      "--base_ref",
+      default="origin/main",
+      help="Git ref to diff against (default: origin/main)",
+  )
+  args = parser.parse_args(argv[1:])
+  allowed_rules = set(args.allowed_rules)
+
+  diff = get_diff(args.base_ref)
+  changed = changed_files_from_diff(diff)
+  if not changed:
+    logging.info("No files changed.")
+    sys.exit(0)
+
+  packages = find_packages(changed)
+  if not packages:
+    logging.info("No Bazel packages affected.")
+    sys.exit(0)
+
+  changed_basenames, build_files_changed = _group_changed_files_by_package(
+      changed, packages
+  )
+  targets = find_affected_targets(
+      packages, allowed_rules, changed_basenames, build_files_changed
+  )
+  if not targets:
+    logging.info("No targets of allowed types found in affected packages.")
+    sys.exit(0)
+
+  logging.info("Found %d target(s) to check:", len(targets))
+  for t in targets:
+    logging.info("  %s", t)
+
+  # Write targets to stdout, one per line, for consumption by bant.
+  sys.stdout.write("\n".join(targets) + "\n")
+
+
+if __name__ == "__main__":
+  main(sys.argv)

third_party/xla/build_tools/lint/check_dwyu_test.py

@@ -0,0 +1,140 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+import unittest
+
+from build_tools.lint.check_dwyu import extract_targets
+
+ALLOWED_RULES = {"cc_library", "xla_test", "xla_cc_test"}
+
+
+class ExtractTargetsTest(unittest.TestCase):
+
+  def test_cc_library(self):
+    build = """\
+cc_library(
+    name = "foo",
+    srcs = ["foo.cc"],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "foo")
+    self.assertEqual(result[0][1], {"foo.cc"})
+
+  def test_xla_test(self):
+    build = """\
+xla_test(
+    name = "bar_test",
+    srcs = ["bar_test.cc"],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "bar_test")
+    self.assertEqual(result[0][1], {"bar_test.cc"})
+
+  def test_xla_cc_test(self):
+    build = """\
+xla_cc_test(
+    name = "baz_test",
+    srcs = ["baz_test.cc"],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "baz_test")
+    self.assertEqual(result[0][1], {"baz_test.cc"})
+
+  def test_skips_non_allowed_rules(self):
+    build = """\
+cc_binary(
+    name = "my_binary",
+    srcs = ["main.cc"],
+)
+
+py_library(
+    name = "my_lib",
+    srcs = ["lib.py"],
+)
+"""
+    self.assertEqual(extract_targets(build, ALLOWED_RULES), [])
+
+  def test_mixed_rules(self):
+    build = """\
+cc_library(
+    name = "lib",
+    srcs = ["lib.cc"],
+)
+
+cc_binary(
+    name = "bin",
+    srcs = ["main.cc"],
+)
+
+xla_test(
+    name = "lib_test",
+    srcs = ["lib_test.cc"],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 2)
+    self.assertEqual(result[0][0], "lib")
+    self.assertEqual(result[1][0], "lib_test")
+
+  def test_empty_build_file(self):
+    self.assertEqual(extract_targets("", ALLOWED_RULES), [])
+
+  def test_name_on_same_line(self):
+    build = 'cc_library(name = "inline_lib", srcs = ["a.cc"])\n'
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "inline_lib")
+    self.assertEqual(result[0][1], {"a.cc"})
+
+  def test_srcs_and_hdrs(self):
+    build = """\
+cc_library(
+    name = "mylib",
+    srcs = ["mylib.cc"],
+    hdrs = ["mylib.h"],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "mylib")
+    self.assertEqual(result[0][1], {"mylib.cc", "mylib.h"})
+
+  def test_multiple_srcs(self):
+    build = """\
+cc_library(
+    name = "multi",
+    srcs = [
+        "a.cc",
+        "b.cc",
+    ],
+    hdrs = [
+        "a.h",
+        "b.h",
+    ],
+)
+"""
+    result = extract_targets(build, ALLOWED_RULES)
+    self.assertEqual(len(result), 1)
+    self.assertEqual(result[0][0], "multi")
+    self.assertEqual(result[0][1], {"a.cc", "b.cc", "a.h", "b.h"})
+
+
+if __name__ == "__main__":
+  unittest.main()

third_party/xla/build_tools/lint/check_header_guards.py

@@ -0,0 +1,126 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Checks that newly added .h files have standard header guards."""
+
+import argparse
+from collections.abc import Sequence
+import logging
+import re
+import subprocess
+import sys
+
+# Compiled regular expressions for header guard checks.
+_IFNDEF_RE = re.compile(r"^#ifndef\s+([A-Z0-9_]+_H_)\b", re.MULTILINE)
+_DEFINE_RE = re.compile(r"^#define\s+([A-Z0-9_]+_H_)\b", re.MULTILINE)
+
+
+def get_added_header_files() -> list[str]:
+  """Gets the list of newly added or copied header files.
+
+  This function compares the current HEAD with 'origin/main' and returns a list
+  of '.h' files that have been added or copied in the current branch.
+
+  Returns:
+    A list of strings, where each string is the path to a new header file.
+    Returns an empty list if there's an error running the git command.
+  """
+  cmd = [
+      "git",
+      "diff",
+      "--name-only",
+      "--diff-filter=ARC",
+      "origin/main",
+      "HEAD",
+  ]
+  try:
+    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+    return [f for f in result.stdout.splitlines() if f.endswith(".h")]
+  except subprocess.CalledProcessError as e:
+    logging.error("Error running git diff: %s", e)
+    logging.error(e.stderr)
+    # Returns empty list to avoid failing the build over git issues.
+    return []
+
+
+def check_file(path: str) -> tuple[bool, str]:
+  """Checks if a given file has standard C++ header guards.
+
+  Args:
+    path: The path to the header file to check (str).
+
+  Returns:
+    A tuple (is_valid, error_message), where is_valid is True if the file has
+    a valid header guard and False otherwise, and error_message is a string
+    describing the error if is_valid is False.
+  """
+  with open(path, "r", encoding="utf-8") as f:
+    content = f.read()
+  # Check for #ifndef and #define at the top of the file.
+  ifndef_match = _IFNDEF_RE.search(content)
+  define_match = _DEFINE_RE.search(content)
+
+  if not ifndef_match or not define_match:
+    return (
+        False,
+        ("Missing or malformed #ifndef or #define guard (must follow"
+         " PROJECT_PATH_FILE_H_ convention)"),
+    )
+
+  # Ensure they use the same guard name
+  guard_name = ifndef_match.group(1)
+  if define_match.group(1) != guard_name:
+    return (
+        False,
+        f"Mismatched guard name: {guard_name} vs {define_match.group(1)}",
+    )
+
+  # Check for #endif with the guard comment
+  # It usually looks like `#endif  // XLA_PATH_H_`
+  endif_re = re.compile(
+      r"^#endif\s+//\s*" + re.escape(guard_name) + r"\b", re.MULTILINE
+  )
+  if not endif_re.search(content):
+    return False, f"Missing or malformed #endif comment for {guard_name}"
+
+  return True, ""
+
+
+def main(argv: Sequence[str]) -> None:
+  parser = argparse.ArgumentParser(
+      description=(
+          "Checks that newly added .h files have standard header guards."
+      )
+  )
+  parser.parse_args(argv[1:])
+
+  logging.basicConfig(level=logging.INFO)
+
+  headers = get_added_header_files()
+
+  failed = []
+  for h in headers:
+    is_valid, error_msg = check_file(h)
+    if not is_valid:
+      failed.append((h, error_msg))
+
+  if failed:
+    for f, msg in failed:
+      logging.error("Error in %s: %s", f, msg)
+    sys.exit(1)
+  logging.info("Header guard check passed or no new headers found!")
+
+
+if __name__ == "__main__":
+  main(sys.argv)

third_party/xla/build_tools/lint/check_header_guards_test.py

@@ -0,0 +1,131 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+from unittest import mock
+
+from absl.testing import absltest
+
+from build_tools.lint import check_header_guards
+
+
+class CheckHeaderGuardsTest(absltest.TestCase):
+
+  def _create_file(self, content):
+    return self.create_tempfile("test.h", content).full_path
+
+  def test_valid_guard(self):
+    content = """
+#ifndef XLA_TEST_H_
+#define XLA_TEST_H_
+
+void foo();
+
+#endif  // XLA_TEST_H_
+"""
+    path = self._create_file(content)
+    is_valid, _ = check_header_guards.check_file(path)
+    self.assertTrue(is_valid)
+
+  def test_missing_ifndef(self):
+    content = """
+#define XLA_TEST_H_
+
+void foo();
+
+#endif  // XLA_TEST_H_
+"""
+    path = self._create_file(content)
+    is_valid, msg = check_header_guards.check_file(path)
+    self.assertFalse(is_valid)
+    self.assertIn("Missing or malformed #ifndef", msg)
+
+  def test_mismatched_guard_name(self):
+    content = """
+#ifndef XLA_TEST_H_
+#define XLA_WRONG_H_
+
+void foo();
+
+#endif  // XLA_TEST_H_
+"""
+    path = self._create_file(content)
+    is_valid, msg = check_header_guards.check_file(path)
+    self.assertFalse(is_valid)
+    self.assertIn("Mismatched guard name", msg)
+
+  def test_missing_endif_comment(self):
+    content = """
+#ifndef XLA_TEST_H_
+#define XLA_TEST_H_
+
+void foo();
+
+#endif
+"""
+    path = self._create_file(content)
+    is_valid, msg = check_header_guards.check_file(path)
+    self.assertFalse(is_valid)
+    self.assertIn("Missing or malformed #endif", msg)
+
+  def test_guard_not_at_start_of_line(self):
+    content = """
+ // #ifndef XLA_TEST_H_
+ // #define XLA_TEST_H_
+
+void foo();
+
+#endif  // XLA_TEST_H_
+"""
+    path = self._create_file(content)
+    is_valid, msg = check_header_guards.check_file(path)
+    self.assertFalse(is_valid)
+    self.assertIn("Missing or malformed #ifndef", msg)
+
+  def test_endif_not_at_start_of_line(self):
+    content = """
+#ifndef XLA_TEST_H_
+#define XLA_TEST_H_
+
+void foo();
+
+ // #endif  // XLA_TEST_H_
+"""
+    path = self._create_file(content)
+    is_valid, msg = check_header_guards.check_file(path)
+    self.assertFalse(is_valid)
+    self.assertIn("Missing or malformed #endif", msg)
+
+  def test_get_added_header_files(self):
+    with mock.patch.object(check_header_guards.subprocess, "run") as mock_run:
+      mock_run.return_value.stdout = "foo.h\nbar.cc\nbaz.h\n"
+      added_files = check_header_guards.get_added_header_files()
+      self.assertEqual(added_files, ["foo.h", "baz.h"])
+      mock_run.assert_called_once_with(
+          [
+              "git",
+              "diff",
+              "--name-only",
+              "--diff-filter=ARC",
+              "origin/main",
+              "HEAD",
+          ],
+          capture_output=True,
+          text=True,
+          check=True,
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()

third_party/xla/build_tools/lint/diff_parser.py

@@ -0,0 +1,116 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Parses diffs into hunks."""
+import dataclasses
+import re
+import subprocess
+from typing import Generator, Iterable, TypeVar
+
+_T = TypeVar("_T")
+
+
+@dataclasses.dataclass(frozen=True)
+class Hunk:
+  """Represents a hunk of a diff."""
+
+  file: str
+  start: int
+  length: int
+  lines: list[str]
+
+  def added_lines(self) -> Generator[tuple[int, str], None, None]:
+    current_line_no = self.start
+    for line in self.lines:
+      if line.startswith("+"):
+        yield current_line_no, line[1:]  # elide leading '+'
+        current_line_no += 1
+      elif line.startswith("-"):
+        continue
+      else:
+        current_line_no += 1
+
+
+def batch(
+    iterable: Iterable[_T], n: int
+) -> Generator[tuple[_T, ...], None, None]:
+  """Splits an iterable into chunks of size n.
+
+  TODO(ddunleavy): once python 3.12 is available, use itertools.batch.
+
+  Arguments:
+    iterable: the iterable to batch.
+    n: the number of elements in each batch.
+
+  Yields:
+    A tuple of length n of the type that the iterable produces.
+  """
+  iterator = iter(iterable)
+  while True:
+    try:
+      # Unnecessary list here, but a generator won't raise StopIteration,
+      # instead it will raise RuntimeError: "generator raises StopIteration".
+      # I'd rather have a list comprehension in place of a generator expression
+      # than catch RuntimeError and have to inspect the payload to verify it's
+      # the one I want to be catching.
+      yield tuple([next(iterator) for _ in range(n)])
+    except StopIteration:
+      return
+
+
+def parse_hunks(diff: str) -> list[Hunk]:
+  """Parses a diff into hunks.
+
+  Arguments:
+    diff: The raw output of git diff.
+
+  Returns:
+    A list of Hunks.
+  """
+  diff_pattern = (
+      r"diff --git a/.* b/(.*)\n"  # capture file name
+      r"(?:\w+ file mode \d+\n)?"  # maybe 'new file mode 100644' or similar
+      r"index .*\n"
+      r"--- .*\n"
+      r"\+\+\+ .*\n"
+  )
+
+  # capture line number and length from header
+  hunk_header_pattern = r"@@ -\d+,\d+ \+(\d+),(\d+) @@.*\n"
+
+  # ignore initial empty match
+  raw_per_file_hunks = re.split(diff_pattern, diff)[1:]
+
+  parsed_hunks = []
+
+  for file, raw_hunks in batch(raw_per_file_hunks, 2):
+    # ignore initial empty match
+    hunks = re.split(hunk_header_pattern, raw_hunks, re.MULTILINE)[1:]
+    for start, length, body in batch(hunks, 3):
+      lines = body.split("\n")
+      lines = lines if lines[-1] else lines[:-1]  # trim empty line
+      parsed_hunks.append(Hunk(file, int(start), int(length), lines))
+
+  return parsed_hunks
+
+
+def get_git_diff_stdout() -> str:
+  """Run git diff with appropriate arguments and capture stdout as a str."""
+  proc = subprocess.run(
+      ["git", "diff", "origin/main", "HEAD"],
+      capture_output=True,
+      check=True,
+      text=True,
+  )
+  return proc.stdout

third_party/xla/build_tools/lint/diff_parser_test.py

@@ -0,0 +1,129 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from absl.testing import absltest
+
+from build_tools import test_utils
+from build_tools.lint import diff_parser
+
+
+class ParseDiffTest(absltest.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+
+    testdata = test_utils.xla_src_root() / "build_tools" / "lint" / "testdata"
+
+    with (testdata / "bad_cc.diff").open() as f:
+      cls.bad_cc_diff = f.read()
+
+    with (testdata / "important_cc.diff").open() as f:
+      cls.important_cc_diff = f.read()
+
+    with (testdata / "crosstool.diff").open() as f:
+      cls.crosstool_diff = f.read()
+
+  def test_parse_important_cc_diff(self):
+    hunks = diff_parser.parse_hunks(self.important_cc_diff)
+    self.assertLen(hunks, 1)
+    [hunk] = hunks
+    self.assertEqual(hunk.file, "src/important.cc")
+    self.assertEqual(hunk.start, 1)
+    self.assertEqual(hunk.length, 3)
+
+    expected_lines = [
+        "+// Here we care if we find prohibited regexes.",
+        "+std::unique_ptr<int> add(int a, int b) {",
+        "+  return std::make_unique<int>(a + b);",
+        "+}",
+    ]
+
+    self.assertEqual(hunk.lines, expected_lines)
+
+  def test_parse_bad_cc_diff(self):
+    hunks = diff_parser.parse_hunks(self.bad_cc_diff)
+    self.assertLen(hunks, 2)
+
+    bad_cc_hunk, important_cc_hunk = hunks
+
+    # check bad_cc_hunk
+    self.assertEqual(bad_cc_hunk.file, "src/dir/bad.cc")
+    self.assertEqual(bad_cc_hunk.start, 1)
+    self.assertEqual(bad_cc_hunk.length, 7)
+
+    expected_lines = [
+        "+// This code is bad!",
+        "+",
+        "+using Make_Unique = std::make_unique; // OK",
+        "+",
+        "+std::unique_ptr<int> add(int a, int b) {",
+        "+  return Make_Unique<int>(a + b); // OK. Fixed now!",
+        "+}",
+    ]
+
+    self.assertEqual(bad_cc_hunk.lines, expected_lines)
+
+    # check important_cc_hunk
+    self.assertEqual(important_cc_hunk.file, "src/important.cc")
+    self.assertEqual(important_cc_hunk.start, 1)
+    self.assertEqual(important_cc_hunk.length, 5)
+
+    expected_lines = [
+        "+// Here we care if we find prohibited regexes.",
+        "+",
+        "+std::unique_ptr<int> add(int a, int b) {",
+        "+  return std::make_unique<int>(a + b);",
+        "+}",
+    ]
+
+    self.assertEqual(important_cc_hunk.lines, expected_lines)
+
+  def test_parse_crosstool_diff(self):
+    hunks = diff_parser.parse_hunks(self.crosstool_diff)
+
+    self.assertLen(hunks, 3)
+
+    small_hunk, big_hunk, literal_cc_hunk = hunks
+
+    self.assertEqual(
+        small_hunk.file,
+        "third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl",
+    )
+    self.assertEqual(small_hunk.start, 24)
+    self.assertEqual(small_hunk.length, 7)
+
+    self.assertEqual(
+        big_hunk.file, "third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl"
+    )
+    self.assertEqual(big_hunk.start, 300)
+    self.assertEqual(big_hunk.length, 45)
+
+    self.assertEqual(literal_cc_hunk.file, "xla/literal.cc")
+    self.assertEqual(literal_cc_hunk.start, 47)
+    self.assertEqual(literal_cc_hunk.length, 7)
+
+  def test_added_lines(self):
+    hunks = diff_parser.parse_hunks(self.crosstool_diff)
+    small_hunk, big_hunk, literal_cc_hunk = hunks
+
+    line_numbers = lambda hunk: [line_no for line_no, _ in hunk.added_lines()]
+
+    self.assertEqual(line_numbers(small_hunk), [27])
+    self.assertEqual(line_numbers(big_hunk), list(range(303, 342)))
+    self.assertEqual(line_numbers(literal_cc_hunk), [50])
+
+
+if __name__ == "__main__":
+  absltest.main()

third_party/xla/build_tools/lint/generate_compile_commands.py

@@ -0,0 +1,133 @@
+# Copyright 2023 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+r"""Produces a `compile_commands.json` from the output of `bazel aquery`.
+
+This tool requires that a build has been completed for all targets in the
+query (e.g., for the example usage below `bazel build //xla/...`). This is due
+to generated files like proto headers and files generated via tablegen. So if
+LSP or other tools get out of date, it may be necessary to rebuild or regenerate
+`compile_commands.json`, or both.
+
+Example usage:
+  bazel aquery "mnemonic(CppCompile, //xla/...)" --output=jsonproto | \
+      python3 build_tools/lint/generate_compile_commands.py
+"""
+import dataclasses
+import json
+import logging
+import pathlib
+import sys
+from typing import Any
+
+_JSONDict = dict[Any, Any]  # Approximates parsed JSON
+
+_DISALLOWED_ARGS = frozenset(["-fno-canonical-system-headers"])
+_XLA_SRC_ROOT = pathlib.Path(__file__).absolute().parent.parent.parent
+
+
+@dataclasses.dataclass
+class CompileCommand:
+  """Represents a compilation command with options on a specific file."""
+
+  file: str
+  arguments: list[str]
+
+  @classmethod
+  def from_args_list(cls, args_list: list[str]) -> "CompileCommand":
+    """Alternative constructor which uses the args_list from `bazel aquery`.
+
+    This collects arguments and the file being run on from the output of
+    `bazel aquery`. Also filters out arguments which break clang-tidy.
+
+    Arguments:
+      args_list: List of arguments generated by `bazel aquery`
+
+    Returns:
+      The corresponding ClangTidyCommand.
+    """
+    cc_file = None
+    filtered_args = []
+
+    for arg in args_list:
+      if arg in _DISALLOWED_ARGS:
+        continue
+
+      if arg.endswith(".cc"):
+        cc_file = arg
+
+      # Split generated commands, because otherwise they get wrapped
+      # into "command with spaces" when passed to clangd, and clangd
+      # can't parse them correctly.
+      for s in arg.split(" "):
+        filtered_args.append(s)
+
+    return cls(cc_file, filtered_args)
+
+  def to_dumpable_json(self, directory: str) -> _JSONDict:
+    return {
+        "directory": directory,
+        "file": self.file,
+        "arguments": self.arguments,
+    }
+
+
+def extract_compile_commands(
+    parsed_aquery_output: _JSONDict,
+) -> list[CompileCommand]:
+  """Gathers compile commands to run from `bazel aquery` JSON output.
+
+  Arguments:
+    parsed_aquery_output: Parsed JSON representing the output of `bazel aquery
+      --output=jsonproto`.
+
+  Returns:
+    The list of CompileCommands that should be executed.
+  """
+  actions = parsed_aquery_output["actions"]
+
+  commands = []
+  for action in actions:
+    command = CompileCommand.from_args_list(action["arguments"])
+    commands.append(command)
+  return commands
+
+
+def main():
+  # Setup logging
+  logging.basicConfig()
+  logging.getLogger().setLevel(logging.INFO)
+
+  # Setup external symlink if necessary so headers can be found in include paths
+  if not (external := _XLA_SRC_ROOT / "external").exists():
+    logging.info("Symlinking `xla/bazel-xla/external` to `xla/external`")
+    external.symlink_to(_XLA_SRC_ROOT / "bazel-xla" / "external")
+
+  logging.info("Reading `bazel aquery` output from stdin...")
+  parsed_aquery_output = json.loads(sys.stdin.read())
+
+  commands = extract_compile_commands(parsed_aquery_output)
+
+  with (_XLA_SRC_ROOT / "compile_commands.json").open("w") as f:
+    json.dump(
+        [
+            command.to_dumpable_json(directory=str(_XLA_SRC_ROOT))
+            for command in commands
+        ],
+        f,
+    )
+
+
+if __name__ == "__main__":
+  main()

third_party/xla/build_tools/lint/generate_compile_commands_test.py

@@ -0,0 +1,54 @@
+# Copyright 2024 The OpenXLA Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+from absl.testing import absltest
+
+from build_tools.lint import generate_compile_commands
+
+CompileCommand = generate_compile_commands.CompileCommand
+
+
+class CompileCommandsTest(absltest.TestCase):
+
+  def test_command_from_args_list(self):
+    arguments = [
+        "/usr/bin/gcc",
+        "-DTEST_DEFINE",
+        "-fstack-protector",
+        "-c",
+        "xla/compiler.cc",
+        "-o",
+        "bazel-out/k8-opt/bin/xla/_objs/compiler/compiler.pic.o",
+    ]
+
+    command = CompileCommand.from_args_list(arguments)
+
+    self.assertEqual(command.file, "xla/compiler.cc")
+    self.assertEqual(command.arguments, arguments)
+
+  def test_command_from_args_list_with_disallowed_option(self):
+    arguments = [
+        "/usr/bin/gcc",
+        "-DTEST_DEFINE",
+        "-fno-canonical-system-headers",
+        "-c",
+        "xla/compiler.cc",
+        "-o",
+        "bazel-out/k8-opt/bin/xla/_objs/compiler/compiler.pic.o",
+    ]
+
+    command = CompileCommand.from_args_list(arguments)
+
+    self.assertEqual(command.file, "xla/compiler.cc")
+    self.assertEqual(command.arguments, arguments[0:2] + arguments[3:])

third_party/xla/build_tools/lint/tags.py

@@ -0,0 +1,156 @@
+# Copyright 2024 The OpenXLA Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Asserts all tags in XLA are documented.
+
+`bazel query //xla/... --output=build` is read from stdin, and then we check
+all tags are present in the `_TAGS_TO_DOCUMENTATION_MAP`. Ideally we would parse
+using
+https://github.com/bazelbuild/bazel/blob/master/src/main/protobuf/build.proto
+but this is not possible due to XLA's old protobuf version. So we parse by hand
+instead.
+"""
+
+import logging
+import sys
+from typing import Set
+
+_TAGS_TO_DOCUMENTATION_MAP = {
+    # Tags that Bazel recognizes
+    "local": "https://bazel.build/reference/be/common-definitions",
+    "manual": "https://bazel.build/reference/be/common-definitions",
+    "exclusive-if-local": "https://bazel.build/reference/be/common-definitions",
+    "large": "Conventional tag for `test_suites` of large tests",
+    "__PYTHON_RULES_MIGRATION_DO_NOT_USE_WILL_BREAK__": "Internal bazel tag",
+    # Various disable tags (currently recognized by OpenXLA CI)
+    "no_oss": "Test is disabled on OpenXLA CI.",
+    "no_mac": "Disabled on MacOS.",
+    "no_windows": "Disabled on Windows.",
+    "no_mac_arm64": "Disabled on ARM MacOS.",
+    "not_run:arm": (
+        "Not ran on ARM. Currently OpenXLA CI doesn't make a distinction and"
+        " doesn't build things tagged with this either."
+    ),
+    # Various disable tags (currently *unrecognized* by OpenXLA CI)
+    "notap": "Internal tag which disables the test. Not used on OpenXLA CI.",
+    "nosan": "Disabled under all sanitizers. Not used on OpenXLA CI.",
+    "noasan": "Disabled under asan. Not used on OpenXLA CI.",
+    "nomsan": "Disabled under msan. Not used on OpenXLA CI.",
+    "notsan": "Disabled under tsan. Not used on OpenXLA CI",
+    "nobuilder": "Not built internally.",
+    "nofixdeps": "Internal tag. Disables build_cleaner.",
+    "nozapfhahn": "Internal tag. Disables gathering coverage",
+    "optonly": "Should only be tested with -c opt",
+    "nodebug": "Should not be tested in debug builds.",
+    "config-cuda-only": (
+        "Meaningless in OSS as all GPU tests are built with `--config=cuda`"
+    ),
+    # GPU tags
+    "requires-gpu": (
+        "Test requires GPU to execute. Fallback if neither CUDA nor ROCm is"
+        " specified."
+    ),
+    "requires-gpu-amd": "Test requires AMD GPU to execute",
+    "requires-gpu-intel": "Test requires Intel GPU to execute",
+    "requires-gpu-nvidia": "Test requires NVIDIA GPU to execute",
+    "requires-gpu-nvidia:2": "Test needs 2 NVIDIA GPUs to run",
+    "requires-gpu-sm60-only": "Requires exactly sm60.",
+    "requires-gpu-sm70-only": "Requires exactly sm70.",
+    "requires-gpu-sm80-only": "Requires exactly sm80.",
+    "requires-gpu-sm90-only": "Requires exactly sm90.",
+    "requires-gpu-sm100-only": "Requires exactly sm100.",
+    "requires-gpu-sm103-only": "Requires exactly sm103.",
+    "requires-gpu-sm120-only": "Requires exactly sm120.",
+    "full": "Test requires a full GPU, not a partitioned one. No effect in"
+            " OSS.",
+    "gpu": "Catch-all tag for targets that should be built/tested on GPU CI",
+    "cpu": "Catch-all tag for targets that should be built/tested on CPU CI.",
+    "cuda-only": "Targets that require the CUDA backend to be enabled.",
+    "rocm-only": "Targets that require the ROCm backend to be enabled.",
+    "oneapi-only": "Targets that require the oneAPI backend to be enabled.",
+    "no-oneapi": "Targets that are not configured for the oneAPI backend.",
+    # Below tags are generated by `xla_test`.
+    "broken": "Test will be marked with other tags to disable in `xla_test`.",
+    "xla_interpreter": "Uses interpreter backend.",
+    "xla_cpu": "Uses CPU backend.",
+    "xla_amdgpu_any": "Uses ROCm backend.",
+    "xla_nvgpu_any": "Uses NVIDIA GPU backend.",
+    "xla_intelgpu_any": "Uses Intel GPU backend.",
+    # Below tags are emitted alongside `requires-gpu-x` tags, which is what the
+    # CI actually follows. So we may not execute on an A100, and instead use an
+    # L4. These tags are taken literally internally.
+    "xla_p100": "Runs on a p100.",
+    "xla_v100": "Runs on a v100.",
+    "xla_a100": "Runs on an a100.",
+    "xla_h100": "Runs on an h100.",
+    "xla_b200": "Runs on a b200.",
+    "xla_gb200": "Runs on a gb200.",
+    "xla_gb300": "Runs on a gb300.",
+    "xla_rtx6000pro": "Runs on an rtx6000pro.",
+    "xla_device_p100": "Runs on a p100.",
+    "xla_device_v100": "Runs on a v100.",
+    "xla_device_a100": "Runs on an a100.",
+    "xla_device_h100": "Runs on an h100.",
+    "xla_device_b200": "Runs on a b200.",
+    "xla_device_gb200": "Runs on a gb200.",
+    "xla_device_gb300": "Runs on a gb300.",
+    "xla_device_rtx6000pro": "Runs on an rtx6000pro.",
+    # Below tags are consumed by `xla_test`.
+    "test_migrated_to_hlo_runner_pjrt": (
+        "Adds the appropriate `xla/tests:pjrt_$BACKEND_client_registry` to the"
+        " annotated `xla_test` target. Adding this tag does not synthesize"
+        " additional targets."
+    ),
+    "pjrt_migration_candidate": (
+        "Tags the target as a PJRT migration candidate."
+    ),
+    "multi_gpu": "Used by `xla_test` to signal that multiple GPUs are needed.",
+}
+
+
+def get_tags_from_line(line: str) -> Set[str]:
+  if line.strip().startswith("tags = "):
+    tags_list = line[10:-3]  # "tag1", "tag2"
+    if tags_list.strip():
+      # Remove extraneous quotes, tags like `-broken` used in test_suites,
+      # and split on ", "
+      return {tag.strip('-"') for tag in tags_list.split(", ")}
+
+  return set()
+
+
+def main():
+  logging.basicConfig()
+  logging.getLogger().setLevel(logging.INFO)
+
+  tags = set.union(*(get_tags_from_line(line) for line in sys.stdin))
+
+  logging.info(str(tags))
+
+  if undocumented_tags := tags - _TAGS_TO_DOCUMENTATION_MAP.keys():
+    raise ValueError(
+        f"Tag(s) {undocumented_tags} are undocumented! Please document them in"
+        " `build_tools/lint/tags.py`."
+    )
+
+  if unused_but_documented_tags := _TAGS_TO_DOCUMENTATION_MAP.keys() - tags:
+    logging.info(
+        "The following tags are documented but unused: %s. Do we expect they'll"
+        " be used in the future?",
+        str(unused_but_documented_tags),
+    )
+
+
+if __name__ == "__main__":
+  main()