From ae62e15f8711eb47f3107dc93291f8708e31c93a Mon Sep 17 00:00:00 2001
From: Simon Layton <slayton58@gmail.com>
Date: Wed, 4 Jan 2017 15:06:20 -0500
Subject: [PATCH] Added MPI operators to cmake

---
 caffe2/CMakeLists.txt     |  2 +-
 caffe2/mpi/CMakeLists.txt | 28 ++++++++++++++++++++++++++++
 caffe2/mpi/mpi_python.cc  |  2 +-
 cmake/Dependencies.cmake  |  5 +++++
 4 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 caffe2/mpi/CMakeLists.txt

diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
index d825236900be..9509255fc74b 100644
--- a/caffe2/CMakeLists.txt
+++ b/caffe2/CMakeLists.txt
@@ -33,7 +33,7 @@ add_subdirectory(db)
 # add_subdirectory(experiments) # note, we may remove this folder at some point
 add_subdirectory(image)
 add_subdirectory(mkl)
-# add_subdirectory(mpi)
+add_subdirectory(mpi)
 add_subdirectory(operators)
 # add_subdirectory(proto) # note, proto is processed above already.
 # add_subdirectory(python)
diff --git a/caffe2/mpi/CMakeLists.txt b/caffe2/mpi/CMakeLists.txt
new file mode 100644
index 000000000000..9f938c7b807e
--- /dev/null
+++ b/caffe2/mpi/CMakeLists.txt
@@ -0,0 +1,28 @@
+if(USE_MPI)
+    set(Caffe2_MPI_CPU_SRC
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_common.cc"
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_ops.cc"
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_python.cc"
+    )
+    set(Caffe2_MPI_GPU_SRC
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_ops_gpu.cc"
+    )
+    set(Caffe2_MPI_CPU_TEST_SRC
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_test.cc"
+    )
+    set(Caffe2_MPI_GPU_TEST_SRC
+        "${CMAKE_CURRENT_SOURCE_DIR}/mpi_gpu_test.cc"
+    )
+
+    # pass up to calling script
+    set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} ${Caffe2_MPI_CPU_SRC})
+    set(Caffe2_CPU_SRCS ${Caffe2_CPU_SRCS} PARENT_SCOPE)
+    set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} ${Caffe2_MPI_GPU_SRC})
+    set(Caffe2_GPU_SRCS ${Caffe2_GPU_SRCS} PARENT_SCOPE)
+    set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} ${Caffe2_MPI_CPU_TEST_SRC})
+    set(Caffe2_CPU_TEST_SRCS ${Caffe2_CPU_TEST_SRCS} PARENT_SCOPE)
+    set(Caffe2_GPU_TEST_SRCS ${Caffe2_GPU_TEST_SRCS} ${Caffe2_MPI_GPU_TEST_SRC})
+    set(Caffe2_GPU_TEST_SRCS ${Caffe2_GPU_TEST_SRCS} PARENT_SCOPE)
+else()
+    message(STATUS "MPI operators skipped due to no MPI support")
+endif()
diff --git a/caffe2/mpi/mpi_python.cc b/caffe2/mpi/mpi_python.cc
index c811f2e30602..425e91bd5c96 100644
--- a/caffe2/mpi/mpi_python.cc
+++ b/caffe2/mpi/mpi_python.cc
@@ -1,7 +1,7 @@
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 
-#include "caffe2/caffe2/mpi/mpi_common.h"
+#include "caffe2/mpi/mpi_common.h"
 
 namespace caffe2 {
 
diff --git a/cmake/Dependencies.cmake b/cmake/Dependencies.cmake
index 62176ad358b8..a8b803337a60 100644
--- a/cmake/Dependencies.cmake
+++ b/cmake/Dependencies.cmake
@@ -116,6 +116,11 @@ include_directories(SYSTEM ${CMAKE_SOURCE_DIR}/third_party/pybind11/include)
 if(USE_MPI)
   find_package(MPI)
   if(MPI_CXX_FOUND)
+    message(STATUS "MPI support found")
+    message(STATUS "MPI compile flags: " ${MPI_CXX_COMPILE_FLAGS})
+    message(STATUS "MPI include path: " ${MPI_CXX_INCLUDE_PATH})
+    message(STATUS "MPI LINK flags path: " ${MPI_CXX_LINK_FLAGS})
+    message(STATUS "MPI libraries: " ${MPI_CXX_LIBRARIES})
     include_directories(SYSTEM ${MPI_CXX_INCLUDE_PATH})
     list(APPEND Caffe2_LINKER_LIBS ${MPI_CXX_LIBRARIES})
     set(CMAKE_EXE_LINKER_FLAGS ${MPI_CXX_LINK_FLAGS})