From 17d426d81acb42c50d94e65c23ac21dceccae09c Mon Sep 17 00:00:00 2001
From: stgpetrovic <stgpetrovic@gmail.com>
Date: Tue, 2 May 2023 19:22:33 +0200
Subject: [PATCH] Support coverage when compiling for CUDA. (#4962)

* Support coverage when compiling for CUDA.

Passing `--config=cuda` to bazel makes it use a nvc/gcc compiler
infrastructure in `@local_config_cuda` that is immune to the
`bazel coverage` intervention in coverage flags. Setting those
flags manually solves the problem.

* Only cover repository files, not external dependencies.

* Improve coverage handling

* Fix copying coverage when only testing
---
 .bazelrc            |  4 ++++
 .circleci/common.sh | 38 +++++++++++++++++++++++---------------
 2 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/.bazelrc b/.bazelrc
index 3aac2b06a7a..d19d8faf75d 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -55,6 +55,10 @@ build:cuda --@org_tensorflow//tensorflow/compiler/xla/python:enable_gpu=true
 build:cuda --define=xla_python_enable_gpu=true
 build:cuda --cxxopt=-DXLA_CUDA=1
 
+# Coverage with cuda/gcc/nvcc requires manually setting coverage flags.
+coverage:cuda --per_file_copt=-external/.*,-test/.*@--coverage
+coverage:cuda --linkopt=-lgcov
+
 build:acl --define==build_with_acl=true
 
 build:nonccl --define=no_nccl_support=true
diff --git a/.circleci/common.sh b/.circleci/common.sh
index c39768396b2..7428a0032f3 100755
--- a/.circleci/common.sh
+++ b/.circleci/common.sh
@@ -157,25 +157,33 @@ function run_torch_xla_tests() {
       fi
     fi
 
-    pushd test/cpp
-      echo "Running C++ Tests on PJRT"
-      EXTRA_ARGS=""
+    echo "Running C++ Tests on PJRT"
+    EXTRA_ARGS=""
+    if [ "$USE_COVERAGE" != "0" ]; then
+	    EXTRA_ARGS="-C"
+    fi
+    if [ ! -z "$GCLOUD_SERVICE_KEY_FILE" ]; then
+	    EXTRA_ARGS="$EXTRA_ARGS -R"
+    fi
+    if [ -x "$(command -v nvidia-smi)" ]; then
+      PJRT_DEVICE=GPU test/cpp/run_tests.sh $EXTRA_ARGS
       if [ "$USE_COVERAGE" != "0" ]; then
-	      EXTRA_ARGS="-C"
+        cp $XLA_DIR/bazel-out/_coverage/_coverage_report.dat /tmp/cov1.dat
       fi
-      if [ ! -z "$GCLOUD_SERVICE_KEY_FILE" ]; then
-	      EXTRA_ARGS="$EXTRA_ARGS -R"
-      fi
-      if [ -x "$(command -v nvidia-smi)" ]; then
-        PJRT_DEVICE=GPU ./run_tests.sh $EXTRA_ARGS
-        PJRT_DEVICE=GPU ./run_tests.sh -X early_sync -F AtenXlaTensorTest.TestEarlySyncLiveTensors -L"" $EXTRA_ARGS
-      else
-        PJRT_DEVICE=CPU ./run_tests.sh $EXTRA_ARGS
+      PJRT_DEVICE=GPU test/cpp/run_tests.sh -X early_sync -F AtenXlaTensorTest.TestEarlySyncLiveTensors -L"" $EXTRA_ARGS
+      if [ "$USE_COVERAGE" != "0" ]; then
+        cp $XLA_DIR/bazel-out/_coverage/_coverage_report.dat /tmp/cov2.dat
+        lcov --add-tracefile /tmp/cov1.dat -a /tmp/cov2.dat -o /tmp/merged.dat
       fi
+    else
+      PJRT_DEVICE=CPU test/cpp/run_tests.sh $EXTRA_ARGS
       if [ "$USE_COVERAGE" != "0" ]; then
-        genhtml $XLA_DIR/bazel-out/_coverage/_coverage_report.dat -o ~/htmlcov/cpp/cpp_lcov.info
-        mv $XLA_DIR/bazel-out/_coverage/_coverage_report.dat ~/htmlcov/cpp_lcov.info
+        cp $XLA_DIR/bazel-out/_coverage/_coverage_report.dat /tmp/merged.dat
       fi
-    popd
+    fi
+    if [ "$USE_COVERAGE" != "0" ]; then
+      genhtml /tmp/merged.dat -o ~/htmlcov/cpp/cpp_lcov.info
+      mv /tmp/merged.dat ~/htmlcov/cpp_lcov.info
+    fi
   popd
 }