diff --git a/include/nvexec/stream/common.cuh b/include/nvexec/stream/common.cuh index df60275f9..28a149704 100644 --- a/include/nvexec/stream/common.cuh +++ b/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; + } } } diff --git a/test/nvexec/CMakeLists.txt b/test/nvexec/CMakeLists.txt index d46a551e5..4e45c8234 100644 --- a/test/nvexec/CMakeLists.txt +++ b/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 diff --git a/test/nvexec/stream_pool.cpp b/test/nvexec/stream_pool.cpp new file mode 100644 index 000000000..dac246e29 --- /dev/null +++ b/test/nvexec/stream_pool.cpp @@ -0,0 +1,55 @@ +#include + +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 + +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