#!/usr/bin/env ksh93
# Copyright (C) 1994-2021 Altair Engineering, Inc.
# For more information, contact Altair at www.altair.com.
#
# This file is part of both the OpenPBS software ("OpenPBS")
# and the PBS Professional ("PBS Pro") software.
#
# Open Source License Information:
#
# OpenPBS is free software. You can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the
# Free Software Foundation, either version 3 of the License, or (at your
# option) any later version.
#
# OpenPBS is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public
# License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Commercial License Information:
#
# PBS Pro is commercially licensed software that shares a common core with
# the OpenPBS software.  For a copy of the commercial license terms and
# conditions, go to: (http://www.pbspro.com/agreement.html) or contact the
# Altair Legal Department.
#
# Altair's dual-license business model allows companies, individuals, and
# organizations to create proprietary derivative works of OpenPBS and
# distribute them - whether embedded or bundled with other software -
# under a commercial license agreement.
#
# Use of Altair's trademarks, including but not limited to "PBS™",
# "OpenPBS®", "PBS Professional®", and "PBS Pro™" and Altair's logos is
# subject to Altair's trademark licensing policies.

#
#
# This script is run on each node of an MPI job by poe.  It uses the
# MP_CHILD value to figure out what adapter information to put into
# MP_MPI_NETWORK or MP_LAPI_NETWORK.  If the value of PBS_EUILIB is not
# "ip" and $PBS_HPS_JOBKEY has a value, the setup is done to use the HPS.
# Similarly, if the value of PBS_EUILIB is not "ip" and $PBS_AIXIB_JOBKEY
# has a value, the setup is done to use the InfiniBand.  The function
# extract_wins takes the window list and extracts the window information
# for the child being run.  The PBS_EUILIB variable is used instead of
# MP_EUILIB because poe may change MP_EUILIB.
#
# All arguments are passed to pbs_attach to run the MPI program.
#

if [ $# -eq 1 ] && [ $1 = "--version" ]; then
   echo pbs_version = 23.06.06
   exit 0
fi

extract_wins() {
	# Extract our list of switch windows from full list
	# arguments:
	#	number of windows per process
	#	window list
	#	split windows between MPI and LAPI

	winsperproc=$1
	winlist=$2
	split=${3:-0}

	(( start = MP_CHILD * winsperproc ))
	oldifs="$IFS"
	IFS=":"
	set -A all_windows $winlist
	set -A windows1
	set -A windows2

	i=0
	while (( i < winsperproc )) ;do
		# use windows2 to save the second half of the windows list
		# if it is being split
		if (( $split && $i >= $split )) ;then
			windows2[$i]="${all_windows[(( start + i ))]}"
		else
			windows1[$i]="${all_windows[(( start + i ))]}"
		fi
		(( i = i + 1 ))
	done
	data1="${windows1[*]}"
	data2="${windows2[*]}"
	IFS="$oldifs"
}

if [ "XX$PBS_EUILIB" != "XXip" -a "${PBS_HPS_JOBKEY:-XX}" != "XX" ]; then
	export MP_PARTITION=$PBS_HPS_JOBKEY

	extract_wins $PBS_HPS_ADAPTERS $PBS_HPS_WINDOWS
	if [ "$MP_MSG_API" = "MPI" ]; then
		unset MP_CHILD_INET_ADDR
		export MP_MPI_NETWORK=@$PBS_HPS_ADAPTERS:$data1
	fi

	if [ "$MP_MSG_API" = "LAPI" ]; then
		unset MP_LAPI_INET_ADDR
		export MP_LAPI_NETWORK=@$PBS_HPS_ADAPTERS:$data1
	fi
elif [ "XX$PBS_EUILIB" != "XXip" -a "${PBS_AIXIB_JOBKEY:-XX}" != "XX" ]; then
	export MP_PARTITION=$PBS_AIXIB_JOBKEY

	# The lapi keywords are only genereated in lower case by poe.
	# The mpi keyword will be "MPI" if the user doesn't set it and
	# all lower case if the user sets it.
	if [ "$MP_MSG_API" = "MPI" -o "$MP_MSG_API" = "mpi" ]; then
		extract_wins $PBS_AIXIB_NETWORKS $PBS_AIXIB_WINDOWS
		unset MP_CHILD_INET_ADDR
		export MP_MPI_NETWORK=@$PBS_AIXIB_NETWORKS:$data1
	elif [ "$MP_MSG_API" = "lapi" -o "$MP_MSG_API" = "mpi_lapi" ]; then
		extract_wins $PBS_AIXIB_NETWORKS $PBS_AIXIB_WINDOWS
		unset MP_LAPI_INET_ADDR
		export MP_LAPI_NETWORK=@$PBS_AIXIB_NETWORKS:$data1
	elif [ "$MP_MSG_API" = "mpi,lapi" ]; then
		(( wins = PBS_AIXIB_NETWORKS / 2 ))
		extract_wins $PBS_AIXIB_NETWORKS $PBS_AIXIB_WINDOWS $wins
		unset MP_CHILD_INET_ADDR
		unset MP_LAPI_INET_ADDR
		export MP_MPI_NETWORK=@$wins:$data1
		export MP_LAPI_NETWORK=@$wins:$data2
	fi
fi

exec pbs_attach -s -j $PBS_JOBID "$@"