From 3bd0df8e12f3f8c4bf3e30f22b1d32dfe60f329c Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Tue, 30 Mar 2021 16:07:39 +0800
Subject: [PATCH] fix(rocm): enable var_releaser for rocm

GitOrigin-RevId: a42185aa6f3a21b79901a87facdc8c4bc0063557
---
 src/core/impl/graph/var_node_mem_mgr.cpp | 8 ++++----
 src/core/impl/graph/var_node_mem_mgr.h   | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/core/impl/graph/var_node_mem_mgr.cpp b/src/core/impl/graph/var_node_mem_mgr.cpp
index 6b77e68c2..7bb130a0b 100644
--- a/src/core/impl/graph/var_node_mem_mgr.cpp
+++ b/src/core/impl/graph/var_node_mem_mgr.cpp
@@ -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();
     };
diff --git a/src/core/impl/graph/var_node_mem_mgr.h b/src/core/impl/graph/var_node_mem_mgr.h
index 3f414b38e..be69da170 100644
--- a/src/core/impl/graph/var_node_mem_mgr.h
+++ b/src/core/impl/graph/var_node_mem_mgr.h
@@ -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;
-- 
GitLab