#!/bin/bash

# enable bash debugging
KMDBSCRIPT="${KMDBSCRIPT:=0}"
if [ $KMDBSCRIPT == "1" ]
then
  set -x
fi

# check number of arguments
if [ "$#" -eq 0 ]; then
  echo "$0 is NOT recommended to be directly used" >&2
  exit 1
fi

# enable bundle dumping
KMDUMPBUNDLE="${KMDUMPBUNDLE:=0}"

BINDIR=$(dirname $0)
LINK=$BINDIR/llvm-link
OPT=$BINDIR/opt
CLAMP_DEVICE=$BINDIR/clamp-device
CLAMP_EMBED=$BINDIR/clamp-embed
CLANG_OFFLOAD_BUNDLER=$BINDIR/clang-offload-bundler

################
# Verbose flag
################

VERBOSE=0

################
# GPU targets
# This array could be populated via 3 methods:
# 1) clang driver adds multiple --amdgpu-target=<GPU target> options
# 2) HCC_AMDGPU_TARGET env var, multiple targets are delimited by ",".
#    Example:
#
#    export HCC_AMDGPU_TARGET=gfx701,gfx801,gfx802,gfx803
#    # ISA for Hawaii(gfx701), Carrizo(gfx801), Tonga(gfx802) and Fiji(gfx803)
#    # would be produced
#    hcc `hcc-config --cxxflags --ldflags` foo.cpp
#
# 3) CMake AMDGPU_TARGET config value. This is the last resort options.
################

AMDGPU_TARGET_ARRAY=()

################
# link
################

LINK_KERNEL_ARGS=()
LINK_HOST_ARGS=()
LINK_CPU_ARG=()
LINK_OTHER_ARGS=()

TEMP_DIR=`mktemp -d`

# a file which contains the list of __cxxamp_serialize symbols in each CPU object file
CXXAMP_SERIALIZE_SYMBOL_FILE=$TEMP_DIR/symbol.txt
touch $CXXAMP_SERIALIZE_SYMBOL_FILE

# find object file
_find_object() {
  local FILE=$1
  local ret=${FILE%.o}

  if [ -e $FILE ]; then
    local file_output=`file $FILE | grep 'ELF 64-bit LSB  relocatable, x86-64'`
    local readelf_output=`readelf -h $FILE | grep 'Relocatable file'`

    if [ ! -z "$file_output" ] && [ ! -z "$readelf_output" ]; then
      # remove postfix
      ret=${FILE%.*}
    fi
  fi

  echo $ret
}

ARGS="$@"
STATIC_LIB_LIST=()
TEMP_AR_DIRS=()

# detect the verbose flags before doing anything
if [[ "$ARGS" =~ --verbose ]]; then
  VERBOSE=1
fi
#very verbose
#VERBOSE=2



# gather a list of library search paths
LIB_SEARCH_PATHS=()
for ARG in "$@"
do
  # matches -L<path>
  if [[ "$ARG" =~ ^-L.* ]]; then
    REAL_PATH="$(readlink -f "${ARG:2}")"
    if [ $VERBOSE == 2 ]; then
      echo "add library path: ${ARG:2}, canonical path: $REAL_PATH"
    fi
    LIB_SEARCH_PATHS+=( "$REAL_PATH" )
  fi
done

# gather input arguments from linker command files
INPUT_ARGUMENTS=()
for ARG in "$@"
do
  # matches @<path>
  if [[ "$ARG" =~ ^@.* ]]; then
    REAL_PATH="$(readlink -f "${ARG:1}")"
    if [ $VERBOSE == 2 ]; then
      echo "add linker command path: ${ARG:1}, canonical path: $REAL_PATH"
    fi

    # read from linker command file
    IFS=$'\n' read -d '' -r -a LINES < "$REAL_PATH"
    for LINE in ${LINES[@]}
    do
      if [ $VERBOSE == 2 ]; then
        echo "add linker command: $LINE"
      fi
      INPUT_ARGUMENTS+=($LINE)
    done
  else
    INPUT_ARGUMENTS+=("$ARG")
  fi
done

for ARG in "${INPUT_ARGUMENTS[@]}"
do

  case $ARG in
    ######################
    # Parse AMDGPU target
    ######################
    --amdgpu-target=*)
    AMDGPU_TARGET_ARRAY+=("${ARG#*=}")
    continue
    ;;
    ################################################
    # Parse dump options and export them for called
    # scripts e.g. clamp-device
    ################################################
    -dump-isa)
    export KMDUMPISA=1
    continue
    ;;
    -dump-llvm)
    export KMDUMPLLVM=1
    continue
    ;;
    -dump-dir=*)
    export KMDUMPDIR="${ARG#*=}"
    continue
    ;;
  esac

  #####################################
  # detect object or static library
  #####################################

  OBJS_TO_PROCESS=()


  if [[ "$ARG" =~ [^[:space:]]+\.cpu$ ]]; then

    cp "$ARG" $TEMP_DIR/kernel_cpu.o
    LINK_CPU_ARG+=( "$TEMP_DIR/kernel_cpu.o" )

  elif [[ "$ARG" =~ [^[:space:]]+\.o$ ]]; then
    # detected a .o file
    if [ $VERBOSE == 2 ]; then
      echo "detect object file to process further: $ARG"
    fi

    OBJS_TO_PROCESS+=( "$ARG" )
  elif [[ "$ARG" =~ ^-l[^[:space:]]+$ ]] || [[ "$ARG" =~ [^[:space:]]+.a$ ]] || [[ "$ARG" =~ [^[:space:]]+.lo$ ]]; then

    # proccess a static library

    DETECTED_STATIC_LIBRARY=""

    # detected whether it's an -l<library naem> option
    if [[ "$ARG" =~ ^-l[^[:space:]]+$ ]]; then

      # expand the option into a library name
      STATIC_LIB_NAME="lib${ARG:2}.a"

      if [ $VERBOSE == 2 ]; then
        echo "looking for static library $STATIC_LIB_NAME"
      fi

      # look for the static library in the library search paths
      for LIB_PATH in "${LIB_SEARCH_PATHS[@]}"
      do
        FULL_LIB_PATH="$LIB_PATH"
        FULL_LIB_PATH+="/"
        FULL_LIB_PATH+="$STATIC_LIB_NAME"
        FULL_LIB_PATH=$(readlink -f "$FULL_LIB_PATH")
        if [ $VERBOSE == 2 ]; then
          echo "trying to detect $FULL_LIB_PATH"
        fi
        if [ -f "$FULL_LIB_PATH" ]; then
          if [ $VERBOSE == 2 ]; then
            echo "$FULL_LIB_PATH detected"
          fi
          DETECTED_STATIC_LIBRARY="$FULL_LIB_PATH";
          break;
        fi
      done
    else
      # this is .a or .lo static library file specified at the commad line
      if [ -f "$ARG" ]; then
        FULL_LIB_PATH=$(readlink -f "$ARG")
        if [ $VERBOSE == 2 ]; then
          echo "use .a / .lo specified at: $FULL_LIB_PATH"
        fi
        DETECTED_STATIC_LIBRARY="$FULL_LIB_PATH"
      fi
    fi  # if [[ "$ARG" =~ ^-l[^[:space:]]+$ ]]; then


    # check for duplicated static library options
    if [[ $DETECTED_STATIC_LIBRARY != "" ]]; then
      for LIB in "${STATIC_LIB_LIST[@]}"
      do
        if [[ $LIB == $DETECTED_STATIC_LIBRARY ]]; then
          # this library has already been looked at, skip it
          DETECTED_STATIC_LIBRARY=""
          break;
        fi
      done
      if [[ $DETECTED_STATIC_LIBRARY != "" ]]; then
        STATIC_LIB_LIST+=( "$DETECTED_STATIC_LIBRARY" )
      fi
    fi


    KERNEL_UNDETECTED="1"
    if [[ $DETECTED_STATIC_LIBRARY != "" ]]; then

      # we found a static library library
      if [ $VERBOSE == 2 ]; then
        echo "processing static library $DETECTED_STATIC_LIBRARY"
      fi

      # detect whether the objects in the static library contain a .kernel section
      KERNEL_UNDETECTED=`objdump -t "$DETECTED_STATIC_LIBRARY" | grep -q "\.kernel"; echo $?`
      if [[ $KERNEL_UNDETECTED == "0" ]]; then

        # .kernel section detected, extract the objects from the archieve

        if [ $VERBOSE == 2 ]; then
          echo "kernel detected in $DETECTED_STATIC_LIBRARY"
        fi

        CURRENT_DIR=$PWD
        # extract the archive
        FILE=`basename $DETECTED_STATIC_LIBRARY`
        AR_TEMP_DIR="$TEMP_DIR"
        AR_TEMP_DIR+="/"
        AR_TEMP_DIR+="$FILE"

        if [ $VERBOSE == 2 ]; then
          echo "creating temp dir: $AR_TEMP_DIR"
        fi

        mkdir -p "$AR_TEMP_DIR"
        TEMP_AR_DIRS+=( "$AR_TEMP_DIR" )
        cd "$AR_TEMP_DIR"
        `ar x "$DETECTED_STATIC_LIBRARY"`

        cd "$CURRENT_DIR"

        # store all the extract objects to process further
        OBJS_TO_PROCESS=($(ls "$AR_TEMP_DIR"/*.o))

      fi  # if [[ $KERNEL_UNDETECTED == "0" ]]; then
    fi  # if [[ $DETECTED_STATIC_LIBRARY != "" ]]; then
  elif [ -f "$ARG" ]; then
    # an object file but doesn't have an .o extension??
    file_output=`file "$ARG" | grep 'ELF 64-bit LSB  relocatable, x86-64'`
    readelf_output=`readelf -h "$ARG" | grep 'Relocatable file'`
    if [ ! -z "$file_output" ] && [ ! -z "$readelf_output" ]; then
      OBJS_TO_PROCESS+=( "$ARG" )
    fi
  fi


  # no objects to further process, pass the original args down to the host linker
  if [ ${#OBJS_TO_PROCESS[@]} == 0 ]; then
    # no objects to further process, pass the original args down to the host linker
    if [ $VERBOSE == 2 ]; then
      echo "passing down link args: $ARG"
    fi
    LINK_OTHER_ARGS+=( "$ARG" )
    continue
  fi

  # processs the objects we put aside
  for OBJ in "${OBJS_TO_PROCESS[@]}"
  do
    if [ $VERBOSE == 2 ]; then
      echo "processing $OBJ"
    fi

    # detect whether the objects in the static library contain a .kernel section
    KERNEL_UNDETECTED=`objdump -t "$OBJ" | grep -q "\.kernel"; echo $?`
    if [[ $KERNEL_UNDETECTED == "0" ]]; then

      FILE=`basename "$OBJ"` # remove path
      FILENAME="${FILE%.*}"
      KERNEL_FILE="$TEMP_DIR/$FILENAME.kernel.bc"
      HOST_FILE="$TEMP_DIR/$FILENAME.host.o"

      # extract kernel section
      objcopy -O binary -j .kernel "$OBJ" "$KERNEL_FILE"

      # extract host section
      objcopy -R .kernel "$OBJ" "$HOST_FILE"

      # strip all symbols specified in symbol.txt from $HOST_FILE
      objcopy @$CXXAMP_SERIALIZE_SYMBOL_FILE "$HOST_FILE" "$HOST_FILE.new" 2> /dev/null
      if [ -f "$HOST_FILE.new" ]; then
        mv "$HOST_FILE.new" "$HOST_FILE"
      fi

      # find cxxamp_serialize symbols and save them into symbol.txt
      objdump -t "$HOST_FILE" -j .text 2> /dev/null | grep "g.*__cxxamp_serialize" | awk '{print "-L"$6}' >> $CXXAMP_SERIALIZE_SYMBOL_FILE

      LINK_KERNEL_ARGS+=( "$KERNEL_FILE" )
      LINK_HOST_ARGS+=( "$HOST_FILE" )
    else
      LINK_OTHER_ARGS+=( "$OBJ" )
    fi


  done # for OBJ in "${OBJS_TO_PROCESS[@]}"
done

# AMDGPU_TARGET_ARRAY could be overridden by HCC_AMDGPU_TARGET env variable
if [ -n "$HCC_AMDGPU_TARGET" ]; then
  # tokenize HCC_AMDGPU_TARGET
  IFS=','
  HCC_AMDGPU_TARGET_TOKENS=($HCC_AMDGPU_TARGET)
  unset IFS
  # parse individual token, override AMDGPU_TARGET_ARRAY
  iter=0
  for HCC_AMDGPU_TARGET_TOKEN in ${HCC_AMDGPU_TARGET_TOKENS[@]}; do
    AMDGPU_TARGET_ARRAY[$iter]=$HCC_AMDGPU_TARGET_TOKEN
    ((iter++))
  done
fi

# In case neither Clang Driver nor HCC_AMDGPU_TARGET env var specify target,
# use default value in CMake configuration
if [ ${#AMDGPU_TARGET_ARRAY[@]} == 0 ]; then
  AMDGPU_TARGET_ARRAY=(auto)
fi

if [ $VERBOSE != 0 ]; then
  echo "AMDGPU target array: "${AMDGPU_TARGET_ARRAY[@]}
  echo ""
fi

if [ $VERBOSE != 0 ]; then
  echo "new kernel args: ${LINK_KERNEL_ARGS[@]}"
  echo ""
  echo "new host args: ${LINK_HOST_ARGS[@]}"
  echo ""
  echo "new other args: ${LINK_OTHER_ARGS[@]}"
  echo ""
fi

# linker return value
ret=0

# only do kernel lowering if there are objects given
if [ ${#LINK_KERNEL_ARGS[@]} != 0 ]; then

  # combine kernel sections together
  $LINK "${LINK_KERNEL_ARGS[@]}" | $OPT -always-inline - -o $TEMP_DIR/kernel.bc

  # lower to GCN ISA
  if [ $VERBOSE == 1 ]; then
    echo "Generating AMD GCN kernel"
  fi

  # touch an empty object for host part, to accomodate rule required by
  # clang-offload-bundler
  touch $TEMP_DIR/__empty.o

  # invoke clang-offload-bundler to create kernel bundle
  CLANG_OFFLOAD_BUNDLER_INPUTS_ARGS="-inputs=$TEMP_DIR/__empty.o"
  CLANG_OFFLOAD_BUNDLER_TARGETS_ARGS="-targets=host-x86_64-unknown-linux"

  # for each GPU target, lower to GCN ISA in HSACO format
  for AMDGPU_TARGET in ${AMDGPU_TARGET_ARRAY[@]}; do
    $CLAMP_DEVICE $TEMP_DIR/kernel.bc $TEMP_DIR/kernel-$AMDGPU_TARGET.hsaco --amdgpu-target=$AMDGPU_TARGET

    # error handling
    ret=$?
    if [ $ret != 0 ]; then
      exit $ret
    fi

    # augment arguments to clang-offload-bundler
    CLANG_OFFLOAD_BUNDLER_INPUTS_ARGS+=",$TEMP_DIR/kernel-$AMDGPU_TARGET.hsaco"
    CLANG_OFFLOAD_BUNDLER_TARGETS_ARGS+=",hcc-amdgcn--amdhsa-$AMDGPU_TARGET"
  done

  # invoke clang-offload-bundler
  $CLANG_OFFLOAD_BUNDLER -type=o $CLANG_OFFLOAD_BUNDLER_INPUTS_ARGS $CLANG_OFFLOAD_BUNDLER_TARGETS_ARGS -outputs=$TEMP_DIR/kernel.bundle

  # error handling
  ret=$?
  if [ $ret != 0 ]; then
    exit $ret
  fi

  if [ $KMDUMPBUNDLE == "1" ]; then
    cp $TEMP_DIR/kernel.bundle ./dump.bundle
  fi

  # build a new kernel object
  pushd . > /dev/null
  cd $TEMP_DIR
  $CLAMP_EMBED kernel.bundle kernel_hsa.o
  popd > /dev/null

  # link everything together
  ld --allow-multiple-definition -L`hcc -print-resource-dir`/lib $TEMP_DIR/kernel_hsa.o "${LINK_HOST_ARGS[@]}" "${LINK_CPU_ARG[@]}" "${LINK_OTHER_ARGS[@]}"
  ret=$?

fi # if [ -n $LINK_KERNEL_ARGS ];

# remove temp files
if [ -e $TEMP_DIR/kernel_hsa.o ]; then
  rm $TEMP_DIR/kernel_hsa.o
fi

if [ -e $TEMP_DIR/kernel_cpu.o ]; then
  rm $TEMP_DIR/kernel_cpu.o
fi

if [ -e $TEMP_DIR/__empty.o ]; then
  rm $TEMP_DIR/__empty.o
fi

if [ -e $TEMP_DIR/kernel.bundle ]; then
  rm $TEMP_DIR/kernel.bundle
fi

rm -f $TEMP_DIR/kernel-*.hsaco

if [ -e $TEMP_DIR/kernel.bc ]; then
  rm $TEMP_DIR/kernel.bc
fi

if [ ${#LINK_KERNEL_ARGS[@]} != 0 ]; then
  rm "${LINK_KERNEL_ARGS[@]}" # individual kernels
fi

if [ -n "$LINK_HOST_ARGS" ]; then
  rm "${LINK_HOST_ARGS[@]}" # individual host codes
fi

if [ -e $CXXAMP_SERIALIZE_SYMBOL_FILE ]; then
  rm $CXXAMP_SERIALIZE_SYMBOL_FILE # __cxxamp_serialize symbols
fi

for TD in "${TEMP_AR_DIRS[@]}"
do
  rm -rf $TD
done

if [ -d $TEMP_DIR ]; then
  rm -f $TEMP_DIR/*
  rmdir $TEMP_DIR
fi

# return value
exit $ret
