#!/bin/bash
#
# hipfc: Wrapper to call fortran compiler with the hipfort interface
#
PROGVERSION=0.2-0

# Copyright (c) 2020 Advanced Micro Devices, Inc. All rights reserved.
# [MITx11 License]

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

function usage(){
/bin/cat 2>&1 <<"EOF" 

   hipfc: Wrapper to call fortran compiler with hipfort
          This script also call hipcc for non fortran files.

   Usage:  hipfc [ options ] input-files

   Options without values:
    -g          Generate debug information
    -version    Display version of hipfc then exit
    -v          Verbose, just print commands
    -vv         Very verbose, pass -v to commands
    -n          Dryrun, do nothing, show commands that would execute
    -h          Print this help message
    -k          Keep temporary files
    -c          Compile to object code only
    -lrocsparse Link to rocsparse library
    -lrocblas   Link to rocblas library
    -lrocfft    Link to rocfft library
    -lhipblas   Link to hipblas library

   Options with values:       Defaults  
    -hipfort-compiler <bin>   $HIPFORT_COMPILER or /usr/bin/gfortran
    -hipfort          <path>  $HIPFORT   or /opt/rocm-4.0.0/hipfort
    -cuda-path        <path>  $CUDA_PATH or /usr/local/cuda
    -rocm-dir         <path>  $ROCM_DIR  or /opt/rocm
    -I         <include dir>  Provide one directory per -I option
    -O         <LLVM opt>     LLVM optimization level
    -o         <outfilename>  Default=a.out
    -t         <tdir>         Temporary directory for intermediate files
                              Default=/tmp/hipfc-tmp-$$
    --offload-arch <gputype>  Default=<value returned by mygpu utility>

   Examples:
    hipfc myapp.f -o myapp             /* Compile FORTRAN to create myapp  */
    hipfc myapp.f -o myapp -lrocblas   /* Compile FORTRAN with rocblas     */
    hipfc myapp.f kernels.cpp -o myapp /* Compile FORTRAN and c++ and link */

   Instead of these command line options:
        -hipfort, -cuda-path, -hipfort-compiler -rocm-dir
      you may alternatively set these environment variables, respectively:
        HIPFORT,  CUDA_PATH,  HIPFORT_COMPILER, ROCM_DIR
      Command line options take precedence over environment variables. 

   Copyright (c) 2020 ADVANCED MICRO DEVICES, INC.

EOF
   exit 0 
}

DEADRC=12

#  Utility Functions
function do_err(){
   if [ $NEWTMPDIR ] ; then 
      if [ $KEEPTDIR ] ; then 
         cp -rp $TMPDIR $OUTDIR
         [ $VERBOSE ] && echo "#Info:  Temp files copied to $OUTDIR/$TMPNAME"
      fi
      rm -rf $TMPDIR
   else 
      if [ $KEEPTDIR ] ; then 
         [ $VERBOSE ] && echo "#Info:  Temp files kept in $TMPDIR"
      fi 
   fi
   [ $VV ] && echo "#Info:  Done"
   exit $1
}

function version(){
   echo $PROGVERSION
   exit 0
}

function runcmd(){
   THISCMD=$1
   if [ $DRYRUN ] ; then
      echo "$THISCMD"
   else 
      [ $VERBOSE ] && echo "$THISCMD"
      $THISCMD
      rc=$?
      if [ $rc != 0 ] ; then 
         echo "ERROR:  The following command failed with return code $rc."
         echo "        $THISCMD"
         do_err $rc
      fi
   fi
}

function getdname(){
   local __DIRN=`dirname "$1"`
   if [ "$__DIRN" = "." ] ; then 
      __DIRN=$PWD; 
   else
      if [ ${__DIRN:0:1} != "/" ] ; then 
         if [ ${__DIRN:0:2} == ".." ] ; then 
               __DIRN=`dirname $PWD`/${__DIRN:3}
         else
            if [ ${__DIRN:0:1} = "." ] ; then 
               __DIRN=$PWD/${__DIRN:2}
            else
               __DIRN=$PWD/$__DIRN
            fi
         fi
      fi
   fi
   echo $__DIRN
}

#  --------  The main code starts here -----
INCLUDES=""
PASSTHRUARGS=""
INPUTFILES=""
#  Argument processing
while [ $# -gt 0 ] ; do 
   case "$1" in 
      -q)               QUIET=true;;
      --quiet)          QUIET=true;;
      -k) 		KEEPTDIR=true;; 
      -n) 		DRYRUN=true;; 
      -c) 		GEN_OBJECT_ONLY=true;; 
      -g) 		GEN_DEBUG=true;; 
      -noshared) 	NOSHARED=true;;
      -cuopts) 		CUOPTS=$2; shift ;; 
      -I) 		INCLUDES="$INCLUDES -I $2"; shift ;; 
      -O) 		FORTOPT=$2; shift ;; 
      -O3) 		FORTOPT=3 ;; 
      -O2) 		FORTOPT=2 ;; 
      -O1) 		FORTOPT=1 ;; 
      -O0) 		FORTOPT=0 ;; 
      -o) 		OUTFILE=$2; shift ;; 
      -t)		TMPDIR=$2; shift ;; 
      -lrocsparse)      ROCSPARSE=true ; PASSTHRUARGS+=" $1" ;;
      -lrocblas)        ROCBLAS=true ; PASSTHRUARGS+=" $1" ;;
      -lrocfft)         ROCFFT=true ; PASSTHRUARGS+=" $1" ;;
      -lhipblas)        HIPBLAS=true ; PASSTHRUARGS+=" $1" ;;
      --offload-arch)   HIPFORT_GPU=$2; shift ;;
      -triple)          TARGET_TRIPLE=$2; shift ;;
      -cuda-path)       CUDA_PATH=$2; shift ;;
      -rocm-dir)        ROCM_DIR=$2; shift ;;
      -hipfort)         HIPFORT=$2; shift ;;
      -hipfort-compiler) HIPFORT_COMPILER=$2; shift ;;
      -h) 	        usage ;;
      -help) 	        usage ;;
      --help) 	        usage ;;
      -version) 	version ;;
      --version) 	version ;;
      -v) 		VERBOSE=true;;
      -vv) 		VV=true;;
      --) 		shift ;;
      *)
        dash=${1:0:1}
	if [ $dash == "-" ] ; then
	   PASSTHRUARGS+=" $1"
        else
	   INPUTFILES+=" $1"
        fi
   esac
   shift
done

fcount=0
for __input_file in `echo $INPUTFILES` ; do
   fcount=$(( fcount + 1 ))
   if [ $fcount == 1 ] ; then
      FIRST_INPUT_FILE_NAME=$__input_file
   fi
   if [ ! -e "$__input_file" ] ; then
      echo "ERROR:  The file $__input_file does not exist."
      exit $DEADRC
   fi
done
if [ -z "$FIRST_INPUT_FILE_NAME" ]  ; then
   echo "ERROR:  No File specified."
   exit $DEADRC
fi

cdir=$(getdname $0)
[ ! -L "$cdir/hipfc" ] || cdir=$(getdname `readlink "$cdir/hipfc"`)
HOW_CALLED=${0##*/}

CUDA_PATH=${CUDA_PATH:-/usr/local/cuda}
ROCM_DIR=${ROCM_DIR:-/opt/rocm}

HIPFORT=${HIPFORT:-/opt/rocm-4.0.0/hipfort}
if [ ! -d $HIPFORT ] ; then
   HIPFORT="$ROCM_DIR/hipfort"
fi
if [ ! -d $HIPFORT ] ; then
   HIPFORT="/tmp/hipfort"
fi
if [ ! -d $HIPFORT ] ; then
   echo "ERROR: HIPFORT installation not found at $HIPFORT"
   echo "       Please install HIPFORT or set environment variable HIPFORT"
   exit 1
fi

HIPFORT_COMPILER=${HIPFORT_COMPILER:-/usr/bin/gfortran}
if [ ! -f $HIPFORT_COMPILER ] ; then
   HIPFORT_COMPILER=`which $HIPFORT_COMPILER`
fi
if [ ! -f $HIPFORT_COMPILER ] ; then
   echo "ERROR: HIPFORT_COMPILER not found at $HIPFORT_COMPILER"
   echo "       Please install $HIPFORT_COMPILER"
   exit 1
fi

# Determine which gfx processor to use, default to Vega (gfx900)
if [ ! $HIPFORT_GPU ] ; then 
   # Use the mygpu in pair with this script, no the pre-installed one.
   HIPFORT_GPU=`$cdir/mygpu -d gfx900`
   if [ "$HIPFORT_GPU" == "" ] ; then 
      HIPFORT_GPU="gfx900"
   fi
fi

# Future versions of ROCm will install the device library in ROCM_DIR/amdgcn/bitcode
if [ -d $ROCM_DIR/amdgcn/bitcode ] ; then
    DEVICE_LIB_PATH=$ROCM_DIR/amdgcn/bitcode
else
    DEVICE_LIB_PATH=$ROCM_DIR/lib
fi

#  Handle all differences between amdgcn and nvidia here
if [ "${HIPFORT_GPU:0:3}" == "sm_" ] ; then 
   TARGET_ARCH="nvptx"
   TARGET_TRIPLE=${TARGET_TRIPLE:-nvptx64-nvidia-cuda}
   TARGET_LIBS="-L$CUDA_PATH/targets/x86_64-linux/lib -lcudart"
   HIPCC_ENV="HIP_PLATFORM=nvcc"
   HIPCC_OPTS="--gpu-architecture=$HIPFORT_GPU"
   # fixme: add test for minimum cuda version here
else 
   if [ ! -f $ROCM_DIR/bin/hipconfig ] ; then
      echo "ERROR: hipconfig not found at $ROCM_DIR/bin/hipconfig"
      echo "       Please install ROCm hip"
      exit 1
   fi
   hipversion=`$ROCM_DIR/bin/hipconfig -v`
   hipver=`echo $hipversion | cut -d"." -f1`
   hiprel=`echo $hipversion | cut -d"." -f2`
   if [ $hipver -lt 3 ] || ( [ $hipver == 3 ] && [ $hiprel -lt 5 ] ) ; then
      echo "ERROR: Minimum required version of hip is in ROCm 3.5"
      echo "       Please update your hip installation"
      exit 1
   fi
   TARGET_TRIPLE=${TARGET_TRIPLE:-amdgcn-amd-amdhsa}
   TARGET_ARCH="amdgcn"
   TARGET_LIBS="-L$ROCM_DIR/lib -lamdhip64 -Wl,-rpath=$ROCM_DIR/lib "
   HIPCC_ENV="HIP_PLATFORM=hcc DEVICE_LIB_PATH=$DEVICE_LIB_PATH HIP_CLANG_PATH=$ROCM_DIR/llvm/bin"
   HIPCC_OPTS="-fno-gpu-rdc -fPIC --offload-arch=$HIPFORT_GPU"
fi

FORTOPT=${FORTOPT:-2}

if [ $VV ]  ; then 
   VERBOSE=true
fi

RUNDATE=`date`

# Parse FIRST_INPUT_FILE_NAME for filetype, directory, and filename
INPUT_FTYPE=${FIRST_INPUT_FILE_NAME##*\.}
INDIR=$(getdname $FIRST_INPUT_FILE_NAME)
FILENAME=${FIRST_INPUT_FILE_NAME##*/}
# FNAME has the filetype extension removed, used for naming intermediate filenames
FNAME=${FILENAME%.*}

if [ -z $OUTFILE ] ; then 
#  Output file not specified so use input directory
   OUTDIR=$INDIR
   if [ $GEN_OBJECT_ONLY ] ; then
      OUTFILE=${FNAME}.o
   else
      OUTFILE="a.out"
   fi
else 
#  Use the specified OUTFILE
   OUTDIR=$(getdname $OUTFILE)
   OUTFILE=${OUTFILE##*/}
fi 

sdir=$(getdname $0)
[ ! -L "$sdir/hipfc" ] || sdir=$(getdname `readlink "$sdir/hipfc"`)
ROCC_DIR=$sdir

TMPNAME="hipfc-tmp-$$"
TMPDIR=${TMPDIR:-/tmp/$TMPNAME}
if [ -d $TMPDIR ] ; then 
   KEEPTDIR=true
else 
   if [ $DRYRUN ] ; then
      echo "mkdir -p $TMPDIR"
   else
      mkdir -p $TMPDIR
      NEWTMPDIR=true
   fi
fi

# Be sure not to delete the output directory
if [ $TMPDIR == $OUTDIR ] ; then 
   KEEPTDIR=true
fi
if [ ! -d $TMPDIR ] && [ ! $DRYRUN ] ; then 
   echo "ERROR:  Directory $TMPDIR does not exist or could not be created"
   exit $DEADRC
fi 
if [ ! -d $OUTDIR ] && [ ! $DRYRUN ]  ; then 
   echo "ERROR:  The output directory $OUTDIR does not exist"
   exit $DEADRC
fi 

UNAMEP=`uname -p`
HOST_TARGET="$UNAMEP-pc-linux-gnu"
if [ "$UNAMEP" == "ppc64le" ] ; then 
  HOST_TARGET="ppc64le-linux-gnu"
fi

FCARGS="-cpp -I$HIPFORT/include/$TARGET_ARCH"
if [ $GEN_OBJECT_ONLY ]  ; then 
   FCARGS=" -c $FCARGS"
   HIPCC_OPTS=" -c $HIPCC_OPTS"
   LINKOPTS=""
else
   LINKOPTS="$TARGET_LIBS -lstdc++ -L$HIPFORT/lib -lhipfort-$TARGET_ARCH"

   if [ $ROCSPARSE ] ; then 
     if [ ! -d $ROCM_DIR/rocsparse/lib ] ; then 
        echo "ERROR:  Directory $ROCM_DIR/rocsparse/lib does not exist"
        echo "        Please install rocsparse"
        exit $DEADRC
     fi
     LINKOPTS="$LINKOPTS -L$ROCM_DIR/rocsparse/lib -lrocsparse"
   fi
   if [ $ROCBLAS ] ; then 
     if [ ! -d $ROCM_DIR/rocblas/lib ] ; then 
        echo "ERROR:  Directory $ROCM_DIR/rocblas/lib does not exist"
        echo "        Please install rocblas"
        exit $DEADRC
     fi
     LINKOPTS="$LINKOPTS -L$ROCM_DIR/rocblas/lib -lrocblas"
   fi
   if [ $HIPBLAS ] ; then 
     if [ ! -d $ROCM_DIR/hipblas/lib ] ; then 
        echo "ERROR:  Directory $ROCM_DIR/hipblas/lib does not exist"
        echo "        Please install hipblas"
        exit $DEADRC
     fi
     LINKOPTS="$LINKOPTS -L$ROCM_DIR/hipblas/lib -lhipblas"
   fi
   if [ $ROCFFT ] ; then 
     if [ ! -d $ROCM_DIR/rocfft/lib ] ; then 
        echo "ERROR:  Directory $ROCM_DIR/rocfft/lib does not exist"
        echo "        Please install rocfft"
        exit $DEADRC
     fi
     LINKOPTS="$LINKOPTS -L$ROCM_DIR/rocfft/lib -lrocfft"
   fi
fi

# Separate hipcc inputs from fortran inputs
__INPUTS=""
__HIPCC_INPUTS=""
for __input_file in `echo $INPUTFILES` ; do
  _ftype=${__input_file##*\.}
  if [ "$_ftype" == "hip" ] ; then
     __HIPCC_INPUTS+="$__input_file"
  elif [ "$_ftype" == "cpp" ] ; then
     __HIPCC_INPUTS+=" $__input_file"
  elif [ "$_ftype" == "hpp" ] ; then
     __HIPCC_INPUTS+=" $__input_file"
  elif [ "$_ftype" == "c" ] ; then
     __HIPCC_INPUTS+=" $__input_file"
  elif [ "$_ftype" == "h" ] ; then
     __HIPCC_INPUTS+=" $__input_file"
  else
     __INPUTS+=" $__input_file"
  fi
done

if [ $GEN_DEBUG ]  ; then
   FCARGS=" -g $FCARGS"
   HIPCC_OPTS=" -g $HIPCC_OPTS"
fi
if [ $VV ]  ; then 
   FCARGS=" -v $FCARGS"
   HIPCC_OPTS=" -v $HIPCC_OPTS"
fi

#  Print Header block
if [ $VV ] ; then 
   echo "#   "
   echo "#Info:  HIPFORT Version:	$PROGVERSION" 
   echo "#Info:  HIPFORT Compiler:	$HIPFORT_COMPILER"
   echo "#Info:  HIPFORT Path:		$HIPFORT"
   echo "#Info:  ROCM Path:		$ROCM_DIR"
   echo "#Info:  CUDA Path:		$CUDA_PATH"
   echo "#Info:  How called:		$HOW_CALLED"
   echo "#Info:  Target GPU:		$HIPFORT_GPU"
   echo "#Info:  Input files:		$INPUTFILES"
   echo "#Info:  Output file:		$OUTDIR/$OUTFILE"
   echo "#Info:  Passthru args:		$PASSTHRUARGS"
   echo "#Info:  Run date:		$RUNDATE"
   [ $KEEPTDIR ] &&  echo "#Info:  Temp dir:	$TMPDIR" 
   echo "#   "
fi 

rc=0
if [ "$__HIPCC_INPUTS" != "" ] ; then
   if [ "$__INPUTS" == "" ] ; then
      __HIPCC_OUTFILE="$OUTDIR/$OUTFILE"
      __HIPCC_LINKOPTS=$LINKOPTS
   else
      HIPCC_OPTS+=" -c"
      __HIPCC_LINKOPTS=""
      __HIPCC_OUTFILE="$TMPDIR/hipcc.o"
      __INPUTS+=" $__HIPCC_OUTFILE"
   fi
   if [ ! -f $ROCM_DIR/bin/hipcc ] ; then
      echo "ERROR:  hipcc compiler not found at $ROCM_DIR/bin/hipcc"
      echo "        Please install hip"
      exit $DEADRC
   fi
   [ $VERBOSE ] && echo "export $HIPCC_ENV"
   export $HIPCC_ENV
   runcmd "$ROCM_DIR/bin/hipcc $HIPCC_OPTS $PASSTHRUARGS $__HIPCC_INPUTS $__HIPCC_LINKOPTS -o $__HIPCC_OUTFILE"
fi

if [ "$__INPUTS" != "" ] ; then
   runcmd "$HIPFORT_COMPILER $FCARGS $__INPUTS $PASSTHRUARGS $LINKOPTS -o $OUTDIR/$OUTFILE"
fi

# cleanup
do_err 0
exit 0
