SWDEV-457316 - Return invalid val from add mem free node if corresponding mem alloc node is missing.

Change-Id: Ib0c346a439fc38ebfd106bcbdf75bd10bfd2f090


[ROCm/clr commit: cf0320a0b9]

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<hip::Graph*>(graph),
+  auto hgraph = reinterpret_cast<hip::Graph*>(graph);
+  auto status = ihipGraphAddNode(node, hgraph,
                        reinterpret_cast<hip::GraphNode* const*>(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<hipGraphNode_t>(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<hip::GraphMemFreeNode*>(n)->GetParams(&param);
-        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<void*>::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<hip::Graph*>(graph),

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<GraphNode*> capturedNodes_;
   bool graphInstantiated_;
-
+  std::unordered_set<void*> 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); }