From 116fbe124074d915f0319baefe907f3b868fb242 Mon Sep 17 00:00:00 2001
From: Adam <szadamx@gmail.com>
Date: Thu, 26 Oct 2023 10:05:50 +0200
Subject: [PATCH] [UR][Tests] Add a CMake option to limit the test devices
 count

and test platforms count in CTS.
---
 test/conformance/CMakeLists.txt         |  4 +++-
 test/conformance/README.md              |  9 ++++++++-
 test/conformance/cts_exe.py             |  6 +++++-
 test/conformance/source/environment.cpp | 17 +++++++++++++++++
 4 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/test/conformance/CMakeLists.txt b/test/conformance/CMakeLists.txt
index 5a898846ec..9085968f51 100644
--- a/test/conformance/CMakeLists.txt
+++ b/test/conformance/CMakeLists.txt
@@ -4,6 +4,8 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 set(UR_CONFORMANCE_TEST_DIR ${CMAKE_CURRENT_SOURCE_DIR})
+option(UR_TEST_DEVICES_COUNT "Count of devices on which CTS will be run" 1)
+option(UR_TEST_PLATFORMS_COUNT "Count of platform on which CTS will be run" 1)
 
 function(add_test_adapter name adapter)
     set(TEST_TARGET_NAME test-${name})
@@ -12,7 +14,7 @@ function(add_test_adapter name adapter)
     add_test(NAME ${TEST_NAME}
         COMMAND ${CMAKE_COMMAND}
         -D TEST_FILE=${Python3_EXECUTABLE}
-        -D TEST_ARGS="${UR_CONFORMANCE_TEST_DIR}/cts_exe.py --test_command ${CMAKE_BINARY_DIR}/bin/${TEST_TARGET_NAME}"
+        -D TEST_ARGS="${UR_CONFORMANCE_TEST_DIR}/cts_exe.py --test_command ${CMAKE_BINARY_DIR}/bin/${TEST_TARGET_NAME} --test_devices_count=${UR_TEST_DEVICES_COUNT} --test_platforms_count=${UR_TEST_PLATFORMS_COUNT}"
         -D MODE=stdout
         -D MATCH_FILE=${CMAKE_CURRENT_SOURCE_DIR}/${name}_${adapter}.match
         -P ${PROJECT_SOURCE_DIR}/cmake/match.cmake
diff --git a/test/conformance/README.md b/test/conformance/README.md
index db90fc759b..9e78337151 100644
--- a/test/conformance/README.md
+++ b/test/conformance/README.md
@@ -8,4 +8,11 @@ In the future, when all bugs are fixed, and the tests pass,
 this solution will no longer be necessary.
 When you fix any test, the match file must be updated
 Empty match files indicate that there are no failing tests
-in a particular group for the corresponding adapter.
\ No newline at end of file
+in a particular group for the corresponding adapter.
+
+## How to limit the test devices count
+
+To limit how many devices you want to run the CTS on,
+use CMake option UR_TEST_DEVICES_COUNT. If you want to run
+the tests on all available devices, set 0.
+The default value is 1.
\ No newline at end of file
diff --git a/test/conformance/cts_exe.py b/test/conformance/cts_exe.py
index ce3ca00a20..3efcea0d0d 100644
--- a/test/conformance/cts_exe.py
+++ b/test/conformance/cts_exe.py
@@ -20,9 +20,13 @@
 
     parser = ArgumentParser()
     parser.add_argument("--test_command", help="Ctest test case")
+    parser.add_argument("--test_devices_count", type=str, help="Number of devices on which tests will be run")
+    parser.add_argument("--test_platforms_count", type=str, help="Number of devices on which tests will be run")
     args = parser.parse_args()
     
-    result = subprocess.Popen([args.test_command, '--gtest_brief=1'], stdout = subprocess.PIPE, text = True)  # nosec B603
+    result = subprocess.Popen([args.test_command, '--gtest_brief=1', f'--devices_count={args.test_devices_count}',
+                               f'--platforms_count={args.test_platforms_count}'],
+                               stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)  # nosec B603
 
     pat = re.compile(r'\[( )*FAILED( )*\]')
     output_list = []
diff --git a/test/conformance/source/environment.cpp b/test/conformance/source/environment.cpp
index 5fe7d9b352..f263dda1e3 100644
--- a/test/conformance/source/environment.cpp
+++ b/test/conformance/source/environment.cpp
@@ -3,6 +3,7 @@
 // See LICENSE.TXT
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+#include <algorithm>
 #include <cstring>
 #include <fstream>
 
@@ -210,6 +211,22 @@ DevicesEnvironment::DevicesEnvironment(int argc, char **argv)
         error = "Could not find any devices associated with the platform";
         return;
     }
+    // Get the argument (test_devices_count) to limit test devices count.
+    u_long count_set = 0;
+    for (int i = 1; i < argc; ++i) {
+        if (std::strcmp(argv[i], "--test_devices_count") == 0 && i + 1 < argc) {
+            count_set = std::strtoul(argv[i + 1], nullptr, 10);
+            break;
+        }
+    }
+    // In case, the count_set is "0", the variable count will not be changed.
+    // The CTS will run on all devices.
+    if (count_set > (std::numeric_limits<uint32_t>::max)()) {
+        error = "Invalid test_devices_count argument";
+        return;
+    } else if (count_set > 0) {
+        count = (std::min)(count, static_cast<uint32_t>(count_set));
+    }
     devices.resize(count);
     if (urDeviceGet(platform, UR_DEVICE_TYPE_ALL, count, devices.data(),
                     nullptr)) {