build _ms_mpi with mpi_initializer.cc

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

mindspore/ccsrc/device/cpu/mpi/mpi_interface.cc

@@ -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 <pybind11/operators.h>
-#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

mindspore/ccsrc/device/gpu/mpi/mpi_initializer.cc

@@ -17,6 +17,7 @@
 #include "device/gpu/mpi/mpi_initializer.h"
 
 #include <mpi.h>
+#include <pybind11/operators.h>
 #include <iostream>
 
 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