Introduce MPI allreduce in a new contrib project.

configure

@@ -63,6 +63,7 @@ if is_windows; then
   TF_NEED_HDFS=0
   TF_NEED_JEMALLOC=0
   TF_NEED_OPENCL=0
+  TF_NEED_MPI=0
 fi
 
 while [ "$TF_NEED_JEMALLOC" == "" ]; do
@@ -112,6 +113,68 @@ else
   sed -i -e "s/WITH_GCP_SUPPORT = True/WITH_GCP_SUPPORT = False/" tensorflow/core/platform/default/build_config.bzl
 fi
 
+while [ "$TF_NEED_MPI" == "" ]; do
+  read -p "Do you wish to build TensorFlow with "\
+"Message Passing Interface (MPI) support? [y/N] " INPUT
+  case $INPUT in
+    [Yy]* ) echo "MPI support will be enabled for "\
+"TensorFlow"; TF_NEED_MPI=1;;
+    [Nn]* ) echo "No MPI support will be enabled for "\
+"TensorFlow"; TF_NEED_MPI=0;;
+    "" ) echo "No Hadoop File System support will be enabled for "\
+"TensorFlow"; TF_NEED_MPI=0;;
+    * ) echo "Invalid selection: " $INPUT;;
+  esac
+done
+
+while true; do
+  if [ "$TF_NEED_MPI" == "0" ]; then
+    break;
+  fi
+
+  fromuser=""
+  if [ -z "$MPI_PATH" ]; then
+    default_mpi_path=$(dirname $(dirname $(which mpirun)) || dirname $(dirname $(which mpiexec))  || true)
+    read -p "Please specify the location of MPI. [Default is $default_mpi_path]: " MPI_PATH
+    fromuser="1"
+    if [ -z "$MPI_PATH" ]; then
+      MPI_PATH=$default_mpi_path
+    fi
+  fi
+  if [ -e "$MPI_PATH/include" -a -e "$MPI_PATH/lib" ]; then
+    break
+  fi
+  echo "Invalid MPI path. ${MPI_PATH}/include or ${MPI_PATH}/lib cannot be found" 1>&2
+  if [ -z "$fromuser" ]; then
+    exit 1
+  fi
+  MPI_PATH=""
+  # Retry
+done
+
+if [ "$TF_NEED_MPI" == "1" ]; then
+  # Symlink necessary parts of MPI into the third party directory.
+  ln -sf "${MPI_PATH}/include/mpi.h" third_party/mpi/mpi.h
+
+  if [ -e "${MPI_PATH}/include/mpi_portable_platform.h" ]; then
+    ln -sf "${MPI_PATH}/include/mpi_portable_platform.h" third_party/mpi/mpi_portable_platform.h
+  fi
+
+  if [ -e "${MPI_PATH}/lib/libmpi.so" ]; then
+    ln -sf "${MPI_PATH}/lib/libmpi.so" third_party/mpi/libmpi.so
+  fi
+
+  if [ -e "${MPI_PATH}/lib/libmpi.dylib" ]; then
+    ln -sf "${MPI_PATH}/lib/libmpi.dylib" third_party/mpi/libmpi.dylib
+  fi
+
+  # Update Bazel build configuration.
+  sed -i -e "s/WITH_MPI_SUPPORT = False/WITH_MPI_SUPPORT = True/" tensorflow/contrib/BUILD
+else
+  # Update Bazel build configuration.
+  sed -i -e "s/WITH_MPI_SUPPORT = True/WITH_MPI_SUPPORT = False/" tensorflow/contrib/BUILD
+fi
+
 while [ "$TF_NEED_HDFS" == "" ]; do
   read -p "Do you wish to build TensorFlow with "\
 "Hadoop File System support? [y/N] " INPUT

tensorflow/contrib/BUILD

@@ -1,6 +1,9 @@
 # Description:
 #   contains parts of TensorFlow that are experimental or unstable and which are not supported.
 
+# The configure script may change this to True.
+WITH_MPI_SUPPORT = True
+
 licenses(["notice"])  # Apache 2.0
 
 exports_files(["LICENSE"])
@@ -56,7 +59,7 @@ py_library(
         "//tensorflow/contrib/tfprof",
         "//tensorflow/contrib/training:training_py",
         "//tensorflow/contrib/util:util_py",
-    ],
+    ] + ["//tensorflow/contrib/mpi:mpi_ops_py"] if WITH_MPI_SUPPORT else [],
 )
 
 cc_library(

tensorflow/contrib/mpi/BUILD

@@ -0,0 +1,66 @@
+# Ops that communicate with other processes via MPI.
+
+licenses(["notice"])  # Apache 2.0
+
+package(default_visibility = ["//tensorflow:__subpackages__"])
+
+load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
+load("//tensorflow:tensorflow.bzl", "tf_copts")
+load("//tensorflow:tensorflow.bzl", "tf_custom_op_library")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
+load("//tensorflow/core:platform/default/build_config.bzl", "tf_proto_library")
+
+tf_custom_op_library(
+    name = "mpi.so",
+    srcs = ["mpi_ops.cc", "ring.cc", "ring.h"],
+    gpu_srcs = ["ring.cu.cc", "ring.h"],
+    deps = [
+        "//third_party/mpi:mpi",
+        ":mpi_message_proto_cc",
+    ],
+)
+
+tf_py_test(
+    name = "mpi_ops_test",
+    srcs = ["mpi_ops_test.py"],
+    additional_deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/python:platform",
+    ],
+    data = [
+        ":mpi.so",
+    ],
+    tags = ["manual"],
+)
+
+py_library(
+    name = "mpi_ops_py",
+    srcs = [
+        "__init__.py",
+        "mpi_ops.py",
+    ],
+    data = [
+        ":mpi.so",
+    ],
+    srcs_version = "PY2AND3",
+    visibility = ["//visibility:public"],
+)
+
+filegroup(
+    name = "all_files",
+    srcs = glob(
+        ["**/*"],
+        exclude = [
+            "**/METADATA",
+            "**/OWNERS",
+        ],
+    ),
+    visibility = ["//tensorflow:__subpackages__"],
+)
+
+tf_proto_library(
+    name = "mpi_message_proto",
+    srcs = ["mpi_message.proto"],
+    cc_api_version = 2,
+    visibility = ["//visibility:public"],
+)

tensorflow/contrib/mpi/README.md

@@ -0,0 +1,5 @@
+# MPI TensorFlow integration
+
+Tensorflow MPI integration allows communicating between different TensorFlow
+processes using MPI. This enables training across multiple nodes and GPUs
+using high-speed interconnects.