[XLA] Add coverage support to XLA lit tests.

PiperOrigin-RevId: 646882531
tensorflow · Jun 26, 2024 · 1aeebeb · 1aeebeb
1 parent b572b20
commit 1aeebeb
Showing 1 changed file with 19 additions and 4 deletions.
diff --git a/third_party/xla/xla/lit.cfg.py b/third_party/xla/xla/lit.cfg.py
@@ -29,8 +29,24 @@
 config.test_format = lit.formats.ShTest(execute_external=True)
 
 
-# Passthrough XLA_FLAGS.
-config.environment["XLA_FLAGS"] = os.environ.get("XLA_FLAGS", "")
+for env in [
+    # Passthrough XLA_FLAGS.
+    "XLA_FLAGS",
+    # Propagate environment variables used by 'bazel coverage'.
+    # These are exported by tools/coverage/collect_coverage.sh
+    "BULK_COVERAGE_RUN",
+    "COVERAGE",
+    "COVERAGE_DIR",
+    "COVERAGE_MANIFEST",
+    "LLVM_PROFILE_FILE",
+    "LLVM_COVERAGE_FILE",
+    "GCOV_PREFIX",
+    "GCOV_PREFIX_STRIP",
+]:
+  value = os.environ.get(env)
+  if value:
+    config.environment[env] = value
+
 
 # Use the most preferred temp directory.
 config.test_exec_root = (
@@ -45,6 +61,5 @@
 
 # Include additional substitutions that may be defined via params
 config.substitutions.extend(
-    ("%%{%s}" % key, val)
-    for key, val in lit_config.params.items()
+    ("%%{%s}" % key, val) for key, val in lit_config.params.items()
 )