chore: import upstream snapshot with attribution

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()

third_party/xla/build_tools/lint/testdata/bad_cc.diff

@@ -0,0 +1,24 @@
+diff --git a/src/dir/bad.cc b/src/dir/bad.cc
+new file mode 100644
+index 0000000..2508a76
+--- /dev/null
++++ b/src/dir/bad.cc
+@@ -0,0 +1,7 @@
++// 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!
++}
+diff --git a/src/important.cc b/src/important.cc
+new file mode 100644
+index 0000000..dc06702
+--- /dev/null
++++ b/src/important.cc
+@@ -0,0 +1,5 @@
++// Here we care if we find prohibited regexes.
++
++std::unique_ptr<int> add(int a, int b) {
++  return std::make_unique<int>(a + b);
++}

third_party/xla/build_tools/lint/testdata/crosstool.diff

@@ -0,0 +1,71 @@
+diff --git a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
+index ffa305c77..8ad22502a 100644
+--- a/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
++++ b/third_party/gpus/crosstool/cc_toolchain_config.bzl.tpl
+@@ -24,7 +24,7 @@ def all_assembly_actions():
+     ]
+ 
+ def all_compile_actions():
+-    return [
++    retur [
+         ACTION_NAMES.assemble,
+         ACTION_NAMES.c_compile,
+         ACTION_NAMES.cpp_compile,
+@@ -300,6 +300,45 @@ def _features(cpu, compiler, ctx):
+     if cpu in ["local", "darwin"]:
+         return [
+             feature(name = "no_legacy_features"),
++            feature(
++             name = "use_module_maps",
++             requires = [feature_set(features = ["module_maps"])],
++             flag_sets = [
++                 flag_set(
++                     actions = all_compile_actions(),
++                     flag_groups = [
++                         flag_group(
++                             flags = [
++                                 "-fmodule-name=%{module_name}",
++                                 "-fmodule-map-file=%{module_map_file}",
++                             ],
++                         ),
++                     ],
++                 ),
++             ],
++           ),
++           feature(name = "module_maps", enabled = True),
++           feature(
++               name = "layering_check",
++               implies = ["use_module_maps"],
++               flag_sets = [
++                   flag_set(
++                       actions = all_compile_actions(),
++                       flag_groups = [
++                           flag_group(flags = [
++                               "-fmodules-strict-decluse"
++                               "-Wprivate-header",
++                           ]),
++                           flag_group(
++                               iterate_over = "dependent_module_map_files",
++                               flags = [
++                                   "-fmodule-map-file=%{dependent_module_map_files}",
++                               ],
++                           ),
++                       ],
++                   ),
++               ],
++           ),
+             feature(
+                 name = "all_compile_flags",
+                 enabled = True,
+diff --git a/xla/literal.cc b/xla/literal.cc
+index 2cef23917..d362f8af9 100644
+--- a/xla/literal.cc
++++ b/xla/literal.cc
+@@ -47,6 +47,7 @@ limitations under the License.
+ #include "tsl/platform/errors.h"
+ #include "tsl/platform/float8.h"
+ #include "tsl/platform/logging.h"
++#include "tsl/platform/platform.h"
+ #include "tsl/platform/mem.h"
+ #include "tsl/util/byte_swap_array.h"
+ 

third_party/xla/build_tools/lint/testdata/important_cc.diff

@@ -0,0 +1,10 @@
+diff --git a/src/important.cc b/src/important.cc
+new file mode 100644
+index 0000000..9c68461
+--- /dev/null
++++ b/src/important.cc
+@@ -0,0 +1,3 @@
++// Here we care if we find prohibited regexes.
++std::unique_ptr<int> add(int a, int b) {
++  return std::make_unique<int>(a + b);
++}

third_party/xla/build_tools/pjrt_wheels/BUILD.bazel

@@ -0,0 +1,172 @@
+load("@cuda_cudart//:version.bzl", cuda_major_version = "VERSION")
+load("@nightly_timestamp//:timestamp.bzl", "XLA_NIGHTLY_TIMESTAMP")
+load("@rc_number//:rc_number.bzl", "XLA_RC_NUMBER")
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_test.bzl", "cc_test")
+load("@rules_python//python:packaging.bzl", "py_wheel")
+
+# This ensures we can only build plugins for selected CUDA versions.
+cuda_label = "cuda" + cuda_major_version if cuda_major_version else "null"
+
+# If we're building a nightly, append .devYYYYMMDD to the version
+# If we're not, the timestamp is an empty string
+# If we're building a release candidate, append rc# to the version
+# If we're not, the rc number is an empty string
+# Note that PEP 440 supports both of these at the same time, but
+# our CI will never set both.
+wheel_version = "0.0.0" + XLA_RC_NUMBER + XLA_NIGHTLY_TIMESTAMP
+
+wheel_platform = select({
+    "//conditions:default": "manylinux_2_27_x86_64",
+    "//xla/tsl:linux_aarch64": "manylinux_2_27_aarch64",
+})
+
+# Aliases and filegroups don't move files within bazel-out, so we need to use genrules to place them
+# in the correct directory structure for the wheel.
+genrule(
+    name = "init_file_" + cuda_label,
+    srcs = ["__init__.py"],
+    outs = ["xla_plugins/xla_" + cuda_label + "_pjrt/__init__.py"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "init_file_cpu",
+    srcs = ["__init__.py"],
+    outs = ["xla_plugins/xla_cpu_pjrt/__init__.py"],
+    cmd = "cp $< $@",
+)
+
+# The wheels have the same C API header, but we need to give each their own copy
+genrule(
+    name = "pjrt_c_api_hdr_cuda",
+    srcs = ["//xla/pjrt/c:pjrt_c_api.h"],
+    outs = ["xla_plugins/xla_" + cuda_label + "_pjrt/include/pjrt_c_api.h"],
+    cmd = "cp $< $@",
+)
+
+genrule(
+    name = "pjrt_c_api_hdr_cpu",
+    srcs = ["//xla/pjrt/c:pjrt_c_api.h"],
+    outs = ["xla_plugins/xla_cpu_pjrt/include/pjrt_c_api.h"],
+    cmd = "cp $< $@",
+)
+
+# GPU-specific files
+cc_binary(
+    name = "xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    linkopts = [
+        "-Wl,--version-script,$(location :pjrt_symbols.lds)",
+        "-Wl,--no-undefined",
+    ],
+    linkshared = True,
+    deps = [
+        ":pjrt_symbols.lds",
+        "//xla/pjrt/c:pjrt_c_api_gpu",
+    ],
+)
+
+py_wheel(
+    name = "xla_" + cuda_label + "_pjrt",
+    author = "The OpenXLA Authors",
+    classifiers = ["Development Status :: 1 - Planning"],
+    distribution = "xla_" + cuda_label + "_pjrt",
+    entry_points = {
+        "xla_plugins": [
+            "xla_" + cuda_label + "_pjrt = xla_plugins.xla_" + cuda_label + "_pjrt\n",
+        ],
+    },
+    platform = wheel_platform,
+    python_tag = "py3",
+    strip_path_prefixes = ["build_tools/pjrt_wheels"],
+    summary = "XLA PJRT Plugin",
+    version = wheel_version,
+    deps = [
+        ":init_file_" + cuda_label,
+        ":pjrt_c_api_hdr_cuda",
+        ":xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    ],
+)
+
+alias(
+    name = "xla_cuda_pjrt.dist",
+    actual = ":xla_" + cuda_label + "_pjrt.dist",
+)
+
+# CPU-specific files
+cc_binary(
+    name = "xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    linkopts = [
+        "-Wl,--version-script,$(location :pjrt_symbols.lds)",
+        "-Wl,--no-undefined",
+    ],
+    linkshared = True,
+    deps = [
+        ":pjrt_symbols.lds",
+        "//xla/pjrt/c:pjrt_c_api_cpu",
+    ],
+)
+
+py_wheel(
+    name = "xla_cpu_pjrt",
+    author = "The OpenXLA Authors",
+    classifiers = ["Development Status :: 1 - Planning"],
+    distribution = "xla_cpu_pjrt",
+    entry_points = {
+        "xla_plugins": [
+            "xla_cpu_pjrt = xla_plugins.xla_cpu_pjrt\n",
+        ],
+    },
+    platform = wheel_platform,
+    python_tag = "py3",
+    strip_path_prefixes = ["build_tools/pjrt_wheels"],
+    summary = "XLA PJRT Plugin",
+    version = wheel_version,
+    deps = [
+        ":init_file_cpu",
+        ":pjrt_c_api_hdr_cpu",
+        ":xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    ],
+)
+
+# Tests
+cc_test(
+    name = "cpu_smoke_test",
+    srcs = ["smoke_test.cc"],
+    data = [":xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so"],
+    env = {
+        "PJRT_PLUGIN_PATH": "build_tools/pjrt_wheels/xla_plugins/xla_cpu_pjrt/xla_cpu_pjrt.so",
+    },
+    linkopts = ["-ldl"],
+    tags = ["no_windows"],
+    deps = ["//xla/pjrt/c:pjrt_c_api_hdrs"],
+)
+
+cc_test(
+    name = cuda_label + "_smoke_test",
+    srcs = ["smoke_test.cc"],
+    data = [":xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so"],
+    env = {
+        "PJRT_PLUGIN_PATH": "build_tools/pjrt_wheels/xla_plugins/xla_" + cuda_label + "_pjrt/xla_gpu_pjrt.so",
+    },
+    linkopts = ["-ldl"],
+    tags = ["no_windows"],
+    deps = ["//xla/pjrt/c:pjrt_c_api_hdrs"],
+)
+
+# The CPU and CUDA test suites are run in CI's build script
+test_suite(
+    name = "cpu_test_suite",
+    tags = ["no_windows"],
+    tests = [
+        ":cpu_smoke_test",
+    ],
+)
+
+test_suite(
+    name = "cuda_test_suite",
+    tags = ["no_windows"],
+    tests = [
+        ":" + cuda_label + "_smoke_test",
+    ],
+)

third_party/xla/build_tools/pjrt_wheels/__init__.py

@@ -0,0 +1 @@
+# This is currently just here to mark the directory as a module.

third_party/xla/build_tools/pjrt_wheels/nightly.bzl

@@ -0,0 +1,35 @@
+"""If we're building a nightly, we use this to pass a timestamp for the wheel version."""
+
+def _nightly_timestamp_impl(rctx):
+    timestamp_val = rctx.getenv("XLA_NIGHTLY_TIMESTAMP", "")  # Default to ""
+
+    # Smuggle the value via a new .bzl file
+    if timestamp_val:
+        rctx.file(
+            "timestamp.bzl",
+            content = 'XLA_NIGHTLY_TIMESTAMP = ".dev{}"'.format(timestamp_val),
+        )
+    else:
+        rctx.file(
+            "timestamp.bzl",
+            content = 'XLA_NIGHTLY_TIMESTAMP = ""',
+        )
+
+    # Create a BUILD file to make timestamp.bzl addressable
+    rctx.file("BUILD.bazel", content = "")
+
+nightly_timestamp_repo = repository_rule(
+    implementation = _nightly_timestamp_impl,
+    environ = ["XLA_NIGHTLY_TIMESTAMP"],
+)
+
+# bzlmod implementation
+def _nightly_timestamp_ext_impl(mctx):  # @unused
+    nightly_timestamp_repo(
+        name = "nightly_timestamp",
+    )
+
+nightly_timestamp_repo_bzlmod = module_extension(
+    implementation = _nightly_timestamp_ext_impl,
+    environ = ["XLA_NIGHTLY_TIMESTAMP"],
+)

third_party/xla/build_tools/pjrt_wheels/pjrt_symbols.lds

@@ -0,0 +1,9 @@
+VERS_1.0 {
+  global:
+    extern "C" {
+      GetPjrtApi;
+    };
+
+  local:
+    *;
+};

third_party/xla/build_tools/pjrt_wheels/release_candidate.bzl

@@ -0,0 +1,35 @@
+"""If we're building a release candidate, we use this to pass a rc number for the wheel version."""
+
+def _rc_number_impl(rctx):
+    rc_number = rctx.getenv("XLA_RC_NUMBER", "")  # Default to ""
+
+    # Smuggle the value via a new .bzl file
+    if rc_number:
+        rctx.file(
+            "rc_number.bzl",
+            content = 'XLA_RC_NUMBER = "{}"'.format(rc_number),
+        )
+    else:
+        rctx.file(
+            "rc_number.bzl",
+            content = 'XLA_RC_NUMBER = ""',
+        )
+
+    # Create a BUILD file to make timestamp.bzl addressable
+    rctx.file("BUILD.bazel", content = "")
+
+rc_number_repo = repository_rule(
+    implementation = _rc_number_impl,
+    environ = ["XLA_RC_NUMBER"],
+)
+
+# bzlmod implementation
+def _rc_number_ext_impl(mctx):  # @unused
+    rc_number_repo(
+        name = "rc_number",
+    )
+
+rc_number_repo_bzlmod = module_extension(
+    implementation = _rc_number_ext_impl,
+    environ = ["XLA_RC_NUMBER"],
+)

third_party/xla/build_tools/pjrt_wheels/smoke_test.cc

@@ -0,0 +1,57 @@
+/* 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.
+==============================================================================*/
+
+#include <dlfcn.h>
+
+#include <iostream>
+
+#include "xla/pjrt/c/pjrt_c_api.h"
+
+typedef const PJRT_Api* (*GetPjrtApi_Func)();
+
+int main() {
+  // 1. Open the shared object
+  const char* so_path = std::getenv("PJRT_PLUGIN_PATH");
+  std::cout << "so_path: " << so_path << std::endl;
+  void* handle = dlopen(so_path, RTLD_LAZY);
+  if (!handle) {
+    std::cerr << "Error: Could not open shared object." << std::endl;
+    std::cerr << "Reason: " << dlerror() << std::endl;
+    return 1;
+  }
+
+  // 2. Load the symbol (the function)
+  GetPjrtApi_Func get_pjrt_api = (GetPjrtApi_Func)dlsym(handle, "GetPjrtApi");
+  const char* dlsym_error = dlerror();
+  if (dlsym_error) {
+    std::cerr << "Error: Could not find symbol 'GetPjrtApi'." << std::endl;
+    std::cerr << "Reason: " << dlsym_error << std::endl;
+    dlclose(handle);
+    return 1;
+  }
+
+  // 3. Call the function
+  std::cout << "Successfully loaded symbol. Calling GetPjrtApi()..."
+            << std::endl;
+  const PJRT_Api* api = get_pjrt_api();
+  if (api) {
+    std::cout << "Success! Received PjrtApi struct pointer." << std::endl;
+  } else {
+    std::cerr << "Error: GetPjrtApi() returned a null pointer." << std::endl;
+    return 1;
+  }
+
+  return 0;
+}

third_party/xla/build_tools/rocm/BUILD

@@ -0,0 +1,87 @@
+# 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.
+# ============================================================================
+
+load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
+load("@rules_shell//shell:sh_binary.bzl", "sh_binary")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+string_flag(
+    name = "sanitizer",
+    build_setting_default = "none",
+    values = [
+        "none",
+        "asan",
+        "tsan",
+    ],
+)
+
+config_setting(
+    name = "asan",
+    flag_values = {"//build_tools/rocm:sanitizer": "asan"},
+)
+
+config_setting(
+    name = "tsan",
+    flag_values = {"//build_tools/rocm:sanitizer": "tsan"},
+)
+
+filegroup(
+    name = "sanitizer_ignore_lists",
+    srcs = select({
+        ":asan": [
+            "asan_ignore_list.txt",
+            "lsan_ignore_list.txt",
+        ],
+        ":tsan": ["tsan_ignore_list.txt"],
+        "//conditions:default": [],
+    }),
+    visibility = ["//visibility:public"],
+)
+
+genrule(
+    name = "san_wrapper_script",
+    srcs = [":sanitizer_ignore_lists"],
+    outs = ["san_wrapper.sh"],
+    cmd = """
+      echo '#!/bin/bash' > $@
+      echo 'exec "$$@"' >> $@
+      chmod +x $@
+    """,
+)
+
+# this wrapper ensures the test target
+# take into account any changes in the ignore list files
+sh_binary(
+    name = "sanitizer_wrapper",
+    srcs = [":san_wrapper_script"],
+    data = [":sanitizer_ignore_lists"],
+    tags = ["manual"],
+    visibility = ["//visibility:public"],
+)
+
+sh_binary(
+    name = "parallel_gpu_execute",
+    srcs = ["parallel_gpu_execute.sh"],
+    data = [
+        ":sanitizer_ignore_lists",
+        "@local_config_rocm//rocm:rocminfo",
+    ],
+    tags = ["manual"],
+    visibility = ["//visibility:public"],
+)

third_party/xla/build_tools/rocm/asan_ignore_list.txt

@@ -0,0 +1,2 @@
+interceptor_via_lib:libhsa-runtime64.so
+interceptor_via_lib:libamdhip64.so

third_party/xla/build_tools/rocm/lsan_ignore_list.txt

@@ -0,0 +1,10 @@
+leak:libhsa-runtime64.so
+leak:libstdc++.so
+leak:libamdhip64.so
+
+# RCCL leaks in librccl (not XLA), seen by single-rank comm tests. Shipped
+# librccl is stripped (frames lack function names), so we suppress by lib name.
+#  - ncclGroupJob in ncclGroupEndInternal: fixed in ROCm 7.2.4
+#    (https://github.com/ROCm/rocm-systems/pull/3198).
+#  - lsaRankList in ncclDevrInitOnce: still present on develop (2.30.4); fix pending.
+leak:librccl.so

third_party/xla/build_tools/rocm/parallel_gpu_execute.sh

@@ -0,0 +1,92 @@
+#!/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.
+
+ROCMINFO=$(find -L "${TEST_SRCDIR:-.}" -name "rocminfo" -path "*/bin/rocminfo" | head -n 1)
+TF_GPU_COUNT=$($ROCMINFO | grep "Name: *gfx*" | wc -l)
+TF_TESTS_PER_GPU=${TF_TESTS_PER_GPU:-8}
+
+# There are certain tests in xla that do not require any gpu in order to be executed
+# here we allow executing these tests on a machine without gpu support.
+# if there are no GPUs on that system e.g rbe default pool then execute the test without lock
+if [[ $TF_GPU_COUNT == 0 ]];then
+    echo "Execute with no GPU support"
+    exec "$@"
+fi
+
+# 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/rocm/rocm_tag_filters.sh

@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+# Copyright 2024 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.
+#
+# ==============================================================================
+
+TAG_FILTERS=(
+    -no_gpu
+    -requires-gpu-intel
+    -requires-gpu-nvidia
+    -requires-gpu-cuda
+    -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
+)
+
+echo $(
+    IFS=,
+    echo "${TAG_FILTERS[*]}"
+)

third_party/xla/build_tools/rocm/rocm_xla.bazelrc

@@ -0,0 +1,108 @@
+# Test-related settings.
+
+build:rocm_dev --remote_upload_local_results=false
+build:rocm_dev --remote_cache="grpcs://wardite.cluster.engflow.com"
+
+build:rocm_rbe --remote_cache="grpcs://wardite.cluster.engflow.com"
+build:rocm_rbe --repo_env=REMOTE_GPU_TESTING=1
+build:rocm_rbe --bes_backend="grpcs://wardite.cluster.engflow.com"
+build:rocm_rbe --bes_results_url="https://wardite.cluster.engflow.com/invocation/"
+build:rocm_rbe --host_platform="@local_config_rocm//rocm:linux_x64"
+build:rocm_rbe --extra_execution_platforms="@local_config_rocm//rocm:linux_x64"
+build:rocm_rbe --platforms="@local_config_rocm//rocm:linux_x64"
+build:rocm_rbe --bes_timeout=600s
+build:rocm_rbe --tls_client_certificate="/data/ci-cert.crt"
+build:rocm_rbe --tls_client_key="/data/ci-cert.key"
+build:rocm_rbe --spawn_strategy=remote,local
+build:rocm_rbe --grpc_keepalive_time=30s
+build:rocm_rbe --remote_cache_compression
+
+test:rocm_rbe --jobs=30
+test:rocm_rbe --remote_executor=grpcs://wardite.cluster.engflow.com
+test:rocm_rbe --remote_timeout=3600
+test:rocm_rbe --strategy=TestRunner=remote,local
+
+# Dynamic execution config for ROCm RBE - builds locally, tests use hybrid mode
+build:rocm_rbe_dynamic --config=rocm_rbe
+build:rocm_rbe_dynamic --spawn_strategy=local
+test:rocm_rbe_dynamic --experimental_spawn_scheduler
+test:rocm_rbe_dynamic --strategy=TestRunner=dynamic
+test:rocm_rbe_dynamic --dynamic_mode=default
+test:rocm_rbe_dynamic --dynamic_local_strategy=worker,standalone,local
+test:rocm_rbe_dynamic --dynamic_remote_strategy=remote
+test:rocm_rbe_dynamic --experimental_local_execution_delay=1000
+test:rocm_rbe_dynamic --local_resources=cpu=HOST_CPUS*0.5
+
+build:tsan --strip=never
+build:tsan --copt -fsanitize=thread
+build:tsan --copt -g
+build:tsan --copt -fno-omit-frame-pointer
+build:tsan --linkopt -fsanitize=thread
+build:tsan --linkopt -g
+build:tsan --//build_tools/rocm:sanitizer=tsan
+build:tsan --test_env=TSAN_OPTIONS=suppressions=build_tools/rocm/tsan_ignore_list.txt::history_size=7:ignore_noninstrumented_modules=1
+build:tsan --run_under=//build_tools/rocm:sanitizer_wrapper
+
+build:asan --test_env=ASAN_OPTIONS=suppressions=build_tools/rocm/asan_ignore_list.txt:use_sigaltstack=0
+build:asan --test_env=LSAN_OPTIONS=suppressions=build_tools/rocm/lsan_ignore_list.txt:use_sigaltstack=0
+build:asan --//build_tools/rocm:sanitizer=asan
+build:asan --run_under=//build_tools/rocm:sanitizer_wrapper
+
+build:ci_single_gpu --run_under=//build_tools/rocm:parallel_gpu_execute
+build:ci_single_gpu --flaky_test_attempts=3
+
+build:ci_multi_gpu --action_env=XLA_FLAGS="--xla_gpu_force_compilation_parallelism=16"
+build:ci_multi_gpu --action_env=NCCL_MAX_NCHANNELS=1
+build:ci_multi_gpu --run_under=//build_tools/rocm:sanitizer_wrapper
+build:ci_multi_gpu --test_sharding_strategy=disabled
+build:ci_multi_gpu --flaky_test_attempts=3
+build:ci_multi_gpu --experimental_guard_against_concurrent_changes
+build:ci_multi_gpu --test_env=HIP_VISIBLE_DEVICES=0,1,2,3
+
+build:ci_rocm_cpu --run_under=//build_tools/rocm:sanitizer_wrapper
+build:ci_rocm_cpu --local_test_jobs=200
+build:ci_rocm_cpu --strategy=TestRunner=local
+
+test:xla_sgpu -- \
+//xla/... \
+-//xla/tests:iota_test_amdgpu_any \
+-//xla/backends/gpu/collectives:gpu_clique_key_test \
+-//xla/service:collective_ops_utils_test \
+-//xla/service:collective_pipeliner_test \
+-//xla/service:collective_permute_cycle_test \
+-//xla/service:batched_gather_scatter_normalizer_test \
+-//xla/service:all_reduce_simplifier_test \
+-//xla/service:all_gather_simplifier_test \
+-//xla/service:reduce_scatter_decomposer_test \
+-//xla/service:reduce_scatter_reassociate_test \
+-//xla/service:reduce_scatter_combiner_test \
+-//xla/service:scatter_simplifier_test \
+-//xla/service:sharding_propagation_test \
+-//xla/service:sharding_remover_test \
+-//xla/service:p2p_schedule_preparation_test \
+-//xla/pjrt/distributed:topology_util_test \
+-//xla/pjrt/distributed:client_server_test
+
+test:xla_mgpu -- \
+//xla/tests:collective_ops_e2e_test \
+//xla/tests:collective_ops_test \
+//xla/tests:collective_pipeline_parallelism_test \
+//xla/tests:replicated_io_feed_test \
+//xla/backends/gpu/collectives:gpu_clique_key_test \
+//xla/backends/gpu/runtime:all_reduce_test \
+//xla/service:collective_ops_utils_test \
+//xla/service:collective_pipeliner_test \
+//xla/service:collective_permute_cycle_test \
+//xla/service:batched_gather_scatter_normalizer_test \
+//xla/service:all_reduce_simplifier_test \
+//xla/service:all_gather_simplifier_test \
+//xla/service:reduce_scatter_decomposer_test \
+//xla/service:reduce_scatter_reassociate_test \
+//xla/service:reduce_scatter_combiner_test \
+//xla/service:scatter_simplifier_test \
+//xla/service:sharding_propagation_test \
+//xla/service:sharding_remover_test \
+//xla/service:p2p_schedule_preparation_test \
+//xla/tools/multihost_hlo_runner:functional_hlo_runner_test \
+//xla/pjrt/distributed:topology_util_test \
+//xla/pjrt/distributed:client_server_test 

third_party/xla/build_tools/rocm/rocm_xla_ci.bazelrc

@@ -0,0 +1,3 @@
+# CI related imports
+try-import /usertools/rocm.bazelrc
+try-import %workspace%/build_tools/rocm/rocm_xla.bazelrc

third_party/xla/build_tools/rocm/run_xla.sh

@@ -0,0 +1,68 @@
+#!/usr/bin/env bash
+# Copyright 2024 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 script runs XLA unit tests on ROCm platform by selecting tests that are
+# tagged with requires-gpu-amd
+
+set -e
+set -x
+
+N_BUILD_JOBS=$(grep -c ^processor /proc/cpuinfo)
+# If rocm-smi exists locally (it should) use it to find
+# out how many GPUs we have to test with.
+rocm-smi -i
+STATUS=$?
+if [ $STATUS -ne 0 ]; then TF_GPU_COUNT=1; else
+   TF_GPU_COUNT=$(rocm-smi -i|grep 'Device ID' |grep 'GPU' |wc -l)
+fi
+TF_TESTS_PER_GPU=1
+N_TEST_JOBS=$(expr ${TF_GPU_COUNT} \* ${TF_TESTS_PER_GPU})
+amdgpuname=(`rocminfo | grep gfx | head -n 1`)
+AMD_GPU_GFX_ID=${amdgpuname[1]}
+echo ""
+echo "Bazel will use ${N_BUILD_JOBS} concurrent build job(s) and ${N_TEST_JOBS} concurrent test job(s) for gpu ${AMD_GPU_GFX_ID}."
+echo ""
+
+export PYTHON_BIN_PATH=`which python3`
+export TF_NEED_ROCM=1
+export ROCM_PATH=/opt/rocm
+
+SCRIPT_DIR=$(realpath $(dirname $0))
+TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh),gpu,-multi_gpu,-multi_gpu_h100,requires-gpu-amd,,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial
+
+if [ ! -d /tf/pkg ]; then
+	mkdir -p /tf/pkg
+fi
+
+bazel --bazelrc=build_tools/rocm/rocm_xla.bazelrc test \
+    --config=rocm_ci \
+    --config=xla_sgpu \
+    --build_tag_filters=$TAG_FILTERS \
+    --test_tag_filters=$TAG_FILTERS \
+    --profile=/tf/pkg/profile.json.gz \
+    --test_timeout=920,2400,7200,9600 \
+    --test_sharding_strategy=disabled \
+    --test_output=errors \
+    --flaky_test_attempts=3 \
+    --keep_going \
+    --local_test_jobs=${N_TEST_JOBS} \
+    --test_env=TF_TESTS_PER_GPU=$TF_TESTS_PER_GPU \
+    --test_env=TF_GPU_COUNT=$TF_GPU_COUNT \
+    --action_env=TF_ROCM_AMDGPU_TARGETS=${AMD_GPU_GFX_ID} \
+    --action_env=XLA_FLAGS="--xla_gpu_force_compilation_parallelism=16" \
+    --repo_env="ROCM_PATH=$ROCM_PATH" \
+    --run_under=//build_tools/ci:parallel_gpu_execute

third_party/xla/build_tools/rocm/run_xla_ci_build.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+# Copyright 2025 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.
+#
+# ==============================================================================
+
+set -e
+set -x
+
+SCRIPT_DIR=$(realpath $(dirname $0))
+TAG_FILTERS=$($SCRIPT_DIR/rocm_tag_filters.sh)
+
+mkdir -p /tf/pkg
+
+for arg in "$@"; do
+    if [[ "$arg" == "--config=ci_multi_gpu" ]]; then
+        TAG_FILTERS="${TAG_FILTERS},multi_gpu"
+    fi
+    if [[ "$arg" == "--config=ci_single_gpu" ]]; then
+        TAG_FILTERS="${TAG_FILTERS},requires-gpu-rocm,requires-gpu-amd,-multi_gpu"
+    fi
+    if [[ "$arg" == "--config=ci_rocm_cpu" ]]; then
+        TAG_FILTERS="${TAG_FILTERS},gpu,-requires-gpu-rocm,-requires-gpu-amd"
+    fi
+done
+
+SCRIPT_DIR=$(dirname $0)
+bazel --bazelrc="$SCRIPT_DIR/rocm_xla_ci.bazelrc" test \
+    --build_tag_filters=$TAG_FILTERS \
+    --test_tag_filters=$TAG_FILTERS \
+    --profile=/tf/pkg/profile.json.gz \
+    --nokeep_going \
+    --test_env=TF_TESTS_PER_GPU=1 \
+    --action_env=XLA_FLAGS="--xla_gpu_force_compilation_parallelism=16" \
+    --test_output=errors \
+    --run_under=//build_tools/rocm:parallel_gpu_execute \
+    "$@"

third_party/xla/build_tools/rocm/run_xla_multi_gpu.sh

@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+# Copyright 2024 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 is a rocm specific script housed under `build_tools/rocm`
+# It runs following distributed tests which require more >= 4 gpus and these tests
+# are skipped currently in the CI due to tag selection. These tests are tagged either as manual or with oss
+# ```
+# //xla/tests:collective_ops_e2e_test_gpu_amd_any
+# //xla/tests:collective_ops_test_gpu_amd_any
+# //xla/tests:replicated_io_feed_test_gpu_amd_any
+# //xla/tools/multihost_hlo_runner:functional_hlo_runner_test_gpu_amd_any
+# //xla/pjrt/distributed:topology_util_test
+# //xla/pjrt/distributed:client_server_test
+# ```
+# Also these tests do not use `--run_under=//build_tools/ci:parallel_gpu_execute` with bazel which
+# locks down individual gpus thus making multi gpu tests impossible to run
+
+set -e
+set -x
+
+N_BUILD_JOBS=$(grep -c ^processor /proc/cpuinfo)
+# If rocm-smi exists locally (it should) use it to find
+# out how many GPUs we have to test with.
+rocm-smi -i
+STATUS=$?
+if [ $STATUS -ne 0 ]; then TF_GPU_COUNT=1; else
+   TF_GPU_COUNT=$(rocm-smi -i|grep 'Device ID' |grep 'GPU' |wc -l)
+fi
+if [[ $TF_GPU_COUNT -lt 4 ]]; then
+    echo "Found only ${TF_GPU_COUNT} gpus, multi-gpu tests need atleast 4 gpus."
+    exit
+fi
+
+TF_TESTS_PER_GPU=1
+N_TEST_JOBS=$(expr ${TF_GPU_COUNT} \* ${TF_TESTS_PER_GPU})
+amdgpuname=(`rocminfo | grep gfx | head -n 1`)
+AMD_GPU_GFX_ID=${amdgpuname[1]}
+echo ""
+echo "Bazel will use ${N_BUILD_JOBS} concurrent build job(s) and ${N_TEST_JOBS} concurrent test job(s) for gpu ${AMD_GPU_GFX_ID}."
+echo ""
+
+export PYTHON_BIN_PATH=`which python3`
+export TF_NEED_ROCM=1
+export ROCM_PATH=/opt/rocm/
+
+SCRIPT_DIR=$(realpath $(dirname $0))
+TAG_FILTERS="$($SCRIPT_DIR/rocm_tag_filters.sh)"
+
+if [ ! -d /tf/pkg ]; then
+	mkdir -p /tf/pkg
+fi
+
+bazel --bazelrc=build_tools/rocm/rocm_xla.bazelrc test \
+    --config=rocm_ci \
+    --config=xla_mgpu \
+    --build_tag_filters=${TAG_FILTERS} \
+    --test_tag_filters=${TAG_FILTERS} \
+    --profile=/tf/pkg/profile.json.gz \
+    --test_timeout=920,2400,7200,9600 \
+    --test_sharding_strategy=disabled \
+    --test_output=errors \
+    --flaky_test_attempts=3 \
+    --keep_going \
+    --local_test_jobs=${N_TEST_JOBS} \
+    --test_env=TF_TESTS_PER_GPU=$TF_TESTS_PER_GPU \
+    --test_env=TF_GPU_COUNT=$TF_GPU_COUNT \
+    --action_env=TF_ROCM_AMDGPU_TARGETS=${AMD_GPU_GFX_ID} \
+    --action_env=XLA_FLAGS=--xla_gpu_force_compilation_parallelism=16 \
+    --action_env=NCCL_MAX_NCHANNELS=1 \
+    --repo_env="ROCM_PATH=$ROCM_PATH"

third_party/xla/build_tools/rocm/tsan_ignore_list.txt

@@ -0,0 +1,40 @@
+race:libhsa-runtime64.so
+race:libamdhip64.so
+race:hipStreamSynchronize
+race:libhipblaslt.so
+race:libamd_comgr.so
+race:librccl.so
+
+# ROCr's os_thread destructor intentionally calls pthread_detach (not
+# pthread_join) on worker threads that are still running at teardown, so
+# they can exit on their own without blocking shutdown.
+thread:rocr::os::os_thread::os_thread
+thread:libhsa-runtime64.so
+
+# Abseil reference counting (DropRef / RefCount init)
+race:tsl::ReferenceCounted
+race:absl::lts_*::Mutex
+race:absl::lts_*::CondVar
+
+# XLA GPU RawSEDeviceMemory RCReference reuse
+race:xla::RawSEDeviceMemory
+race:xla::LocalDeviceState::ThenRelease
+
+# To be fixed
+race:xla::PjRtStreamExecutorClient::~PjRtStreamExecutorClient
+race:xla::GpuAsyncHostToDeviceTransferManager::TransferRawDataToSubBuffer
+race:xla::LiteralBase::Piece::DeallocateBuffers
+race:xla::CommonPjRtLoadedExecutable::ExecuteHelperOnSingleDevice
+race:xla::LocalDeviceState::AllocateAndRecordEvent
+race:xla::HloRunner::TransferLiteralsFromDevice
+race:xla::MutableLiteralBase::~MutableLiteralBase
+race:xla::MutableLiteralBase::PopulateR1<int>
+race:xla::gpu::GpuCompiler::CompileSingleModule
+race:xla::LiteralBase::Piece::Storage::Storage
+race:xla::LocalClient::TransferFromOutfeedLocal
+race:llvm::cl::opt_storage<bool, false, false>::setValue<int>
+
+race:xla::gpu::(anonymous namespace)::RecoverExp2Pattern::initStaticsIfNeeded*
+race:lld::lldMain
+race:llvm::*
+race:xla::gpu::GpuExecutable::ExecuteAsyncOnStream

third_party/xla/build_tools/sycl/ci_docker_test.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+# Copyright 2026 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.
+# ==============================================================================
+set -euo pipefail
+
+# Running inside GH Actions job container already
+  wget -qO - https://repositories.intel.com/gpu/intel-graphics.key | \
+    gpg --yes --dearmor --output /usr/share/keyrings/intel-graphics.gpg
+
+  if ! command -v lsb_release >/dev/null 2>&1; then
+    apt-get update && apt-get install -y lsb-release
+  fi
+
+  echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu $(lsb_release -cs) unified" | \
+   tee /etc/apt/sources.list.d/intel-gpu-$(lsb_release -cs).list
+
+  apt-get update && \
+    apt-get install -y --no-install-recommends \
+      intel-opencl-icd intel-ocloc libze1 libze-dev xpu-smi && \
+    apt-get clean && rm -rf /var/lib/apt/lists/*
+
+  bash build_tools/sycl/ci_test_xla.sh

third_party/xla/build_tools/sycl/ci_test_xla.sh

@@ -0,0 +1,57 @@
+#!/usr/bin/env bash
+# Copyright 2024 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 script builds and executes tests. It can be run only on a system that
+# has an Intel GPU with the appropriate driver installed.
+# TEST_TARGETS="//xla/stream_executor/sycl/..." build_tools/sycl/ci_test_xla.sh
+
+set -euo pipefail
+
+no_of_gpus=$(ls /dev/dri/ | fgrep render | wc -l)
+if [[ "${no_of_gpus}" -eq 0 ]]; then
+  echo "unknown number of gpus."
+  exit 1
+fi
+
+local_test_jobs=$((no_of_gpus * 8))
+
+if [[ ${local_test_jobs} -gt 64 ]]; then
+  local_test_jobs=64
+  echo "local_test_jobs ${local_test_jobs} too high, using 64"
+fi
+
+TEST_TARGETS="${TEST_TARGETS:-\
+  //xla/stream_executor/... \
+  //xla/backends/gpu/codegen/emitters/tests/... \
+  //xla/codegen/emitters/tests/... \
+  -//xla/backends/gpu/codegen/emitters/tests:transpose/packed_transpose_s4.hlo.test \
+  -//xla/codegen/emitters/tests:loop/s8_to_s2.hlo.test \
+  -//xla/backends/gpu/codegen/emitters/tests:transpose/multiple_roots_mixed_rank.hlo.test \
+  -//xla/backends/gpu/codegen/emitters/tests:transpose/multiple_roots_one_shmem_transpose.hlo.test \
+  -//xla/backends/gpu/codegen/emitters/tests:transpose/packed_transpose_bf16.hlo.test \
+  -//xla/backends/gpu/codegen/emitters/tests:transpose/packed_transpose_two_heroes.hlo.test \
+  -//xla/codegen/emitters/tests:loop/broadcast_constant.hlo.test\
+}"
+
+echo "TEST_TARGETS=${TEST_TARGETS}"
+
+bazel test \
+  --config=sycl_hermetic --verbose_failures -c opt \
+  --test_timeout=900 --flaky_test_attempts=2 --keep_going --test_keep_going \
+  --build_tag_filters=gpu,oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi \
+  --test_tag_filters=gpu,oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi \
+  -- \
+  ${TEST_TARGETS}

third_party/xla/build_tools/test_utils.py

@@ -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.
+# ============================================================================
+"""Test utils for python tests in XLA."""
+import os
+import pathlib
+
+
+def xla_src_root() -> pathlib.Path:
+  """Gets the path to the root of the XLA source tree."""
+  is_oss = "BAZEL_TEST" in os.environ
+  test_srcdir = os.environ["TEST_SRCDIR"]
+  test_workspace = os.environ["TEST_WORKSPACE"]
+  if is_oss:
+    return pathlib.Path(test_srcdir) / test_workspace
+  else:
+    return pathlib.Path(test_srcdir) / test_workspace / "third_party" / "xla"