multi-gpu-proxy.sh

#!/bin/bash
# This file sets the environment variable CUDA_VISIBLE_DEVICES to the MPI local rank to enable multi-GPU usage of this benchmark. Note that this disables any GPU distribution handling by he batch scheduler.
# Background: Most/some batch schedulers set CUDA_VISIBLE_DEVICES to all available GPUs on a node. In that case, the Arbor benchmark would only use the first entry in the list, probably GPU#0. This script changes that.
# -Andreas Herten, Nov 2018

_verbose=1

localrank=$CUDA_VISIBLE_DEVICES

if [[ -n "$OMPI_COMM_WORLD_NODE_RANK" ]]; then
	# OPENMPI
	if [[ "$_verbose" == "1" ]]; then echo "OMPI_COMM_WORLD_NODE_RANK is set to $OMPI_COMM_WORLD_NODE_RANK"; fi
	localrank=$OMPI_COMM_WORLD_NODE_RANK
elif [[ -n "$MV2_COMM_WORLD_LOCAL_RANK" ]]; then
	# MVAPICH2
	if [[ "$_verbose" == "1" ]]; then echo "MV2_COMM_WORLD_LOCAL_RANK is set to $MV2_COMM_WORLD_LOCAL_RANK"; fi
	localrank=$MV2_COMM_WORLD_LOCAL_RANK
elif [[ -n "$MPI_LOCALRANKID" ]]; then
	# MVAPICH2 (also)
	if [[ "$_verbose" == "1" ]]; then echo "MPI_LOCALRANKID is set to $MPI_LOCALRANKID"
	localrank=$MPI_LOCALRANKID; fi
fi

export CUDA_VISIBLE_DEVICES=$localrank
if [[ "$_verbose" == "1" ]]; then echo "CUDA_VISIBLE_DEVICES is set to $CUDA_VISIBLE_DEVICES"; fi
$@