diff --git a/hipamd/src/hip_graph.cpp b/hipamd/src/hip_graph.cpp index 75d0f3d737..642a4bd76a 100644 --- a/hipamd/src/hip_graph.cpp +++ b/hipamd/src/hip_graph.cpp @@ -1215,7 +1215,8 @@ hipError_t hipGraphAddChildGraphNode(hipGraphNode_t* pGraphNode, hipGraph_t grap HIP_RETURN(status); } -hipError_t ihipGraphInstantiate(hipGraphExec_t* pGraphExec, hipGraph_t graph) { +hipError_t ihipGraphInstantiate(hipGraphExec_t* pGraphExec, hipGraph_t graph, + uint64_t flags = 0) { if (pGraphExec == nullptr || graph == nullptr) { HIP_RETURN(hipErrorInvalidValue); } @@ -1232,7 +1233,8 @@ hipError_t ihipGraphInstantiate(hipGraphExec_t* pGraphExec, hipGraph_t graph) { clonedGraph->LevelOrder(levelOrder); clonedGraph->GetUserObjs(graphExeUserObj); *pGraphExec = - new hipGraphExec(levelOrder, parallelLists, nodeWaitLists, clonedNodes, graphExeUserObj); + new hipGraphExec(levelOrder, parallelLists, nodeWaitLists, clonedNodes, + graphExeUserObj, flags); if (*pGraphExec != nullptr) { return (*pGraphExec)->Init(); } else { @@ -1247,7 +1249,7 @@ hipError_t hipGraphInstantiate(hipGraphExec_t* pGraphExec, hipGraph_t graph, } hipError_t hipGraphInstantiateWithFlags(hipGraphExec_t* pGraphExec, hipGraph_t graph, - unsigned long long flags) { + unsigned long long flags = 0) { HIP_INIT_API(hipGraphInstantiateWithFlags, pGraphExec, graph, flags); if (pGraphExec == nullptr || graph == nullptr) { HIP_RETURN(hipErrorInvalidValue); diff --git a/hipamd/src/hip_graph_internal.cpp b/hipamd/src/hip_graph_internal.cpp index 3ce74cab31..ab7fc65d92 100644 --- a/hipamd/src/hip_graph_internal.cpp +++ b/hipamd/src/hip_graph_internal.cpp @@ -796,6 +796,11 @@ hipError_t hipGraphExec::Run(hipStream_t stream) { if (queue == nullptr) { return hipErrorInvalidResourceHandle; } + if (flags_ == hipGraphInstantiateFlagAutoFreeOnLaunch) { + if (!levelOrder_.empty()) { + levelOrder_[0]->GetParentGraph()->FreeAllMemory(); + } + } UpdateQueue(parallelLists_, queue, this); std::vector rootCommands; amd::Command* endCommand = nullptr; diff --git a/hipamd/src/hip_graph_internal.hpp b/hipamd/src/hip_graph_internal.hpp index dfebf11917..536e0387fd 100644 --- a/hipamd/src/hip_graph_internal.hpp +++ b/hipamd/src/hip_graph_internal.hpp @@ -530,6 +530,10 @@ struct ihipGraph { } return false; } + + void FreeAllMemory() { + mem_pool_->FreeAllMemory(); + } }; struct hipGraphExec { @@ -544,19 +548,21 @@ struct hipGraphExec { static std::unordered_set graphExecSet_; std::unordered_set graphExeUserObj_; static amd::Monitor graphExecSetLock_; - + uint64_t flags_ = 0; public: hipGraphExec(std::vector& levelOrder, std::vector>& lists, std::unordered_map>& nodeWaitLists, std::unordered_map& clonedNodes, - std::unordered_set& userObjs) + std::unordered_set& userObjs, + uint64_t flags = 0) : parallelLists_(lists), levelOrder_(levelOrder), nodeWaitLists_(nodeWaitLists), clonedNodes_(clonedNodes), lastEnqueuedCommand_(nullptr), graphExeUserObj_(userObjs), - currentQueueIndex_(0) { + currentQueueIndex_(0), + flags_(flags) { amd::ScopedLock lock(graphExecSetLock_); graphExecSet_.insert(this); } diff --git a/hipamd/src/hip_mempool_impl.cpp b/hipamd/src/hip_mempool_impl.cpp index ddec8f49b8..2606688ff0 100644 --- a/hipamd/src/hip_mempool_impl.cpp +++ b/hipamd/src/hip_mempool_impl.cpp @@ -397,4 +397,10 @@ void MemoryPool::GetAccess(hip::Device* device, hipMemAccessFlags* flags) { } } +void MemoryPool::FreeAllMemory(hip::Stream* stream) { + while (!busy_heap_.Allocations().empty()) { + FreeMemory(busy_heap_.Allocations().begin()->first, stream); + } +} + } diff --git a/hipamd/src/hip_mempool_impl.hpp b/hipamd/src/hip_mempool_impl.hpp index e42bc7ebf7..9d176b1710 100644 --- a/hipamd/src/hip_mempool_impl.hpp +++ b/hipamd/src/hip_mempool_impl.hpp @@ -136,7 +136,7 @@ public: bool IsActiveMemory(amd::Memory* memory) const { return (allocations_.find(memory) != allocations_.end()); } - + const auto& Allocations() { return allocations_; } private: Heap() = delete; Heap(const Heap&) = delete; @@ -217,7 +217,7 @@ public: bool EventDependencies() const { return (state_.event_dependencies_) ? true : false; } bool Opportunistic() const { return (state_.opportunistic_) ? true : false; } bool InternalDependencies() const { return (state_.internal_dependencies_) ? true : false; } - + void FreeAllMemory(hip::Stream* stream = nullptr); private: MemoryPool() = delete; MemoryPool(const MemoryPool&) = delete;