Do not pool failed CUDA streams

Author: fallintoplace

include/nvexec/stream/common.cuh

@@ -358,7 +358,13 @@ namespace nv::execution
       {
         if (!borrows_stream)
         {
-          std::tie(own_stream_, status_) = context_.borrow_stream();
+          cudaStream_t stream{};
+          std::tie(stream, status_) = context_.borrow_stream();
+
+          if (status_ == cudaSuccess)
+          {
+            own_stream_ = stream;
+          }
         }
       }
 

test/nvexec/CMakeLists.txt

@@ -47,6 +47,7 @@ set(nvexec_test_sources
 
 if(NOT (CMAKE_CXX_COMPILER_ID STREQUAL "NVHPC"))
   set_source_files_properties(${nvexec_test_sources} PROPERTIES LANGUAGE CUDA)
+  set_source_files_properties(stream_pool.cpp PROPERTIES LANGUAGE CUDA)
 endif()
 
 add_executable(test.nvexec ${nvexec_test_sources})
@@ -61,6 +62,18 @@ target_link_libraries(test.nvexec STDEXEC::nvexec stdexec_executable_flags
 
 catch_discover_tests(test.nvexec PROPERTIES TIMEOUT 30)
 
+add_executable(test.nvexec.stream_pool stream_pool.cpp test_main.cpp)
+set_target_properties(
+  test.nvexec.stream_pool
+  PROPERTIES CXX_STANDARD 20
+             CXX_STANDARD_REQUIRED ON
+             CXX_EXTENSIONS OFF)
+target_include_directories(test.nvexec.stream_pool PRIVATE ${CMAKE_CURRENT_LIST_DIR}/..)
+target_link_libraries(test.nvexec.stream_pool STDEXEC::nvexec stdexec_executable_flags
+                      Catch2::Catch2WithMain nvexec_executable_flags)
+
+catch_discover_tests(test.nvexec.stream_pool PROPERTIES TIMEOUT 30)
+
 icm_add_build_failure_test(
   NAME
   when_all_fail

test/nvexec/stream_pool.cpp

@@ -0,0 +1,55 @@
+#include <cuda_runtime_api.h>
+
+static int test_stream_pool_create_calls{};
+static int test_stream_pool_destroy_calls{};
+
+static cudaError_t test_stream_pool_cudaStreamCreate(cudaStream_t* stream) noexcept
+{
+  ++test_stream_pool_create_calls;
+  *stream = nullptr;
+  return cudaErrorMemoryAllocation;
+}
+
+static cudaError_t test_stream_pool_cudaStreamDestroy(cudaStream_t) noexcept
+{
+  ++test_stream_pool_destroy_calls;
+  return cudaSuccess;
+}
+
+#define cudaStreamCreate test_stream_pool_cudaStreamCreate
+#define cudaStreamDestroy test_stream_pool_cudaStreamDestroy
+#include "nvexec/stream/common.cuh"
+#undef cudaStreamDestroy
+#undef cudaStreamCreate
+
+#include <test_common/catch2.hpp>
+
+namespace
+{
+  TEST_CASE("stream provider does not pool a failed stream", "[cuda][stream][stream_pool]")
+  {
+    test_stream_pool_create_calls  = 0;
+    test_stream_pool_destroy_calls = 0;
+
+    {
+      nvexec::_strm::stream_pools_t stream_pools;
+      nvexec::_strm::context        context{nullptr, nullptr, &stream_pools, nullptr};
+
+      {
+        nvexec::_strm::stream_provider provider{context};
+        REQUIRE(provider.status_ == cudaErrorMemoryAllocation);
+        CHECK_FALSE(provider.own_stream_.has_value());
+      }
+
+      {
+        nvexec::_strm::stream_provider provider{context};
+        CHECK(provider.status_ == cudaErrorMemoryAllocation);
+        CHECK_FALSE(provider.own_stream_.has_value());
+      }
+
+      CHECK(test_stream_pool_create_calls == 2);
+    }
+
+    CHECK(test_stream_pool_destroy_calls == 0);
+  }
+}  // namespace