#!/bin/bash

# hc-kernel-assemble kernel-bitcode kernel-object

# enable bash debugging
KMDBSCRIPT="${KMDBSCRIPT:=0}"
KMDBGRIDLAUNCH="${KMDBGRIDLAUNCH:=0}"

# dump the LLVM bitcode
KMDUMPLLVM="${KMDUMPLLVM:=0}"

if [ $KMDBSCRIPT == "1" ]
then
  set -x
fi

BINDIR=$(dirname "$0")
CLANG=$BINDIR/hcc
LLVM_LINK=$BINDIR/llvm-link
OPT=$BINDIR/opt
LLVM_AS=$BINDIR/llvm-as
LLVM_DIS=$BINDIR/llvm-dis
CLAMP_ASM=$BINDIR/clamp-assemble
LIBPATH=$BINDIR

CXXFLAGS="-std=c++amp -I$BINDIR/../../include -fPIC"

# Add additional flags if using libc++ for C++ runtime
USE_LIBCXX=OFF
if [ "$USE_LIBCXX" == "ON" ]; then
 CXXFLAGS+=" -stdlib=libc++ -I" 
fi

# Set additional CXXFLAGS based on CMAKE_BUILD_TYPE
shopt -s nocasematch
CMAKE_BUILD_TYPE="RELWITHDEBINFO"
if [ $KMDBGRIDLAUNCH == "1" ]; then
    CMAKE_BUILD_TYPE="debug";
fi
case $CMAKE_BUILD_TYPE in
  release)
    CXXFLAGS=$CXXFLAGS" -O3"
    ;;
  relwithdebinfo)
    CXXFLAGS=$CXXFLAGS" -O2 -g"
    ;;
  minsizerel)
    CXXFLAGS=$CXXFLAGS" -Os"
    ;;
  debug)
    CXXFLAGS=$CXXFLAGS" -g"
    ;;
  *)
    CXXFLAGS=$CXXFLAGS
esac

if [ "$#" -ne 2 ]; then
  echo "Usage: $0 kernel-bitcode object" >&2
  exit 1
fi

# inject a new GPU IR to replace the IR from the FE
HC_REPLACE_GPU_FE_LLVM="${HC_REPLACE_GPU_FE_LLVM:=0}"
KERNEL_INPUT=$1
if [ $HC_REPLACE_GPU_FE_LLVM != "0" ]; then
  KERNEL_INPUT=$HC_REPLACE_GPU_FE_LLVM
  echo "Replacing GPU input IR from clang with $KERNEL_INPUT" 
fi

if [ ! -f "$KERNEL_INPUT" ]; then
  echo "kernel-bitcode $KERNEL_INPUT is not valid" >&2
  exit 1
fi

CO="-c -o"

TEMP_DIR=`mktemp -d`
BASENAME=`basename "$2"`
TEMP_NAME="$TEMP_DIR/$BASENAME"

# hip-kernel-assemble goes after hip-host-assemble, so attempt to link object from host
if [ -f "$2" ]; then
  mv "$2" "$TEMP_DIR/$BASENAME.tmp.o"
fi

$LLVM_DIS "$KERNEL_INPUT" -o "$TEMP_NAME.ll"
if [ $KMDUMPLLVM == "1" ]; then
  cp "$TEMP_NAME.ll" ./dump.kernel_input.ll
fi
$OPT -load $LIBPATH/LLVMDirectFuncCall.so -redirect < "$TEMP_NAME.ll" 2>"$TEMP_NAME.kernel_redirect.ll" >/dev/null
if [ $KMDUMPLLVM == "1" ]; then
  cp "$TEMP_NAME.kernel_redirect.ll" ./dump.kernel_redirect.ll
fi
if [[ -s "$TEMP_NAME.kernel_redirect.ll" ]]; then
  $OPT -load $LIBPATH/LLVMWrapperGen.so -gensrc < "$TEMP_NAME.ll" 2>"$TEMP_NAME.camp.cpp" >/dev/null
  if [ $KMDUMPLLVM == "1" ]; then
    cp "$TEMP_NAME.camp.cpp" ./dump.kernel_camp.cpp
  fi
  $LLVM_AS "$TEMP_NAME.kernel_redirect.ll" -o "$TEMP_NAME.kernel_redirect.bc"
  $CLANG $CXXFLAGS "$TEMP_NAME.camp.cpp" $CO "$TEMP_NAME.camp.s" -emit-llvm
  $CLANG $CXXFLAGS "$TEMP_NAME.camp.cpp" $CO "$TEMP_NAME.camp.o"
  objcopy -R .kernel "$TEMP_NAME.camp.o"
  $LLVM_LINK "$TEMP_NAME.kernel_redirect.bc" "$TEMP_NAME.camp.s" -o "$TEMP_NAME.link.bc"
  $CLAMP_ASM "$TEMP_NAME.link.bc" "$TEMP_NAME.camp.o"
else
  ln -s "$KERNEL_INPUT" "$1.bc"
  $CLAMP_ASM "$1.bc" "$TEMP_NAME.camp.o"
fi

if [ -f "$TEMP_DIR/$BASENAME.tmp.o" ]; then
  ld -r --allow-multiple-definition "$TEMP_DIR/$BASENAME.tmp.o" "$TEMP_NAME.camp.o" -o "$2"
else
  mv "$TEMP_NAME.camp.o" "$2"
fi

if [ $KMDBGRIDLAUNCH == "1" ]; then
  echo "db: KMDBGRIDLAUNCH=1 and temps saved in $TEMP_DIR"
else
   rm -rf "$TEMP_DIR"
fi
