From 623477fdebc5d93e5cd0026f5c19fb42cdbf421a Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 29 Oct 2021 15:30:25 +0800
Subject: [PATCH] refactor(megbrain): combine mc40 and mc20 compnode and change
 mc20 to ax

GitOrigin-RevId: bcef8c6f040b97de3b287f9cbc07cc06aea6fa17
---
 src/core/impl/comp_node/comp_node.cpp | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/core/impl/comp_node/comp_node.cpp b/src/core/impl/comp_node/comp_node.cpp
index af2d13df..7a33033f 100644
--- a/src/core/impl/comp_node/comp_node.cpp
+++ b/src/core/impl/comp_node/comp_node.cpp
@@ -154,11 +154,17 @@ CompNode::Locator CompNode::Locator::parse(const std::string& id) {
 
     // parse dev_type
             if (ptr[0] == 'a') {
-        if (strncmp(ptr, "atlas", 5)) {
-            err();
-        }
-        dev_type = DeviceType::ATLAS;
-        ptr += 5;
+#if MGB_MC20
+        if (strncmp(ptr, "ax", 2) == 0) {
+            dev_type = DeviceType::AX;
+            ptr += 2;
+        } else {
+#endif
+            if (strncmp(ptr, "atlas", 5)) {
+                err();
+            }
+            dev_type = DeviceType::ATLAS;
+            ptr += 5;
     }
     else if (ptr[0] == 'r') {
         if (strncmp(ptr, "rocm", 4)) {
@@ -173,11 +179,11 @@ CompNode::Locator CompNode::Locator::parse(const std::string& id) {
         dev_type = DeviceType::CAMBRICON;
         ptr += 9;
     } else if (ptr[0] == 'm') {
-                if (strncmp(ptr, "multithread", 11)) {
-                    err();
-                }
-                dev_type = DeviceType::MULTITHREAD;
-                ptr += 11;
+            if (strncmp(ptr, "multithread", 11)) {
+                err();
+            }
+            dev_type = DeviceType::MULTITHREAD;
+            ptr += 11;
     }
     else {
         if (ptr[1] != 'p' || ptr[2] != 'u') {