diff --git a/projects/clr/hipamd/src/hip_graph.cpp b/projects/clr/hipamd/src/hip_graph.cpp index 1a8c020c97..ee5b94ce23 100644 --- a/projects/clr/hipamd/src/hip_graph.cpp +++ b/projects/clr/hipamd/src/hip_graph.cpp @@ -2516,13 +2516,15 @@ hipError_t hipGraphAddMemAllocNode(hipGraphNode_t* pGraphNode, hipGraph_t graph, pNodeParams->dptr = nullptr; auto mem_alloc_node = new hip::GraphMemAllocNode(pNodeParams); hip::GraphNode* node = mem_alloc_node; - auto status = - ihipGraphAddNode(node, reinterpret_cast(graph), + auto hgraph = reinterpret_cast(graph); + auto status = ihipGraphAddNode(node, hgraph, reinterpret_cast(pDependencies), numDependencies); // The address must be provided during the node creation time pNodeParams->dptr = (HIP_MEM_POOL_USE_VM) ? mem_alloc_node->ReserveAddress() : mem_alloc_node->Execute(); *pGraphNode = reinterpret_cast(node); + amd::ScopedLock lock(hip::Graph::graphSetLock_); + hgraph->memAllocNodePtrs_.insert(pNodeParams->dptr); HIP_RETURN(status); } @@ -2572,18 +2574,26 @@ hipError_t hipGraphAddMemFreeNode(hipGraphNode_t* pGraphNode, hipGraph_t graph, dev_ptr == nullptr) { HIP_RETURN(hipErrorInvalidValue); } + // memAllocNodePtrs_ stores only local to graph alloc dptrs whose free node is not added. + // and so we need to traverse all graphs of graphSet_ and below cases are handled. + // 1) Free node cannot be added twice to the same graph + // 2) Free node if it part of another graph cannot be added to this graph hip::GraphNode* pNode; - for (auto it : hip::Graph::graphSet_) { - for (auto n : it->vertices_) { - if (n->GetType() == hipGraphNodeTypeMemFree) { - void* param; - reinterpret_cast(n)->GetParams(¶m); - if (param == dev_ptr) { - HIP_RETURN(hipErrorInvalidValue); - } + bool bGraphFound = false; + { + amd::ScopedLock lock(hip::Graph::graphSetLock_); + for (auto itGraph : hip::Graph::graphSet_) { + std::unordered_set::iterator itDevPtr = itGraph->memAllocNodePtrs_.find(dev_ptr); + if (itDevPtr != itGraph->memAllocNodePtrs_.end()) { + bGraphFound = true; + itGraph->memAllocNodePtrs_.erase(itDevPtr); + break; } } } + if (bGraphFound == false) { + HIP_RETURN(hipErrorInvalidValue); + } auto status = ihipGraphAddMemFreeNode(&pNode, reinterpret_cast(graph), diff --git a/projects/clr/hipamd/src/hip_graph_internal.hpp b/projects/clr/hipamd/src/hip_graph_internal.hpp index 2cbefb6238..bd56360e9e 100644 --- a/projects/clr/hipamd/src/hip_graph_internal.hpp +++ b/projects/clr/hipamd/src/hip_graph_internal.hpp @@ -394,7 +394,7 @@ struct Graph { hip::MemoryPool* mem_pool_; //!< Memory pool, associated with this graph std::unordered_set capturedNodes_; bool graphInstantiated_; - + std::unordered_set memAllocNodePtrs_; public: Graph(hip::Device* device, const Graph* original = nullptr) : pOriginalGraph_(original) @@ -419,7 +419,7 @@ struct Graph { if (mem_pool_ != nullptr) { mem_pool_->release(); } - + memAllocNodePtrs_.clear(); } void AddManualNodeDuringCapture(GraphNode* node) { capturedNodes_.insert(node); }