From 078c02bda579243d04adfd0fa762fd24833cf895 Mon Sep 17 00:00:00 2001
From: Allison Vacanti <alliepiper16@gmail.com>
Date: Thu, 16 Jul 2020 12:12:14 -0400
Subject: [PATCH] add tests for legacy and per-thread default streams

---
 testing/CMakeLists.txt               |  8 ++++++++
 testing/cuda/stream_legacy.cu        | 21 +++++++++++++++++++++
 testing/cuda/stream_per_thread.cmake |  6 ++++++
 testing/cuda/stream_per_thread.cu    | 21 +++++++++++++++++++++
 testing/cuda/stream_per_thread.mk    |  1 +
 5 files changed, 57 insertions(+)
 create mode 100644 testing/cuda/stream_legacy.cu
 create mode 100644 testing/cuda/stream_per_thread.cmake
 create mode 100644 testing/cuda/stream_per_thread.cu
 create mode 100644 testing/cuda/stream_per_thread.mk

diff --git a/testing/CMakeLists.txt b/testing/CMakeLists.txt
index 67e25af45d..fdfc04e97b 100644
--- a/testing/CMakeLists.txt
+++ b/testing/CMakeLists.txt
@@ -111,6 +111,14 @@ function(thrust_add_test target_name_var test_name test_src thrust_target)
   if (("OMP" IN_LIST config_systems) OR ("TBB" IN_LIST config_systems))
     set_tests_properties(${test_target} PROPERTIES RUN_SERIAL ON)
   endif()
+
+  # Check for per-test script. Script will be included in the current scope
+  # to allow custom property modifications.
+  get_filename_component(test_cmake_script "${test_src}" NAME_WLE)
+  set(test_cmake_script "${CMAKE_CURRENT_LIST_DIR}/${test_cmake_script}.cmake")
+  if (EXISTS "${test_cmake_script}")
+    include("${test_cmake_script}")
+  endif()
 endfunction()
 
 file(GLOB test_srcs
diff --git a/testing/cuda/stream_legacy.cu b/testing/cuda/stream_legacy.cu
new file mode 100644
index 0000000000..51c82a096d
--- /dev/null
+++ b/testing/cuda/stream_legacy.cu
@@ -0,0 +1,21 @@
+#include <unittest/unittest.h>
+#include <thrust/execution_policy.h>
+#include <thrust/system/cuda/detail/util.h>
+
+#include <thread>
+
+void verify_stream()
+{
+  auto exec = thrust::device;
+  auto stream = thrust::cuda_cub::stream(exec);
+  ASSERT_EQUAL(stream, cudaStreamLegacy);
+}
+
+void TestLegacyDefaultStream()
+{
+  verify_stream();
+
+  std::thread t(verify_stream);
+  t.join();
+}
+DECLARE_UNITTEST(TestLegacyDefaultStream);
diff --git a/testing/cuda/stream_per_thread.cmake b/testing/cuda/stream_per_thread.cmake
new file mode 100644
index 0000000000..5325a07025
--- /dev/null
+++ b/testing/cuda/stream_per_thread.cmake
@@ -0,0 +1,6 @@
+# This test should always use per-thread streams on NVCC.
+# NVC++ does not have an equivalent option.
+set_target_properties(${test_target} PROPERTIES
+  COMPILE_OPTIONS
+    $<$<AND:$<COMPILE_LANGUAGE:CUDA>,$<CUDA_COMPILER_ID:NVIDIA>>:--default-stream=per-thread>
+)
diff --git a/testing/cuda/stream_per_thread.cu b/testing/cuda/stream_per_thread.cu
new file mode 100644
index 0000000000..ef126e78a5
--- /dev/null
+++ b/testing/cuda/stream_per_thread.cu
@@ -0,0 +1,21 @@
+#include <unittest/unittest.h>
+#include <thrust/execution_policy.h>
+#include <thrust/system/cuda/detail/util.h>
+
+#include <thread>
+
+void verify_stream()
+{
+  auto exec = thrust::device;
+  auto stream = thrust::cuda_cub::stream(exec);
+  ASSERT_EQUAL(stream, cudaStreamPerThread);
+}
+
+void TestPerThreadDefaultStream()
+{
+  verify_stream();
+
+  std::thread t(verify_stream);
+  t.join();
+}
+DECLARE_UNITTEST(TestPerThreadDefaultStream);
diff --git a/testing/cuda/stream_per_thread.mk b/testing/cuda/stream_per_thread.mk
new file mode 100644
index 0000000000..da9adfe1b8
--- /dev/null
+++ b/testing/cuda/stream_per_thread.mk
@@ -0,0 +1 @@
+CUDACC_FLAGS += --default-stream per-thread