From 954711ed8f0d1eba3e7708e786a1dadf7e9a18ce Mon Sep 17 00:00:00 2001
From: Shilei Tian <tianshilei1992@gmail.com>
Date: Sun, 18 Jul 2021 13:16:03 -0400
Subject: [PATCH] [OpenMP][Offloading] Add a CMake argument
 LIBOMPTARGET_LIT_ARGS to control behavior of libomptarget lit test

By default, `lit` uses all threads to invoke tests, which  can easily cause out
of memory on GPUs because most of OpenMP offloading test usually take about 1GB
GPU memory, but a typical GPU only has 4-8GB memory. This patch introduce a
CMake argument `LIBOMPTARGET_LIT_ARGS` to allow users to control the behavior of
`libomptarget` tests, similar to `LLVM_LIT_ARGS`.

Reviewed By: JonChesterfield

Differential Revision: https://reviews.llvm.org/D106236
---
 openmp/README.rst                       |  6 ++++++
 openmp/libomptarget/test/CMakeLists.txt | 13 +++++++++++--
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/openmp/README.rst b/openmp/README.rst
index bd8071f43f95..ffa49e4d2a49 100644
--- a/openmp/README.rst
+++ b/openmp/README.rst
@@ -300,6 +300,12 @@ Options for ``NVPTX device RTL``
  **LIBOMPTARGET_NVPTX_DEBUG** = ``OFF|ON``
   Enable printing of debug messages from the NVPTX device RTL.
 
+**LIBOMPTARGET_LIT_ARGS** = ``""``
+  Arguments given to lit. ``make check-libomptarget`` and
+  ``make check-libomptarget-*`` are affected. For example, use
+  ``LIBOMPTARGET_LIT_ARGS="-j4"`` to force ``lit`` to start only four parallel
+  jobs instead of by default the number of threads in the system.
+
 Example Usages of CMake
 =======================
 
diff --git a/openmp/libomptarget/test/CMakeLists.txt b/openmp/libomptarget/test/CMakeLists.txt
index ddc05e82ef9e..9b6de8ba08df 100644
--- a/openmp/libomptarget/test/CMakeLists.txt
+++ b/openmp/libomptarget/test/CMakeLists.txt
@@ -15,7 +15,11 @@ endif()
 string(REGEX MATCHALL "([^\ ]+\ |[^\ ]+$)" SYSTEM_TARGETS "${LIBOMPTARGET_SYSTEM_TARGETS}")
 foreach(CURRENT_TARGET IN LISTS SYSTEM_TARGETS)
   string(STRIP "${CURRENT_TARGET}" CURRENT_TARGET)
-  add_openmp_testsuite(check-libomptarget-${CURRENT_TARGET} "Running libomptarget tests" ${CMAKE_CURRENT_BINARY_DIR}/${CURRENT_TARGET} DEPENDS omptarget omp ${LIBOMPTARGET_TESTED_PLUGINS})
+  add_openmp_testsuite(check-libomptarget-${CURRENT_TARGET}
+    "Running libomptarget tests"
+    ${CMAKE_CURRENT_BINARY_DIR}/${CURRENT_TARGET}
+    DEPENDS omptarget omp ${LIBOMPTARGET_TESTED_PLUGINS}
+    ARGS ${LIBOMPTARGET_LIT_ARGS})
   list(APPEND LIBOMPTARGET_LIT_TESTSUITES ${CMAKE_CURRENT_BINARY_DIR}/${CURRENT_TARGET})
 
   # Configure the lit.site.cfg.in file
@@ -24,4 +28,9 @@ foreach(CURRENT_TARGET IN LISTS SYSTEM_TARGETS)
   configure_file(lit.site.cfg.in ${CURRENT_TARGET}/lit.site.cfg @ONLY)
 endforeach()
 
-add_openmp_testsuite(check-libomptarget "Running libomptarget tests" ${LIBOMPTARGET_LIT_TESTSUITES} EXCLUDE_FROM_CHECK_ALL DEPENDS omptarget omp ${LIBOMPTARGET_TESTED_PLUGINS})
\ No newline at end of file
+add_openmp_testsuite(check-libomptarget
+  "Running libomptarget tests"
+  ${LIBOMPTARGET_LIT_TESTSUITES}
+  EXCLUDE_FROM_CHECK_ALL
+  DEPENDS omptarget omp ${LIBOMPTARGET_TESTED_PLUGINS}
+  ARGS ${LIBOMPTARGET_LIT_ARGS})