#! /bin/sh
#
# Uniquely map GPU devices to tasks in an MPI job 
# based on the GPU assignment from Torque.
#
# Also works for serial jobs.
#
# It would be better to put this into mpiexec, though.
#
# In your Torque job script, you would do this:
#
# export PBS_GPUS=`gpus=; while read gpu ; do gpus="$gpus,$gpu"; done < $PBS_GPUFILE; echo $gpus | sed -e "s%^,%%"`
# mpiexec /path/to/this/script <your_mpi_job> <your_mpi_job_arguments> 
#

if [ -z "$PBS_GPUS" ] ; then
    echo "PBS_GPUS undefined - exiting"
    exit 1
fi

MYID=
if   [ ! -z "$PMI_ID" ] ; then
    # mvapich/mvapich2 mpirun
    MYID=$PMI_ID
elif [ ! -z "$OMPI_COMM_WORLD_LOCAL_RANK" ] ; then
    # OpenMPI
    MYID=$OMPI_COMM_WORLD_LOCAL_RANK
elif [ ! -z "$PBS_VNODENUM" ] ; then
    # OSU mpiexec - we appear to be a little "lucky" that this
    #   works for us - OSU mpiexec gives a unique PBS_VNODENUM
    #   for each task (OpenMPI does not)
    # Serial case - this works for serial, too, but it is kind
    #   of a sledgehammer, since the case is trivial - users 
    #   can continue to parse PBS_GPUFILE themselves - we add no 
    #   real value here.  But it is nice that a consistent means
    #   of extracting GPU assigment for serial and MPI cases exists
    MYID=$PBS_VNODENUM
fi

myhostname=`hostname -s`

NUMGPUS=0
for i in `echo $PBS_GPUS | sed -e "s%,% %g"`
do
  echo $i | grep -e "^${myhostname}-" >/dev/null 2>&1
  if [ $? -eq 0 ] ; then
      NUMGPUS=`expr $NUMGPUS + 1`
  fi
done

unset CUDA_VISIBLE_DEVICES
N=`expr \( $MYID % $NUMGPUS \) + 1`
count=0
for i in `echo $PBS_GPUS | sed -e "s%,% %g"`
do
  echo $i | grep -e "^${myhostname}-" >/dev/null 2>&1
  if [ $? -eq 0 ] ; then
      export CUDA_VISIBLE_DEVICES=`echo $i |sed -e "s@.*-gpu@@g"`
      count=`expr $count + 1`
      if [ $count -eq $N ] ; then
	  break
      fi
  fi
done

$*