Prepare stream testing job for static cudart

Rather than using `LD_PRELOAD` to swap in our redirects at runtime,
rewrite the ELF symbols at compile time.

Author: KyleFromNVIDIA

cpp/CMakeLists.txt

@@ -61,16 +61,8 @@ option(CUDA_WARNINGS_AS_ERRORS "Enable -Werror=all-warnings for all CUDA compila
 # cudart can be statically linked or dynamically linked. The python ecosystem wants dynamic linking
 option(CUDA_STATIC_RUNTIME "Statically link the CUDA runtime" OFF)
 
-set(DEFAULT_CUDF_BUILD_STREAMS_TEST_UTIL ON)
-
-if(CUDA_STATIC_RUNTIME OR NOT BUILD_SHARED_LIBS)
-  set(DEFAULT_CUDF_BUILD_STREAMS_TEST_UTIL OFF)
-endif()
-
-option(
-  CUDF_BUILD_STREAMS_TEST_UTIL
-  "Whether to build the utilities for stream testing contained in libcudf"
-  ${DEFAULT_CUDF_BUILD_STREAMS_TEST_UTIL}
+option(CUDF_BUILD_STREAMS_TEST_UTIL
+       "Whether to build the utilities for stream testing contained in libcudf" OFF
 )
 mark_as_advanced(CUDF_BUILD_STREAMS_TEST_UTIL)
 option(CUDF_CLANG_TIDY "Enable clang-tidy during compilation" OFF)
@@ -1080,8 +1072,12 @@ if(CUDF_BUILD_TESTUTIL)
   )
 
   target_link_libraries(
-    cudftestutil INTERFACE cuco::cuco Threads::Threads cudf cudftest_default_stream
-                           $<TARGET_NAME_IF_EXISTS:conda_env>
+    cudftestutil
+    INTERFACE cuco::cuco
+              Threads::Threads
+              cudf
+              $<$<NOT:$<BOOL:${CUDF_BUILD_STREAMS_TEST_UTIL}>>:cudftest_default_stream>
+              $<TARGET_NAME_IF_EXISTS:conda_env>
   )
 
   target_include_directories(
@@ -1135,19 +1131,12 @@ endif()
 
 # * build cudf_identify_stream_usage --------------------------------------------------------------
 if(CUDF_BUILD_STREAMS_TEST_UTIL)
-  if(CUDA_STATIC_RUNTIME)
-    message(
-      FATAL_ERROR
-        "Stream identification cannot be used with a static CUDA runtime. Please set CUDA_STATIC_RUNTIME=OFF or CUDF_BUILD_STREAMS_TEST_UTIL=OFF."
-    )
-  endif()
-
   # Libraries for stream-related testing. We build the library twice, one with STREAM_MODE_TESTING
   # on and one with it set to off. Each test will then be configured to use the appropriate library
   # depending via ctest and whether it has been updated to expose public stream APIs.
   foreach(_mode cudf testing)
     set(_tgt "cudf_identify_stream_usage_mode_${_mode}")
-    add_library(${_tgt} SHARED tests/utilities/identify_stream_usage.cpp)
+    add_library(${_tgt} STATIC tests/utilities/identify_stream_usage.cpp)
 
     if(CUDF_USE_PER_THREAD_DEFAULT_STREAM)
       target_compile_definitions(
@@ -1168,7 +1157,11 @@ if(CUDF_BUILD_STREAMS_TEST_UTIL)
       ${_tgt} PRIVATE "$<BUILD_INTERFACE:$<$<COMPILE_LANGUAGE:CXX>:${CUDF_CXX_FLAGS}>>"
     )
     target_include_directories(${_tgt} PRIVATE "$<BUILD_INTERFACE:${CUDF_SOURCE_DIR}/include>")
-    target_link_libraries(${_tgt} PUBLIC CUDA::cudart rmm::rmm)
+    target_link_libraries(
+      ${_tgt}
+      PUBLIC $<COMPILE_ONLY:CUDA::cudart> rmm::rmm
+      PRIVATE cudf
+    )
 
     rapids_cuda_set_runtime(${_tgt} USE_STATIC ${CUDA_STATIC_RUNTIME})
     add_library(cudf::${_tgt} ALIAS ${_tgt})
@@ -1177,6 +1170,28 @@ if(CUDF_BUILD_STREAMS_TEST_UTIL)
       target_compile_definitions(${_tgt} PUBLIC STREAM_MODE_TESTING)
     endif()
   endforeach()
+
+  find_package(Python3 REQUIRED COMPONENTS Interpreter)
+
+  # Redirect CUDA calls for a target to use our stream testing versions
+  function(fixup_target_for_stream_testing tgt)
+    set(fixup_script "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/cmake/fixup_target_for_stream_testing.py")
+    add_custom_command(
+      TARGET ${tgt}
+      PRE_LINK
+      COMMAND Python3::Interpreter "${fixup_script}" $<TARGET_OBJECTS:${tgt}>
+      VERBATIM COMMAND_EXPAND_LISTS
+      COMMENT "Redirecting CUDA calls for ${tgt}"
+    )
+    set_property(
+      TARGET ${tgt}
+      PROPERTY LINK_DEPENDS "${fixup_script}"
+      APPEND
+    )
+  endfunction()
+
+  fixup_target_for_stream_testing(cudf)
+  target_compile_definitions(cudf PRIVATE CUDF_BUILD_STREAMS_TEST_UTIL)
 endif()
 
 # ##################################################################################################

cpp/cmake/fixup_target_for_stream_testing.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+
+import lief
+
+# This list must be kept in sync with cpp/tests/utilities/identify_stream_usage.cpp
+SYMBOLS_TO_REWRITE = [
+    "cudaEventRecord",
+    "cudaEventRecordWithFlags",
+    "cudaLaunchKernel",
+    "__cudaLaunchKernel",
+    "__cudaLaunchKernel_ptsz",
+    "cudaMemPrefetchAsync",
+    "cudaMemcpy2DAsync",
+    "cudaMemcpy2DFromArrayAsync",
+    "cudaMemcpy2DToArrayAsync",
+    "cudaMemcpy3DAsync",
+    "cudaMemcpy3DPeerAsync",
+    "cudaMemcpyAsync",
+    "cudaMemcpyFromSymbolAsync",
+    "cudaMemcpyToSymbolAsync",
+    "cudaMemset2DAsync",
+    "cudaMemset3DAsync",
+    "cudaMemsetAsync",
+    "cudaFreeAsync",
+    "cudaMallocAsync",
+    "cudaMallocFromPoolAsync",
+]
+
+
+for filename in sys.argv[1:]:
+    elf = lief.ELF.parse(filename)
+
+    for symbol_name in SYMBOLS_TO_REWRITE:
+        if symbol := elf.get_symbol(symbol_name):
+            symbol.name = f"cudf_{symbol.name}"
+
+    elf.write(filename)

cpp/include/cudf_test/default_stream.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -17,12 +17,11 @@ namespace test {
  *
  * The standard behavior of this function is to return cudf's default stream
  * (cudf::get_default_stream). This function is primarily provided as an
- * overload target for preload libraries (via LD_PRELOAD) so that the default
- * stream used for tests may be modified for tracking purposes. All tests of
- * public APIs that accept streams should pass `cudf::test::get_default_stream`
- * as the stream argument so that a preload library changing the behavior of
- * this function will trigger those tests to run on a different stream than
- * `cudf::get_default_stream`.
+ * overload target so that the default stream used for tests may be modified
+ * for tracking purposes. All tests of public APIs that accept streams should
+ * pass `cudf::test::get_default_stream` as the stream argument so that a
+ * library changing the behavior of this function will trigger those tests to
+ * run on a different stream than `cudf::get_default_stream`.
  *
  * @return The default stream to use for tests.
  */

cpp/src/utilities/default_stream.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -17,6 +17,18 @@ rmm::cuda_stream_view const default_stream_value{rmm::cuda_stream_per_thread};
 rmm::cuda_stream_view const default_stream_value{};
 #endif
 
+rmm::cuda_stream_view const get_default_stream()
+{
+  static auto const default_stream = []() {
+    if (std::getenv("CUDF_PER_THREAD_STREAM") != nullptr) {
+      return rmm::cuda_stream_per_thread;
+    } else {
+      return detail::default_stream_value;
+    }
+  }();
+  return default_stream;
+}
+
 }  // namespace detail
 
 /**

cpp/tests/CMakeLists.txt

@@ -71,11 +71,13 @@ function(ConfigureTest CMAKE_TEST_NAME)
 
   if(CUDF_BUILD_STREAMS_TEST_UTIL)
     set_tests_properties(
-      ${CMAKE_TEST_NAME}
-      PROPERTIES
-        ENVIRONMENT
-        "GTEST_CUDF_STREAM_MODE=new_${_CUDF_TEST_STREAM_MODE}_default;LD_PRELOAD=$<TARGET_FILE:cudf_identify_stream_usage_mode_${_CUDF_TEST_STREAM_MODE}>"
+      ${CMAKE_TEST_NAME} PROPERTIES ENVIRONMENT
+                                    "GTEST_CUDF_STREAM_MODE=new_${_CUDF_TEST_STREAM_MODE}_default"
     )
+    target_link_libraries(
+      ${CMAKE_TEST_NAME} PRIVATE "cudf_identify_stream_usage_mode_${_CUDF_TEST_STREAM_MODE}"
+    )
+    fixup_target_for_stream_testing(${CMAKE_TEST_NAME})
   endif()
 endfunction()
 

cpp/tests/utilities/identify_stream_usage.cpp

@@ -19,9 +19,9 @@
 #include <string>
 #include <unordered_map>
 
-// This file is compiled into a separate library that is dynamically loaded with LD_PRELOAD at
-// runtime to libcudf to override some stream-related symbols in libcudf. The goal of such a library
-// is to verify if the stream/stream pool is being correctly forwarded between API calls.
+// This file is compiled into a separate library that is statically linked to tests to
+// override some stream-related symbols in libcudf. The goal of such a library is to
+// verify if the stream/stream pool is being correctly forwarded between API calls.
 //
 // We control whether to override cudf::test::get_default_stream or
 // cudf::get_default_stream with a compile-time flag. The behaviour of tests
@@ -44,19 +44,46 @@
 
 namespace cudf {
 
-#ifdef STREAM_MODE_TESTING
-namespace test {
+namespace detail {
+
+#if defined(CUDF_USE_PER_THREAD_DEFAULT_STREAM)
+rmm::cuda_stream_view const default_stream_value{rmm::cuda_stream_per_thread};
+#else
+rmm::cuda_stream_view const default_stream_value{};
 #endif
 
+}  // namespace detail
+
 rmm::cuda_stream_view const get_default_stream()
 {
+#ifdef STREAM_MODE_TESTING
+  static auto const default_stream = []() {
+    if (std::getenv("CUDF_PER_THREAD_STREAM") != nullptr) {
+      return rmm::cuda_stream_per_thread;
+    } else {
+      return detail::default_stream_value;
+    }
+  }();
+  return default_stream;
+#else
   static rmm::cuda_stream stream{};
   return stream;
+#endif
 }
 
+namespace test {
+
+rmm::cuda_stream_view const get_default_stream()
+{
 #ifdef STREAM_MODE_TESTING
-}  // namespace test
+  static rmm::cuda_stream stream{};
+  return stream;
+#else
+  return cudf::get_default_stream();
 #endif
+}
+
+}  // namespace test
 
 #ifdef STREAM_MODE_TESTING
 namespace detail {
@@ -119,15 +146,6 @@ void check_stream_and_error(cudaStream_t stream)
   }
 }
 
-/**
- * @brief Container for CUDA APIs that have been overloaded using DEFINE_OVERLOAD.
- *
- * This variable must be initialized before everything else.
- *
- * @see find_originals for a description of the priorities
- */
-__attribute__((init_priority(1001))) std::unordered_map<std::string, void*> originals;
-
 /**
  * @brief Macro for generating functions to override existing CUDA functions.
  *
@@ -145,15 +163,12 @@ __attribute__((init_priority(1001))) std::unordered_map<std::string, void*> orig
  * @param signature The function signature (must include names, not just types).
  * @parameter arguments The function arguments (names only, no types).
  */
-#define DEFINE_OVERLOAD(function, signature, arguments)     \
-  using function##_t = cudaError_t (*)(signature);          \
-                                                            \
-  cudaError_t function(signature)                           \
-  {                                                         \
-    check_stream_and_error(stream);                         \
-    return ((function##_t)originals[#function])(arguments); \
-  }                                                         \
-  __attribute__((constructor(1002))) void queue_##function() { originals[#function] = nullptr; }
+#define DEFINE_OVERLOAD(function, signature, arguments) \
+  extern "C" cudaError_t cudf_##function(signature)     \
+  {                                                     \
+    check_stream_and_error(stream);                     \
+    return function(arguments);                         \
+  }
 
 /**
  * @brief Helper macro to define macro arguments that contain a comma.
@@ -177,6 +192,8 @@ __attribute__((init_priority(1001))) std::unordered_map<std::string, void*> orig
    - https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__INTEROP.html#group__CUDART__INTEROP
    - https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__GRAPH.html#group__CUDART__GRAPH
    - https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__HIGHLEVEL.html#group__CUDART__HIGHLEVEL
+
+   This list must be kept in sync with cpp/cmake/fixup_target_for_stream_testing.py
  */
 // clang-format on
 
@@ -335,18 +352,3 @@ DEFINE_OVERLOAD(cudaMallocAsync,
 DEFINE_OVERLOAD(cudaMallocFromPoolAsync,
                 ARG(void** ptr, size_t size, cudaMemPool_t memPool, cudaStream_t stream),
                 ARG(ptr, size, memPool, stream));
-
-/**
- * @brief Function to collect all the original CUDA symbols corresponding to overloaded functions.
- *
- * Note on priorities:
- * - `originals` must be initialized first, so it is 1001.
- * - The function names must be added to originals next in the macro, so those are 1002.
- * - Finally, this function actually finds the original symbols so it is 1003.
- */
-__attribute__((constructor(1003))) void find_originals()
-{
-  for (auto it : originals) {
-    originals[it.first] = dlsym(RTLD_NEXT, it.first.data());
-  }
-}

dependencies.yaml

@@ -154,6 +154,7 @@ files:
       - build_base
       - cuda
       - cuda_version
+      - stream_test_dependencies
   docs:
     output: none
     includes:
@@ -1343,3 +1344,9 @@ dependencies:
       - output_types: [conda, requirements, pyproject]
         packages:
           - structlog
+  stream_test_dependencies:
+    common:
+      - output_types: conda
+        packages:
+          - python
+          - py-lief