SWDEV-353281 - Initial support of memalloc in graph

Add memory allocation support in graph. Current implementation uses
cache from mempool  to hold the allocations which belong to the graph.
Also the resource tracking is disabled at this moment because mempool
operates with hip::Stream objects, but graph has execution with
amd::HostQueue objects.

Change-Id: I54fe3250126d24f5a26ada975f37d429bb4ef17b


[ROCm/clr commit: eef47ca24a]

projects/clr/hipamd/include/hip/amd_detail/hip_prof_str.h

@@ -3792,9 +3792,19 @@ typedef struct hip_api_data_s {
 };
 // hipGraphAddMemAllocNode[('hipGraphNode_t*', 'pGraphNode'), ('hipGraph_t', 'graph'), ('const hipGraphNode_t*', 'pDependencies'), ('size_t', 'numDependencies'), ('hipMemAllocNodeParams*', 'pNodeParams')]
 #define INIT_hipGraphAddMemAllocNode_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipGraphAddMemAllocNode.pGraphNode = (hipGraphNode_t*)pGraphNode; \
+  cb_data.args.hipGraphAddMemAllocNode.graph = (hipGraph_t)graph; \
+  cb_data.args.hipGraphAddMemAllocNode.pDependencies = (const hipGraphNode_t*)pDependencies; \
+  cb_data.args.hipGraphAddMemAllocNode.numDependencies = (size_t)numDependencies; \
+  cb_data.args.hipGraphAddMemAllocNode.pNodeParams = (hipMemAllocNodeParams*)pNodeParams; \
 };
 // hipGraphAddMemFreeNode[('hipGraphNode_t*', 'pGraphNode'), ('hipGraph_t', 'graph'), ('const hipGraphNode_t*', 'pDependencies'), ('size_t', 'numDependencies'), ('void*', 'dev_ptr')]
 #define INIT_hipGraphAddMemFreeNode_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipGraphAddMemFreeNode.pGraphNode = (hipGraphNode_t*)pGraphNode; \
+  cb_data.args.hipGraphAddMemFreeNode.graph = (hipGraph_t)graph; \
+  cb_data.args.hipGraphAddMemFreeNode.pDependencies = (const hipGraphNode_t*)pDependencies; \
+  cb_data.args.hipGraphAddMemFreeNode.numDependencies = (size_t)numDependencies; \
+  cb_data.args.hipGraphAddMemFreeNode.dev_ptr = (void*)dev_ptr; \
 };
 // hipGraphAddMemcpyNode[('hipGraphNode_t*', 'pGraphNode'), ('hipGraph_t', 'graph'), ('const hipGraphNode_t*', 'pDependencies'), ('size_t', 'numDependencies'), ('const hipMemcpy3DParms*', 'pCopyParams')]
 #define INIT_hipGraphAddMemcpyNode_CB_ARGS_DATA(cb_data) { \
@@ -4055,9 +4065,13 @@ typedef struct hip_api_data_s {
 };
 // hipGraphMemAllocNodeGetParams[('hipGraphNode_t', 'node'), ('hipMemAllocNodeParams*', 'pNodeParams')]
 #define INIT_hipGraphMemAllocNodeGetParams_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipGraphMemAllocNodeGetParams.node = (hipGraphNode_t)node; \
+  cb_data.args.hipGraphMemAllocNodeGetParams.pNodeParams = (hipMemAllocNodeParams*)pNodeParams; \
 };
 // hipGraphMemFreeNodeGetParams[('hipGraphNode_t', 'node'), ('void*', 'dev_ptr')]
 #define INIT_hipGraphMemFreeNodeGetParams_CB_ARGS_DATA(cb_data) { \
+  cb_data.args.hipGraphMemFreeNodeGetParams.node = (hipGraphNode_t)node; \
+  cb_data.args.hipGraphMemFreeNodeGetParams.dev_ptr = (void*)dev_ptr; \
 };
 // hipGraphMemcpyNodeGetParams[('hipGraphNode_t', 'node'), ('hipMemcpy3DParms*', 'pNodeParams')]
 #define INIT_hipGraphMemcpyNodeGetParams_CB_ARGS_DATA(cb_data) { \

projects/clr/hipamd/include/hip/nvidia_detail/nvidia_hip_runtime_api.h

@@ -1340,6 +1340,10 @@ typedef struct cudaHostNodeParams hipHostNodeParams;
 typedef struct cudaKernelNodeParams hipKernelNodeParams;
 typedef struct cudaMemsetParams hipMemsetParams;
 
+#if CUDA_VERSION >= CUDA_11040
+typedef struct cudaMemAllocNodeParams hipMemAllocNodeParams;
+#endif
+
 typedef enum cudaGraphExecUpdateResult hipGraphExecUpdateResult;
 #define hipGraphExecUpdateSuccess cudaGraphExecUpdateSuccess
 #define hipGraphExecUpdateError cudaGraphExecUpdateError
@@ -3208,6 +3212,30 @@ inline static hipError_t hipGraphInstantiateWithFlags(hipGraphExec_t* pGraphExec
                                                       unsigned long long flags) {
     return hipCUDAErrorTohipError(cudaGraphInstantiateWithFlags(pGraphExec, graph, flags));
 }
+
+inline hipError_t hipGraphAddMemAllocNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
+                                          const hipGraphNode_t* pDependencies,
+                                          size_t numDependencies,
+                                          hipMemAllocNodeParams* pNodeParams) {
+    return hipCUDAErrorTohipError(cudaGraphAddMemAllocNode(
+        pGraphNode, graph, pDependencies, numDependencies, pNodeParams));
+}
+
+inline hipError_t hipGraphMemAllocNodeGetParams(hipGraphNode_t node,
+                                                hipMemAllocNodeParams* pNodeParams) {
+    return hipCUDAErrorTohipError(cudaGraphMemAllocNodeGetParams(node, pNodeParams));
+}
+
+inline hipError_t hipGraphAddMemFreeNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
+                                         const hipGraphNode_t* pDependencies,
+                                         size_t numDependencies, void* dev_ptr) {
+    return hipCUDAErrorTohipError(cudaGraphAddMemFreeNode(
+        pGraphNode, graph, pDependencies, numDependencies, dev_ptr));
+}
+
+inline hipError_t hipGraphMemFreeNodeGetParams(hipGraphNode_t node, void* dev_ptr) {
+    return hipCUDAErrorTohipError(cudaGraphMemFreeNodeGetParams(node, dev_ptr));
+}
 #endif
 inline static hipError_t hipGraphLaunch(hipGraphExec_t graphExec, hipStream_t stream) {
     return hipCUDAErrorTohipError(cudaGraphLaunch(graphExec, stream));

projects/clr/hipamd/src/amdhip.def

@@ -436,3 +436,7 @@ hipGraphKernelNodeCopyAttributes
 hipGraphNodeGetEnabled
 hipGraphNodeSetEnabled
 hipGraphUpload
+hipGraphAddMemAllocNode
+hipGraphMemAllocNodeGetParams
+hipGraphAddMemFreeNode
+hipGraphMemFreeNodeGetParams

projects/clr/hipamd/src/hip_graph.cpp

@@ -18,12 +18,14 @@
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  THE SOFTWARE. */
 
+#include "top.hpp"
 #include "hip_graph_internal.hpp"
 #include "platform/command.hpp"
 #include "hip_conversions.hpp"
 #include "hip_platform.hpp"
 #include "hip_event.hpp"
-#include "top.hpp"
+#include "hip_mempool_impl.hpp"
+
 
 std::vector<hip::Stream*> g_captureStreams;
 amd::Monitor g_captureStreamsLock{"StreamCaptureGlobalList"};
@@ -832,6 +834,59 @@ hipError_t capturehipLaunchHostFunc(hipStream_t& stream, hipHostFn_t& fn, void*&
   return hipSuccess;
 }
 
+// ================================================================================================
+hipError_t capturehipMallocAsync(hipStream_t stream, hipMemPool_t mem_pool,
+                                 size_t size, void** dev_ptr) {
+  auto s = reinterpret_cast<hip::Stream*>(stream);
+  auto mpool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+
+  hipMemAllocNodeParams node_params{};
+
+  node_params.poolProps.allocType = hipMemAllocationTypePinned;
+  node_params.poolProps.location.id = mpool->Device()->deviceId();
+  node_params.poolProps.location.type = hipMemLocationTypeDevice;
+
+  std::vector<hipMemAccessDesc> descs;
+  for (const auto device : g_devices ) {
+    hipMemLocation  location{hipMemLocationTypeDevice, device->deviceId()};
+    hipMemAccessFlags flags{};
+    mpool->GetAccess(device, &flags);
+    descs.push_back({location, flags});
+  }
+
+  node_params.accessDescs = &descs[0];
+  node_params.accessDescCount = descs.size();
+  node_params.bytesize = size;
+
+  auto mem_alloc_node = new hipGraphMemAllocNode(&node_params);
+  auto status = ihipGraphAddNode(mem_alloc_node, s->GetCaptureGraph(),
+      s->GetLastCapturedNodes().data(), s->GetLastCapturedNodes().size());
+  if (status != hipSuccess) {
+    return status;
+  }
+  // Execute the node during capture, so runtime can return a valid device pointer
+  *dev_ptr = mem_alloc_node->Execute(s);
+  s->SetLastCapturedNode(mem_alloc_node);
+
+  return hipSuccess;
+}
+
+// ================================================================================================
+hipError_t capturehipFreeAsync(hipStream_t stream, void* dev_ptr) {
+  hip::Stream* s = reinterpret_cast<hip::Stream*>(stream);
+  auto mem_free_node = new hipGraphMemFreeNode(dev_ptr);
+  auto status = ihipGraphAddNode(mem_free_node, s->GetCaptureGraph(),
+      s->GetLastCapturedNodes().data(), s->GetLastCapturedNodes().size());
+  if (status != hipSuccess) {
+    return status;
+  }
+  // Execute the node during capture, so runtime can release memory into cache
+  mem_free_node->Execute(s);
+  s->SetLastCapturedNode(mem_free_node);
+  return hipSuccess;
+}
+
+// ================================================================================================
 hipError_t hipStreamIsCapturing_common(hipStream_t stream, hipStreamCaptureStatus* pCaptureStatus) {
   if (pCaptureStatus == nullptr) {
     return hipErrorInvalidValue;
@@ -893,7 +948,7 @@ hipError_t hipStreamBeginCapture_common(hipStream_t stream, hipStreamCaptureMode
     return hipErrorIllegalState;
   }
 
-  s->SetCaptureGraph(new ihipGraph());
+  s->SetCaptureGraph(new ihipGraph(s->GetDevice()));
   s->SetCaptureId();
   s->SetCaptureMode(mode);
   s->SetOriginStream();
@@ -1006,7 +1061,7 @@ hipError_t hipGraphCreate(hipGraph_t* pGraph, unsigned int flags) {
   if ((pGraph == nullptr) || (flags != 0)) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  *pGraph = new ihipGraph();
+  *pGraph = new ihipGraph(hip::getCurrentDevice());
   HIP_RETURN(hipSuccess);
 }
 
@@ -1023,7 +1078,6 @@ hipError_t hipGraphDestroy(hipGraph_t graph) {
   HIP_RETURN(hipSuccess);
 }
 
-
 hipError_t hipGraphAddKernelNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
                                  const hipGraphNode_t* pDependencies, size_t numDependencies,
                                  const hipKernelNodeParams* pNodeParams) {
@@ -2064,6 +2118,61 @@ hipError_t hipGraphExecUpdate(hipGraphExec_t hGraphExec, hipGraph_t hGraph,
   HIP_RETURN(hipSuccess);
 }
 
+// ================================================================================================
+hipError_t hipGraphAddMemAllocNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
+    const hipGraphNode_t* pDependencies, size_t numDependencies,
+    hipMemAllocNodeParams* pNodeParams) {
+  HIP_INIT_API(hipGraphAddMemAllocNode, pGraphNode, graph,
+      pDependencies, numDependencies, pNodeParams);
+  if (pGraphNode == nullptr || graph == nullptr ||
+      (numDependencies > 0 && pDependencies == nullptr) || pNodeParams == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  // Clear the pointer to allocated memory because it may contain stale/uninitialized data
+  pNodeParams->dptr = nullptr;
+  auto mem_alloc_node = new hipGraphMemAllocNode(pNodeParams);
+  *pGraphNode = mem_alloc_node;
+  auto status = ihipGraphAddNode(*pGraphNode, graph, pDependencies, numDependencies);
+  // The address must be provided during the node creation time
+  pNodeParams->dptr = mem_alloc_node->Execute();
+  HIP_RETURN(status);
+}
+
+// ================================================================================================
+hipError_t hipGraphMemAllocNodeGetParams(hipGraphNode_t node, hipMemAllocNodeParams* pNodeParams) {
+  HIP_INIT_API(hipGraphMemAllocNodeGetParams, node, pNodeParams);
+  if (node == nullptr || pNodeParams == nullptr || !hipGraphNode::isNodeValid(node)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  reinterpret_cast<hipGraphMemAllocNode*>(node)->GetParams(pNodeParams);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipGraphAddMemFreeNode(hipGraphNode_t* pGraphNode, hipGraph_t graph,
+    const hipGraphNode_t* pDependencies, size_t numDependencies, void* dev_ptr) {
+  HIP_INIT_API(hipGraphAddMemFreeNode, pGraphNode, graph, pDependencies, numDependencies, dev_ptr);
+  if (pGraphNode == nullptr || graph == nullptr ||
+      (numDependencies > 0 && pDependencies == nullptr) || dev_ptr == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto mem_free_node = new hipGraphMemFreeNode(dev_ptr);
+  *pGraphNode = mem_free_node;
+  auto status = ihipGraphAddNode(*pGraphNode, graph, pDependencies, numDependencies);
+  HIP_RETURN(status);
+}
+
+// ================================================================================================
+hipError_t hipGraphMemFreeNodeGetParams(hipGraphNode_t node, void* dev_ptr) {
+  HIP_INIT_API(hipGraphMemFreeNodeGetParams, node, dev_ptr);
+  if (node == nullptr || dev_ptr == nullptr || !hipGraphNode::isNodeValid(node)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  reinterpret_cast<hipGraphMemFreeNode*>(node)->GetParams(reinterpret_cast<void**>(dev_ptr));
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
 hipError_t hipDeviceGetGraphMemAttribute(int device, hipGraphMemAttributeType attr, void* value) {
   HIP_INIT_API(hipDeviceGetGraphMemAttribute, device, attr, value);
   if ((static_cast<size_t>(device) >= g_devices.size()) || device < 0 || value == nullptr) {

projects/clr/hipamd/src/hip_graph_capture.hpp

@@ -105,3 +105,7 @@ hipError_t capturehipEventRecord(hipStream_t& stream, hipEvent_t& event);
 hipError_t capturehipStreamWaitEvent(hipEvent_t& event, hipStream_t& stream, unsigned int& flags);
 
 hipError_t capturehipLaunchHostFunc(hipStream_t& stream, hipHostFn_t& fn, void*& userData);
+
+hipError_t capturehipMallocAsync(hipStream_t stream, hipMemPool_t mem_pool, size_t size, void** dev_ptr);
+
+hipError_t capturehipFreeAsync(hipStream_t stream, void* dev_ptr);

projects/clr/hipamd/src/hip_graph_internal.cpp

@@ -636,13 +636,8 @@ void ihipGraph::LevelOrder(std::vector<Node>& levelOrder) {
   }
 }
 
-const ihipGraph* ihipGraph::getOriginalGraph() const { return pOriginalGraph_; }
-void ihipGraph::setOriginalGraph(const ihipGraph* pOriginalGraph) {
-  pOriginalGraph_ = pOriginalGraph;
-}
-
 ihipGraph* ihipGraph::clone(std::unordered_map<Node, Node>& clonedNodes) const {
-  ihipGraph* newGraph = new ihipGraph();
+  ihipGraph* newGraph = new ihipGraph(device_, this);
   for (auto entry : vertices_) {
     hipGraphNode* node = entry->clone();
     node->SetParentGraph(newGraph);
@@ -668,7 +663,6 @@ ihipGraph* ihipGraph::clone(std::unordered_map<Node, Node>& clonedNodes) const {
     }
     clonedNodes[node]->SetDependencies(clonedDependencies);
   }
-  newGraph->setOriginalGraph(this);
   return newGraph;
 }
 

projects/clr/hipamd/src/hip_graph_internal.hpp

@@ -32,6 +32,7 @@
 #include "hip_graph_helper.hpp"
 #include "hip_event.hpp"
 #include "hip_platform.hpp"
+#include "hip_mempool_impl.hpp"
 
 typedef hipGraphNode* Node;
 hipError_t FillCommands(std::vector<std::vector<Node>>& parallelLists,
@@ -377,11 +378,26 @@ struct ihipGraph {
   std::unordered_set<hipUserObject*> graphUserObj_;
   unsigned int id_;
   static int nextID;
+  hip::Device* device_;       //!< HIP device object
+  hip::MemoryPool* mem_pool_; //!< Memory pool, associated with this graph
 
  public:
-  ihipGraph() : id_(nextID++) {
+  ihipGraph(hip::Device* device, const ihipGraph* original = nullptr)
+      : pOriginalGraph_(original)
+      , id_(nextID++)
+      , device_(device) {
     amd::ScopedLock lock(graphSetLock_);
     graphSet_.insert(this);
+    if (original == nullptr) {
+      // Create memory pool, associated with the graph
+      mem_pool_ = new hip::MemoryPool(device);
+      uint64_t max_size = std::numeric_limits<uint64_t>::max();
+      // Note: the call for the threshold is always successful
+      auto error = mem_pool_->SetAttribute(hipMemPoolAttrReleaseThreshold, &max_size);
+    } else {
+      mem_pool_ = original->mem_pool_;
+      mem_pool_->retain();
+    }
   };
 
   ~ihipGraph() {
@@ -393,6 +409,10 @@ struct ihipGraph {
     for (auto userobj : graphUserObj_) {
       userobj->release();
     }
+    if (mem_pool_ != nullptr) {
+      mem_pool_->release();
+    }
+
   };
 
   /// Return graph unique ID
@@ -417,7 +437,7 @@ struct ihipGraph {
   /// returns all the edges in the graph
   std::vector<std::pair<Node, Node>> GetEdges() const;
   // returns the original graph ptr if cloned
-  const ihipGraph* getOriginalGraph() const;
+  const ihipGraph* getOriginalGraph() const { return pOriginalGraph_; }
   // Add user obj resource to graph
   void addUserObjGraph(hipUserObject* pUserObj) {
     amd::ScopedLock lock(graphSetLock_);
@@ -432,8 +452,6 @@ struct ihipGraph {
   }
   // Delete user obj resource from graph
   void RemoveUserObjGraph(hipUserObject* pUserObj) { graphUserObj_.erase(pUserObj); }
-  // saves the original graph ptr if cloned
-  void setOriginalGraph(const ihipGraph* pOriginalGraph);
 
   void GetRunListUtil(Node v, std::unordered_map<Node, bool>& visited,
                       std::vector<Node>& singleList, std::vector<std::vector<Node>>& parallelLists,
@@ -464,6 +482,31 @@ struct ihipGraph {
       node->GenerateDOT(fout, flag);
     }
   }
+
+  void* AllocateMemory(size_t size, hip::Stream* stream, void* dptr) const {
+    auto ptr = mem_pool_->AllocateMemory(size, stream, dptr);
+    return ptr;
+  }
+
+  void FreeMemory(void* dev_ptr, hip::Stream* stream) const {
+    size_t offset = 0;
+    auto memory = getMemoryObject(dev_ptr, offset);
+    if (memory != nullptr) {
+      auto device_id = memory->getUserData().deviceId;
+      if (!g_devices[device_id]->FreeMemory(memory, stream)) {
+        LogError("Memory didn't belong to any pool!");
+      }
+    }
+  }
+
+  bool ProbeMemory(void* dev_ptr) const {
+    size_t offset = 0;
+    auto memory = getMemoryObject(dev_ptr, offset);
+    if (memory != nullptr) {
+      return mem_pool_->IsBusyMemory(memory);
+    }
+    return false;
+  }
 };
 
 struct hipGraphExec {
@@ -804,6 +847,8 @@ class hipGraphKernelNode : public hipGraphNode {
     if (status != hipSuccess) {
       ClPrint(amd::LOG_ERROR, amd::LOG_CODE, "[hipGraph] Failed to allocate memory to copy params");
     }
+    memset(&kernelAttr_, 0, sizeof(kernelAttr_));
+    kernelAttrInUse_ = 0;
     status = CopyAttr(&rhs);
     if (status != hipSuccess) {
       ClPrint(amd::LOG_ERROR, amd::LOG_CODE, "[hipGraph] Failed to during copy attrs");
@@ -1753,3 +1798,80 @@ class hipGraphEmptyNode : public hipGraphNode {
     return hipSuccess;
   }
 };
+
+class hipGraphMemAllocNode : public hipGraphNode {
+  hipMemAllocNodeParams node_params_; // Node parameters for memory allocation
+
+ public:
+  hipGraphMemAllocNode(const hipMemAllocNodeParams* node_params)
+      : hipGraphNode(hipGraphNodeTypeEmpty, "solid", "rectangle", "MEM_ALLOC") {
+        node_params_ = *node_params;
+      }
+  ~hipGraphMemAllocNode() {}
+
+  hipGraphNode* clone() const {
+    return new hipGraphMemAllocNode(static_cast<hipGraphMemAllocNode const&>(*this));
+  }
+
+  virtual hipError_t CreateCommand(amd::HostQueue* queue) {
+    auto error = hipGraphNode::CreateCommand(queue);
+    // Note: memory pool can work with hip::Streams only. It can't accept amd::HostQueue.
+    // Resource tracking is disabled!
+    auto ptr = Execute();
+    return error;
+  }
+
+  void* Execute(hip::Stream* stream = nullptr) {
+    auto graph = GetParentGraph();
+    if (graph != nullptr) {
+      // The node creation requires to return a valid address, however FreeNode can't
+      // free memory on creation because it doesn't have any execution point yet. Thus
+      // the code below makes sure memory won't be recreated on the first execution of the graph
+      if ((node_params_.dptr == nullptr) || !graph->ProbeMemory(node_params_.dptr)) {
+        auto dptr = graph->AllocateMemory(node_params_.bytesize, stream, node_params_.dptr);
+        if ((node_params_.dptr != nullptr) && (node_params_.dptr != dptr)) {
+          LogPrintfError("Ptr mismatch in graph mem alloc %p != %p", node_params_.dptr, dptr);
+        }
+        node_params_.dptr = dptr;
+      }
+    }
+    return node_params_.dptr;
+  }
+
+  void GetParams(hipMemAllocNodeParams* params) const {
+    std::memcpy(params, &node_params_, sizeof(hipMemAllocNodeParams));
+  }
+};
+
+class hipGraphMemFreeNode : public hipGraphNode {
+  void* device_ptr_;    // Device pointer of the freed memory
+
+ public:
+  hipGraphMemFreeNode(void* dptr)
+    : hipGraphNode(hipGraphNodeTypeEmpty, "solid", "rectangle", "MEM_FREE")
+    , device_ptr_(dptr) {}
+  ~hipGraphMemFreeNode() {}
+
+  hipGraphNode* clone() const {
+    return new hipGraphMemFreeNode(static_cast<hipGraphMemFreeNode const&>(*this));
+  }
+
+  virtual hipError_t CreateCommand(amd::HostQueue* queue) {
+    auto error = hipGraphNode::CreateCommand(queue);
+    // Note: memory pool can work with hip::Streams only. It can't accept amd::HostQueue.
+    // Resource tracking is disabled!
+    Execute();
+    return error;
+  }
+
+  void Execute(hip::Stream* stream = nullptr) {
+    auto graph = GetParentGraph();
+    if (graph != nullptr) {
+      graph->FreeMemory(device_ptr_, stream);
+    }
+  }
+
+  void GetParams(void** params) const {
+    *params = device_ptr_;
+  }
+};

projects/clr/hipamd/src/hip_hcc.def.in

@@ -451,3 +451,7 @@ hipGraphKernelNodeCopyAttributes
 hipGraphNodeGetEnabled
 hipGraphNodeSetEnabled
 hipGraphUpload
+hipGraphAddMemAllocNode
+hipGraphMemAllocNodeGetParams
+hipGraphAddMemFreeNode
+hipGraphMemFreeNodeGetParams

projects/clr/hipamd/src/hip_hcc.map.in

@@ -510,6 +510,10 @@ hip_5.5 {
 global:
     hipModuleLaunchCooperativeKernel;
     hipModuleLaunchCooperativeKernelMultiDevice;
+    hipGraphAddMemAllocNode;
+    hipGraphMemAllocNodeGetParams;
+    hipGraphAddMemFreeNode;
+    hipGraphMemFreeNodeGetParams;
 local:
     *;
 } hip_5.3;

projects/clr/hipamd/src/hip_memory.cpp

@@ -108,7 +108,10 @@ hipError_t ihipFree(void *ptr) {
       queue->finish();
     }
     hip::Stream::syncNonBlockingStreams(device_id);
-    amd::SvmBuffer::free(memory_object->getContext(), ptr);
+    // Find out if memory belongs to any memory pool
+    if (!g_devices[device_id]->FreeMemory(memory_object, nullptr)) {
+      amd::SvmBuffer::free(memory_object->getContext(), ptr);
+    }
     return hipSuccess;
   }
   return hipErrorInvalidValue;

projects/clr/hipamd/src/hip_mempool.cpp

@@ -74,7 +74,10 @@ hipError_t hipMallocAsync(void** dev_ptr, size_t size, hipStream_t stream) {
     reinterpret_cast<hip::Stream*>(stream);
   auto device = hip_stream->GetDevice();
   auto mem_pool = device->GetCurrentMemoryPool();
-  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(size, hip_stream);
+
+  STREAM_CAPTURE(hipMallocAsync, stream, reinterpret_cast<hipMemPool_t>(mem_pool), size, dev_ptr);
+
+  *dev_ptr = mem_pool->AllocateMemory(size, hip_stream);
   HIP_RETURN(hipSuccess);
 }
 
@@ -84,15 +87,18 @@ hipError_t hipFreeAsync(void* dev_ptr, hipStream_t stream) {
   if ((dev_ptr == nullptr) || (!hip::isValid(stream))) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+  STREAM_CAPTURE(hipFreeAsync, stream, dev_ptr);
   size_t offset = 0;
   auto memory = getMemoryObject(dev_ptr, offset);
-  auto id = memory->getUserData().deviceId;
-  auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
-    reinterpret_cast<hip::Stream*>(stream);
-  if (!g_devices[id]->FreeMemory(memory, hip_stream)) {
-    //! @todo It's not the most optimal logic. The current implementation has unconditional waits
-    HIP_RETURN(ihipFree(dev_ptr));
-   }
+  if (memory != nullptr) {
+    auto id = memory->getUserData().deviceId;
+    auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
+      reinterpret_cast<hip::Stream*>(stream);
+    if (!g_devices[id]->FreeMemory(memory, hip_stream)) {
+      //! @todo It's not the most optimal logic. The current implementation has unconditional waits
+      HIP_RETURN(ihipFree(dev_ptr));
+    }
+  }
 
   HIP_RETURN(hipSuccess);
 }
@@ -232,9 +238,12 @@ hipError_t hipMallocFromPoolAsync(
   if ((dev_ptr == nullptr) || (size == 0) || (mem_pool == nullptr) || (!hip::isValid(stream))) {
     HIP_RETURN(hipErrorInvalidValue);
   }
+  STREAM_CAPTURE(hipMallocAsync, stream, mem_pool, size, dev_ptr);
+
+  auto mpool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
   auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
     reinterpret_cast<hip::Stream*>(stream);
-  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(size, hip_stream);
+  *dev_ptr = mpool->AllocateMemory(size, hip_stream);
   HIP_RETURN(hipSuccess);
 }
 

projects/clr/hipamd/src/hip_mempool_impl.cpp

@@ -37,11 +37,13 @@ void Heap::AddMemory(amd::Memory* memory, const MemoryTimestamp& ts) {
 }
 
 // ================================================================================================
-amd::Memory* Heap::FindMemory(size_t size, hip::Stream* stream, bool opportunistic) {
+amd::Memory* Heap::FindMemory(size_t size, hip::Stream* stream, bool opportunistic, void* dptr) {
   amd::Memory* memory = nullptr;
   for (auto it = allocations_.begin(); it != allocations_.end();) {
+    bool check_address = (dptr == nullptr) || (it->first->getSvmPtr() == dptr);
     // Check if size can match and it's safe to use this resource
-    if ((it->first->getSize() >= size) && (it->second.IsSafeFind(stream, opportunistic))) {
+    if ((it->first->getSize() >= size) && check_address &&
+       (it->second.IsSafeFind(stream, opportunistic))) {
       memory = it->first;
       total_size_ -= memory->getSize();
       // Remove found allocation from the map
@@ -147,11 +149,11 @@ void Heap::SetAccess(hip::Device* device, bool enable) {
 }
 
 // ================================================================================================
-void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream) {
+void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream, void* dptr) {
   amd::ScopedLock lock(lock_pool_ops_);
 
   void* dev_ptr = nullptr;
-  amd::Memory* memory = free_heap_.FindMemory(size, stream, Opportunistic());
+  amd::Memory* memory = free_heap_.FindMemory(size, stream, Opportunistic(), dptr);
   if (memory == nullptr) {
     amd::Context* context = device_->asContext();
     const auto& dev_info = context->devices()[0]->info();
@@ -203,20 +205,26 @@ bool MemoryPool::FreeMemory(amd::Memory* memory, hip::Stream* stream) {
   amd::ScopedLock lock(lock_pool_ops_);
 
   MemoryTimestamp ts;
-  // Remove memory object fro the busy pool
+  // Remove memory object from the busy pool
   if (!busy_heap_.RemoveMemory(memory, &ts)) {
     // This pool doesn't contain memory
     return false;
   }
-  // The stream of destruction is a safe stream, because the app must handle sync
-  ts.AddSafeStream(stream);
 
-  // Add a marker to the stream to trace availability of this memory
-  Event* e = new hip::Event(0);
-  if (e != nullptr) {
-    if (hipSuccess == e->addMarker(reinterpret_cast<hipStream_t>(stream), nullptr, true)) {
-      ts.SetEvent(e);
+  if (stream != nullptr) {
+    // The stream of destruction is a safe stream, because the app must handle sync
+    ts.AddSafeStream(stream);
+
+    // Add a marker to the stream to trace availability of this memory
+    Event* e = new hip::Event(0);
+    if (e != nullptr) {
+      if (hipSuccess == e->addMarker(reinterpret_cast<hipStream_t>(stream), nullptr, true)) {
+        ts.SetEvent(e);
+      }
     }
+  } else {
+    // Assume a safe release from hipFree() if stream is nullptr
+    ts.SetEvent(nullptr);
   }
   free_heap_.AddMemory(memory, ts);
 

projects/clr/hipamd/src/hip_mempool_impl.hpp

@@ -93,7 +93,7 @@ public:
   void AddMemory(amd::Memory* memory, const MemoryTimestamp& ts);
 
   /// Finds memory object with the specified size
-  amd::Memory* FindMemory(size_t size, hip::Stream* stream, bool opportunistic);
+  amd::Memory* FindMemory(size_t size, hip::Stream* stream, bool opportunistic, void* dptr = nullptr);
 
   /// Removes allocation from the map
   bool RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts = nullptr);
@@ -132,6 +132,11 @@ public:
   std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator EraseAllocaton(
     std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator& it);
 
+  /// Checks if memory belongs to this heap
+  bool IsActiveMemory(amd::Memory* memory) const {
+    return (allocations_.find(memory) != allocations_.end());
+  }
+
 private:
   Heap() = delete;
   Heap(const Heap&) = delete;
@@ -170,11 +175,16 @@ public:
   }
 
   /// The same stream can reuse memory without HIP event validation
-  void* AllocateMemory(size_t size, hip::Stream* stream);
+  void* AllocateMemory(size_t size, hip::Stream* stream, void* dptr = nullptr);
 
   /// Frees memory by placing memory object with HIP event into free_heap_
   bool FreeMemory(amd::Memory* memory, hip::Stream* stream);
 
+  /// Check if memory is active and belongs to the busy heap
+  bool IsBusyMemory(amd::Memory* memory) const {
+    return busy_heap_.IsActiveMemory(memory);
+  }
+
   /// Releases all allocations from free_heap_. It can be called on Stream or Device synchronization
   /// @note The caller must make sure it's safe to release memory
   void ReleaseFreedMemory(hip::Stream* stream = nullptr);