#!/bin/bash
#===============================================================================
# Copyright 2001-2023 Intel Corporation.
#
# This software and the related documents are Intel copyrighted  materials,  and
# your use of  them is  governed by the  express license  under which  they were
# provided to you (License).  Unless the License provides otherwise, you may not
# use, modify, copy, publish, distribute,  disclose or transmit this software or
# the related documents without Intel's prior written permission.
#
# This software and the related documents  are provided as  is,  with no express
# or implied  warranties,  other  than those  that are  expressly stated  in the
# License.
#===============================================================================

# For Mvapich
if [ -n "${MPIRUN_RANK}" ]
    then PMI_RANK=${MPIRUN_RANK}
fi

# For OpenMPI
if [ -n "${OMPI_COMM_WORLD_RANK}" ]
    then PMI_RANK=${OMPI_COMM_WORLD_RANK}
fi

LOCAL_RANK=$((PMI_RANK % MPI_PER_NODE))

export HPL_HOST_NODE=$((LOCAL_RANK * NUMA_PER_MPI))-$(((LOCAL_RANK + 1) * NUMA_PER_MPI - 1))

if [ -z ${USE_HPL_GPU} ]; then
    echo RANK=${PMI_RANK}, NODE=${HPL_HOST_NODE}
    ./${HPL_EXE} "$@"
else
    # Device information
    if [ ${PMI_RANK} -eq 0 ]; then
        numactl -H
        clinfo
    fi

    if [ -z ${HPL_NUMSTACK} ]; then
        export HPL_NUMSTACK=2
    fi

    DEV_START=$((HPL_NUMSTACK * HPL_NUMDEV / MPI_PER_NODE * LOCAL_RANK))
    DEV_END=$((HPL_NUMSTACK * HPL_NUMDEV / MPI_PER_NODE * (LOCAL_RANK + 1)))

    for ((d = ${DEV_START}; d < ${DEV_END}; d++)); do
        if [ $d -eq ${DEV_START} ]; then
            HPL_DEVICE=
        else
            HPL_DEVICE=${HPL_DEVICE},
        fi
        export HPL_DEVICE=${HPL_DEVICE}:$((d / HPL_NUMSTACK)).$((d % HPL_NUMSTACK))
    done

    if [ $((LOCAL_RANK * NUMA_PER_MPI)) -eq $(((LOCAL_RANK + 1) * NUMA_PER_MPI - 1)) ]; then
        NUMA_RANK=$((LOCAL_RANK * NUMA_PER_MPI))
    else
        NUMA_RANK=${HPL_HOST_NODE}
    fi

    echo RANK=${PMI_RANK}, NODE=${HPL_HOST_NODE}, HPL_DEVICE=${HPL_DEVICE}
    ./${HPL_EXE} "$@"
fi