From a558d4a253f22665354321e58c0dff5e4cbe88ed Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Mon, 7 Sep 2020 18:37:19 +0800
Subject: [PATCH] fix(mgb/atlas): remove unnessary setdevice

GitOrigin-RevId: 7492de1f017c676dc74c79b55b13f3428bec7a51
---
 src/core/include/megbrain/comp_node_env.h | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/core/include/megbrain/comp_node_env.h b/src/core/include/megbrain/comp_node_env.h
index 9e46450fc..ee5817f33 100644
--- a/src/core/include/megbrain/comp_node_env.h
+++ b/src/core/include/megbrain/comp_node_env.h
@@ -45,7 +45,7 @@
 #endif //MGB_CUDA
 
 #if MGB_ATLAS
-#include "acl/acl.h"
+#include "megcore_atlas.h"
 #include <atomic>
 
 #if MGB_ENABLE_LOGGING
@@ -378,7 +378,16 @@ public:
 
         void activate() const {
             init();
-            MGB_ATLAS_CHECK(aclrtSetDevice(device));
+            int32_t device_id = -1;
+            auto err = aclrtGetDevice(&device_id);
+            if (err == ACL_ERROR_INVALID_DEVICE || device != device_id) {
+                MGB_ATLAS_CHECK(aclrtSetDevice(device));
+            } else {
+                MGB_ATLAS_CHECK(err);
+                mgb_assert(err == ACL_ERROR_NONE,
+                           "Failed to invoke aclrtGetDevice, get %s(%d)",
+                           megcore::atlas::get_error_str(err), err);
+            }
         }
     };
 
-- 
GitLab