6 files changed, 24 insertions, 4 deletions
diff --git a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc
index 2e815477ee..2e9544b0b5 100644
--- a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc
+++ b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#ifdef TENSORFLOW_USE_MPI
+
 #include "tensorflow/contrib/mpi/mpi_rendezvous_mgr.h"
 
 #include <chrono>
@@ -310,3 +312,4 @@ void MPIRendezvousMgr::MPIBackgroundThread() {
 }
 
 }  // namespace tensorflow
+#endif  // TENSORFLOW_USE_MPI
diff --git a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h
index 8c908e635f..50fc380496 100644
--- a/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h
+++ b/tensorflow/contrib/mpi/mpi_rendezvous_mgr.h
@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_CONTRIB_MPI_MPI_RENDEZVOUS_MGR_H_
 #define TENSORFLOW_CONTRIB_MPI_MPI_RENDEZVOUS_MGR_H_
 
+#ifdef TENSORFLOW_USE_MPI
+
 #include <queue>
 #include <thread>
 #include <list>
@@ -254,4 +256,5 @@ class MPIRendezvousMgr : public BaseRendezvousMgr {
 };  // MPIRendezvousMgr
 }  // namespace tensorflow
 
+#endif  // TENSORFLOW_USE_MPI
 #endif  // TENSORFLOW_CONTRIB_MPI_MPI_RENDEZVOUS_MGR_H_
diff --git a/tensorflow/contrib/mpi/mpi_utils.cc b/tensorflow/contrib/mpi/mpi_utils.cc
index bea17bae41..b8e7d1a274 100644
--- a/tensorflow/contrib/mpi/mpi_utils.cc
+++ b/tensorflow/contrib/mpi/mpi_utils.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#ifdef TENSORFLOW_USE_MPI
+
 #include "tensorflow/contrib/mpi/mpi_utils.h"
 namespace tensorflow {
 
@@ -66,3 +68,5 @@ void MPIUtils::InitMPI() {
 }
 
 }  // namespace tensorflow
+
+#endif  // TENSORFLOW_USE_MPI
diff --git a/tensorflow/contrib/mpi/mpi_utils.h b/tensorflow/contrib/mpi/mpi_utils.h
index 4ee060efec..45e21f2b25 100644
--- a/tensorflow/contrib/mpi/mpi_utils.h
+++ b/tensorflow/contrib/mpi/mpi_utils.h
@@ -16,6 +16,8 @@ limitations under the License.
 #ifndef TENSORFLOW_CONTRIB_MPI_MPI_UTILS_H_
 #define TENSORFLOW_CONTRIB_MPI_MPI_UTILS_H_
 
+#ifdef TENSORFLOW_USE_MPI
+
 #include <string>
 #include <map>
 #include <vector>
@@ -54,4 +56,5 @@ class MPIUtils {
 };
 }  // namespace tensorflow
 
+#endif  // TENSORFLOW_USE_MPI
 #endif  // TENSORFLOW_CONTRIB_MPI_MPI_UTILS_H_
diff --git a/third_party/mpi/BUILD b/third_party/mpi/BUILD
index ec09b045b7..ff3f437e92 100644
--- a/third_party/mpi/BUILD
+++ b/third_party/mpi/BUILD
@@ -13,12 +13,13 @@ filegroup(
 )
 
 load("//third_party/mpi:mpi.bzl", "mpi_hdr")
+load("//third_party/mpi:mpi.bzl", "if_mpi")
 
 cc_library(
     name = "mpi",
-    srcs = select({
-        "//conditions:default": ["libmpi.so"],
-    }),
-    hdrs = mpi_hdr(),
+    srcs = if_mpi([
+        "libmpi.so",
+    ]),
+    hdrs = if_mpi(mpi_hdr()),
     visibility = ["//visibility:public"],
 )
diff --git a/third_party/mpi/mpi.bzl b/third_party/mpi/mpi.bzl
index 1f7d2f4102..38ce91c4d0 100644
--- a/third_party/mpi/mpi.bzl
+++ b/third_party/mpi/mpi.bzl
@@ -9,3 +9,9 @@ def mpi_hdr():
     else:
         hdrs = ["mpi.h",  "mpio.h", "mpicxx.h"]        #When using MVAPICH
     return hdrs
+
+def if_mpi(if_true, if_false = []):
+    return select({
+        "//tensorflow:with_mpi_support": if_true,
+        "//conditions:default": if_false
+    })