Browse Source

fix(rocm): enable var_releaser for rocm

GitOrigin-RevId: a42185aa6f
release-1.4
Megvii Engine Team 4 years ago
parent
commit
17ef140cc7
2 changed files with 5 additions and 5 deletions
  1. +4
    -4
      src/core/impl/graph/var_node_mem_mgr.cpp
  2. +1
    -1
      src/core/impl/graph/var_node_mem_mgr.h

+ 4
- 4
src/core/impl/graph/var_node_mem_mgr.cpp View File

@@ -125,7 +125,7 @@ StaticDeviceMemoryManager::make_default_impl() {
#endif // MGB_THREAD_SAFE

/* ==================== AsyncVarReleaser ==================== */
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON || MGB_ROCM
class VarNodeMemManager::AsyncVarReleaser {
struct WaiterParam {
CompNode cn;
@@ -248,7 +248,7 @@ bool VarNodeMemManager::ImpureMemPlanManager::check_need_realloc() {
VarNodeMemManager::VarNodeMemManager(ComputingGraphImpl* graph)
: m_owner_graph(graph),
m_seq_mem_opt(graph)
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON || MGB_ROCM
,m_asyn_var_releaser(new AsyncVarReleaser)
#endif
{
@@ -256,7 +256,7 @@ VarNodeMemManager::VarNodeMemManager(ComputingGraphImpl* graph)
MGB_MARK_USED_VAR(ev);
// async release is only used for sync between multiple comp nodes, and
// does not wait for device to finish
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON || MGB_ROCM
m_asyn_var_releaser->wait_release_finish();
#endif
m_cpu_async_release_barrier.wait_zero();
@@ -298,7 +298,7 @@ VarNodeMemManager::VarNodeMemManager(ComputingGraphImpl* graph)
on_comp_seq_error);

#if MGB_ENABLE_VAR_DEV_MEM_DEFRAGMENTER && \
(MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON )
(MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON || MGB_ROCM)
auto on_mem_defrag_start = [this](const event::BeforeMemDefrag&) {
m_asyn_var_releaser->wait_release_finish();
};


+ 1
- 1
src/core/impl/graph/var_node_mem_mgr.h View File

@@ -446,7 +446,7 @@ class VarNodeMemManager {
SyncableCounter m_cpu_async_release_barrier;


#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON
#if MGB_CUDA || MGB_ATLAS || MGB_CAMBRICON || MGB_ROCM
//! release dynamic var on after compnode event finishes
class AsyncVarReleaser;
std::unique_ptr<AsyncVarReleaser> m_asyn_var_releaser;


Loading…
Cancel
Save