blob: 38ce91c4d069fc311d5e7f17a49ff7904c9c67eb (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
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
})
|
