Skip to content
Snippets Groups Projects
set_core_device 1.3 KiB
Newer Older
#!/bin/bash

#set -e
#Numactl='numactl --physcpubind '
Numactl='taskset -c '

if [[ "x${SLURM_HINT}" != *nomultithread* ]]
then
else
HYP_FAC=1
fi

#NB_DEVICE=$( echo ${SLURM_STEP_GPUS//,/ } | wc -w )
NB_DEVICE=$( nvidia-smi -L | grep GPU | wc -l )
[ ${NB_DEVICE} -eq 0 ] && NB_DEVICE=1

export IP=${OMPI_COMM_WORLD_RANK:-${SLURM_PROCID}}
export LIP=${OMPI_COMM_WORLD_LOCAL_RANK:-${SLURM_LOCALID}}
export NP=${OMPI_COMM_WORLD_SIZE:-${SLURM_NTASKS}}
export NN=${OMPI_MCA_orte_num_nodes:-${SLURM_NNODES}}

export NPN=$(( NP / NN ))
#export NB_HYP=${SLURM_CPUS_ON_NODE}
export NB_HYP=${SLURM_JOB_CPUS_PER_NODE/(*)/}
export NB_CORE=$(( ${NB_HYP} / HYP_FAC ))
export NPC=$(( NB_CORE / NPN ))
CORE=$(( LIP * NPC ))

export HALF=$(( 1+ ( NPN -1 ) / 2 ))
export SOC=$(( 1 *( LIP / HALF ) ))
export RANK_SOC=$(( LIP % HALF ))
export CORE_IMP=$(( SOC + 2*NPC*RANK_SOC ))

export SURBOOK=$(( NPN / NB_DEVICE ))

if [ ${SURBOOK} == 0 ]
then
export ACC_DEVICE_NUM=$(( LIP * 2 ))
else
export ACC_DEVICE_NUM=$(( LIP / SURBOOK ))
fi

[[ "x${BIND_VERBOSE}" != "x" ]] && echo IP=${IP} LIP=${LIP} NP=${NP} NN=${NN} NPN=${NPN} NPC=${NPC} HOST=`hostname` NB_CORE=${NB_CORE} CORE=${CORE} CORE_IMP=${CORE_IMP} HALF=$HALF SOC=$SOC RS=${RANK_SOC} ND=${NB_DEVICE} CD=${ACC_DEVICE_NUM}