From 9f2af2099c2c2360bc72f5df58e81406e1a5c43f Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Thu, 24 Sep 2020 14:45:18 +0800 Subject: [PATCH] feat(mgb): add enflame comp node GitOrigin-RevId: 478c8538aa890dddf4e6ca95d1c4bb8a8b49ed8e --- dnn/src/common/handle.cpp | 1 + dnn/src/common/megcore/common/computing_context.cpp | 1 + src/core/impl/comp_node/comp_node.cpp | 5 +++-- src/core/impl/comp_node/cpu/comp_node.cpp | 13 ++++++------- src/core/impl/comp_node_env.cpp | 4 ++++ src/core/impl/exception.cpp | 1 + src/core/include/megbrain/comp_node.h | 2 +- src/core/include/megbrain/comp_node_env.h | 3 +++ src/core/include/megbrain/exception.h | 5 +++++ src/core/test/comp_node.cpp | 2 ++ 10 files changed, 27 insertions(+), 10 deletions(-) diff --git a/dnn/src/common/handle.cpp b/dnn/src/common/handle.cpp index 44fb13ba3..17e7edfae 100644 --- a/dnn/src/common/handle.cpp +++ b/dnn/src/common/handle.cpp @@ -31,6 +31,7 @@ #endif + #if MEGDNN_WITH_CUDA #include "src/cuda/handle.h" #endif diff --git a/dnn/src/common/megcore/common/computing_context.cpp b/dnn/src/common/megcore/common/computing_context.cpp index f1d3a76b3..df191d75b 100644 --- a/dnn/src/common/megcore/common/computing_context.cpp +++ b/dnn/src/common/megcore/common/computing_context.cpp @@ -18,6 +18,7 @@ #endif + #if MEGDNN_WITH_ROCM #include "src/rocm/megcore/computing_context.hpp" #endif diff --git a/src/core/impl/comp_node/comp_node.cpp b/src/core/impl/comp_node/comp_node.cpp index 725eb3832..47b9d4e07 100644 --- a/src/core/impl/comp_node/comp_node.cpp +++ b/src/core/impl/comp_node/comp_node.cpp @@ -182,7 +182,8 @@ CompNode::Locator CompNode::Locator::parse(const std::string &id) { } dev_type = DeviceType::MULTITHREAD; ptr += 11; - } else { + } + else { if (ptr[1] != 'p' || ptr[2] != 'u') { err(); } @@ -237,7 +238,7 @@ CompNode::Locator CompNode::Locator::parse(const std::string &id) { //! num_steam store the nr_thread std::swap(num_dev, num_stream); } - + return {dev_type, num_dev, {num_stream}}; } diff --git a/src/core/impl/comp_node/cpu/comp_node.cpp b/src/core/impl/comp_node/cpu/comp_node.cpp index 68cbba901..d5d3bf4b4 100644 --- a/src/core/impl/comp_node/cpu/comp_node.cpp +++ b/src/core/impl/comp_node/cpu/comp_node.cpp @@ -1021,13 +1021,12 @@ void CpuCompNode::CpuDispatchableBase::EventImpl::do_device_wait_by( { auto type = cn_impl->env().property().type; - mgb_throw_if( - type != CompNode::DeviceType::CPU && - type != CompNode::DeviceType::CUDA - && type != CompNode::DeviceType::ATLAS && - type != CompNode::DeviceType::CAMBRICON, - MegBrainError, - "currently CPU can only wait for CPU, CUDA, ATLAS, CAMBRICON" + mgb_throw_if(type != CompNode::DeviceType::CPU + && type != CompNode::DeviceType::CUDA + && type != CompNode::DeviceType::ATLAS + , + MegBrainError, + "currently CPU can only wait for CPU, CUDA, ATLAS" ); } diff --git a/src/core/impl/comp_node_env.cpp b/src/core/impl/comp_node_env.cpp index f58f0f3c6..3ca26df43 100644 --- a/src/core/impl/comp_node_env.cpp +++ b/src/core/impl/comp_node_env.cpp @@ -36,6 +36,7 @@ #endif + using namespace mgb; /* =================== MegDNNHandle =================== */ @@ -232,6 +233,7 @@ void CompNodeEnv::init_cuda_async(int dev, CompNode comp_node, } #endif + #if MGB_ATLAS void mgb::_on_atlas_error(const char* expr, int err, const char* file, @@ -421,6 +423,7 @@ void CompNodeEnv::fini() { MGB_CUDA_CHECK(cudaStreamDestroy(m_cuda_env.stream)); } #endif + #if MGB_ROCM if (m_property.type == DeviceType::ROCM) { m_rocm_env.activate(); @@ -440,6 +443,7 @@ void CompNodeEnv::fini() { MGB_ATLAS_CHECK(aclrtDestroyStream(m_atlas_env.stream)); } #endif + } #if MGB_ENABLE_COMP_NODE_ASYNC_INIT diff --git a/src/core/impl/exception.cpp b/src/core/impl/exception.cpp index 2c3e42c33..7d990a180 100644 --- a/src/core/impl/exception.cpp +++ b/src/core/impl/exception.cpp @@ -73,6 +73,7 @@ std::string CudaError::get_cuda_extra_info() { #endif } + AtlasError::AtlasError(const std::string &msg): SystemError(msg) { diff --git a/src/core/include/megbrain/comp_node.h b/src/core/include/megbrain/comp_node.h index 265f8e677..647f74662 100644 --- a/src/core/include/megbrain/comp_node.h +++ b/src/core/include/megbrain/comp_node.h @@ -82,7 +82,7 @@ class CompNode { CAMBRICON = 3, ROCM = 8, ATLAS = 9, - MULTITHREAD, + MULTITHREAD = 11, MAX_DEVICE_ID, }; static constexpr size_t NR_DEVICE_TYPE = diff --git a/src/core/include/megbrain/comp_node_env.h b/src/core/include/megbrain/comp_node_env.h index 960e1b720..a850b6cf5 100644 --- a/src/core/include/megbrain/comp_node_env.h +++ b/src/core/include/megbrain/comp_node_env.h @@ -63,6 +63,7 @@ #endif //MGB_ENABLE_LOGGING #endif //MGB_CUDA + #if MGB_ATLAS #include "megcore_atlas.h" #include @@ -205,6 +206,7 @@ namespace mgb { #endif + #if MGB_ROCM [[noreturn]] void _on_hip_error(const char* expr, hipError_t err, const char* file, const char* func, int line); @@ -369,6 +371,7 @@ public: const ContinuationCtx& cont); #endif + #if MGB_ATLAS struct AtlasEnv { int device = -1; diff --git a/src/core/include/megbrain/exception.h b/src/core/include/megbrain/exception.h index d2d1eda08..a58bd8433 100644 --- a/src/core/include/megbrain/exception.h +++ b/src/core/include/megbrain/exception.h @@ -139,6 +139,11 @@ public: CudaError(const std::string& msg); }; +class EnFlameError final : public SystemError { +public: + EnFlameError(const std::string& msg); +}; + class AtlasError final: public SystemError { public: AtlasError(const std::string& msg); diff --git a/src/core/test/comp_node.cpp b/src/core/test/comp_node.cpp index 9b9b9fccb..663d80b6d 100644 --- a/src/core/test/comp_node.cpp +++ b/src/core/test/comp_node.cpp @@ -166,6 +166,7 @@ TEST(TestCompNode, Load) { ASSERT_NE(atlas0, atlas1); #endif + } TEST(TestCompNode, FreeAfterFinalize) { @@ -754,6 +755,7 @@ TEST(TestCompNodeCambricon, P2PCopy) { #endif #endif // MGB_CAMBRICON + #if MGB_ATLAS TEST(TestCompNodeAtlas, D2DCopy) { -- GitLab