Added link libraries when MPI is enabled

CMakeLists.txt

@@ -244,7 +244,12 @@ if (MAGMADNN_ENABLE_MPI)
 
   find_package(MPI REQUIRED)
 
-  message(STATUS "MPI c++ libraries: ${MPI_CXX_LIBRARIES}")
+  add_definitions(-D MAGMADNN_HAVE_MPI)
+  set(MAGMADNN_HAVE_MPI TRUE)
+
+  message(STATUS "MPI C++ libraries: ${MPI_CXX_LIBRARIES}")
+
+  set(LIBS ${LIBS} ${MPI_CXX_LIBRARIES})
 
 endif()
 

cmake/build_helpers.cmake

@@ -14,6 +14,10 @@ function(magmadnn_default_includes name)
     target_include_directories("${name}" PRIVATE ${MAGMA_INCLUDE_DIRS}) 
   endif ()
 
+  if (MAGMADNN_ENABLE_MPI)
+    target_include_directories("${name}" PRIVATE ${MPI_CXX_INCLUDE_DIRS})
+  endif ()
+
   if (MAGMADNN_ENABLE_MKLDNN)
     target_include_directories("${name}" PRIVATE ${MKLDNN_INCLUDE_DIRS}) 
   endif ()
@@ -41,6 +45,9 @@ function(magmadnn_add_example tests_driver)
   if (MAGMADNN_ENABLE_MKLDNN)
     target_include_directories("${tests_driver_name}" PRIVATE ${MKLDNN_INCLUDE_DIRS}) 
   endif ()
+  if (MAGMADNN_ENABLE_MPI)
+    target_include_directories("${tests_driver_name}" PRIVATE ${MPI_CXX_INCLUDE_DIRS})
+  endif ()
   target_link_libraries(${tests_driver_name} PRIVATE magmadnn)
   target_link_libraries(${tests_driver_name} PRIVATE ${LIBS})
 

include/magmadnn/config.h.in

@@ -18,3 +18,6 @@
 
 // Has MagmaDNN been compiled with MKLDNN 
 #cmakedefine MAGMADNN_HAVE_MKLDNN
+
+// Has MagmaDNN been compiled with MPI 
+#cmakedefine MAGMADNN_HAVE_MPI

include/magmadnn/init_finalize.h

@@ -15,6 +15,10 @@
 #include "magma.h"
 #endif
 
+#if defined(MAGMADNN_HAVE_MPI)
+#include <mpi.h>
+#endif
+
 namespace magmadnn {
 
 /** Should be called at the start of every program.

src/init_finalize.cpp

@@ -38,9 +38,9 @@ magmadnn_error_t magmadnn_init() {
     internal::MAGMADNN_SETTINGS->n_devices = 1; /* TODO : read in number of devices */
 
     int rank = 0;
-    #if defined(_HAS_MPI_)
+#if defined(MAGMADNN_HAVE_MPI)
     MPI_Comm_rank(MPI_COMM_WORLD, &rank);
-    #endif
+#endif
     int num_devices;
 
     // query number of devices