add OpenMPI support besides MLSL (#2353)

* quick fix to add openmpi as default

* add finalize to distributed class & use unit test

* use intel mlsl github link

* apply style

* address a few comments

* fix test

* update nbench cmake

* remove extras

* fix a bug

* add counter to finalize and cleanup

* test ci error

* address mlsl ci error

* update flag names, as mentioned in pr comment

* revert back the link to mlsl repo and tag

* add flag to finalize

* apply style

* debug with info

* delete when flag is true

* add distributed setup class works, tests pass

* fix style

* remove extra instance

* disable the test due to a bug

* change flag to ompi

* remove the dependency of setting NGRAPH_DISTRIBUTED_ENABLE flag

* cleanup

* change extern to static

* remove the option NGRAPH_DISTRIBUTED_ENABLE setting this flag

* formatting

* update flags not catched by ci

* make unique pointer

* remove unused bool, fix clang error

Author: sasadep

CMakeLists.txt

@@ -111,19 +111,31 @@ option(NGRAPH_INTERPRETER_ENABLE "Control the building of the INTERPRETER backen
 option(NGRAPH_NOP_ENABLE "Control the building of the NOP backend" TRUE)
 option(NGRAPH_GPUH_ENABLE "Control the building of the Hybrid GPU backend" FALSE)
 option(NGRAPH_GENERIC_CPU_ENABLE "Enable build nGraph for generic CPU backend" FALSE)
-option(NGRAPH_DISTRIBUTED_ENABLE "Add distributed mode to the CPU backend" FALSE)
 option(NGRAPH_DEBUG_ENABLE "Enable output for NGRAPH_DEBUG statements" FALSE)
 option(NGRAPH_ONNX_IMPORT_ENABLE "Enable ONNX importer" FALSE)
 option(NGRAPH_DEX_ONLY "Build CPU DEX without codegen" FALSE)
 option(NGRAPH_CODE_COVERAGE_ENABLE "Enable code coverage data collection" FALSE)
 option(NGRAPH_LIB_VERSIONING_ENABLE "Enable shared library versioning" FALSE)
 option(NGRAPH_PYTHON_BUILD_ENABLE "Enable build nGraph python package wheel" FALSE)
+option(NGRAPH_DISTRIBUTED_MLSL_ENABLE "Add distributed MLSL mode for CPU only backend" FALSE)
+option(NGRAPH_DISTRIBUTED_OMPI_ENABLE "Add distributed Open-MPI mode for all backend" FALSE)
 option(NGRAPH_PLAIDML_ENABLE "Enable the PlaidML backend" ${PLAIDML_FOUND})
 
 if (NGRAPH_GPUH_ENABLE)
     set(NGRAPH_GPU_ENABLE TRUE)
 endif()
 
+if (NGRAPH_DISTRIBUTED_MLSL_ENABLE AND NGRAPH_DISTRIBUTED_OMPI_ENABLE)
+    message(FATAL_ERROR 
+                "Does not support the use of two distributed libraries simultaneously.\n"
+                "If CPU only backend recommend Intel MLSL by setting NGRAPH_DISTRIBUTED_MLSL_ENABLE flag to true.\n"
+                "For all other backends use OpenMPI by setting NGRAPH_DISTRIBUTED_OMPI_ENABLE flag to true.\n")
+elseif(NGRAPH_DISTRIBUTED_MLSL_ENABLE  OR NGRAPH_DISTRIBUTED_OMPI_ENABLE) 
+    set(NGRAPH_DISTRIBUTED_ENABLE TRUE)
+else()
+    set(NGRAPH_DISTRIBUTED_ENABLE FALSE)
+endif()
+
 if (NGRAPH_ONNX_IMPORT_ENABLE)
     option(NGRAPH_USE_SYSTEM_PROTOBUF "Use system provided Protobuf shared object" FALSE)
     option(NGRAPH_ONNXIFI_ENABLE "Enable ONNX Interface for Framework Integration" TRUE)
@@ -138,7 +150,8 @@ message(STATUS "NGRAPH_INTERPRETER_ENABLE:    ${NGRAPH_INTERPRETER_ENABLE}")
 message(STATUS "NGRAPH_NOP_ENABLE:            ${NGRAPH_NOP_ENABLE}")
 message(STATUS "NGRAPH_GPUH_ENABLE:           ${NGRAPH_GPUH_ENABLE}")
 message(STATUS "NGRAPH_GENERIC_CPU_ENABLE:    ${NGRAPH_GENERIC_CPU_ENABLE}")
-message(STATUS "NGRAPH_DISTRIBUTED_ENABLE:    ${NGRAPH_DISTRIBUTED_ENABLE}")
+message(STATUS "NGRAPH_DISTRIBUTED_MLSL_ENABLE:    ${NGRAPH_DISTRIBUTED_MLSL_ENABLE}")
+message(STATUS "NGRAPH_DISTRIBUTED_OMPI_ENABLE:    ${NGRAPH_DISTRIBUTED_OMPI_ENABLE}")
 message(STATUS "NGRAPH_DEBUG_ENABLE:          ${NGRAPH_DEBUG_ENABLE}")
 message(STATUS "NGRAPH_ONNX_IMPORT_ENABLE:    ${NGRAPH_ONNX_IMPORT_ENABLE}")
 message(STATUS "NGRAPH_DEX_ONLY:              ${NGRAPH_DEX_ONLY}")
@@ -260,6 +273,15 @@ if (NGRAPH_PLAIDML_ENABLE)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNGRAPH_PlaidML_ENABLE")
 endif()
 
+if (NGRAPH_DISTRIBUTED_ENABLE)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNGRAPH_DISTRIBUTED_ENABLE")
+    if (NGRAPH_DISTRIBUTED_MLSL_ENABLE)
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNGRAPH_DISTRIBUTED_MLSL_ENABLE")
+    elseif (NGRAPH_DISTRIBUTED_OMPI_ENABLE)
+        set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DNGRAPH_DISTRIBUTED_OMPI_ENABLE")
+    endif()
+endif()
+
 if (NOT DEFINED NGRAPH_TBB_ENABLE)
     set(NGRAPH_TBB_ENABLE ${NGRAPH_CPU_ENABLE})
 endif()
@@ -336,11 +358,11 @@ if (WIN32 OR APPLE)
 else()
     include(cmake/external_tbb.cmake)
 endif()
-if (NGRAPH_DISTRIBUTED_ENABLE)
+
+if (NGRAPH_DISTRIBUTED_MLSL_ENABLE)
     include(cmake/external_mlsl.cmake)
 endif()
 
-
 if (NGRAPH_HALIDE)
     message(WARNING "Halide build system integration is currently using an older LLVM release \
                      and is not expected to work across most build environments. Consider \

doc/examples/mnist_mlp/CMakeLists.txt

@@ -19,7 +19,7 @@ add_dependencies(mnist_mlp ngraph cpu_backend)
 target_link_libraries(mnist_mlp ngraph cpu_backend)
 if (NGRAPH_DISTRIBUTED_ENABLE)
     add_executable(dist_mnist_mlp mnist_loader.cpp dist_mnist_mlp.cpp)
-    target_compile_definitions(dist_mnist_mlp PRIVATE NGRAPH_DISTRIBUTED)
+    target_compile_definitions(dist_mnist_mlp PRIVATE NGRAPH_DISTRIBUTED_ENABLE)
     target_include_directories(dist_mnist_mlp SYSTEM PRIVATE libmlsl)
     target_link_libraries(dist_mnist_mlp ngraph cpu_backend libmlsl)
 endif()

src/ngraph/CMakeLists.txt

@@ -168,6 +168,7 @@ set (SRC
     placement.cpp
     cpio.cpp
     )
+
 if(NGRAPH_DISTRIBUTED_ENABLE)
     list(APPEND SRC distributed.cpp)
 endif()
@@ -178,9 +179,16 @@ add_library(ngraph SHARED ${SRC})
 
 if(NGRAPH_DISTRIBUTED_ENABLE)
     target_sources(ngraph PRIVATE distributed.cpp)
-    target_compile_definitions(ngraph PRIVATE NGRAPH_DISTRIBUTED)
-    target_include_directories(ngraph SYSTEM PRIVATE libmlsl)
-    target_link_libraries(ngraph PRIVATE libmlsl)
+    if(NGRAPH_DISTRIBUTED_MLSL_ENABLE)
+        target_include_directories(ngraph SYSTEM PRIVATE libmlsl)
+        target_link_libraries(ngraph PRIVATE libmlsl)
+    elseif(NGRAPH_DISTRIBUTED_OMPI_ENABLE)
+        find_package(MPI REQUIRED)
+        target_include_directories(ngraph SYSTEM PRIVATE ${MPI_C_INCLUDE_PATH} ${MPI_CXX_INCLUDE_PATH})
+        target_link_libraries(ngraph PRIVATE ${MPI_C_LIBRARIES} ${MPI_CXX_LIBRARIES})
+    else()
+        message(FATAL_ERROR "Distributed Library not supported/mentioned")
+    endif()
 endif()
 
 add_subdirectory(frontend)

src/ngraph/codegen/CMakeLists.txt

@@ -47,9 +47,15 @@ if ((NGRAPH_GPU_ENABLE OR NGRAPH_CPU_ENABLE) AND NOT NGRAPH_DEX_ONLY)
     list(APPEND HEADER_SEARCH_DEFINES NGRAPH_HEADERS_PATH="${NGRAPH_INCLUDE_PATH}")
 
     if(NGRAPH_DISTRIBUTED_ENABLE)
-        get_target_property(MLSL_INCLUDE_DIR libmlsl INTERFACE_INCLUDE_DIRECTORIES)
-        list(APPEND HEADER_SEARCH_DEFINES MLSL_HEADER_PATH="${MLSL_INCLUDE_DIR}")
-        add_definitions(-DNGRAPH_DISTRIBUTED)
+        if (NGRAPH_DISTRIBUTED_MLSL_ENABLE)
+            get_target_property(MLSL_INCLUDE_DIR libmlsl INTERFACE_INCLUDE_DIRECTORIES)
+            list(APPEND HEADER_SEARCH_DEFINES MLSL_HEADER_PATH="${MLSL_INCLUDE_DIR}")
+        elseif(NGRAPH_DISTRIBUTED_OMPI_ENABLE)
+            find_package(MPI REQUIRED)
+            add_definitions(-DMPI_HEADER_PATH="${MPI_PATH}")
+        else()
+            message(FATAL_ERROR "Distributed Library not supported/mentioned")
+        endif()
     endif()
 
     if(NGRAPH_GPU_ENABLE)

src/ngraph/codegen/compiler.cpp

@@ -472,8 +472,14 @@ void codegen::CompilerCore::configure_search_path()
     add_header_search_path(CUDNN_HEADER_PATHS);
 #endif
 
-#ifdef NGRAPH_DISTRIBUTED
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
     add_header_search_path(MLSL_HEADER_PATH);
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+    add_header_search_path(MPI_HEADER_PATH);
+#else
+    throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
 #endif
 }
 

src/ngraph/distributed.cpp

@@ -14,37 +14,90 @@
 // limitations under the License.
 //*****************************************************************************
 
-#ifdef NGRAPH_DISTRIBUTED
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
 
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
 #include <mlsl.hpp>
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+#include <mpi.h>
+#endif
 
 #include "ngraph/distributed.hpp"
+#include "ngraph/log.hpp"
 
 using namespace ngraph;
 
 ngraph::Distributed::Distributed()
 {
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
     if (!MLSL::Environment::GetEnv().IsInitialized())
     {
         MLSL::Environment::GetEnv().Init(nullptr, nullptr);
+        this_init_comm = true;
+    }
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+    int flag = 0;
+    MPI_Initialized(&flag);
+    if (!flag)
+    {
+        MPI_Init(NULL, NULL);
+        this_init_comm = true;
     }
+#else
+    throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
 }
 
 ngraph::Distributed::~Distributed()
 {
+    if (this_init_comm == true)
+    {
+        finalize();
+    }
+}
+
+void ngraph::Distributed::finalize()
+{
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
     if (MLSL::Environment::GetEnv().IsInitialized())
     {
         MLSL::Environment::GetEnv().Finalize();
     }
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+    int flag = 0;
+    MPI_Initialized(&flag);
+    if (flag)
+    {
+        MPI_Finalize();
+    }
+#else
+    throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
 }
 
-size_t ngraph::Distributed::get_size() const
+int ngraph::Distributed::get_size() const
 {
-    return MLSL::Environment::GetEnv().GetProcessCount();
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+    return static_cast<int>(MLSL::Environment::GetEnv().GetProcessCount());
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+    int size;
+    MPI_Comm_size(MPI_COMM_WORLD, &size);
+    return size;
+#else
+    throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
 }
 
-size_t ngraph::Distributed::get_rank() const
+int ngraph::Distributed::get_rank() const
 {
-    return MLSL::Environment::GetEnv().GetProcessIdx();
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
+    return static_cast<int>(MLSL::Environment::GetEnv().GetProcessIdx());
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+    int rank;
+    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
+    return rank;
+#else
+    throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
 }
 #endif

src/ngraph/distributed.hpp

@@ -25,7 +25,11 @@ namespace ngraph
     public:
         Distributed();
         ~Distributed();
-        size_t get_size() const;
-        size_t get_rank() const;
+        int get_size() const;
+        int get_rank() const;
+
+    private:
+        bool this_init_comm;
+        void finalize();
     };
 }

src/ngraph/runtime/cpu/CMakeLists.txt

@@ -202,9 +202,18 @@ if (NGRAPH_CPU_ENABLE)
     target_compile_definitions(cpu_backend PRIVATE CPU_BACKEND_DLL_EXPORTS)
 
     if(NGRAPH_DISTRIBUTED_ENABLE)
-        target_compile_definitions(cpu_backend PRIVATE NGRAPH_DISTRIBUTED)
-        target_include_directories(cpu_backend SYSTEM PRIVATE libmlsl)
-        target_link_libraries(cpu_backend PRIVATE libmlsl)
+        if(NGRAPH_DISTRIBUTED_MLSL_ENABLE)
+            target_include_directories(cpu_backend SYSTEM PRIVATE libmlsl)
+            target_link_libraries(cpu_backend PRIVATE libmlsl)
+        elseif(NGRAPH_DISTRIBUTED_OMPI_ENABLE)
+            find_package(MPI REQUIRED)
+            target_include_directories(cpu_backend
+                SYSTEM PRIVATE ${MPI_C_INCLUDE_PATH} ${MPI_CXX_INCLUDE_PATH})
+            target_link_libraries(cpu_backend
+                PRIVATE ${MPI_C_LIBRARIES} ${MPI_CXX_LIBRARIES})
+        else()
+            message(FATAL_ERROR "Distributed Library not supported/mentioned")
+        endif()
     endif()
 
     add_dependencies(cpu_backend ext_mkldnn ext_eigen)

src/ngraph/runtime/cpu/builder/allreduce.cpp

@@ -13,9 +13,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 //*****************************************************************************
-#ifdef NGRAPH_DISTRIBUTED
+#ifdef NGRAPH_DISTRIBUTED_ENABLE
 
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
 #include <mlsl.hpp>
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+#include <mpi.h>
+#endif
 
 #include "ngraph/op/allreduce.hpp"
 #include "ngraph/runtime/cpu/cpu_builder.hpp"
@@ -37,6 +41,8 @@ namespace ngraph
                 auto& arg_tensor = external_function->get_tensor_data(args[0].get_name());
                 auto& out_tensor = external_function->get_tensor_data(out[0].get_name());
                 auto count = static_cast<int>(out[0].get_size());
+
+#ifdef NGRAPH_DISTRIBUTED_MLSL_ENABLE
                 auto data_type = MLSL::DT_FLOAT;
 
                 if (args[0].get_element_type() == element::f32)
@@ -54,7 +60,26 @@ namespace ngraph
                         arg_tensor, out_tensor, count, data_type, MLSL::RT_SUM, MLSL::GT_DATA);
                     ctx->mlsl_env->Wait(req);
                 };
+#elif NGRAPH_DISTRIBUTED_OMPI_ENABLE
+                auto data_type = MPI_FLOAT;
 
+                if (args[0].get_element_type() == element::f32)
+                {
+                    data_type = MPI_FLOAT;
+                }
+                else if (args[0].get_element_type() == element::f64)
+                {
+                    data_type = MPI_DOUBLE;
+                }
+
+                auto functor = [&, count, data_type](CPURuntimeContext* ctx,
+                                                     CPUExecutionContext* ectx) {
+                    MPI_Allreduce(
+                        arg_tensor, out_tensor, count, data_type, MPI_SUM, MPI_COMM_WORLD);
+                };
+#else
+                throw ngraph_error("Distributed Library not supported/mentioned");
+#endif
                 functors.emplace_back(functor);
             }
 

src/ngraph/runtime/cpu/cpu_builder.cpp

@@ -103,6 +103,11 @@
 #include "ngraph/type/element_type.hpp"
 #include "ngraph/util.hpp"
 
+#ifdef NGRAPH_DISTRIBUTED_OMPI_ENABLE
+#include <mpi.h>
+#include "ngraph/op/allreduce.hpp"
+#endif
+
 using namespace std;
 using namespace ngraph;