#!/bin/bash
#
# This shell script (nvcc_wrapper) wraps both the host compiler and
# NVCC, if you are building legacy C or C++ code with CUDA enabled.
# The script remedies some differences between the interface of NVCC
# and that of the host compiler, in particular for linking.
# It also means that a legacy code doesn't need separate .cu files;
# it can just use .cpp files.
#
# Default settings: change those according to your machine.  For
# example, you may have have two different wrappers with either icpc
# or g++ as their back-end compiler.  The defaults can be overwritten
# by using the usual arguments (e.g., -arch=sm_30 -ccbin icpc).

default_arch="sm_35"
#default_arch="sm_50"

#
# The default C++ compiler.
#
host_compiler=${NVCC_WRAPPER_DEFAULT_COMPILER:-"g++"}
#host_compiler="icpc"
#host_compiler="/usr/local/gcc/4.8.3/bin/g++"
#host_compiler="/usr/local/gcc/4.9.1/bin/g++"

#
# Internal variables
#

# C++ files
cpp_files=""

# Host compiler arguments
xcompiler_args=""

# Cuda (NVCC) only arguments
cuda_args=""

# Arguments for both NVCC and Host compiler
shared_args=""

# Argument -c
compile_arg=""

# Argument -o <obj>
output_arg=""

# Linker arguments
xlinker_args=""

# Object files passable to NVCC
object_files=""

# Link objects for the host linker only
object_files_xlinker=""

# Shared libraries with version numbers are not handled correctly by NVCC
shared_versioned_libraries_host=""
shared_versioned_libraries=""

# Does the User set the architecture 
arch_set=0

# Does the user overwrite the host compiler
ccbin_set=0

#Error code of compilation
error_code=0

# Do a dry run without actually compiling
dry_run=0

# Skip NVCC compilation and use host compiler directly
host_only=0
host_only_args=""

# Just run version on host compiler
get_host_version=0

# Enable workaround for CUDA 6.5 for pragma ident 
replace_pragma_ident=0

# Mark first host compiler argument
first_xcompiler_arg=1

temp_dir=${TMPDIR:-/tmp}

# Check if we have an optimization argument already
optimization_applied=0

# Check if we have -std=c++X  or --std=c++X already
stdcxx_applied=0

# Run nvcc a second time to generate dependencies if needed
depfile_separate=0
depfile_output_arg=""
depfile_target_arg=""

# Option to remove duplicate libraries and object files
remove_duplicate_link_files=0

#echo "Arguments: $# $@"

while [ $# -gt 0 ]
do
  case $1 in
  #show the executed command
  --show|--nvcc-wrapper-show)
    dry_run=1
    ;;
  #run host compilation only
  --host-only)
    host_only=1
    ;;
  #get the host version only
  --host-version)
    get_host_version=1
    ;;
  #replace '#pragma ident' with '#ident' this is needed to compile OpenMPI due to a configure script bug and a non standardized behaviour of pragma with macros
  --replace-pragma-ident)
    replace_pragma_ident=1
    ;;
  #remove duplicate link files
  --remove-duplicate-link-files)
    remove_duplicate_link_files=1
    ;;
  #handle source files to be compiled as cuda files
  *.cpp|*.cxx|*.cc|*.C|*.c++|*.cu)
    cpp_files="$cpp_files $1"
    ;;
   # Ensure we only have one optimization flag because NVCC doesn't allow muliple
  -O*)
    if [ $optimization_applied -eq 1 ]; then
       echo "nvcc_wrapper - *warning* you have set multiple optimization flags (-O*), only the first is used because nvcc can only accept a single optimization setting."
    else
       shared_args="$shared_args $1"
       optimization_applied=1
    fi
    ;;
  #Handle shared args (valid for both nvcc and the host compiler)
  -D*)
    unescape_commas=`echo "$1" | sed -e 's/\\\,/,/g'`
    arg=`printf "%q" $unescape_commas`
    shared_args="$shared_args $arg"
    ;;
  -I*|-L*|-l*|-g|--help|--version|-E|-M|-shared|-w)
    shared_args="$shared_args $1"
    ;;
  #Handle compilation argument
  -c)
    compile_arg="$1"
    ;;
  #Handle output argument
  -o)
    output_arg="$output_arg $1 $2"
    shift
    ;;
  # Handle depfile arguments.  We map them to a separate call to nvcc.
  -MD|-MMD)
    depfile_separate=1
    host_only_args="$host_only_args $1"
    ;;
  -MF)
    depfile_output_arg="-o $2"
    host_only_args="$host_only_args $1 $2"
    shift
    ;;
  -MT)
    depfile_target_arg="$1 $2"
    host_only_args="$host_only_args $1 $2"
    shift
    ;;
  #Handle known nvcc args
  --dryrun|--verbose|--keep|--keep-dir*|-G|--relocatable-device-code*|-lineinfo|-expt-extended-lambda|--resource-usage|-Xptxas*)
    cuda_args="$cuda_args $1"
    ;;
  #Handle more known nvcc args
  --expt-extended-lambda|--expt-relaxed-constexpr)
    cuda_args="$cuda_args $1"
    ;;
  #Handle known nvcc args that have an argument
  -rdc|-maxrregcount|--default-stream)
    cuda_args="$cuda_args $1 $2"
    shift
    ;;
  -rdc=*|-maxrregcount*|--maxrregcount*)
    cuda_args="$cuda_args $1"
    ;;
  #Handle c++11
  --std=c++11|-std=c++11|--std=c++14|-std=c++14|--std=c++1y|-std=c++1y|--std=c++17|-std=c++17|--std=c++1z|-std=c++1z)
    if [ $stdcxx_applied -eq 1 ]; then
       echo "nvcc_wrapper - *warning* you have set multiple optimization flags (-std=c++1* or --std=c++1*), only the first is used because nvcc can only accept a single std setting"
    else
       shared_args="$shared_args $1"
       stdcxx_applied=1
    fi
    ;;

  #strip of -std=c++98 due to nvcc warnings and Tribits will place both -std=c++11 and -std=c++98
  -std=c++98|--std=c++98)
    ;;
  #strip of pedantic because it produces endless warnings about #LINE added by the preprocessor
  -pedantic|-Wpedantic|-ansi)
    ;;
  #strip of -Woverloaded-virtual to avoid "cc1: warning: command line option ‘-Woverloaded-virtual’ is valid for C++/ObjC++ but not for C"
  -Woverloaded-virtual)
    ;;
  #strip -Xcompiler because we add it
  -Xcompiler)
    if [ $first_xcompiler_arg -eq 1 ]; then
      xcompiler_args="$2"
      first_xcompiler_arg=0
    else
      xcompiler_args="$xcompiler_args,$2"
    fi
    shift
    ;;
  #strip of "-x cu" because we add that
  -x)
    if [[ $2 != "cu" ]]; then
      if [ $first_xcompiler_arg -eq 1 ]; then
        xcompiler_args="-x,$2"
        first_xcompiler_arg=0
      else
        xcompiler_args="$xcompiler_args,-x,$2"
      fi
    fi
    shift
    ;;
  #Handle -+ (same as -x c++, specifically used for xl compilers, but mutually exclusive with -x. So replace it with -x c++)
  -+)
    if [ $first_xcompiler_arg -eq 1 ]; then
      xcompiler_args="-x,c++"
      first_xcompiler_arg=0
    else
      xcompiler_args="$xcompiler_args,-x,c++"
    fi
    ;;
  #Handle -ccbin (if its not set we can set it to a default value)
  -ccbin)
    cuda_args="$cuda_args $1 $2"
    ccbin_set=1
    host_compiler=$2
    shift
    ;;

  #Handle -arch argument (if its not set use a default) this is the version with = sign
  -arch*|-gencode*)
    cuda_args="$cuda_args $1"
    arch_set=1
    ;;
  #Handle -code argument (if its not set use a default) this is the version with = sign
  -code*)
    cuda_args="$cuda_args $1"
    ;;
  #Handle -arch argument (if its not set use a default) this is the version without = sign
  -arch|-gencode)
    cuda_args="$cuda_args $1 $2"
    arch_set=1
    shift
    ;;
  #Handle -code argument (if its not set use a default) this is the version without = sign
  -code)
    cuda_args="$cuda_args $1 $2"
    shift
    ;;
  #Handle -Xcudafe argument
  -Xcudafe)
    cuda_args="$cuda_args -Xcudafe $2"
    shift
    ;;
  #Handle -Xlinker argument
  -Xlinker)
    xlinker_args="$xlinker_args -Xlinker $2"
    shift
    ;;
  #Handle args that should be sent to the linker
  -Wl,*)
    xlinker_args="$xlinker_args -Xlinker ${1:4:${#1}}"
    host_linker_args="$host_linker_args ${1:4:${#1}}"
    ;;
  #Handle object files: -x cu applies to all input files, so give them to linker, except if only linking
  *.a|*.so|*.o|*.obj)
    object_files="$object_files $1"
    object_files_xlinker="$object_files_xlinker -Xlinker $1"
    ;;
  #Handle object files which always need to use "-Xlinker": -x cu applies to all input files, so give them to linker, except if only linking
  @*|*.dylib)
    object_files="$object_files -Xlinker $1"
    object_files_xlinker="$object_files_xlinker -Xlinker $1"
    ;;
  #Handle shared libraries with *.so.* names which nvcc can't do.
  *.so.*)
    shared_versioned_libraries_host="$shared_versioned_libraries_host $1"
    shared_versioned_libraries="$shared_versioned_libraries -Xlinker $1"
  ;;
  #All other args are sent to the host compiler
  *)
    if [ $first_xcompiler_arg -eq 1 ]; then
      xcompiler_args=$1
      first_xcompiler_arg=0
    else 
      xcompiler_args="$xcompiler_args,$1"
    fi
    ;;
  esac

  shift
done

#Check if nvcc exists
if [ $host_only -ne 1 ]; then
  var=$(which nvcc )
  if [ $? -gt 0 ]; then
    echo "Could not find nvcc in PATH"
    exit $?
  fi
fi


# Only print host compiler version
if [ $get_host_version -eq 1 ]; then
  $host_compiler --version
  exit
fi

#Remove duplicate object files
if [ $remove_duplicate_link_files -eq 1 ]; then
for obj in $object_files
do
  object_files_reverse="$obj $object_files_reverse"
done

object_files_reverse_clean=""
for obj in $object_files_reverse
do
  exists=false
  for obj2 in $object_files_reverse_clean
  do
    if [ "$obj" == "$obj2" ]
    then
      exists=true
      echo "Exists: $obj"
    fi
  done
  if [ "$exists" == "false" ]
  then
    object_files_reverse_clean="$object_files_reverse_clean $obj"
  fi
done

object_files=""
for obj in $object_files_reverse_clean
do
  object_files="$obj $object_files"
done
fi

#Add default host compiler if necessary
if [ $ccbin_set -ne 1 ]; then
  cuda_args="$cuda_args -ccbin $host_compiler"
fi

#Add architecture command
if [ $arch_set -ne 1 ]; then
  cuda_args="$cuda_args -arch=$default_arch"
fi

#Compose compilation command
nvcc_command="nvcc $cuda_args $shared_args $xlinker_args $shared_versioned_libraries"
if [ $first_xcompiler_arg -eq 0 ]; then
  nvcc_command="$nvcc_command -Xcompiler $xcompiler_args"
fi

#Compose host only command
host_command="$host_compiler $shared_args $host_only_args $compile_arg $output_arg $xcompiler_args $host_linker_args $shared_versioned_libraries_host"

#nvcc does not accept '#pragma ident SOME_MACRO_STRING' but it does accept '#ident SOME_MACRO_STRING'
if [ $replace_pragma_ident -eq 1 ]; then
  cpp_files2=""
  for file in $cpp_files
  do
    var=`grep pragma ${file} | grep ident | grep "#"`
    if [ "${#var}" -gt 0 ]
    then
      sed 's/#[\ \t]*pragma[\ \t]*ident/#ident/g' $file > $temp_dir/nvcc_wrapper_tmp_$file
      cpp_files2="$cpp_files2 $temp_dir/nvcc_wrapper_tmp_$file"
    else
      cpp_files2="$cpp_files2 $file"
    fi
  done
  cpp_files=$cpp_files2
  #echo $cpp_files
fi

if [ "$cpp_files" ]; then
  nvcc_command="$nvcc_command $object_files_xlinker -x cu $cpp_files"
else
  nvcc_command="$nvcc_command $object_files"
fi

if [ "$cpp_files" ]; then
  host_command="$host_command $object_files $cpp_files"
else
  host_command="$host_command $object_files"
fi

if [ $depfile_separate -eq 1 ]; then
  # run nvcc a second time to generate dependencies (without compiling)
  nvcc_depfile_command="$nvcc_command -M $depfile_target_arg $depfile_output_arg"
else
  nvcc_depfile_command=""
fi

nvcc_command="$nvcc_command $compile_arg $output_arg"

#Print command for dryrun
if [ $dry_run -eq 1 ]; then
  if [ $host_only -eq 1 ]; then
    echo $host_command
  elif [ -n "$nvcc_depfile_command" ]; then
    echo $nvcc_command "&&" $nvcc_depfile_command
  else
    echo $nvcc_command
  fi
  exit 0
fi

#Run compilation command
if [ $host_only -eq 1 ]; then
  if [ "$NVCC_WRAPPER_SHOW_COMMANDS_BEING_RUN" == "1" ] ; then
    echo "$host_command"
  fi
  $host_command
elif [ -n "$nvcc_depfile_command" ]; then
  if [ "$NVCC_WRAPPER_SHOW_COMMANDS_BEING_RUN" == "1" ] ; then
    echo "$nvcc_command && $nvcc_depfile_command"
  fi
  $nvcc_command && $nvcc_depfile_command
else
  if [ "$NVCC_WRAPPER_SHOW_COMMANDS_BEING_RUN" == "1" ] ; then
    echo "$nvcc_command"
  fi
  $nvcc_command
fi
error_code=$?

#Report error code
exit $error_code
