diff --git a/mindspore/ccsrc/device/CMakeLists.txt b/mindspore/ccsrc/device/CMakeLists.txt index 951b6a1f9c..652c04d4cd 100644 --- a/mindspore/ccsrc/device/CMakeLists.txt +++ b/mindspore/ccsrc/device/CMakeLists.txt @@ -15,7 +15,6 @@ endif () if (ENABLE_CPU) file(GLOB_RECURSE CPU_SRC_LIST RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "cpu/*.cc") list(REMOVE_ITEM CPU_SRC_LIST "cpu/mpi/mpi_adapter.cc") - list(REMOVE_ITEM CPU_SRC_LIST "cpu/mpi/mpi_interface.cc") endif () if (ENABLE_MPI) @@ -26,10 +25,10 @@ if (ENABLE_MPI) add_library(mpi_adapter SHARED ${MPI_SRC_LIST}) target_link_libraries(mpi_adapter PRIVATE mindspore::ompi) - set_property(SOURCE "cpu/mpi/mpi_interface.cc" + set_property(SOURCE "gpu/mpi/mpi_initializer.cc" PROPERTY COMPILE_DEFINITIONS SUBMODULE_ID=mindspore::SubModuleId::SM_DEVICE) - pybind11_add_module(_ms_mpi "cpu/mpi/mpi_interface.cc") - target_link_libraries(_ms_mpi PRIVATE mindspore::pybind11_module mpi_adapter) + pybind11_add_module(_ms_mpi "gpu/mpi/mpi_initializer.cc") + target_link_libraries(_ms_mpi PRIVATE mindspore::pybind11_module mindspore::ompi) endif () # gpu diff --git a/mindspore/ccsrc/device/cpu/mpi/mpi_interface.cc b/mindspore/ccsrc/device/cpu/mpi/mpi_interface.cc deleted file mode 100644 index c047c9ae4e..0000000000 --- a/mindspore/ccsrc/device/cpu/mpi/mpi_interface.cc +++ /dev/null @@ -1,33 +0,0 @@ -/** - * Copyright 2020 Huawei Technologies Co., Ltd - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include -#include "device/cpu/mpi/mpi_adapter.h" - -namespace mindspore { -namespace device { -namespace cpu { -int get_rank_id() { return MPIAdapter::Instance()->GetRankId(); } - -int get_rank_size() { return MPIAdapter::Instance()->GetRankSize(); } - -PYBIND11_MODULE(_ms_mpi, mpi_interface) { - mpi_interface.doc() = "mindspore mpi python wrapper"; - mpi_interface.def("get_rank_id", &get_rank_id, "get rank id"); - mpi_interface.def("get_rank_size", &get_rank_size, "get rank size"); -} -} // namespace cpu -} // namespace device -} // namespace mindspore diff --git a/mindspore/ccsrc/device/gpu/mpi/mpi_initializer.cc b/mindspore/ccsrc/device/gpu/mpi/mpi_initializer.cc index 9baec161e5..bcad74e5b5 100644 --- a/mindspore/ccsrc/device/gpu/mpi/mpi_initializer.cc +++ b/mindspore/ccsrc/device/gpu/mpi/mpi_initializer.cc @@ -17,6 +17,7 @@ #include "device/gpu/mpi/mpi_initializer.h" #include +#include #include namespace mindspore { @@ -53,6 +54,12 @@ MPIInitializer &MPIInitializer::GetInstance() { int MPIInitializer::get_rank_id() { return MPIInitializer::GetInstance().rank_id_; } int MPIInitializer::get_rank_size() { return MPIInitializer::GetInstance().rank_size_; } + +PYBIND11_MODULE(_ms_mpi, mpi_initializer) { + mpi_initializer.doc() = "mindspore mpi python wrapper"; + mpi_initializer.def("get_rank_id", &MPIInitializer::get_rank_id, "get rank id"); + mpi_initializer.def("get_rank_size", &MPIInitializer::get_rank_size, "get rank size"); +} } // namespace gpu } // namespace device } // namespace mindspore