diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp index edd05ea13d..89bbe898de 100644 --- a/hipamd/src/hip_graph.cpp +++ b/hipamd/src/hip_graph.cpp @@ -24,7 +24,10 @@ #include "hip_platform.hpp" #include "hip_event.hpp" -thread_local std::vector g_captureStreams; +std::vector g_captureStreams; +amd::Monitor g_captureStreamsLock{"StreamCaptureGlobalList"}; +thread_local std::vector l_captureStreams; +thread_local hipStreamCaptureMode l_streamCaptureMode{hipStreamCaptureModeGlobal}; inline hipError_t ihipGraphAddNode(hipGraphNode_t graphNode, hipGraph_t graph, const hipGraphNode_t* pDependencies, size_t numDependencies) { @@ -685,9 +688,6 @@ hipError_t capturehipStreamWaitEvent(hipEvent_t& event, hipStream_t& stream, uns s->SetCaptureMode(reinterpret_cast(e->GetCaptureStream())->GetCaptureMode()); s->SetParentStream(e->GetCaptureStream()); s->SetParallelCaptureStream(stream); - } else { - assert(std::find(g_captureStreams.begin(), g_captureStreams.end(), stream) != - g_captureStreams.end() && "capturing stream should be present"); } s->AddCrossCapturedNode(e->GetNodesPrevToRecorded()); return hipSuccess; @@ -731,13 +731,12 @@ hipError_t hipThreadExchangeStreamCaptureMode(hipStreamCaptureMode* mode) { if (mode == nullptr || *mode < hipStreamCaptureModeGlobal || - *mode > hipStreamCaptureModeRelaxed || - g_captureStreams.size() == 0) { + *mode > hipStreamCaptureModeRelaxed) { HIP_RETURN(hipErrorInvalidValue); } - hipStreamCaptureMode oldMode = reinterpret_cast(g_captureStreams[0])->GetCaptureMode(); - reinterpret_cast(g_captureStreams[0])->SetCaptureMode(*mode); + auto oldMode = l_streamCaptureMode; + l_streamCaptureMode = *mode; *mode = oldMode; HIP_RETURN_DURATION(hipSuccess); @@ -752,14 +751,21 @@ hipError_t hipStreamBeginCapture(hipStream_t stream, hipStreamCaptureMode mode) // capture cannot be initiated on legacy stream // It can be initiated if the stream is not already in capture mode if (stream == nullptr || - (mode < hipStreamCaptureModeGlobal || mode > hipStreamCaptureModeRelaxed) || + mode < hipStreamCaptureModeGlobal || + mode > hipStreamCaptureModeRelaxed || s->GetCaptureStatus() == hipStreamCaptureStatusActive) { HIP_RETURN(hipErrorInvalidValue); } s->SetCaptureGraph(new ihipGraph()); s->SetCaptureMode(mode); s->SetOriginStream(); - g_captureStreams.push_back(stream); + if (mode != hipStreamCaptureModeRelaxed) { + l_captureStreams.push_back(s); + } + if (mode == hipStreamCaptureModeGlobal) { + amd::ScopedLock lock(g_captureStreamsLock); + g_captureStreams.push_back(s); + } HIP_RETURN_DURATION(hipSuccess); } @@ -775,10 +781,16 @@ hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* pGraph) { } // If mode is not hipStreamCaptureModeRelaxed, hipStreamEndCapture must be called on the stream // from the same thread - const auto& it = std::find(g_captureStreams.begin(), g_captureStreams.end(), stream); - if (s->GetCaptureMode() != hipStreamCaptureModeRelaxed && - it == g_captureStreams.end()) { - HIP_RETURN(hipErrorStreamCaptureWrongThread); + const auto& it = std::find(l_captureStreams.begin(), l_captureStreams.end(), s); + if (s->GetCaptureMode() != hipStreamCaptureModeRelaxed) { + if (it == l_captureStreams.end()) { + HIP_RETURN(hipErrorStreamCaptureWrongThread); + } + l_captureStreams.erase(it); + } + if (s->GetCaptureMode() == hipStreamCaptureModeGlobal) { + amd::ScopedLock lock(g_captureStreamsLock); + g_captureStreams.erase(std::find(g_captureStreams.begin(), g_captureStreams.end(), s)); } // If capture was invalidated, due to a violation of the rules of stream capture if (s->GetCaptureStatus() == hipStreamCaptureStatusInvalidated) { @@ -804,8 +816,6 @@ hipError_t hipStreamEndCapture(hipStream_t stream, hipGraph_t* pGraph) { } } *pGraph = s->GetCaptureGraph(); - // erase the stream and move the elements to the erased spot - g_captureStreams.erase(it); // end capture on all streams/events part of graph capture HIP_RETURN_DURATION(s->EndCapture()); } diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp index 606fa02f29..2fe5381b35 100644 --- a/hipamd/src/hip_internal.hpp +++ b/hipamd/src/hip_internal.hpp @@ -153,13 +153,19 @@ static amd::Monitor g_hipInitlock{"hipInit lock"}; } \ } while (0); -extern thread_local std::vector g_captureStreams; - #define CHECK_STREAM_CAPTURE_SUPPORTED() \ - for (auto stream : g_captureStreams) { \ - if (reinterpret_cast(stream)->GetCaptureMode() != hipStreamCaptureModeRelaxed) { \ + if (l_streamCaptureMode == hipStreamCaptureModeThreadLocal) { \ + if (l_captureStreams.size() != 0) { \ HIP_RETURN(hipErrorStreamCaptureUnsupported); \ } \ + } else if (l_streamCaptureMode == hipStreamCaptureModeGlobal) { \ + if (l_captureStreams.size() != 0) { \ + HIP_RETURN(hipErrorStreamCaptureUnsupported); \ + } \ + amd::ScopedLock lock(g_captureStreamsLock); \ + if (g_captureStreams.size() != 0) { \ + HIP_RETURN(hipErrorStreamCaptureUnsupported); \ + } \ } // Sync APIs cannot be called when stream capture is active @@ -210,7 +216,7 @@ namespace hip { hipGraph_t pCaptureGraph_; /// Based on mode stream capture places restrictions on API calls that can be made within or /// concurrently - hipStreamCaptureMode captureMode_; + hipStreamCaptureMode captureMode_{hipStreamCaptureModeGlobal}; bool originStream_; /// Origin sream has no parent. Parent stream for the derived captured streams with event /// dependencies @@ -451,4 +457,9 @@ constexpr bool kNewDevProg = false; constexpr bool kMarkerDisableFlush = true; //!< Avoids command batch flush in ROCclr +extern std::vector g_captureStreams; +extern amd::Monitor g_captureStreamsLock; +extern thread_local std::vector l_captureStreams; +extern thread_local hipStreamCaptureMode l_streamCaptureMode; + #endif // HIP_SRC_HIP_INTERNAL_H