#!/bin/bash 

LANG="C"
export LANG

eval `tau-config`


echoIfVerbose () {
    if [ $verbose = "true" ] ; then
	echo -e "$1"
    fi
}

usage()
{
    echo ""
    echo "Usage: tau_exec [options] [--] <exe> <exe options>"
    echo ""
# Common options first
    echo "Options:"

    echo "        -v            verbose mode"
    echo "        -qsub         Use qsub mode (BG/P only, see below)"
    echo "        -io           track I/O"
    echo "        -memory       track memory allocation/deallocation"
    echo "        -memory_debug enable memory debugger"
    echo "        -cuda         track GPU events via CUDA"
    echo "        -cupti        track GPU events via CUPTI (Also see env. variable TAU_CUPTI_API)"
    echo "        -opencl       track GPU events via OpenCL"
    echo "        -armci        track ARMCI events via PARMCI"
    echo "        -ebs          enable Event-based sampling"
    echo "        -ebs_period=<count> sampling period (default 1000)"
    echo "        -ebs_source=<counter> counter (default itimer)"
    options=`tau-config --list-options`
    echo "        -T <$options> : specify TAU option"
    echo "        -loadlib=<file.so>   : specify additional load library"
    echo "        -XrunTAUsh-<options> : specify TAU library directly"
    echo "        -gdb          run program in the gdb debugger"
    echo ""
    echo "Notes:"
    echo "	Defaults if unspecified: -T MPI"
    echo "	MPI is assumed unless SERIAL is specified"
    echo ""
    echo "Example:"
    echo "    mpirun -np 2 tau_exec -io ./ring"
    echo "Example - event-based sampling with samples taken every 1,000,000 FP instructions"
    echo "    mpirun -np 8 tau_exec -ebs -ebs_period=1000000 -ebs_source=PAPI_FP_INS ./ring"
		echo "Examples - GPU:"
		echo "    tau_exec -T serial,cupti -cupti ./matmult (Preferred for CUDA 4.1 or later)"
		echo "    tau_exec -T serial -cuda ./matmult (Preferred for CUDA 4.0 or earlier)"
		echo "    tau_exec -T serial -opencl (OPENCL)"
    echo ""
    echo "qsub mode (IBM BG/P only):"
    echo "    Original:"
    echo "      qsub -n 1 --mode smp -t 10 ./a.out"
    echo "    With TAU:"
    echo "      tau_exec -qsub -io -memory -- qsub -n 1 --mode smp -t 10 ./a.out"
    echo ""
    echo "Memory Debugging:"
    echo "    -memory option:"
    echo "      Tracks heap allocation/deallocation and memory leaks."
    echo "    -memory_debug option:"
    echo "      Detects memory leaks, checks for invalid alignment, and checks for"
    echo "      array overflow.  This is exactly like setting TAU_TRACK_MEMORY_LEAKS=1"
    echo "      and TAU_MEMDBG_PROTECT_ABOVE=1 and running with -memory"
    echo ""
    exit
}

if [ $# = 0 ] ; then
    usage
fi


dryrun=false
processT=false
TauOptions=""
TauOptionsExclude=""
verbose=false
binding_specified=""
binding_options=""
track_io=false
track_memory=false
memory_debug=false
track_cuda=false
track_cupti=false
track_opencl=false
track_armci=false
tau_use_ebs=false
tau_ebs_period=""
tau_ebs_source=""
qsub_mode=false
TAU_PAPI_DEFAULT_DOMAIN=PAPI_DOM_USER
extraloadlibs=""
track_pthread=false
track_gomp=false
use_gdb=false

for arg in "$@" ; do
  # Thanks to Bernd Mohr for the following that handles quotes and spaces (see configure for explanation)
  modarg=`echo "x$arg" | sed -e 's/^x//' -e 's/"/\\\"/g' -e s,\',%@%\',g -e 's/%@%/\\\/g' -e 's/ /\\\ /g'`

  if [ "$processT" = true ] ; then
      binding_options=`echo $binding_options $arg | sed -e 's/,/ /g' | tr '[A-Z]' '[a-z]'`
      processT="false"
      shift
  else
      case $arg in 
	  -v|-d|-verbose|--verbose)
	      verbose=true
	      shift
	      ;;
	  -h|-help|--help)
	      usage
	      ;;
	  -io)
	      track_io=true
	      shift
	      ;;
	  -memory)
	      track_memory=true
	      shift
	      ;;
	  -memory_debug)
	      memory_debug=true
	      shift
	      ;;
	  -cuda)
	      track_cuda=true
	      shift
	      ;;
	  -cupti)
	      track_cupti=true
	      shift
	      ;;
	  -opencl)
	      track_opencl=true
	      shift
	      ;;
	  -gomp)
	      track_gomp=true
	      shift
	      ;;
	  -armci)
	      track_armci=true
	      shift
	      ;;
	  -ebs)
	      tau_use_ebs=true
	      shift
	      ;;
	  -ebs_period=*)
	      tau_use_ebs=true
	      myarg=`echo $arg | sed 's/-ebs_period=//'`
	      tau_ebs_period="$myarg"
	      shift
	      ;;
	  -ebs_source=*)
	      tau_use_ebs=true
	      myarg=`echo $arg | sed 's/-ebs_source=//'`
	      tau_ebs_source="$myarg"
	      shift
	      ;;
	  -qsub)
	      qsub_mode=true
	      shift
	      ;;
	  -s)
	      dryrun=true
	      shift
	      ;;
	  -gdb)
	      use_gdb=true
	      shift
	      ;;
	  -V)
	      echo '$Id: tau_exec,v 1.19 2010/06/09 18:11:25 amorris Exp $';
	      exit 0;
	      ;;
	  -T)
	      processT=true
	      shift
	      ;;
	  -tau:*)
              binding_options="$binding_options `echo $arg | sed -e 's/-tau://' -e 's/,/ /g'`"
	      shift
              ;;
	  -loadlib=*)
	      myarg=`echo $arg | sed 's/-loadlib=//'`
              extraloadlibs="$extraloadlibs:$myarg"
	      shift
              ;;
	  -XrunTAU-*)
	      myarg=`echo $arg | sed 's/-XrunTAU-//'`
	      binding_specified="shared-$myarg"
	      shift
	      ;;
	  -XrunTAUsh-*)
	      myarg=`echo $arg | sed 's/-XrunTAUsh-//'`
	      binding_specified="shared-$myarg"
	      shift
	      ;;
	  --)
	      shift
	      break
	      ;;
	  -*)
	      echo "Unknown option: $arg" >&2
	      exit 1
# First non-option signifies end of options. This would be much easier with getopt()
	      ;;
	  *)
	      break
	      ;;
      esac
  fi
done


# choose TAU library
new_binding_options=""
if [ "x$binding_options" != "x" ]; then
    for i in $binding_options ; do
      case $i in 
	  *)
	      new_binding_options="$new_binding_options $i"
	      ;;
      esac
    done
fi
binding_options="$new_binding_options"


if [ "x$binding_specified" = "x" ] ; then
    if [ "x$binding_options" = "x" ]; then
	binding_options=$DEFAULT_BINDING
    else
        # Add MPI by default
	serial=`echo $binding_options | grep serial`
	if [ $? != 0 ] ; then
           # add mpi if shmem is not specified
	    shmem=`echo $binding_options | grep shmem`
	    if [ $? != 0 ] ; then
	        binding_options="$binding_options mpi"
            fi
	fi
    fi
    theBinding=`tau-config --binding $binding_options`
    if [ $? != 0 ] ; then
	exit 1
    fi
else
    theBinding=$binding_specified
fi

if [ $verbose = "true" ] ; then
    echo ""
    echo "Program to run : $@"
    echo ""
fi


if [ `uname -s ` = Darwin ]; then
  apple=1
  TAU_SHLIBX=.dylib
else
  apple=0
  TAU_SHLIBX=.so
fi

if [ "x$LD_LIBRARY_PATH" = "x" ] ; then
  TAUEX_LD_LIBRARY_PATH=$BASEDIR/lib/$theBinding:$BASEDIR/lib
else
  TAUEX_LD_LIBRARY_PATH=$BASEDIR/lib/$theBinding:$BASEDIR/lib:$LD_LIBRARY_PATH
fi
if [ $apple = 1 ]; then
  TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU$TAU_SHLIBX
else
  TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU$TAU_SHLIBX:$LD_PRELOAD
fi


isPthread=`echo $theBinding | grep pthread`
if [ ! "x$isPthread" == "x" -a -r $BASEDIR/lib/$theBinding/libTAU-pthread$TAU_SHLIBX ]; then
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-pthread$TAU_SHLIBX:$TAUEX_LD_PRELOAD
fi

if [ $track_gomp = "true" -a -r $BASEDIR/lib/$theBinding/libTAU-gomp$TAU_SHLIBX ]; then
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-gomp$TAU_SHLIBX:$TAUEX_LD_PRELOAD
fi

if [ $track_io = "true" ] ; then
    # Add the io wrapper library to the LD_PRELOAD list
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-iowrap$TAU_SHLIBX:$TAUEX_LD_PRELOAD
    #use the auditor
    TAUEX_LD_AUDITOR=$BASEDIR/lib/$theBinding/libTAU-dl-auditor$TAU_SHLIBX
    #Needed for the auditor
    export LD_BIND_NOW=1
fi

if [ $track_memory = "true" ] || [ $memory_debug = "true" ] ; then
    # Add the memory wrapper library to the LD_PRELOAD list
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-memorywrap$TAU_SHLIBX:$TAUEX_LD_PRELOAD
    #use the auditor
    TAUEX_LD_AUDITOR=$BASEDIR/lib/$theBinding/libTAU-dl-auditor$TAU_SHLIBX
    #Needed for the auditor
    export LD_BIND_NOW=1
    # Track heap usage
    export TAU_TRACK_HEAP=1
    # Track memory leaks
    export TAU_TRACK_MEMORY_LEAKS=1
fi
if [ $track_cuda = "true" ] ; then
		cupti_exists=`test -f $BASEDIR/lib/$theBinding/libTAU-CUpti.so -a -f $BASEDIR/lib/$theBinding/libTAU-CudaQP.so`
		if [ $? == 0 ] ; then
			#tell the user cupti is available.
			echo "NOTE: CUPTI is available with your TAU configuration use '-cupti' insteed of '-cuda' to get the latest available features."
		fi
		# Add the CUDA wrapper library to the LD_PRELOAD list
		#TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CUDA.so:$TAUEX_LD_PRELOAD
		TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CUDArt.so:$TAUEX_LD_PRELOAD
		TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CudaQP.so:$TAUEX_LD_PRELOAD
		if [ `echo $TAU_METRICS | grep "TIME"` ]; then
			export TAU_METRICS=`echo $TAU_METRICS | sed -e 's/TIME/TAUGPU_TIME/'`
		else
			export TAU_METRICS=TAUGPU_TIME:$TAU_METRICS 
		fi
fi
if [ $track_cupti = "true" ] ; then
		cupti_exists=`test -f $BASEDIR/lib/$theBinding/libTAU-CUpti.so -a -f $BASEDIR/lib/$theBinding/libTAU-CudaQP.so`
		if [ $? != 0 ] ; then
			#tell the user to add cupti to bindings list.
			echo "ERROR: CUPTI library not found. Please ensure that 'cupti' is on the list of bindings specified with the '-T' option."
			exit
		else
			# Add the wrapper library to the LD_PRELOAD list
			TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CUact.so:$TAUEX_LD_PRELOAD
			TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CUpti.so:$TAUEX_LD_PRELOAD
			TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-CudaQP.so:$TAUEX_LD_PRELOAD
			if [ `echo $TAU_METRICS | grep "TIME"` ]; then
				export TAU_METRICS=`echo $TAU_METRICS | sed -e 's/TIME/TAUGPU_TIME/'`
			else
				export TAU_METRICS=TAUGPU_TIME:$TAU_METRICS 
			fi
		fi
		
		#Attempt to find the CUPTI API type:
		IFS=' ' && exe=($@)
		cupti_api_driver=`ldd $exe | grep libcuda.so`
		cupti_api_runtime=`ldd $exe | grep libcudart.so`
		if [ -z "$TAU_CUPTI_API" ]; then
			if [ -n "$cupti_api_driver" ]; then 
					export TAU_CUPTI_API=driver
			else
				if [ -n "$cupti_api_runtime" ]; then
					export TAU_CUPTI_API=runtime
				fi
			fi
		fi
		#echo "api found: $TAU_CUPTI_API"
fi
if [ $track_opencl = "true" ] ; then
    # Add the OpenCL wrapper library to the LD_PRELOAD list
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-OpenCL.so:$TAUEX_LD_PRELOAD
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-OCLci.so:$TAUEX_LD_PRELOAD
		if [ `echo $TAU_METRICS | grep "TIME"` ]; then
			export TAU_METRICS=`echo $TAU_METRICS | sed -e 's/TIME/TAUGPU_TIME/'`
		else
			export TAU_METRICS=$TAU_METRICS:TAUGPU_TIME
		fi
fi

if [ $track_armci = "true" ] ; then
    # Add the ARMCI wrapper library to the LD_PRELOAD list
    TAUEX_LD_PRELOAD=$BASEDIR/lib/$theBinding/libTAU-armciwrap.so:$TAUEX_LD_PRELOAD
fi

if [ $tau_use_ebs = "true" ] ; then
    # Set TAU_SAMPLING=1 environment variable. Deferring the capability
    #    to enable HPCToolkit until later.
    export TAU_SAMPLING=1
    if [ "x$tau_ebs_period" != "x" ] ; then
	export TAU_EBS_PERIOD=$tau_ebs_period
    fi
    if [ "x$tau_ebs_source" != "x" ] ; then
	export TAU_EBS_SOURCE=$tau_ebs_source
    fi
fi

# add libraries specified by -loadlib=<foo.so>
TAUEX_LD_PRELOAD=${TAUEX_LD_PRELOAD}${extraloadlibs}

# remove double colons
TAUEX_LD_PRELOAD=`echo $TAUEX_LD_PRELOAD | sed -e "s/::/:/g" -e "s/:$//"`

if [ $apple = 1 ]; then
  TAU_LDD='otool -L'
else
  TAU_LDD=ldd
fi


if [ $qsub_mode = false ]; then
    prog="$1"
    if [ ! -x "$prog" ] ; then
	prog=`which $prog 2>/dev/null`
    fi
    
    if [ ! -x "$prog" ] ; then
	echo "tau_exec: $1: command not found"
	exit
    fi 

    # always use the basic preload library now
    TAUEX_LD_PRELOAD=$TAUEX_LD_PRELOAD:$BASEDIR/lib/$theBinding/libTAU-preload$TAU_SHLIBX
    
fi


if [ $verbose = "true" ] ; then
    echo "Matching bindings:"
    tau-config --list-matching $binding_options
    echo ""
    echo "Using:"
    echo "$theBinding"
    echo ""
    echo "Configuration:"
    echo ""
    echo "Setting LD_LIBRARY_PATH to $TAUEX_LD_LIBRARY_PATH"
    echo "Setting LD_PRELOAD to $TAUEX_LD_PRELOAD"
    echo "Setting LD_AUDIT to $TAUEX_LD_AUDITOR"
    echo ""
fi



if [ $qsub_mode = true ] ; then

    # gather all TAU_* environment variabls, but skip TAU_OPTIONS since it often has spaces, 
    # and the ACLF staff tells us that it's impossible to pass env vars with spaces through qsub
    tau_vars=`env | grep TAU_ | grep -v TAU_OPTIONS | tr '\n' ':'`
    cmd="$@"
    # don't use the current LD_LIBRARY_PATH or it will screw things up for the backend
    TAUEX_LD_LIBRARY_PATH=$BASEDIR/lib/$theBinding
    envs="LD_PRELOAD=$TAUEX_LD_PRELOAD:LD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH:$tau_vars"

    prevEnv=""
    processEnv=false
    newCmd=""
    for arg in $cmd ; do
        # Thanks to Bernd Mohr for the following that handles quotes and spaces (see configure for explanation)
	modarg=`echo "x$arg" | sed -e 's/^x//' -e 's/"/\\\"/g' -e s,\',%@%\',g -e 's/%@%/\\\/g' -e 's/ /\\\ /g'`
	if [ "$processEnv" = true ] ; then
	    prevEnv="$arg"
	    processEnv=false
	else
	    case $arg in 
	  --env)
	      processEnv=true
	      ;;
	  *)
              newCmd="$newCmd $arg"
	      ;;
	    esac
	fi
    done

    envs="$envs:$prevEnv"
    envs=`echo $envs | sed -e "s/::/:/g" -e "s/:$//"`
    env_option="--env $envs"

    newCmd=`echo $newCmd | sed -e 's/^qsub //'`
    if [ $dryrun = "true" ]; then
	echo qsub $env_option $newCmd
    else
	qsub $env_option $newCmd
    fi
    
else

    if [ $dryrun = "true" ]; then
      if [ $track_memory = "true" ] || [ $memory_debug = "true" ] ; then
        echo "export LD_BIND_NOW=1"
      fi
      if [ $track_io = "true" ] ; then
        echo "export LD_BIND_NOW=1"
      fi
      if [ $apple = 1 ]; then
        echo "export DYLD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH"
        echo "export DYLD_INSERT_LIBRARIES=$TAUEX_LD_PRELOAD"
        echo "export DYLD_FORCE_FLAT_NAMESPACE="
      else
        echo "export LD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH"
        echo "export LD_AUDIT=$TAUEX_LD_AUDITOR"
        echo "export LD_PRELOAD=$TAUEX_LD_PRELOAD"
      fi
      if [ $memory_debug = "true" ] ; then
        echo "export TAU_TRACK_MEMORY_LEAKS=1"
        echo "export TAU_MEMDBG_PROTECT_ABOVE=1"
      fi
      echo "$@"
      exit 0;
    elif [ $use_gdb = "true" ]; then
      echo "" > .gdb_commands
      if [ $track_memory = "true" ] || [ $memory_debug = "true" ] ; then
        echo "set env LD_BIND_NOW=1" >> .gdb_commands
      fi
      if [ $track_io = "true" ] ; then
        echo "set env LD_BIND_NOW=1" >> .gdb_commands 
      fi
      if [ $apple = 1 ]; then
        echo "set env DYLD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH" >> .gdb_commands 
        echo "set env DYLD_INSERT_LIBRARIES=$TAUEX_LD_PRELOAD" >> .gdb_commands    
        echo "set env DYLD_FORCE_FLAT_NAMESPACE=" >> .gdb_commands  
      else
        echo "set env LD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH" >> .gdb_commands  
        echo "set env LD_AUDIT=$TAUEX_LD_AUDITOR" >> .gdb_commands  
        echo "set env LD_PRELOAD=$TAUEX_LD_PRELOAD" >> .gdb_commands   
      fi
      if [ $memory_debug = "true" ] ; then
        echo "set env TAU_TRACK_MEMORY_LEAKS=1" >> .gdb_commands 
        echo "set env TAU_MEMDBG_PROTECT_ABOVE=1">> .gdb_commands 
      fi
      echo "set arg ${*:2}" >> .gdb_commands
      gdb -x .gdb_commands "$1"
      rm -f ./.gdb_commands
      exit 0;
    elif [ $apple = 1 ]; then
      export DYLD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH
      export DYLD_INSERT_LIBRARIES=$TAUEX_LD_PRELOAD
      export DYLD_FORCE_FLAT_NAMESPACE=""
    else
      export LD_LIBRARY_PATH=$TAUEX_LD_LIBRARY_PATH
      export LD_AUDIT=$TAUEX_LD_AUDITOR
      export LD_PRELOAD=$TAUEX_LD_PRELOAD
    fi
    if [ $memory_debug = "true" ] ; then
      export TAU_TRACK_MEMORY_LEAKS=1
      export TAU_MEMDBG_PROTECT_ABOVE=1
    fi
    
    $@
    retval=$?
    unset LD_PRELOAD
    exit $retval
fi

