2 files changed, 42 insertions, 0 deletions
diff --git a/third_party/mpi/BUILD b/third_party/mpi/BUILD
new file mode 100644
index 0000000000..ff3f437e92
--- /dev/null
+++ b/third_party/mpi/BUILD
@@ -0,0 +1,25 @@
+licenses(["restricted"])
+
+filegroup(
+    name = "all_files",
+    srcs = glob(
+        ["**/*"],
+        exclude = [
+            "**/METADATA",
+            "**/OWNERS",
+        ],
+    ),
+    visibility = ["//tensorflow:__subpackages__"],
+)
+
+load("//third_party/mpi:mpi.bzl", "mpi_hdr")
+load("//third_party/mpi:mpi.bzl", "if_mpi")
+
+cc_library(
+    name = "mpi",
+    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
new file mode 100644
index 0000000000..38ce91c4d0
--- /dev/null
+++ b/third_party/mpi/mpi.bzl
@@ -0,0 +1,17 @@
+#OpenMPI and Mvapich/mpich require different headers
+#based on the configuration options return one or the other
+
+def mpi_hdr():
+    MPI_LIB_IS_OPENMPI=True
+    hdrs = []    
+    if MPI_LIB_IS_OPENMPI:
+        hdrs = ["mpi.h", "mpi_portable_platform.h"]   #When using OpenMPI
+    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
+    })