chore: import upstream snapshot with attribution

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);
++}